diff --git a/fern/versions/latest/pages/agent-server/trajectory-telemetry.mdx b/fern/versions/latest/pages/agent-server/trajectory-telemetry.mdx
new file mode 100644
index 0000000000..9a8343e8cb
--- /dev/null
+++ b/fern/versions/latest/pages/agent-server/trajectory-telemetry.mdx
@@ -0,0 +1,96 @@
+---
+title: "Trajectory Telemetry"
+description: "The agent's Response is the trajectory: agent-step structure, per-call token stats, and tool timing on the native contract"
+position: 4
+---
+
+Agent loops run many model calls (**agent steps**) with tool executions between them, but return one flattened `Response` — one item list, one summed usage. That is the right *composition interface* (agents serve the same `/v1/responses` contract as models), but it used to discard the execution structure: which items each model call produced, what each call cost, how long each tool ran. Trajectory telemetry records that structure **on the Response contract itself** — the `NeMoGymResponse` is the trajectory entity, and nothing is stored twice.
+
+
+Terminology: an **agent step** is one interaction with the environment through the model — a single LLM generation plus the orchestration of its tool calls and their outputs (one iteration of an agent loop; `max_steps` counts these). A **turn** is a full cycle of control, from user input until the agent hands control back, containing one or more agent steps. Some harnesses call agent steps "turns" (Claude Code's `num_turns` / `--max-turns` count model calls); the telemetry normalizes that to `num_agent_steps`.
+
+
+## Telemetry on the contract
+
+The extension follows the same pattern the `*ForTraining` item variants use for token IDs (`nemo_gym/openai_utils.py`):
+
+**On the items** — `response.output` is the episode's lossless, execution-ordered item list. Model-produced items are `*WithAgentTelemetry` variants carrying `agent_step_no` (which model call produced them); `function_call_output` items also carry `execution` — independent per-call timing, so parallel tool calls stay independently timed:
+
+```json
+{"type": "function_call_output", "agent_step_no": 1, "call_id": "toolu_…", "output": "…",
+ "execution": {"started_at": "…", "completed_at": "…", "duration_ms": 532.3,
+ "error": null, "extra": {"interrupted": false}}}
+```
+
+Reasoning is a native `reasoning` item, function calls appear in issue order with outputs in arrival order, unresolved calls are kept, mid-episode user messages are plain message items (turn boundaries), and a compaction summary is a `NeMoGymContextBoundaryMessage`. The task's initial prompt is not repeated — it stays in `responses_create_params.input`.
+
+**On the Response** — two optional fields that model servers never populate:
+
+```json
+"generations": [
+ {"agent_step_no": 1, "model": "…", "stop_reason": "tool_use",
+ "response_id": "msg_01…", "request_id": "req_…", "ended_at": "…",
+ "usage": {"input_tokens": 12, "input_tokens_details": {"cached_tokens": 9000}, "…": "…"},
+ "provider_usage": {"cache_read_input_tokens": 9000, "cache_creation_input_tokens": 512, "…": "…"}}
+],
+"agent_telemetry": {"schema_version": "1.0", "agent": "claude_code_agent", "source": "transcript",
+ "session_id": "…", "num_agent_steps": 2, "duration_ms": 8123.0,
+ "total_cost_usd": 0.021, "dropped_records": {}}
+```
+
+`generations[]` carries what is generation-shaped: native per-call `NeMoGymResponseUsage` (cache detail included, deduplicated per API message), the provider's raw usage verbatim (so fields with no native slot are never lost), and provider identity per model call. `agent_telemetry` carries run-level facts: provenance (`source`), `num_agent_steps`, duration/cost, and `dropped_records` — counters for events seen but not represented (e.g. subagent sidechains), so "nothing happened" is distinguishable from "events were dropped".
+
+
+Because the telemetry rides the contract, it survives every hop — verify requests, resources servers, rollout rows — without any side-channel. Plain model-server payloads still validate to the plain item classes: the telemetry variants have required discriminating fields, so nothing changes for existing responses.
+
+
+## Reconstructing what each model call saw
+
+The output is append-only in execution order, so any step's exact model-visible input is derivable from the single stored copy:
+
+```python
+from nemo_gym.trajectory import reconstruct_model_input, to_response_create_params, agent_step_slices
+
+items = reconstruct_model_input(
+ response.output,
+ agent_step_no=3,
+ base_input=responses_create_params.input, # the initial prompt lives only here
+)
+```
+
+This is compaction-aware: when a `NeMoGymContextBoundaryMessage` precedes the step, reconstruction restarts at its summary instead of prepending `base_input` — exactly what the post-compaction model call actually saw. `agent_step_slices()` iterates `(agent_step_no, items)` groups for per-step analysis.
+
+## Capture modes: same contract, two ways to produce it
+
+- **In-process loops** (a `responses()` loop the agent owns): drive `TrajectoryBuilder` as the loop executes — exact boundaries, real per-call usage, *measured* tool timing.
+- **Black-box harnesses** (external CLIs like Claude Code): an adapter reconstructs everything post-hoc from the harness's artifacts. The `claude_code_agent` is the reference adapter: it harvests the session transcript Claude Code writes under its ephemeral `CLAUDE_CONFIG_DIR/projects/` and parses it once.
+
+The `source` field labels provenance (`"transcript"`, `"stream_json"`, `"in_process"`). Telemetry a source cannot observe stays `null` — never approximated: the stream-json fallback carries no timestamps, so tool durations are `null` rather than guessed.
+
+
+For the Claude Code agent, capture is on by default (`capture_trajectory: true`). Set it to `false` to skip transcript harvesting; the response then carries no `generations`/`agent_telemetry`.
+
+
+## Adapting another harness
+
+Parse your harness's artifacts in execution order and drive the builder; it owns the shared semantics (agent-step numbering and tags, model-call deduplication by `response_id`, call/execution correlation, orphan handling):
+
+```python
+from nemo_gym.trajectory import TrajectoryBuilder
+
+builder = TrajectoryBuilder(agent="my_agent", source="in_process")
+builder.start_agent_step(response_id=..., request_id=..., model=..., provider_usage=...)
+builder.add_output_text("...") # native message item, tagged
+builder.add_reasoning("...") # native reasoning item, tagged
+builder.add_tool_call(call_id, name, arguments)
+builder.add_tool_result(call_id, output, started_at=..., completed_at=..., error=..., extra=...)
+builder.add_user_message("...") # mid-episode turns only — never the initial prompt
+builder.add_context_boundary(summary="...")
+builder.set_run_totals(num_agent_steps=..., duration_ms=..., total_cost_usd=..., provider_usage=...)
+
+output_items, generations, agent_telemetry = builder.build()
+```
+
+
+The trajectory is observability-only: it carries no token IDs or logprobs, and re-tokenizing its text does not reproduce the tokens the policy sampled (retokenization drift). For training, token capture belongs at the model server, joined via each generation's `response_id`/`request_id` — the items then upgrade to their `*ForTraining` variants in place.
+
diff --git a/nemo_gym/openai_utils.py b/nemo_gym/openai_utils.py
index db55a5ec2c..cc16ed98f8 100644
--- a/nemo_gym/openai_utils.py
+++ b/nemo_gym/openai_utils.py
@@ -230,7 +230,76 @@ class NeMoGymResponseReasoningItemForTraining(NeMoGymResponseReasoningItem, Toke
}
+########################################
+# Agent telemetry (issue #1867)
+#
+# Agent servers run many model calls (agent steps) with tool executions between them and
+# return one flattened Response. These extensions let the Response itself carry the
+# execution structure — following the ForTraining pattern: item variants add telemetry
+# fields on the items, and the Response gains optional per-generation records. Model
+# servers never populate any of this; plain payloads validate to the plain classes.
+#
+# Terminology: an "agent step" is one model call plus the orchestration of its tool
+# calls (one agent-loop iteration; some providers call these "turns").
+########################################
+
+
+class NeMoGymToolExecutionError(BaseModel):
+ message: str
+ data: Optional[Dict[str, Any]] = None
+
+
+class NeMoGymToolExecution(BaseModel):
+ """Execution telemetry for one tool call, carried on its function_call_output item."""
+
+ started_at: Optional[str] = None
+ completed_at: Optional[str] = None
+ duration_ms: Optional[float] = None
+ error: Optional[NeMoGymToolExecutionError] = None
+ # Provider-specific execution metadata, verbatim (e.g. exit codes, sandbox flags).
+ extra: Optional[Dict[str, Any]] = None
+
+
+class NeMoGymAgentStepTagMixin(BaseModel):
+ # 1-based index of the agent step (model call) that produced/observed this item.
+ # Required on the telemetry variants so plain payloads deterministically validate to
+ # the plain classes in the item union.
+ agent_step_no: int
+
+
+class NeMoGymResponseOutputMessageWithAgentTelemetry(NeMoGymResponseOutputMessage, NeMoGymAgentStepTagMixin):
+ pass
+
+
+class NeMoGymResponseFunctionToolCallWithAgentTelemetry(NeMoGymResponseFunctionToolCall, NeMoGymAgentStepTagMixin):
+ pass
+
+
+class NeMoGymResponseReasoningItemWithAgentTelemetry(NeMoGymResponseReasoningItem, NeMoGymAgentStepTagMixin):
+ pass
+
+
+class NeMoGymFunctionCallOutputWithAgentTelemetry(NeMoGymFunctionCallOutput, NeMoGymAgentStepTagMixin):
+ execution: Optional[NeMoGymToolExecution] = None
+
+
+class NeMoGymContextBoundaryMessage(NeMoGymEasyInputMessage):
+ """A compaction marker: this (summary) message replaced all earlier context.
+
+ ``context_boundary`` is required so plain user messages validate to the plain class.
+ """
+
+ context_boundary: Literal[True]
+
+
NeMoGymResponseInputItem = Union[
+ # Agent-telemetry variants first: their required fields (agent_step_no /
+ # context_boundary) make plain payloads fall through to the plain classes below.
+ NeMoGymResponseOutputMessageWithAgentTelemetry,
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry,
+ NeMoGymFunctionCallOutputWithAgentTelemetry,
+ NeMoGymResponseReasoningItemWithAgentTelemetry,
+ NeMoGymContextBoundaryMessage,
NeMoGymEasyInputMessage,
NeMoGymMessage,
NeMoGymResponseOutputMessage,
@@ -303,9 +372,55 @@ class NeMoGymResponseUsage(ResponseUsage):
output_tokens_details: NeMoGymResponseOutputTokensDetails
+class NeMoGymGeneration(BaseModel):
+ """Per-agent-step telemetry: one record per model call inside an agent's Response.
+
+ Items in ``output`` reference their step via ``agent_step_no`` (on the
+ ``*WithAgentTelemetry`` item variants); this record carries what is
+ generation-shaped rather than item-shaped.
+ """
+
+ agent_step_no: int
+ model: Optional[str] = None
+ stop_reason: Optional[str] = None
+ # Provider identity of this model call (message id / HTTP request id).
+ response_id: Optional[str] = None
+ request_id: Optional[str] = None
+ ended_at: Optional[str] = None
+ # Native per-call token stats; the provider's raw usage dict is preserved verbatim so
+ # fields with no native slot (e.g. cache_creation_input_tokens) are never lost.
+ usage: Optional[NeMoGymResponseUsage] = None
+ provider_usage: Optional[Dict[str, Any]] = None
+
+
+class NeMoGymAgentTelemetry(BaseModel):
+ """Run-level agent telemetry carried on the Response (issue #1867)."""
+
+ schema_version: str = "1.0"
+ agent: Optional[str] = None
+ # Which artifact the telemetry was derived from (e.g. "transcript", "stream_json",
+ # "in_process") — fidelity is data, never a schema fork.
+ source: Optional[str] = None
+ session_id: Optional[str] = None
+ # Agent steps (model calls) as reported by the provider. Some providers call these
+ # "turns" (Claude Code's num_turns / --max-turns count model calls).
+ num_agent_steps: Optional[int] = None
+ duration_ms: Optional[float] = None
+ total_cost_usd: Optional[float] = None
+ # The provider's own end-of-run usage report, verbatim, for cross-checking.
+ provider_usage: Optional[Dict[str, Any]] = None
+ # Events seen but not represented (e.g. {"sidechain": 3}), so consumers can tell
+ # "nothing happened" from "events were dropped".
+ dropped_records: Dict[str, int] = {}
+
+
class NeMoGymResponse(Response):
output: List[NeMoGymResponseOutputItem]
usage: Optional[NeMoGymResponseUsage] = None
+ # Agent telemetry: populated by agent servers (whose Response carries a whole
+ # episode), never by model servers (whose Response is a single generation).
+ generations: Optional[List[NeMoGymGeneration]] = None
+ agent_telemetry: Optional[NeMoGymAgentTelemetry] = None
########################################
diff --git a/nemo_gym/trajectory.py b/nemo_gym/trajectory.py
new file mode 100644
index 0000000000..6b8ba7f299
--- /dev/null
+++ b/nemo_gym/trajectory.py
@@ -0,0 +1,400 @@
+# 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.
+
+"""Agent trajectory support over Gym's native Responses contract (issue #1867).
+
+The trajectory entity **is** the ``NeMoGymResponse``. Its ``output`` item list is the
+episode — lossless and in execution order — and the telemetry rides on the contract
+itself (see ``nemo_gym.openai_utils``), following the same pattern the ``*ForTraining``
+variants use for token IDs:
+
+- **On the items**: model-produced items are the ``*WithAgentTelemetry`` variants
+ carrying ``agent_step_no`` (which model call produced/observed them);
+ ``function_call_output`` items additionally carry ``execution``
+ (:class:`~nemo_gym.openai_utils.NeMoGymToolExecution`: independent per-call timing,
+ errors, provider execution metadata); a compaction summary is a
+ :class:`~nemo_gym.openai_utils.NeMoGymContextBoundaryMessage`.
+- **On the Response**: ``generations`` (one
+ :class:`~nemo_gym.openai_utils.NeMoGymGeneration` per agent step: native per-call
+ usage, provider identity, ``stop_reason``) and ``agent_telemetry`` (run-level:
+ provenance ``source``, ``session_id``, ``num_agent_steps``, duration/cost, verbatim
+ provider usage, ``dropped_records``). Model servers leave both ``None``.
+
+Terminology: an **agent step** is one interaction with the environment through the
+model — a single LLM generation plus the orchestration of its tool calls and their
+outputs (one agent-loop iteration; ``max_steps`` counts these). A **turn** is a full
+cycle of control, from user input until the agent hands control back, containing one or
+more agent steps; turns are derivable from ``role: "user"`` items in the output.
+
+Nothing is stored twice: the task's initial prompt lives only in
+``responses_create_params.input`` (pass it to :func:`reconstruct_model_input` as
+``base_input``), and per-step grouping is the ``agent_step_no`` tag rather than a copy
+of the items. Because the output is append-only in execution order, any step's exact
+model-visible input is derivable — compaction-aware — from the single stored copy.
+
+Capture is contract-agnostic: in-process loops drive :class:`TrajectoryBuilder` as they
+execute (exact boundaries, measured tool timing); black-box harness wrappers reconstruct
+post-hoc from artifacts (see ``responses_api_agents/claude_code_agent/trajectory.py``).
+Fidelity differences surface as data — ``source`` labels provenance, unobservable
+telemetry stays ``None`` (never fabricated), ``dropped_records`` counts events seen but
+not represented.
+"""
+
+from datetime import datetime
+from typing import Any, Iterator, Optional, Union
+
+from nemo_gym.openai_utils import (
+ NeMoGymAgentTelemetry,
+ NeMoGymContextBoundaryMessage,
+ NeMoGymEasyInputMessage,
+ NeMoGymFunctionCallOutputWithAgentTelemetry,
+ NeMoGymGeneration,
+ NeMoGymResponseCreateParamsNonStreaming,
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry,
+ NeMoGymResponseInputItem,
+ NeMoGymResponseInputTokensDetails,
+ NeMoGymResponseOutputMessageWithAgentTelemetry,
+ NeMoGymResponseOutputText,
+ NeMoGymResponseOutputTokensDetails,
+ NeMoGymResponseReasoningItemWithAgentTelemetry,
+ NeMoGymResponseUsage,
+ NeMoGymSummary,
+ NeMoGymToolExecution,
+ NeMoGymToolExecutionError,
+)
+
+
+def zero_usage() -> NeMoGymResponseUsage:
+ return NeMoGymResponseUsage(
+ input_tokens=0,
+ input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=0),
+ output_tokens=0,
+ output_tokens_details=NeMoGymResponseOutputTokensDetails(reasoning_tokens=0),
+ total_tokens=0,
+ )
+
+
+def usage_from_provider(raw: dict[str, Any]) -> NeMoGymResponseUsage:
+ """Map a provider usage dict (Anthropic or OpenAI dialect) onto the native usage model.
+
+ Unreported detail counters default to 0 per the OpenAI contract; the raw dict is
+ preserved verbatim on the agent-step record so nothing is lost.
+ """
+ input_tokens = int(raw.get("input_tokens") or 0)
+ output_tokens = int(raw.get("output_tokens") or 0)
+ input_details = raw.get("input_tokens_details") or {}
+ output_details = raw.get("output_tokens_details") or {}
+ cached = raw.get("cache_read_input_tokens", input_details.get("cached_tokens")) or 0
+ reasoning = raw.get("reasoning_tokens", output_details.get("reasoning_tokens")) or 0
+ total = int(raw.get("total_tokens") or (input_tokens + output_tokens))
+ return NeMoGymResponseUsage(
+ input_tokens=input_tokens,
+ input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=int(cached)),
+ output_tokens=output_tokens,
+ output_tokens_details=NeMoGymResponseOutputTokensDetails(reasoning_tokens=int(reasoning)),
+ total_tokens=total,
+ )
+
+
+def summed_usage(generations: list[NeMoGymGeneration]) -> NeMoGymResponseUsage:
+ """Sum the per-agent-step native usage into run totals."""
+ totals = zero_usage()
+ for generation in generations:
+ if generation.usage is None:
+ continue
+ totals.input_tokens += generation.usage.input_tokens
+ totals.output_tokens += generation.usage.output_tokens
+ totals.total_tokens += generation.usage.total_tokens
+ totals.input_tokens_details.cached_tokens += generation.usage.input_tokens_details.cached_tokens
+ totals.output_tokens_details.reasoning_tokens += generation.usage.output_tokens_details.reasoning_tokens
+ return totals
+
+
+def agent_step_slices(
+ output: list[NeMoGymResponseInputItem],
+) -> Iterator[tuple[int, list[NeMoGymResponseInputItem]]]:
+ """Yield ``(agent_step_no, items)`` for each agent step, from the items' tags."""
+ current_step: Optional[int] = None
+ current_items: list[NeMoGymResponseInputItem] = []
+ for item in output:
+ step = getattr(item, "agent_step_no", None)
+ if step is None:
+ continue
+ if step != current_step:
+ if current_step is not None:
+ yield current_step, current_items
+ current_step, current_items = step, []
+ current_items.append(item)
+ if current_step is not None:
+ yield current_step, current_items
+
+
+def reconstruct_model_input(
+ output: list[NeMoGymResponseInputItem],
+ agent_step_no: Optional[int] = None,
+ base_input: Optional[list[NeMoGymResponseInputItem]] = None,
+) -> list[NeMoGymResponseInputItem]:
+ """Resolve the model-visible input of an agent step from the output item list.
+
+ With ``agent_step_no`` set, returns what that step's model call saw:
+ ``base_input + output[:first item tagged with that step]`` — or, if a
+ ``NeMoGymContextBoundaryMessage`` precedes the step, the context restarts at that
+ boundary (its summary replaced everything earlier, including the base input). With
+ ``agent_step_no=None``, returns the full final context. ``base_input`` is the task's
+ original ``responses_create_params.input``, deliberately stored only there.
+ """
+ if agent_step_no is None:
+ end = len(output)
+ else:
+ # First item belonging to this step or a later one; a step that produced no
+ # items (e.g. the run ended mid-step) saw everything recorded before it.
+ end = next(
+ (i for i, item in enumerate(output) if (getattr(item, "agent_step_no", None) or 0) >= agent_step_no),
+ len(output),
+ )
+
+ boundary_index: Optional[int] = None
+ for index in range(end):
+ if isinstance(output[index], NeMoGymContextBoundaryMessage):
+ boundary_index = index
+ if boundary_index is not None:
+ return list(output[boundary_index:end])
+ return list(base_input or []) + list(output[:end])
+
+
+def to_response_create_params(
+ output: list[NeMoGymResponseInputItem],
+ agent_step_no: Optional[int] = None,
+ base_input: Optional[list[NeMoGymResponseInputItem]] = None,
+ model: Optional[str] = None,
+) -> NeMoGymResponseCreateParamsNonStreaming:
+ """Package a reconstructed model input as native Responses create params."""
+ return NeMoGymResponseCreateParamsNonStreaming(
+ input=reconstruct_model_input(output, agent_step_no=agent_step_no, base_input=base_input),
+ model=model,
+ )
+
+
+def _parse_timestamp(value: Any) -> Optional[datetime]:
+ if not isinstance(value, str):
+ return None
+ try:
+ return datetime.fromisoformat(value)
+ except ValueError:
+ return None
+
+
+class TrajectoryBuilder:
+ """Incremental, agent-agnostic assembly of an agent's Response content and telemetry.
+
+ Call the ``add_*`` methods in execution order (from an in-process loop, or from an
+ adapter replaying a black-box harness's artifacts). The builder appends the
+ telemetry-tagged native items to ``output`` and maintains the per-generation records
+ in lockstep: agent-step numbering, model-call deduplication (``start_agent_step``
+ with the current step's ``response_id`` continues it instead of double counting
+ usage), call/execution correlation with independent per-call timing, orphan
+ handling.
+
+ ``build()`` returns ``(output_items, generations, agent_telemetry)`` — everything an
+ agent server needs to construct its ``NeMoGymResponse``.
+ """
+
+ def __init__(self, agent: str, source: str) -> None:
+ self.output: list[NeMoGymResponseInputItem] = []
+ self.generations: list[NeMoGymGeneration] = []
+ self._telemetry = NeMoGymAgentTelemetry(agent=agent, source=source)
+ self._model: Optional[str] = None
+ # call_id -> (agent_step_no, issue timestamp) for execution timing correlation.
+ self._pending_calls: dict[str, tuple[int, Optional[str]]] = {}
+ self._open = False # whether items may still be appended to the last generation
+
+ def _current_agent_step(self) -> NeMoGymGeneration:
+ if not self.generations or not self._open:
+ raise ValueError("no agent step in progress; call start_agent_step() first")
+ return self.generations[-1]
+
+ def set_session_id(self, session_id: Optional[str]) -> None:
+ if self._telemetry.session_id is None and session_id:
+ self._telemetry.session_id = session_id
+
+ def add_user_message(self, content: str, timestamp: Optional[str] = None) -> None:
+ """Record a mid-episode user message (a turn boundary).
+
+ The task's *initial* input must not be added here — it lives in
+ ``responses_create_params.input`` and is passed to reconstruction as ``base_input``.
+ """
+ self._open = False
+ self.output.append(NeMoGymEasyInputMessage(role="user", content=content))
+
+ def add_context_boundary(self, summary: str = "", timestamp: Optional[str] = None) -> None:
+ """Record a compaction: `summary` is the content that replaced all prior history."""
+ self._open = False
+ self.output.append(NeMoGymContextBoundaryMessage(role="user", content=summary, context_boundary=True))
+
+ def start_agent_step(
+ self,
+ response_id: Optional[str] = None,
+ request_id: Optional[str] = None,
+ model: Optional[str] = None,
+ timestamp: Optional[str] = None,
+ stop_reason: Optional[str] = None,
+ provider_usage: Optional[dict[str, Any]] = None,
+ ) -> NeMoGymGeneration:
+ """Start (or continue) the agent step for one model call.
+
+ Providers may emit one record per content block of the same API message; calling
+ this again with the current step's ``response_id`` returns that generation so
+ content accumulates under its tag without double counting usage.
+ """
+ current = self.generations[-1] if self.generations else None
+ if current is not None and self._open and response_id is not None and current.response_id == response_id:
+ if stop_reason:
+ current.stop_reason = stop_reason
+ return current
+
+ generation = NeMoGymGeneration(
+ agent_step_no=len(self.generations) + 1,
+ model=model,
+ stop_reason=stop_reason,
+ response_id=response_id,
+ request_id=request_id,
+ # A provider record is written when the message completes; the generation's
+ # start is not observable from artifacts, so only the completion time is kept.
+ ended_at=timestamp,
+ usage=usage_from_provider(provider_usage) if provider_usage else None,
+ provider_usage=provider_usage,
+ )
+ self.generations.append(generation)
+ self._open = True
+ if model and self._model is None:
+ self._model = model
+ return generation
+
+ def add_output_text(self, text: str) -> None:
+ generation = self._current_agent_step()
+ block = NeMoGymResponseOutputText(annotations=[], text=text)
+ last_item = self.output[-1] if self.output else None
+ if (
+ isinstance(last_item, NeMoGymResponseOutputMessageWithAgentTelemetry)
+ and last_item.agent_step_no == generation.agent_step_no
+ ):
+ last_item.content.append(block)
+ else:
+ self.output.append(
+ NeMoGymResponseOutputMessageWithAgentTelemetry(
+ id=f"msg-{len(self.output)}",
+ content=[block],
+ agent_step_no=generation.agent_step_no,
+ )
+ )
+
+ def add_reasoning(self, text: str) -> None:
+ generation = self._current_agent_step()
+ self.output.append(
+ NeMoGymResponseReasoningItemWithAgentTelemetry(
+ id=f"rs-{len(self.output)}",
+ summary=[NeMoGymSummary(text=text, type="summary_text")],
+ agent_step_no=generation.agent_step_no,
+ )
+ )
+
+ def add_tool_call(self, call_id: str, name: str, arguments: str) -> None:
+ generation = self._current_agent_step()
+ self.output.append(
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry(
+ call_id=call_id,
+ name=name,
+ arguments=arguments,
+ status="completed",
+ agent_step_no=generation.agent_step_no,
+ )
+ )
+ self._pending_calls[call_id] = (generation.agent_step_no, generation.ended_at)
+
+ def add_tool_result(
+ self,
+ call_id: str,
+ output: str,
+ completed_at: Optional[str] = None,
+ started_at: Optional[str] = None,
+ error: Optional[Union[str, NeMoGymToolExecutionError]] = None,
+ extra: Optional[dict[str, Any]] = None,
+ ) -> None:
+ """Record a tool observation: a ``function_call_output`` item tagged with the
+ issuing agent step and carrying its execution telemetry.
+
+ ``started_at`` defaults to the issuing step's completion timestamp; passing it
+ explicitly (e.g. from a provider execution record) overrides that. Each result
+ carries its own timing, so parallel tool calls stay independently timed. Results
+ with no matching call and no step to attach to are counted as dropped.
+ """
+ step_no, registered_started = self._pending_calls.pop(call_id, (None, None))
+ if step_no is None:
+ if not self.generations:
+ self.count_dropped("orphan_tool_results")
+ return
+ step_no = self.generations[-1].agent_step_no
+
+ started_at = started_at or registered_started
+ duration_ms = None
+ started_dt, completed_dt = _parse_timestamp(started_at), _parse_timestamp(completed_at)
+ if started_dt is not None and completed_dt is not None:
+ duration_ms = (completed_dt - started_dt).total_seconds() * 1000.0
+
+ if isinstance(error, str):
+ error = NeMoGymToolExecutionError(message=error)
+ self.output.append(
+ NeMoGymFunctionCallOutputWithAgentTelemetry(
+ call_id=call_id,
+ output=output,
+ status="completed",
+ agent_step_no=step_no,
+ execution=NeMoGymToolExecution(
+ started_at=started_at,
+ completed_at=completed_at,
+ duration_ms=duration_ms,
+ error=error,
+ extra=extra,
+ ),
+ )
+ )
+
+ def count_dropped(self, kind: str) -> None:
+ self._telemetry.dropped_records[kind] = self._telemetry.dropped_records.get(kind, 0) + 1
+
+ def set_run_totals(
+ self,
+ num_agent_steps: Optional[int] = None,
+ duration_ms: Optional[float] = None,
+ total_cost_usd: Optional[float] = None,
+ provider_usage: Optional[dict[str, Any]] = None,
+ ) -> None:
+ if num_agent_steps is not None:
+ self._telemetry.num_agent_steps = int(num_agent_steps)
+ if duration_ms is not None:
+ self._telemetry.duration_ms = float(duration_ms)
+ if total_cost_usd is not None:
+ self._telemetry.total_cost_usd = float(total_cost_usd)
+ if provider_usage is not None:
+ self._telemetry.provider_usage = provider_usage
+
+ @property
+ def model(self) -> Optional[str]:
+ return self._model
+
+ def build(
+ self,
+ ) -> tuple[list[NeMoGymResponseInputItem], list[NeMoGymGeneration], NeMoGymAgentTelemetry]:
+ return self.output, self.generations, self._telemetry
diff --git a/responses_api_agents/claude_code_agent/README.md b/responses_api_agents/claude_code_agent/README.md
index 573c5b394b..f264e26575 100644
--- a/responses_api_agents/claude_code_agent/README.md
+++ b/responses_api_agents/claude_code_agent/README.md
@@ -22,7 +22,7 @@ anthropic_model_name: Qwen/Qwen3-4B-Instruct-2507
anthropic_base_url: http://localhost:8000
```
-`anthropic_base_url` should not include `/v1`. Claude Code appends `/v1/messages` itself.
+`anthropic_base_url` should not include `/v1` — Claude Code appends `/v1/messages` itself. A trailing `/v1` (a common copy-paste from OpenAI-style configs) is stripped automatically with a warning.
### Launch
@@ -113,12 +113,13 @@ claude_code_agent:
bare: true
mcp_config: null
settings: null
+ capture_trajectory: true
```
- `concurrency`: max simultaneous `run()` calls
- `model`: model name. Full names like `Qwen/Qwen3-4B-Instruct-2507` are kept as-is for local endpoints; the provider prefix is stripped only when `anthropic_base_url` is not set
- `anthropic_api_key`: Anthropic API key, or any non-empty string for local endpoints
-- `anthropic_base_url`: if set, used as `ANTHROPIC_BASE_URL`. Leave null for the real Anthropic API
+- `anthropic_base_url`: if set, used as `ANTHROPIC_BASE_URL` (trailing `/` and `/v1` are stripped — the CLI appends `/v1/messages`). Leave null for the real Anthropic API
- `max_turns`: passed to `--max-turns`. Set to `null` to omit the flag entirely (unlimited turns)
- `timeout`: per-request wall-clock seconds
- `system_prompt`: appended to Claude Code's built-in system prompt via `--append-system-prompt`. The data's system message (if any) is also appended after this.
@@ -130,6 +131,7 @@ claude_code_agent:
- `bare`: when `true` (default), pass `--bare` to skip auto-discovery of hooks, skills, plugins, MCP servers, memory, and CLAUDE.md. Set to `false` to let Claude Code discover those from `CLAUDE_CONFIG_DIR` and the working directory
- `mcp_config`: path to an MCP server config file, passed to `--mcp-config`. Explicit, so it works regardless of `bare`
- `settings`: path to a settings JSON layered into the per-run `CLAUDE_CONFIG_DIR/settings.json`. Top-level keys override the defaults; the `env` block is shallow-merged so telemetry stays disabled unless you override it
+- `capture_trajectory`: when `true` (default), each `run()` result carries a standardized `trajectory` built from the session transcript Claude Code writes on disk (see [Trajectory capture](#trajectory-capture))
For the full set of Claude Code CLI flags see the [CLI reference](https://code.claude.com/docs/en/cli-reference).
@@ -193,6 +195,49 @@ Each rollout result is stamped with a `skills_ref` for provenance and grouping d
The skills path is resolved like `input_jsonl_fpath` (relative paths check the working directory, then the Gym root). For distributed runs the directory must be on storage accessible to the agent process.
+## Trajectory capture
+
+Claude Code persists a complete session transcript — one JSON record per event, with timestamps, request ids, per-model-call token usage, and tool execution metadata — under `$CLAUDE_CONFIG_DIR/projects//.jsonl`. Since each rollout runs with an ephemeral `CLAUDE_CONFIG_DIR`, the agent harvests those artifacts just before cleanup and parses them **once**. The result addresses [NVIDIA-NeMo/Gym#1867](https://github.com/NVIDIA-NeMo/Gym/issues/1867) with **the `NeMoGymResponse` itself as the trajectory entity** — telemetry rides on the contract, nothing is stored twice, and no sidecar object exists:
+
+- **On the items** (`response.output`, lossless and in execution order): model-produced items are `*WithAgentTelemetry` variants carrying `agent_step_no` — which model call produced them (the same extension pattern as the `*ForTraining` token-ID variants). `function_call_output` items additionally carry `execution` (`NeMoGymToolExecution`: independent `started_at`/`completed_at`/`duration_ms` per call, `error`, curated provider metadata). Reasoning is a native `reasoning` item, calls are in issue order with outputs in arrival order, unresolved calls are kept, mid-episode user messages are plain items, and a compaction summary is a `NeMoGymContextBoundaryMessage`. The task's initial prompt is *not* repeated — it stays in `responses_create_params.input`.
+- **On the Response**: `generations` — one `NeMoGymGeneration` per agent step with native per-call `usage` (cache detail included, deduplicated per API message), raw `provider_usage` verbatim, `model`, `stop_reason`, and provider identity (`response_id`/`request_id`); and `agent_telemetry` — run-level provenance (`source`, `session_id`), `num_agent_steps`, `duration_ms`, `total_cost_usd`, provider run totals, and `dropped_records` (events seen but not represented, e.g. subagent sidechains). Model servers leave both `None`.
+
+Terminology: an **agent step** is one interaction with the environment through the model — one LLM generation plus the orchestration of its tool calls and their outputs; a **turn** is a full cycle of control back to the user, containing one or more agent steps. Claude Code's `num_turns`/`--max-turns` count model calls, i.e. **agent steps**; the telemetry records that as `num_agent_steps`.
+
+```json
+"response": {
+ "id": "resp_…", "model": "claude-sonnet-4-6",
+ "output": [
+ {"type": "message", "agent_step_no": 1, "content": [{"type": "output_text", "text": "…"}], "…": "…"},
+ {"type": "function_call", "agent_step_no": 1, "call_id": "toolu_…", "name": "Bash", "…": "…"},
+ {"type": "function_call_output", "agent_step_no": 1, "call_id": "toolu_…", "output": "…",
+ "execution": {"started_at": "…", "completed_at": "…", "duration_ms": 532.3, "error": null, "extra": {"interrupted": false}}},
+ {"type": "message", "agent_step_no": 2, "…": "…"}
+ ],
+ "usage": {"input_tokens": 63, "input_tokens_details": {"cached_tokens": 18478}, "…": "…"},
+ "generations": [
+ {"agent_step_no": 1, "model": "claude-sonnet-4-6", "stop_reason": "tool_use",
+ "response_id": "msg_01…", "request_id": "req_…", "ended_at": "…",
+ "usage": {"input_tokens": 12, "input_tokens_details": {"cached_tokens": 9000}, "…": "…"},
+ "provider_usage": {"input_tokens": 12, "cache_read_input_tokens": 9000, "cache_creation_input_tokens": 512, "…": "…"}},
+ {"agent_step_no": 2, "stop_reason": "end_turn", "…": "…"}
+ ],
+ "agent_telemetry": {"schema_version": "1.0", "agent": "claude_code_agent", "source": "transcript",
+ "session_id": "…", "num_agent_steps": 2, "duration_ms": 8123.0, "total_cost_usd": 0.021,
+ "provider_usage": {"…": "…"}, "dropped_records": {}}
+}
+```
+
+Semantics:
+
+- **Reconstruction**: the model-visible input of agent step *k* is `responses_create_params.input + output[:first item tagged ≥ k]` — `nemo_gym.trajectory.reconstruct_model_input(output, agent_step_no=k, base_input=...)` resolves this, compaction-aware (a `NeMoGymContextBoundaryMessage` restarts the context at its summary). `agent_step_slices()` iterates per-step item groups.
+- **Telemetry survives every hop**: because it rides the contract, re-validating `NeMoGymResponse` anywhere (verify requests, resources servers, rollout rows) preserves it; plain model-server payloads still validate to the plain classes (the telemetry variants have required discriminating fields).
+- **Fallback**: with no transcript, everything is built from stream-json stdout (`source: "stream_json"`) — same structure, no timestamps/request ids; missing telemetry is `null`, never fabricated. On a timeout, the partial transcript surfaces as a partial response.
+
+**Response contract note**: `response.output` differs from this agent's pre-trajectory behavior: reasoning is a native item rather than `` text, calls are in issue order rather than result-completion order, and unresolved calls are kept. Verifiers reading the final assistant message are unaffected.
+
+To adapt another agent harness, parse its artifacts in execution order and drive `nemo_gym.trajectory.TrajectoryBuilder` (`start_agent_step` / `add_output_text` / `add_reasoning` / `add_tool_call` / `add_tool_result` / `add_user_message` / `add_context_boundary` / `set_run_totals`); `build()` returns `(output_items, generations, agent_telemetry)` — everything needed to construct the Response. In-process loops (e.g. `simple_agent`) drive the same builder as the loop executes, with *measured* tool timing.
+
## Limitations
- Eval only for now. Token IDs and logprobs are not wired up yet.
diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py
index dc78f3e613..9795280ea5 100644
--- a/responses_api_agents/claude_code_agent/app.py
+++ b/responses_api_agents/claude_code_agent/app.py
@@ -35,11 +35,11 @@
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 (
+ NeMoGymAgentTelemetry,
NeMoGymEasyInputMessage,
- NeMoGymFunctionCallOutput,
+ NeMoGymGeneration,
NeMoGymResponse,
NeMoGymResponseCreateParamsNonStreaming,
- NeMoGymResponseFunctionToolCall,
NeMoGymResponseInputTokensDetails,
NeMoGymResponseOutputMessage,
NeMoGymResponseOutputText,
@@ -48,133 +48,48 @@
)
from nemo_gym.server_utils import get_response_json, raise_for_status
from nemo_gym.skills import stage_skills
+from nemo_gym.trajectory import summed_usage
from responses_api_agents.claude_code_agent.setup_claude_code import ensure_claude_code
+from responses_api_agents.claude_code_agent.trajectory import build_trajectory, decode_jsonl
LOG = logging.getLogger(__name__)
-def _extract_text(content: list[Any]) -> str:
- return "".join(b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text")
+def _usage_metadata(generations: list[NeMoGymGeneration], telemetry: NeMoGymAgentTelemetry) -> dict:
+ """Derive the response usage metadata from the agent telemetry.
-
-def _extract_thinking(content: list[Any]) -> str:
- parts = []
- for b in content:
- if not isinstance(b, dict):
- continue
- if b.get("type") in ("thinking", "reasoning"):
- parts.append(b.get("thinking") or b.get("text") or "")
- return "\n".join(p for p in parts if p)
+ The provider's end-of-run report is the authoritative total when present (summing it
+ with the per-step usage would double count); otherwise the per-agent-step sums are used.
+ """
+ if telemetry.provider_usage:
+ metadata = {
+ "input_tokens": int(telemetry.provider_usage.get("input_tokens") or 0),
+ "output_tokens": int(telemetry.provider_usage.get("output_tokens") or 0),
+ "cached_tokens": int(telemetry.provider_usage.get("cache_read_input_tokens") or 0),
+ }
+ else:
+ totals = summed_usage(generations)
+ metadata = {
+ "input_tokens": totals.input_tokens,
+ "output_tokens": totals.output_tokens,
+ "cached_tokens": totals.input_tokens_details.cached_tokens,
+ }
+ if telemetry.num_agent_steps is not None:
+ # provider dialect: Claude Code reports agent steps as num_turns; keep the key
+ metadata["num_turns"] = telemetry.num_agent_steps
+ return metadata
def parse_stream_json(stdout: str) -> tuple[list[Any], dict]:
- """Convert claude -p --output-format=stream-json stdout into (output_items, usage)."""
- raw_events: list[dict] = []
- for line in stdout.splitlines():
- line = line.strip()
- if not line:
- continue
- try:
- raw_events.append(json.loads(line))
- except json.JSONDecodeError:
- continue
-
- output_items: list[Any] = []
- pending_calls: dict[str, dict] = {}
- buffered_think: str | None = None
- total_input = 0
- total_output = 0
- num_turns: Optional[int] = None
-
- for event in raw_events:
- etype = event.get("type")
-
- if etype == "result":
- usage = event.get("usage") or {}
- total_input += int(usage.get("input_tokens") or 0)
- total_output += int(usage.get("output_tokens") or 0)
- # Claude Code's authoritative turn counter (what --max-turns bounds).
- if event.get("num_turns") is not None:
- num_turns = int(event["num_turns"])
-
- elif etype == "assistant":
- message = event.get("message", {})
- content = message.get("content") or []
- usage = message.get("usage") or {}
- total_input += int(usage.get("input_tokens") or 0)
- total_output += int(usage.get("output_tokens") or 0)
-
- if not isinstance(content, list):
- content = []
-
- think = _extract_thinking(content)
- if think:
- buffered_think = (buffered_think + "\n" + think) if buffered_think else think
-
- text = _extract_text(content)
- if text:
- if buffered_think:
- text = f"\n{buffered_think}\n\n\n{text}"
- buffered_think = None
- output_items.append(
- NeMoGymResponseOutputMessage(
- id=f"msg-{len(output_items)}",
- content=[NeMoGymResponseOutputText(type="output_text", text=text, annotations=[])],
- role="assistant",
- status="completed",
- type="message",
- )
- )
-
- for block in content:
- if not isinstance(block, dict) or block.get("type") != "tool_use":
- continue
- call_id = block.get("id") or f"call-{uuid4().hex[:8]}"
- input_data = block.get("input") or {}
- arguments = json.dumps(input_data) if isinstance(input_data, dict) else str(input_data)
- pending_calls[call_id] = {"name": block.get("name", ""), "call_id": call_id, "arguments": arguments}
-
- elif etype == "user":
- message = event.get("message", {})
- content = message.get("content") or []
- if not isinstance(content, list):
- continue
-
- for block in content:
- if not isinstance(block, dict) or block.get("type") != "tool_result":
- continue
- tool_id = block.get("tool_use_id", "")
- call_info = pending_calls.pop(tool_id, None)
- if call_info:
- output_items.append(
- NeMoGymResponseFunctionToolCall(
- arguments=call_info["arguments"],
- call_id=tool_id,
- name=call_info["name"],
- type="function_call",
- id=tool_id,
- status="completed",
- )
- )
- result_content = block.get("content") or ""
- if isinstance(result_content, list):
- result_text = _extract_text(result_content)
- else:
- result_text = str(result_content)
- output_items.append(
- NeMoGymFunctionCallOutput(
- type="function_call_output",
- call_id=tool_id,
- output=result_text,
- status="completed",
- )
- )
+ """Convert claude -p --output-format=stream-json stdout into (output_items, usage).
- metadata: dict = {"input_tokens": total_input, "output_tokens": total_output}
- if num_turns is not None:
- metadata["num_turns"] = num_turns
- return output_items, metadata
+ One parse produces both the response's content plane (the native output item list)
+ and the telemetry overlay; this helper returns the content plane plus the usage
+ metadata derived from the overlay, so response and telemetry can never drift apart.
+ """
+ output_items, generations, telemetry = build_trajectory(decode_jsonl(stdout), [])
+ return output_items, _usage_metadata(generations, telemetry)
def _extract_instruction(body_input) -> tuple[str, Optional[str]]:
@@ -233,6 +148,10 @@ class ClaudeCodeAgentConfig(BaseResponsesAPIAgentConfig):
bare: bool = True
mcp_config: Optional[str] = None
settings: Optional[str] = None
+ # When True, the session transcript Claude Code writes under the per-run CLAUDE_CONFIG_DIR
+ # is harvested before cleanup and the response carries the agent telemetry on the contract
+ # itself (telemetry-tagged output items + `generations` + `agent_telemetry`).
+ capture_trajectory: bool = True
class ClaudeCodeAgentRunRequest(BaseRunRequest):
@@ -267,7 +186,14 @@ def _resolve_base_url(self) -> str:
self.config.model_server.name,
)
return self.server_client._build_server_base_url(cfg)
- return self.config.anthropic_base_url or ""
+ base_url = (self.config.anthropic_base_url or "").rstrip("/")
+ if base_url.endswith("/v1"):
+ # Claude Code appends /v1/messages itself; a copy-pasted OpenAI-style URL ending
+ # in /v1 would otherwise hit /v1/v1/messages and fail with a misleading
+ # "issue with the selected model" error.
+ base_url = base_url[: -len("/v1")].rstrip("/")
+ LOG.warning("anthropic_base_url ends with /v1; stripping it and using %s", base_url)
+ return base_url
def _build_settings(self) -> dict[str, Any]:
"""Settings written into the run's CLAUDE_CONFIG_DIR.
@@ -364,14 +290,31 @@ def _build_command(
cmd += ["--", instruction]
return cmd
+ def _collect_transcript_records(self, claude_config_dir: Path) -> list[dict]:
+ """Harvest the session transcript(s) Claude Code wrote under the per-run config dir.
+
+ Claude Code persists every session event (with timestamps, request ids, per-call
+ usage, and tool execution metadata) to ``/projects//*.jsonl``.
+ The per-run dir is removed after each request, so this runs just before cleanup.
+ """
+ records: list[dict] = []
+ try:
+ projects_dir = claude_config_dir / "projects"
+ if projects_dir.is_dir():
+ for transcript in sorted(projects_dir.glob("*/*.jsonl")):
+ records.extend(decode_jsonl(transcript.read_text(errors="replace")))
+ except OSError as exc:
+ LOG.warning("failed to read Claude Code transcript from %s: %s", claude_config_dir, exc)
+ return records
+
async def _run_claude_code(
self,
instruction: str,
system_prompt: Optional[str] = None,
mcp_config: Optional[str] = None,
skills_path: Optional[str] = None,
- ) -> tuple[str, str]:
- """Run claude -p --output-format=stream-json and return (stdout, model_name)."""
+ ) -> tuple[str, str, list[dict]]:
+ """Run claude -p --output-format=stream-json; return (stdout, model_name, transcript_records)."""
base_url = self._resolve_base_url()
# Keep full model name for local/custom endpoints; strip provider prefix for real Anthropic API.
model = self.config.model if base_url else self.config.model.split("/")[-1]
@@ -417,17 +360,22 @@ async def _run_claude_code(
proc.kill()
await proc.communicate()
LOG.warning("claude-code timed out after %ds", self.config.timeout)
- return "", model
+ # The partial transcript is still on disk and is the only record of what
+ # happened before the kill — harvest it for debugging.
+ return "", model, self._maybe_collect_transcript(claude_config_dir)
if proc.returncode not in (0, None):
LOG.warning("claude-code exited %d: %s", proc.returncode, stderr.decode(errors="replace")[:500])
LOG.debug("claude-code stdout (%d chars): %s", len(stdout), stdout[:2000].decode(errors="replace"))
- return stdout.decode(errors="replace"), model
+ return stdout.decode(errors="replace"), model, self._maybe_collect_transcript(claude_config_dir)
finally:
if claude_config_dir is not None:
shutil.rmtree(claude_config_dir, ignore_errors=True)
+ def _maybe_collect_transcript(self, claude_config_dir: Path) -> list[dict]:
+ return self._collect_transcript_records(claude_config_dir) if self.config.capture_trajectory else []
+
def _resources_server_base_url(self) -> str:
cfg = get_first_server_config_dict(
self.server_client.global_config_dict,
@@ -502,13 +450,28 @@ async def _create_response(
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_claude_code(
+ stdout, model_name, transcript_records = await self._run_claude_code(
user_message,
system_prompt=system_prompt,
mcp_config=mcp_config,
skills_path=skills_path,
)
- output_items, usage = parse_stream_json(stdout)
+ # One parse: the response IS the trajectory. The build yields the telemetry-tagged
+ # output items plus the per-generation records and run telemetry that ride on the
+ # Response contract itself. The transcript is the preferred source (timestamps,
+ # request ids, tool metadata); stream-json stdout is the fallback.
+ generations: Optional[list[NeMoGymGeneration]] = None
+ telemetry: Optional[NeMoGymAgentTelemetry] = None
+ try:
+ output_items, built_generations, built_telemetry = build_trajectory(
+ decode_jsonl(stdout), transcript_records
+ )
+ usage = _usage_metadata(built_generations, built_telemetry)
+ if self.config.capture_trajectory:
+ generations, telemetry = built_generations, built_telemetry
+ except Exception as exc:
+ LOG.warning("failed to parse Claude Code artifacts: %s", exc)
+ output_items, usage = [], {"input_tokens": 0, "output_tokens": 0, "cached_tokens": 0}
if not any(
getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant"
@@ -539,11 +502,13 @@ async def _create_response(
parallel_tool_calls=body.parallel_tool_calls,
usage=NeMoGymResponseUsage(
input_tokens=input_tokens,
- input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=0),
+ input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=usage.get("cached_tokens", 0)),
output_tokens=output_tokens,
output_tokens_details=NeMoGymResponseOutputTokensDetails(reasoning_tokens=0),
total_tokens=input_tokens + output_tokens,
),
+ generations=generations,
+ agent_telemetry=telemetry,
)
async def responses(
diff --git a/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml b/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml
index 5e3140bea7..2435e0bc16 100644
--- a/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml
+++ b/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml
@@ -19,3 +19,4 @@ claude_code_agent:
bare: true # set false to auto-discover skills, hooks, plugins, MCP, memory, CLAUDE.md
mcp_config: null # path to an MCP server config file (--mcp-config); works even with bare: true
settings: null # path to a settings JSON layered into the per-run CLAUDE_CONFIG_DIR
+ capture_trajectory: true # attach a standardized trajectory (from Claude Code's transcript) to run() results
diff --git a/responses_api_agents/claude_code_agent/tests/test_app.py b/responses_api_agents/claude_code_agent/tests/test_app.py
index 4c12003b0e..e757e9d4fe 100644
--- a/responses_api_agents/claude_code_agent/tests/test_app.py
+++ b/responses_api_agents/claude_code_agent/tests/test_app.py
@@ -102,6 +102,24 @@ def test_semaphore_initialized(self) -> None:
assert agent.sem._value == 4
+class TestResolveBaseUrl:
+ def test_v1_suffix_stripped(self) -> None:
+ agent = _make_agent(anthropic_base_url="https://inference-api.nvidia.com/v1")
+ assert agent._resolve_base_url() == "https://inference-api.nvidia.com"
+
+ def test_v1_with_trailing_slash_stripped(self) -> None:
+ agent = _make_agent(anthropic_base_url="https://host.example/v1/")
+ assert agent._resolve_base_url() == "https://host.example"
+
+ def test_plain_url_and_trailing_slash_normalized(self) -> None:
+ agent = _make_agent(anthropic_base_url="https://host.example/")
+ assert agent._resolve_base_url() == "https://host.example"
+
+ def test_null_base_url_stays_empty(self) -> None:
+ agent = _make_agent(anthropic_base_url=None)
+ assert agent._resolve_base_url() == ""
+
+
class TestBuildCommand:
def test_default_passes_bare(self) -> None:
agent = _make_agent()
@@ -288,7 +306,7 @@ def _run(self, agent: ClaudeCodeAgent, body: ClaudeCodeAgentRunRequest, run_clau
def test_skills_ref_path_forwarded(self) -> None:
agent = _make_agent()
- run_claude_code = AsyncMock(return_value=("", "claude-sonnet-4-6"))
+ run_claude_code = AsyncMock(return_value=("", "claude-sonnet-4-6", []))
body = ClaudeCodeAgentRunRequest.model_validate(
{
"responses_create_params": {"input": []},
@@ -302,7 +320,7 @@ def test_skills_ref_path_forwarded(self) -> None:
def test_no_skills_ref_forwards_none(self) -> None:
agent = _make_agent()
- run_claude_code = AsyncMock(return_value=("", "claude-sonnet-4-6"))
+ run_claude_code = AsyncMock(return_value=("", "claude-sonnet-4-6", []))
body = ClaudeCodeAgentRunRequest.model_validate({"responses_create_params": {"input": []}})
self._run(agent, body, run_claude_code)
@@ -335,7 +353,7 @@ async def fake_exec(*cmd, **kwargs):
patch("responses_api_agents.claude_code_agent.app.Path.home", return_value=tmp_path),
patch("responses_api_agents.claude_code_agent.app.asyncio.create_subprocess_exec", fake_exec),
):
- stdout, model = asyncio.run(agent._run_claude_code("hello", system_prompt="be terse"))
+ stdout, model, transcript_records = asyncio.run(agent._run_claude_code("hello", system_prompt="be terse"))
assert "claude" in captured["cmd"][0]
assert "--mcp-config" in captured["cmd"]
@@ -415,11 +433,12 @@ async def fake_wait_for(coro, timeout):
patch("responses_api_agents.claude_code_agent.app.asyncio.create_subprocess_exec", fake_exec),
patch("responses_api_agents.claude_code_agent.app.asyncio.wait_for", fake_wait_for),
):
- stdout, model = asyncio.run(agent._run_claude_code("hello"))
+ stdout, model, transcript_records = asyncio.run(agent._run_claude_code("hello"))
assert stdout == ""
assert killed["called"] is True
assert model == "claude-sonnet-4-6"
+ assert transcript_records == []
class TestRolloutMCPConfig:
@@ -530,10 +549,14 @@ async def fake_run_claude_code(instruction, system_prompt=None, mcp_config=None,
captured["mcp_config"] = mcp_config
captured["config_exists_during_run"] = Path(mcp_config).is_file()
captured["config"] = json.loads(Path(mcp_config).read_text())
- return _event(
- "assistant",
- message={"content": [{"type": "text", "text": "The weather in Paris is sunny and 72 F."}]},
- ), "claude-sonnet-4-6"
+ return (
+ _event(
+ "assistant",
+ message={"content": [{"type": "text", "text": "The weather in Paris is sunny and 72 F."}]},
+ ),
+ "claude-sonnet-4-6",
+ [],
+ )
agent.server_client.post.side_effect = fake_post
object.__setattr__(agent, "_run_claude_code", fake_run_claude_code)
@@ -585,7 +608,7 @@ async def fake_run_claude_code(instruction, system_prompt=None, mcp_config=None,
captured["config_token"] = json.loads(Path(mcp_config).read_text())["mcpServers"]["example_mcp_weather"][
"headers"
]["X-NeMo-Gym-Session-Token"]
- return _event("assistant", message={"content": [{"type": "text", "text": "ok"}]}), "claude-sonnet-4-6"
+ return _event("assistant", message={"content": [{"type": "text", "text": "ok"}]}), "claude-sonnet-4-6", []
agent.server_client.post.side_effect = fake_post
object.__setattr__(agent, "_run_claude_code", fake_run_claude_code)
@@ -638,7 +661,7 @@ def _user_tool_result(self, tool_use_id: str, result: str) -> str:
def test_empty(self) -> None:
items, usage = parse_stream_json("")
assert items == []
- assert usage == {"input_tokens": 0, "output_tokens": 0}
+ assert usage == {"input_tokens": 0, "output_tokens": 0, "cached_tokens": 0}
def test_text_message(self) -> None:
line = self._assistant([{"type": "text", "text": "hello"}])
@@ -649,7 +672,8 @@ def test_text_message(self) -> None:
assert usage["input_tokens"] == 10
assert usage["output_tokens"] == 5
- def test_thinking_prepended(self) -> None:
+ def test_thinking_is_a_native_reasoning_item(self) -> None:
+ # The content plane is lossless: reasoning is a native item, not text.
line = self._assistant(
[
{"type": "thinking", "thinking": "let me reason"},
@@ -657,23 +681,17 @@ def test_thinking_prepended(self) -> None:
]
)
items, _ = parse_stream_json(line)
- assert len(items) == 1
- text = items[0].content[0].text
- assert "\nlet me reason\n" in text
- assert "answer" in text
+ assert [type(i).__name__ for i in items] == [
+ "NeMoGymResponseReasoningItemWithAgentTelemetry",
+ "NeMoGymResponseOutputMessageWithAgentTelemetry",
+ ]
+ assert items[0].summary[0].text == "let me reason"
+ assert items[1].content[0].text == "answer"
- def test_thinking_without_text_not_emitted(self) -> None:
+ def test_thinking_without_text_is_kept(self) -> None:
line = self._assistant([{"type": "thinking", "thinking": "just thinking"}])
items, _ = parse_stream_json(line)
- assert items == []
-
- def test_thinking_cleared_after_message(self) -> None:
- l1 = self._assistant([{"type": "thinking", "thinking": "think"}, {"type": "text", "text": "msg1"}])
- l2 = self._assistant([{"type": "text", "text": "msg2"}])
- items, _ = parse_stream_json(f"{l1}\n{l2}")
- assert len(items) == 2
- assert "" in items[0].content[0].text
- assert "" not in items[1].content[0].text
+ assert [type(i).__name__ for i in items] == ["NeMoGymResponseReasoningItemWithAgentTelemetry"]
def test_tool_call_and_result(self) -> None:
assistant_line = self._assistant(
@@ -703,6 +721,34 @@ def test_text_then_tool_call(self) -> None:
assert isinstance(items[1], NeMoGymResponseFunctionToolCall)
assert isinstance(items[2], NeMoGymFunctionCallOutput)
+ def test_unresolved_call_is_kept(self) -> None:
+ # Lossless content plane: a call whose result never arrived stays in the record.
+ line = self._assistant([{"type": "tool_use", "id": "t9", "name": "Bash", "input": {}}])
+ items, _ = parse_stream_json(line)
+ assert [type(i).__name__ for i in items] == ["NeMoGymResponseFunctionToolCallWithAgentTelemetry"]
+
+ def test_parallel_calls_keep_issue_order(self) -> None:
+ assistant_line = self._assistant(
+ [
+ {"type": "tool_use", "id": "a", "name": "Bash", "input": {}},
+ {"type": "tool_use", "id": "b", "name": "Read", "input": {}},
+ ]
+ )
+ # b's result arrives first; calls stay in issue order, outputs in arrival order
+ items, _ = parse_stream_json(
+ f"{assistant_line}\n{self._user_tool_result('b', 'second issued')}\n"
+ f"{self._user_tool_result('a', 'first issued')}"
+ )
+ kinds = [(type(i).__name__, getattr(i, "call_id", None)) for i in items]
+ assert kinds == [
+ ("NeMoGymResponseFunctionToolCallWithAgentTelemetry", "a"),
+ ("NeMoGymResponseFunctionToolCallWithAgentTelemetry", "b"),
+ ("NeMoGymFunctionCallOutputWithAgentTelemetry", "b"),
+ ("NeMoGymFunctionCallOutputWithAgentTelemetry", "a"),
+ ]
+ # issue-order tags: both calls came from the same generation
+ assert all(i.agent_step_no == 1 for i in items)
+
def test_malformed_lines_skipped(self) -> None:
good = self._assistant([{"type": "text", "text": "ok"}])
items, _ = parse_stream_json(f"not-json\n{good}\n{{bad")
@@ -714,6 +760,33 @@ def test_result_event_accumulates_usage(self) -> None:
assert usage["input_tokens"] == 100
assert usage["output_tokens"] == 50
+ def test_result_usage_wins_over_assistant_sums(self) -> None:
+ # The result event's usage is the run total — adding per-message usage on top
+ # of it would double count.
+ assistant = self._assistant([{"type": "text", "text": "hi"}])
+ result = _event("result", usage={"input_tokens": 10, "output_tokens": 5, "cache_read_input_tokens": 7})
+ _, usage = parse_stream_json(f"{assistant}\n{result}")
+ assert usage["input_tokens"] == 10
+ assert usage["output_tokens"] == 5
+ assert usage["cached_tokens"] == 7
+
+ def test_usage_deduped_by_message_id(self) -> None:
+ # One API message can arrive as multiple events (one per content block), each
+ # repeating the same message id and usage — it must count once.
+ shared = {"input_tokens": 10, "output_tokens": 5}
+ e1 = _event("assistant", message={"id": "msg_1", "content": [{"type": "text", "text": "a"}], "usage": shared})
+ e2 = _event(
+ "assistant",
+ message={
+ "id": "msg_1",
+ "content": [{"type": "tool_use", "id": "t1", "name": "Bash", "input": {}}],
+ "usage": shared,
+ },
+ )
+ _, usage = parse_stream_json(f"{e1}\n{e2}")
+ assert usage["input_tokens"] == 10
+ assert usage["output_tokens"] == 5
+
def test_result_event_exposes_num_turns(self) -> None:
result = _event("result", num_turns=9, usage={"input_tokens": 1, "output_tokens": 1})
_, usage = parse_stream_json(result)
@@ -725,6 +798,166 @@ def test_num_turns_absent_when_no_result_event(self) -> None:
assert "num_turns" not in usage
+class TestTrajectoryCapture:
+ def _create(self, agent: ClaudeCodeAgent, run_claude_code: AsyncMock):
+ body = NeMoGymResponseCreateParamsNonStreaming(input="solve it")
+ with patch.object(ClaudeCodeAgent, "_run_claude_code", run_claude_code):
+ return asyncio.run(agent._create_response(body))
+
+ def _stream_stdout(self) -> str:
+ return "\n".join(
+ [
+ _event("system", subtype="init", session_id="sess-1"),
+ _event(
+ "assistant",
+ message={
+ "id": "msg_1",
+ "model": "claude-sonnet-4-6",
+ "content": [{"type": "text", "text": "done"}],
+ "usage": {"input_tokens": 10, "output_tokens": 5, "cache_read_input_tokens": 3},
+ },
+ ),
+ _event(
+ "result",
+ num_turns=1,
+ duration_ms=1234.0,
+ total_cost_usd=0.01,
+ usage={"input_tokens": 10, "output_tokens": 5, "cache_read_input_tokens": 3},
+ ),
+ ]
+ )
+
+ def test_response_carries_the_telemetry(self) -> None:
+ # The response IS the trajectory: items are telemetry-tagged and the Response
+ # carries generations + agent_telemetry on the contract itself.
+ agent = _make_agent()
+ run_claude_code = AsyncMock(return_value=(self._stream_stdout(), "claude-sonnet-4-6", []))
+
+ response = self._create(agent, run_claude_code)
+
+ assert response.agent_telemetry.schema_version == "1.0"
+ assert response.agent_telemetry.source == "stream_json"
+ assert response.agent_telemetry.session_id == "sess-1"
+ assert response.agent_telemetry.num_agent_steps == 1
+ assert response.agent_telemetry.duration_ms == 1234.0
+ assert response.agent_telemetry.total_cost_usd == 0.01
+ (generation,) = response.generations
+ assert generation.agent_step_no == 1
+ assert generation.usage.input_tokens_details.cached_tokens == 3
+ assert response.output[0].agent_step_no == 1
+ assert response.usage.input_tokens == 10
+ assert response.usage.input_tokens_details.cached_tokens == 3
+
+ def test_telemetry_survives_response_revalidation(self) -> None:
+ agent = _make_agent()
+ run_claude_code = AsyncMock(return_value=(self._stream_stdout(), "claude-sonnet-4-6", []))
+ response = self._create(agent, run_claude_code)
+
+ from nemo_gym.openai_utils import NeMoGymResponse
+
+ revalidated = NeMoGymResponse.model_validate(response.model_dump(mode="json"))
+ assert revalidated.agent_telemetry.session_id == "sess-1"
+ assert revalidated.output[0].agent_step_no == 1
+
+ def test_transcript_preferred_over_stream(self) -> None:
+ agent = _make_agent()
+ transcript = [
+ {
+ "type": "user",
+ "uuid": "u1",
+ "timestamp": "2026-07-09T00:00:00.000Z",
+ "sessionId": "sess-t",
+ "message": {"role": "user", "content": "solve it"},
+ },
+ {
+ "type": "assistant",
+ "uuid": "a1",
+ "requestId": "req-1",
+ "timestamp": "2026-07-09T00:00:01.000Z",
+ "sessionId": "sess-t",
+ "message": {
+ "id": "msg_1",
+ "model": "claude-sonnet-4-6",
+ "content": [{"type": "text", "text": "done"}],
+ "usage": {"input_tokens": 10, "output_tokens": 5},
+ },
+ },
+ ]
+ run_claude_code = AsyncMock(return_value=(self._stream_stdout(), "claude-sonnet-4-6", transcript))
+
+ response = self._create(agent, run_claude_code)
+
+ assert response.agent_telemetry.source == "transcript"
+ assert response.agent_telemetry.session_id == "sess-t"
+ # run-level totals still come from the stream-json result event
+ assert response.agent_telemetry.num_agent_steps == 1
+ assert response.agent_telemetry.total_cost_usd == 0.01
+ (generation,) = response.generations
+ assert generation.request_id == "req-1"
+ assert generation.ended_at == "2026-07-09T00:00:01.000Z"
+
+ def test_capture_disabled_yields_plain_response(self) -> None:
+ agent = _make_agent(capture_trajectory=False)
+ run_claude_code = AsyncMock(return_value=(self._stream_stdout(), "claude-sonnet-4-6", []))
+
+ response = self._create(agent, run_claude_code)
+
+ assert response.generations is None
+ assert response.agent_telemetry is None
+ assert response.output # content still parsed
+
+ def test_run_claude_code_harvests_transcript_from_config_dir(self, tmp_path: Path) -> None:
+ agent = _make_agent()
+
+ class FakeProc:
+ returncode = 0
+
+ async def communicate(self):
+ return b'{"type":"result","usage":{"input_tokens":1,"output_tokens":1}}\n', b""
+
+ async def fake_exec(*cmd, **kwargs):
+ # Simulate Claude Code persisting the session transcript under the per-run config dir.
+ project_dir = Path(kwargs["env"]["CLAUDE_CONFIG_DIR"]) / "projects" / "-tmp-cwd"
+ project_dir.mkdir(parents=True)
+ (project_dir / "sess-1.jsonl").write_text(
+ json.dumps({"type": "user", "sessionId": "sess-1", "message": {"role": "user", "content": "hi"}})
+ + "\nnot-json\n"
+ )
+ return FakeProc()
+
+ with (
+ patch("responses_api_agents.claude_code_agent.app.Path.home", return_value=tmp_path),
+ patch("responses_api_agents.claude_code_agent.app.asyncio.create_subprocess_exec", fake_exec),
+ ):
+ _, _, records = asyncio.run(agent._run_claude_code("hello"))
+
+ assert len(records) == 1
+ assert records[0]["sessionId"] == "sess-1"
+
+ def test_capture_disabled_skips_harvest(self, tmp_path: Path) -> None:
+ agent = _make_agent(capture_trajectory=False)
+
+ class FakeProc:
+ returncode = 0
+
+ async def communicate(self):
+ return b"", b""
+
+ async def fake_exec(*cmd, **kwargs):
+ project_dir = Path(kwargs["env"]["CLAUDE_CONFIG_DIR"]) / "projects" / "-tmp-cwd"
+ project_dir.mkdir(parents=True)
+ (project_dir / "sess-1.jsonl").write_text('{"type":"user"}\n')
+ return FakeProc()
+
+ with (
+ patch("responses_api_agents.claude_code_agent.app.Path.home", return_value=tmp_path),
+ patch("responses_api_agents.claude_code_agent.app.asyncio.create_subprocess_exec", fake_exec),
+ ):
+ _, _, records = asyncio.run(agent._run_claude_code("hello"))
+
+ assert records == []
+
+
class TestConfigYaml:
def test_module_parses(self) -> None:
app_path = Path(__file__).resolve().parent.parent / "app.py"
diff --git a/responses_api_agents/claude_code_agent/tests/test_trajectory.py b/responses_api_agents/claude_code_agent/tests/test_trajectory.py
new file mode 100644
index 0000000000..f66232b194
--- /dev/null
+++ b/responses_api_agents/claude_code_agent/tests/test_trajectory.py
@@ -0,0 +1,317 @@
+# 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.
+
+from nemo_gym.openai_utils import (
+ NeMoGymContextBoundaryMessage,
+ NeMoGymEasyInputMessage,
+ NeMoGymFunctionCallOutput,
+ NeMoGymFunctionCallOutputWithAgentTelemetry,
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry,
+ NeMoGymResponseOutputMessageWithAgentTelemetry,
+ NeMoGymResponseReasoningItemWithAgentTelemetry,
+)
+from nemo_gym.trajectory import reconstruct_model_input, summed_usage
+from responses_api_agents.claude_code_agent.trajectory import build_trajectory, decode_jsonl
+
+
+SESSION = "sess-1"
+
+
+def _user_record(text: str, ts: str = "2026-07-09T00:00:00.000Z", **kwargs) -> dict:
+ return {
+ "type": "user",
+ "uuid": "u-user",
+ "timestamp": ts,
+ "sessionId": SESSION,
+ "message": {"role": "user", "content": text},
+ **kwargs,
+ }
+
+
+def _assistant_record(
+ uuid: str,
+ message_id: str,
+ blocks: list,
+ usage: dict | None = None,
+ ts: str = "2026-07-09T00:00:01.000Z",
+ request_id: str = "req-1",
+ stop_reason: str | None = None,
+ **kwargs,
+) -> dict:
+ return {
+ "type": "assistant",
+ "uuid": uuid,
+ "requestId": request_id,
+ "timestamp": ts,
+ "sessionId": SESSION,
+ "message": {
+ "id": message_id,
+ "model": "claude-sonnet-4-6",
+ "role": "assistant",
+ "stop_reason": stop_reason,
+ "content": blocks,
+ "usage": usage or {},
+ },
+ **kwargs,
+ }
+
+
+def _tool_result_record(
+ call_id: str,
+ content: str,
+ ts: str,
+ source_uuid: str = "a1",
+ is_error: bool = False,
+ tool_use_result=None,
+) -> dict:
+ return {
+ "type": "user",
+ "uuid": f"u-{call_id}",
+ "timestamp": ts,
+ "sessionId": SESSION,
+ "sourceToolAssistantUUID": source_uuid,
+ "toolUseResult": tool_use_result,
+ "message": {
+ "role": "user",
+ "content": [{"type": "tool_result", "tool_use_id": call_id, "content": content, "is_error": is_error}],
+ },
+ }
+
+
+USAGE = {"input_tokens": 100, "output_tokens": 20, "cache_read_input_tokens": 60, "cache_creation_input_tokens": 10}
+
+
+class TestTranscriptSource:
+ def _transcript(self) -> list[dict]:
+ return [
+ {"type": "queue-operation", "operation": "enqueue"}, # non-message noise is skipped
+ _user_record("fix the bug"), # initial prompt: lives in create_params, not the content plane
+ # One API message written as two records (text block, then two parallel tool_use
+ # blocks) sharing the same message id and usage.
+ _assistant_record("a1", "msg_1", [{"type": "text", "text": "looking"}], usage=USAGE),
+ _assistant_record(
+ "a1",
+ "msg_1",
+ [
+ {"type": "tool_use", "id": "t1", "name": "Bash", "input": {"command": "ls"}},
+ {"type": "tool_use", "id": "t2", "name": "Read", "input": {"file_path": "/x"}},
+ ],
+ usage=USAGE,
+ stop_reason="tool_use",
+ ),
+ _tool_result_record(
+ "t1",
+ "file.txt",
+ "2026-07-09T00:00:01.500Z",
+ tool_use_result={"interrupted": False, "stdout": "file.txt", "big": "x" * 5000, "nested": {"a": 1}},
+ ),
+ _tool_result_record("t2", "boom", "2026-07-09T00:00:03.000Z", is_error=True),
+ _assistant_record(
+ "a2",
+ "msg_2",
+ [{"type": "thinking", "thinking": "hmm"}, {"type": "text", "text": "fixed"}],
+ usage={"input_tokens": 50, "output_tokens": 5},
+ ts="2026-07-09T00:00:04.000Z",
+ request_id="req-2",
+ stop_reason="end_turn",
+ ),
+ ]
+
+ def test_output_is_native_lossless_and_tagged(self) -> None:
+ output, generations, telemetry = build_trajectory([], self._transcript())
+ assert telemetry.source == "transcript"
+ assert telemetry.session_id == SESSION
+ assert [type(i) for i in output] == [
+ NeMoGymResponseOutputMessageWithAgentTelemetry, # "looking"
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry, # t1, issue order
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry, # t2
+ NeMoGymFunctionCallOutputWithAgentTelemetry, # t1 result, arrival order
+ NeMoGymFunctionCallOutputWithAgentTelemetry, # t2 result
+ NeMoGymResponseReasoningItemWithAgentTelemetry, # native reasoning, not text
+ NeMoGymResponseOutputMessageWithAgentTelemetry, # "fixed"
+ ]
+ assert [i.agent_step_no for i in output] == [1, 1, 1, 1, 1, 2, 2]
+ assert output[1].name == "Bash"
+ assert output[5].summary[0].text == "hmm"
+ assert output[6].content[0].text == "fixed"
+
+ def test_initial_prompt_not_duplicated_into_output(self) -> None:
+ output, _, _ = build_trajectory([], self._transcript())
+ assert not any(type(i) is NeMoGymEasyInputMessage for i in output)
+
+ def test_mid_episode_user_message_is_recorded(self) -> None:
+ records = self._transcript() + [
+ _user_record("and now Berlin", ts="2026-07-09T00:01:00.000Z"),
+ _assistant_record("a3", "msg_3", [{"type": "text", "text": "on it"}], ts="2026-07-09T00:01:01.000Z"),
+ ]
+ output, generations, _ = build_trajectory([], records)
+ user_items = [i for i in output if type(i) is NeMoGymEasyInputMessage]
+ assert [i.content for i in user_items] == ["and now Berlin"]
+ assert output[output.index(user_items[0]) + 1].agent_step_no == generations[-1].agent_step_no
+
+ def test_generation_identity_and_dedupe(self) -> None:
+ _, generations, _ = build_trajectory([], self._transcript())
+ g1, g2 = generations
+ assert (g1.agent_step_no, g2.agent_step_no) == (1, 2)
+ assert g1.stop_reason == "tool_use"
+ assert g1.response_id == "msg_1"
+ assert g1.request_id == "req-1"
+ assert g2.request_id == "req-2"
+
+ def test_usage_counted_once_per_message(self) -> None:
+ _, generations, _ = build_trajectory([], self._transcript())
+ g1 = generations[0]
+ assert g1.usage.input_tokens == 100
+ assert g1.usage.input_tokens_details.cached_tokens == 60
+ assert g1.provider_usage["cache_creation_input_tokens"] == 10
+ totals = summed_usage(generations)
+ assert totals.input_tokens == 150
+ assert totals.output_tokens == 25
+
+ def test_parallel_tool_executions_have_independent_timing(self) -> None:
+ output, _, _ = build_trajectory([], self._transcript())
+ out1, out2 = output[3], output[4]
+ assert (out1.call_id, out2.call_id) == ("t1", "t2")
+ assert out1.execution.started_at == "2026-07-09T00:00:01.000Z"
+ assert out1.execution.duration_ms == 500.0
+ assert out2.execution.duration_ms == 2000.0
+
+ def test_execution_error_and_curated_extra(self) -> None:
+ output, _, _ = build_trajectory([], self._transcript())
+ out1, out2 = output[3], output[4]
+ assert out1.execution.error is None
+ assert out1.execution.extra == {"interrupted": False, "stdout": "file.txt"} # short scalars kept
+ assert out2.execution.error is not None
+ outputs = {i.call_id: i.output for i in output if isinstance(i, NeMoGymFunctionCallOutput)}
+ assert outputs == {"t1": "file.txt", "t2": "boom"}
+
+ def test_sidechain_records_skipped_and_counted(self) -> None:
+ records = self._transcript() + [
+ _assistant_record("a3", "msg_3", [{"type": "text", "text": "sub"}], isSidechain=True)
+ ]
+ _, generations, telemetry = build_trajectory([], records)
+ assert telemetry.dropped_records == {"sidechain": 1}
+ assert len(generations) == 2
+
+ def test_compact_summary_becomes_boundary_with_summary_item(self) -> None:
+ records = [
+ _user_record("start"),
+ _assistant_record("a1", "msg_1", [{"type": "text", "text": "working"}]),
+ _user_record("summary of history", ts="2026-07-09T00:01:00.000Z", isCompactSummary=True),
+ _assistant_record("a2", "msg_2", [{"type": "text", "text": "go on"}], ts="2026-07-09T00:01:01.000Z"),
+ ]
+ output, _, _ = build_trajectory([], records)
+ boundaries = [i for i in output if isinstance(i, NeMoGymContextBoundaryMessage)]
+ assert [b.content for b in boundaries] == ["summary of history"]
+ # post-boundary reconstruction starts at the summary, not the base input
+ items = reconstruct_model_input(
+ output, agent_step_no=2, base_input=[NeMoGymEasyInputMessage(role="user", content="start")]
+ )
+ assert getattr(items[0], "content", None) == "summary of history"
+
+ def test_orphan_observation_with_no_step_counted_as_dropped(self) -> None:
+ records = [_tool_result_record("orphan", "out", "2026-07-09T00:00:02.000Z")]
+ output, _, telemetry = build_trajectory([], records)
+ assert output == []
+ assert telemetry.dropped_records == {"orphan_tool_results": 1}
+
+ def test_meta_user_records_skipped(self) -> None:
+ output, _, _ = build_trajectory([], [_user_record("injected", isMeta=True)])
+ assert output == []
+
+ def test_telemetry_round_trips_through_the_contract(self) -> None:
+ from time import time
+
+ from nemo_gym.openai_utils import NeMoGymResponse
+
+ output, generations, telemetry = build_trajectory([], self._transcript())
+ response = NeMoGymResponse(
+ id="r",
+ created_at=int(time()),
+ model="m",
+ object="response",
+ output=output,
+ tool_choice="auto",
+ tools=[],
+ parallel_tool_calls=True,
+ generations=generations,
+ agent_telemetry=telemetry,
+ )
+ revalidated = NeMoGymResponse.model_validate(response.model_dump(mode="json"))
+ assert revalidated == response
+ assert revalidated.output[3].execution.duration_ms == 500.0
+
+
+class TestStreamJsonFallback:
+ def _events(self) -> list[dict]:
+ return [
+ {"type": "system", "subtype": "init", "session_id": SESSION},
+ {
+ "type": "assistant",
+ "message": {
+ "id": "msg_1",
+ "model": "claude-sonnet-4-6",
+ "content": [{"type": "tool_use", "id": "t1", "name": "Bash", "input": {"command": "ls"}}],
+ "usage": {"input_tokens": 10, "output_tokens": 2},
+ },
+ },
+ {
+ "type": "user",
+ "message": {"content": [{"type": "tool_result", "tool_use_id": "t1", "content": "ok"}]},
+ },
+ {"type": "system", "subtype": "compact_boundary"},
+ {
+ "type": "result",
+ "num_turns": 3,
+ "duration_ms": 42.0,
+ "total_cost_usd": 0.5,
+ "usage": {"input_tokens": 10, "output_tokens": 2},
+ },
+ ]
+
+ def test_fallback_used_when_no_transcript_messages(self) -> None:
+ output, _, telemetry = build_trajectory(self._events(), transcript_records=[{"type": "queue-operation"}])
+ assert telemetry.source == "stream_json"
+ assert telemetry.session_id == SESSION
+ assert [type(i).__name__ for i in output] == [
+ "NeMoGymResponseFunctionToolCallWithAgentTelemetry",
+ "NeMoGymFunctionCallOutputWithAgentTelemetry",
+ "NeMoGymContextBoundaryMessage",
+ ]
+
+ def test_no_timestamps_means_no_fabricated_timing(self) -> None:
+ output, _, _ = build_trajectory(self._events(), [])
+ execution = output[1].execution
+ assert execution.started_at is None
+ assert execution.duration_ms is None
+
+ def test_result_event_totals(self) -> None:
+ _, _, telemetry = build_trajectory(self._events(), [])
+ assert telemetry.num_agent_steps == 3 # Claude Code's num_turns, normalized
+ assert telemetry.duration_ms == 42.0
+ assert telemetry.total_cost_usd == 0.5
+ assert telemetry.provider_usage == {"input_tokens": 10, "output_tokens": 2}
+
+ def test_empty_sources(self) -> None:
+ output, generations, telemetry = build_trajectory([], [])
+ assert output == []
+ assert generations == []
+ assert telemetry.source == "stream_json"
+
+
+class TestDecodeJsonl:
+ def test_skips_blank_and_malformed_lines(self) -> None:
+ text = '\n{"a": 1}\nnot-json\n[1, 2]\n{"b": 2}\n'
+ assert decode_jsonl(text) == [{"a": 1}, {"b": 2}]
diff --git a/responses_api_agents/claude_code_agent/trajectory.py b/responses_api_agents/claude_code_agent/trajectory.py
new file mode 100644
index 0000000000..410bf043d3
--- /dev/null
+++ b/responses_api_agents/claude_code_agent/trajectory.py
@@ -0,0 +1,211 @@
+# 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.
+
+"""Claude Code adapter for Gym's standardized trajectory (nemo_gym.trajectory).
+
+Claude Code writes a complete session transcript (one JSON record per event, with
+timestamps, request ids, per-model-call usage, and tool execution metadata) into
+``$CLAUDE_CONFIG_DIR/projects//.jsonl``. Since the agent runs each
+rollout with an ephemeral ``CLAUDE_CONFIG_DIR``, those artifacts are harvested before the
+directory is removed and fed through this adapter, which parses them and drives the
+generic :class:`~nemo_gym.trajectory.TrajectoryBuilder`. The stream-json stdout events
+are the fallback source when no transcript is available (same message structure and
+token usage, but no timestamps or request ids — that telemetry stays ``None``).
+
+Provider-specific mapping notes:
+
+- Tool timing: an observation's ``started_at`` is the timestamp of the assistant record
+ that issued the ``tool_use`` (linked via ``sourceToolAssistantUUID`` when present),
+ and ``completed_at`` is its ``tool_result`` record's timestamp — each result record
+ has its own, so parallel tool calls keep independent timing.
+- ``toolUseResult`` execution metadata: only short scalars are kept (payloads can embed
+ entire file contents; the model-visible output is already on the native
+ ``function_call_output`` item).
+- Subagent (sidechain) records are skipped and counted under
+ ``dropped_records["sidechain"]``.
+- Compaction: transcript records flagged ``isCompactSummary`` and stream-json
+ ``system/compact_boundary`` events become ``context_boundary`` steps.
+"""
+
+import json
+from typing import Any, Optional
+
+from nemo_gym.openai_utils import NeMoGymAgentTelemetry, NeMoGymGeneration, NeMoGymResponseInputItem
+from nemo_gym.trajectory import TrajectoryBuilder
+
+
+# toolUseResult payloads can embed entire file contents; only short scalars are kept as
+# execution metadata so the trajectory stays bounded.
+_EXTRA_MAX_STR_LEN = 256
+
+
+def decode_jsonl(text: str) -> list[dict]:
+ records = []
+ for line in text.splitlines():
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ record = json.loads(line)
+ except json.JSONDecodeError:
+ continue
+ if isinstance(record, dict):
+ records.append(record)
+ return records
+
+
+def _block_text(content: Any) -> str:
+ if isinstance(content, str):
+ return content
+ if isinstance(content, list):
+ return "".join(b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text")
+ return "" if content is None else str(content)
+
+
+def _thinking_text(block: dict) -> str:
+ return block.get("thinking") or block.get("text") or ""
+
+
+def _curate_extra(tool_use_result: Any) -> Optional[dict[str, Any]]:
+ """Keep only short scalar execution metadata (durations, exit codes, flags) from toolUseResult."""
+ if not isinstance(tool_use_result, dict):
+ return None
+ extra = {}
+ for key, value in tool_use_result.items():
+ if isinstance(value, bool) or isinstance(value, (int, float)):
+ extra[key] = value
+ elif isinstance(value, str) and len(value) <= _EXTRA_MAX_STR_LEN:
+ extra[key] = value
+ return extra or None
+
+
+def build_trajectory(
+ stream_events: list[dict], transcript_records: list[dict]
+) -> tuple[list[NeMoGymResponseInputItem], list[NeMoGymGeneration], NeMoGymAgentTelemetry]:
+ """Build an agent Response's items and telemetry from Claude Code artifacts.
+
+ Returns ``(output_items, generations, agent_telemetry)`` — the episode's native,
+ telemetry-tagged item list (destined for ``NeMoGymResponse.output``) plus the
+ per-generation records and run-level telemetry that ride on the Response itself.
+ Prefers the on-disk transcript (timestamps, request ids, tool execution metadata);
+ falls back to the stream-json stdout events. Run-level totals (`num_agent_steps`,
+ `duration_ms`, `total_cost_usd`, provider usage) always come from the stream-json
+ `result` event when present, since the transcript does not carry them.
+
+ The task's initial prompt is deliberately **not** emitted into the output — it
+ already lives in ``responses_create_params.input`` (pass it to reconstruction as
+ ``base_input``); only mid-episode user messages (later turns) are recorded.
+ """
+ has_transcript_messages = any(r.get("type") in ("assistant", "user") for r in transcript_records)
+ if has_transcript_messages:
+ builder = TrajectoryBuilder(agent="claude_code_agent", source="transcript")
+ _replay_records(builder, transcript_records)
+ else:
+ builder = TrajectoryBuilder(agent="claude_code_agent", source="stream_json")
+ _replay_records(builder, stream_events)
+
+ for event in stream_events:
+ if event.get("type") == "system" and event.get("subtype") == "init":
+ builder.set_session_id(event.get("session_id"))
+ if event.get("type") == "result":
+ builder.set_run_totals(
+ num_agent_steps=event.get("num_turns"),
+ duration_ms=event.get("duration_ms"),
+ total_cost_usd=event.get("total_cost_usd"),
+ provider_usage=event.get("usage") if isinstance(event.get("usage"), dict) else None,
+ )
+ return builder.build()
+
+
+def _replay_records(builder: TrajectoryBuilder, records: list[dict]) -> None:
+ """Drive the builder from transcript records or stream-json events (a transcript
+ record is a stream event plus timestamps/uuids/requestId/toolUseResult)."""
+ # assistant record uuid -> timestamp, for sourceToolAssistantUUID-based start times.
+ assistant_record_ts: dict[str, str] = {}
+ saw_assistant = False
+
+ for record in records:
+ rtype = record.get("type")
+ if record.get("isSidechain"):
+ builder.count_dropped("sidechain")
+ continue
+ if rtype not in ("assistant", "user", "system"):
+ continue
+ builder.set_session_id(record.get("sessionId") or record.get("session_id"))
+ timestamp = record.get("timestamp") if isinstance(record.get("timestamp"), str) else None
+
+ if rtype == "system":
+ # stream-json emits a compaction marker as a system event.
+ if record.get("subtype") == "compact_boundary":
+ builder.add_context_boundary(timestamp=timestamp)
+ continue
+
+ message = record.get("message") or {}
+ content = message.get("content")
+
+ if rtype == "assistant":
+ saw_assistant = True
+ if isinstance(record.get("uuid"), str) and timestamp:
+ assistant_record_ts[record["uuid"]] = timestamp
+ usage = message.get("usage") or {}
+ # The transcript writes one record per content block of the same API message
+ # (identical message id and usage); start_agent_step dedupes on response_id.
+ step = builder.start_agent_step(
+ response_id=message.get("id"),
+ request_id=record.get("requestId"),
+ model=message.get("model"),
+ timestamp=timestamp,
+ stop_reason=message.get("stop_reason"),
+ provider_usage=usage or None,
+ )
+ for block in content if isinstance(content, list) else []:
+ if not isinstance(block, dict):
+ continue
+ if block.get("type") == "text" and block.get("text"):
+ builder.add_output_text(block["text"])
+ elif block.get("type") in ("thinking", "reasoning") and _thinking_text(block):
+ builder.add_reasoning(_thinking_text(block))
+ elif block.get("type") == "tool_use":
+ arguments = block.get("input")
+ builder.add_tool_call(
+ call_id=str(block.get("id") or f"call-{step.step_id}-{len(step.items)}"),
+ name=str(block.get("name") or ""),
+ arguments=json.dumps(arguments) if isinstance(arguments, dict) else str(arguments),
+ )
+
+ elif rtype == "user":
+ if record.get("isMeta"):
+ continue
+ if record.get("isCompactSummary"):
+ builder.add_context_boundary(summary=_block_text(content), timestamp=timestamp)
+ continue
+ blocks = content if isinstance(content, list) else []
+ tool_results = [b for b in blocks if isinstance(b, dict) and b.get("type") == "tool_result"]
+ for block in tool_results:
+ source_uuid = record.get("sourceToolAssistantUUID")
+ builder.add_tool_result(
+ call_id=str(block.get("tool_use_id") or ""),
+ output=_block_text(block.get("content")),
+ completed_at=timestamp,
+ started_at=assistant_record_ts.get(source_uuid) if isinstance(source_uuid, str) else None,
+ error="tool_result flagged is_error" if block.get("is_error") else None,
+ extra=_curate_extra(record.get("toolUseResult")),
+ )
+ if not tool_results:
+ text = _block_text(content)
+ # User text before any assistant record is the task's initial prompt,
+ # already recorded in responses_create_params.input — don't duplicate it.
+ if text and saw_assistant:
+ builder.add_user_message(text, timestamp=timestamp)
diff --git a/tests/unit_tests/test_trajectory.py b/tests/unit_tests/test_trajectory.py
new file mode 100644
index 0000000000..ef4fe6c099
--- /dev/null
+++ b/tests/unit_tests/test_trajectory.py
@@ -0,0 +1,377 @@
+# 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.
+
+from time import time
+
+import pytest
+
+from nemo_gym.openai_utils import (
+ NeMoGymContextBoundaryMessage,
+ NeMoGymEasyInputMessage,
+ NeMoGymFunctionCallOutput,
+ NeMoGymFunctionCallOutputWithAgentTelemetry,
+ NeMoGymResponse,
+ NeMoGymResponseFunctionToolCall,
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry,
+ NeMoGymResponseOutputMessageWithAgentTelemetry,
+ NeMoGymResponseReasoningItemWithAgentTelemetry,
+ NeMoGymToolExecutionError,
+)
+from nemo_gym.trajectory import (
+ TrajectoryBuilder,
+ agent_step_slices,
+ reconstruct_model_input,
+ summed_usage,
+ to_response_create_params,
+ usage_from_provider,
+)
+
+
+ANTHROPIC_USAGE = {
+ "input_tokens": 100,
+ "output_tokens": 20,
+ "cache_read_input_tokens": 60,
+ "cache_creation_input_tokens": 10,
+}
+OPENAI_USAGE = {
+ "input_tokens": 100,
+ "output_tokens": 20,
+ "total_tokens": 120,
+ "input_tokens_details": {"cached_tokens": 60},
+ "output_tokens_details": {"reasoning_tokens": 5},
+}
+
+
+class TestUsageFromProvider:
+ def test_anthropic_dialect(self) -> None:
+ usage = usage_from_provider(ANTHROPIC_USAGE)
+ assert usage.input_tokens == 100
+ assert usage.output_tokens == 20
+ assert usage.total_tokens == 120
+ assert usage.input_tokens_details.cached_tokens == 60
+ assert usage.output_tokens_details.reasoning_tokens == 0
+
+ def test_openai_dialect(self) -> None:
+ usage = usage_from_provider(OPENAI_USAGE)
+ assert usage.input_tokens_details.cached_tokens == 60
+ assert usage.output_tokens_details.reasoning_tokens == 5
+ assert usage.total_tokens == 120
+
+ def test_empty(self) -> None:
+ usage = usage_from_provider({})
+ assert usage.input_tokens == 0
+ assert usage.total_tokens == 0
+
+
+def _two_step_builder() -> TrajectoryBuilder:
+ """Step 1: text + two parallel tool calls; step 2: reasoning + final answer."""
+ builder = TrajectoryBuilder(agent="test_agent", source="unit_test")
+ builder.set_session_id("sess-1")
+ builder.start_agent_step(
+ response_id="msg_1",
+ request_id="req-1",
+ model="test-model",
+ timestamp="2026-07-09T00:00:01.000Z",
+ stop_reason="tool_use",
+ provider_usage=ANTHROPIC_USAGE,
+ )
+ builder.add_output_text("looking")
+ builder.add_tool_call("t1", "Bash", '{"command": "ls"}')
+ builder.add_tool_call("t2", "Read", '{"file_path": "/x"}')
+ builder.add_tool_result("t1", "file.txt", completed_at="2026-07-09T00:00:01.500Z", extra={"interrupted": False})
+ builder.add_tool_result(
+ "t2", "boom", completed_at="2026-07-09T00:00:03.000Z", error="tool_result flagged is_error"
+ )
+ builder.start_agent_step(
+ response_id="msg_2",
+ model="test-model",
+ timestamp="2026-07-09T00:00:04.000Z",
+ stop_reason="end_turn",
+ provider_usage={"input_tokens": 50, "output_tokens": 5},
+ )
+ builder.add_reasoning("hmm")
+ builder.add_output_text("fixed")
+ return builder
+
+
+class TestBuilder:
+ def test_output_is_tagged_native_items_in_execution_order(self) -> None:
+ output, _, _ = _two_step_builder().build()
+ assert [type(i) for i in output] == [
+ NeMoGymResponseOutputMessageWithAgentTelemetry,
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry, # issue order
+ NeMoGymResponseFunctionToolCallWithAgentTelemetry,
+ NeMoGymFunctionCallOutputWithAgentTelemetry, # arrival order
+ NeMoGymFunctionCallOutputWithAgentTelemetry,
+ NeMoGymResponseReasoningItemWithAgentTelemetry,
+ NeMoGymResponseOutputMessageWithAgentTelemetry,
+ ]
+ # every item carries its agent step tag — grouping without any copy
+ assert [i.agent_step_no for i in output] == [1, 1, 1, 1, 1, 2, 2]
+ assert output[0].content[0].text == "looking"
+ assert output[3].output == "file.txt"
+ assert output[5].summary[0].text == "hmm"
+
+ def test_tool_execution_rides_on_the_output_item(self) -> None:
+ output, _, _ = _two_step_builder().build()
+ out_t1, out_t2 = output[3], output[4]
+ assert out_t1.execution.started_at == "2026-07-09T00:00:01.000Z"
+ assert out_t1.execution.duration_ms == 500.0
+ assert out_t1.execution.error is None
+ assert out_t1.execution.extra == {"interrupted": False}
+ assert out_t2.execution.duration_ms == 2000.0
+ assert out_t2.execution.error.message == "tool_result flagged is_error"
+
+ def test_generations_carry_per_call_identity_and_usage(self) -> None:
+ _, generations, _ = _two_step_builder().build()
+ g1, g2 = generations
+ assert (g1.agent_step_no, g2.agent_step_no) == (1, 2)
+ assert g1.response_id == "msg_1"
+ assert g1.request_id == "req-1"
+ assert g1.stop_reason == "tool_use"
+ assert g2.stop_reason == "end_turn"
+ assert g1.usage.input_tokens == 100
+ assert g1.usage.input_tokens_details.cached_tokens == 60
+ # raw provider usage preserved verbatim (incl. fields with no native slot)
+ assert g1.provider_usage["cache_creation_input_tokens"] == 10
+
+ def test_same_response_id_continues_step_without_double_count(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.start_agent_step(response_id="m1", provider_usage={"input_tokens": 10, "output_tokens": 5})
+ builder.add_output_text("part 1")
+ builder.start_agent_step(
+ response_id="m1", provider_usage={"input_tokens": 10, "output_tokens": 5}, stop_reason="end_turn"
+ )
+ builder.add_output_text("part 2")
+ output, generations, _ = builder.build()
+ assert len(generations) == 1
+ assert generations[0].stop_reason == "end_turn"
+ assert summed_usage(generations).input_tokens == 10
+ # text accumulated into one message item under the same tag
+ assert [b.text for b in output[0].content] == ["part 1", "part 2"]
+
+ def test_output_without_step_raises(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ with pytest.raises(ValueError):
+ builder.add_output_text("no step")
+
+ def test_mid_episode_user_message_is_plain_and_closes_the_step(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.start_agent_step(response_id="m1")
+ builder.add_output_text("answer 1")
+ builder.add_user_message("follow-up")
+ builder.start_agent_step(response_id="m1") # same provider id, but step was closed by the user turn
+ builder.add_output_text("answer 2")
+ output, generations, _ = builder.build()
+ assert type(output[1]) is NeMoGymEasyInputMessage # untagged: not produced by a step
+ assert output[1].content == "follow-up"
+ assert [g.agent_step_no for g in generations] == [1, 2]
+ assert output[2].agent_step_no == 2
+
+ def test_no_timestamps_means_no_fabricated_timing(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.start_agent_step(response_id="m1")
+ builder.add_tool_call("t1", "Bash", "{}")
+ builder.add_tool_result("t1", "ok")
+ output, _, _ = builder.build()
+ execution = output[-1].execution
+ assert execution.started_at is None
+ assert execution.completed_at is None
+ assert execution.duration_ms is None
+
+ def test_explicit_started_at_overrides_registration(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.start_agent_step(response_id="m1", timestamp="2026-07-09T00:00:00.000Z")
+ builder.add_tool_call("t1", "Bash", "{}")
+ builder.add_tool_result(
+ "t1", "ok", started_at="2026-07-09T00:00:02.000Z", completed_at="2026-07-09T00:00:03.000Z"
+ )
+ output, _, _ = builder.build()
+ assert output[-1].execution.duration_ms == 1000.0
+
+ def test_orphan_result_with_no_step_counted_as_dropped(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.add_tool_result("orphan", "out")
+ output, _, telemetry = builder.build()
+ assert output == []
+ assert telemetry.dropped_records == {"orphan_tool_results": 1}
+
+ def test_unmatched_result_attaches_to_last_step(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.start_agent_step(response_id="m1")
+ builder.add_tool_call("known", "Bash", "{}")
+ builder.add_tool_result("known", "ok")
+ builder.add_tool_result("orphan", "late")
+ output, _, telemetry = builder.build()
+ assert output[-1].call_id == "orphan"
+ assert output[-1].agent_step_no == 1
+ assert telemetry.dropped_records == {}
+
+ def test_error_can_be_structured(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.start_agent_step(response_id="m1")
+ builder.add_tool_call("t1", "Bash", "{}")
+ builder.add_tool_result("t1", "boom", error=NeMoGymToolExecutionError(message="timeout", data={"signal": 9}))
+ output, _, _ = builder.build()
+ assert output[-1].execution.error.data == {"signal": 9}
+
+ def test_run_totals_and_session(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.set_session_id("sess-1")
+ builder.set_session_id("sess-2") # first one wins
+ builder.count_dropped("sidechain")
+ builder.set_run_totals(
+ num_agent_steps=3, duration_ms=42.0, total_cost_usd=0.5, provider_usage={"input_tokens": 1}
+ )
+ _, _, telemetry = builder.build()
+ assert telemetry.agent == "a"
+ assert telemetry.source == "s"
+ assert telemetry.session_id == "sess-1"
+ assert telemetry.dropped_records == {"sidechain": 1}
+ assert telemetry.num_agent_steps == 3
+ assert telemetry.duration_ms == 42.0
+ assert telemetry.total_cost_usd == 0.5
+ assert telemetry.provider_usage == {"input_tokens": 1}
+
+ def test_summed_usage(self) -> None:
+ _, generations, _ = _two_step_builder().build()
+ totals = summed_usage(generations)
+ assert totals.input_tokens == 150
+ assert totals.output_tokens == 25
+ assert totals.total_tokens == 175
+ assert totals.input_tokens_details.cached_tokens == 60
+
+ def test_agent_step_slices(self) -> None:
+ output, _, _ = _two_step_builder().build()
+ slices = dict(agent_step_slices(output))
+ assert [len(items) for items in slices.values()] == [5, 2]
+ assert slices[2][0].summary[0].text == "hmm"
+
+
+class TestReconstruction:
+ def _episode(self):
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.start_agent_step(response_id="m1", model="test-model")
+ builder.add_output_text("a1")
+ builder.add_context_boundary(summary="summary of q1/a1")
+ builder.add_user_message("q2")
+ builder.start_agent_step(response_id="m2", model="test-model")
+ builder.add_output_text("a2")
+ output, generations, telemetry = builder.build()
+ base_input = [NeMoGymEasyInputMessage(role="user", content="q1")]
+ return output, base_input
+
+ def test_step_1_sees_base_input_only(self) -> None:
+ output, base = self._episode()
+ items = reconstruct_model_input(output, agent_step_no=1, base_input=base)
+ assert [i.content for i in items] == ["q1"]
+
+ def test_post_boundary_step_sees_summary_not_base(self) -> None:
+ output, base = self._episode()
+ items = reconstruct_model_input(output, agent_step_no=2, base_input=base)
+ assert [getattr(i, "content", None) for i in items] == ["summary of q1/a1", "q2"]
+ assert isinstance(items[0], NeMoGymContextBoundaryMessage)
+
+ def test_full_reconstruction(self) -> None:
+ output, base = self._episode()
+ items = reconstruct_model_input(output, base_input=base)
+ assert len(items) == 3 # summary, q2, a2 — the boundary replaced q1 + a1
+ assert items[-1].content[0].text == "a2"
+
+ def test_no_boundary_prepends_base_input(self) -> None:
+ builder = TrajectoryBuilder(agent="a", source="s")
+ builder.start_agent_step(response_id="m1")
+ builder.add_output_text("a1")
+ builder.add_user_message("q2")
+ builder.start_agent_step(response_id="m2")
+ output, _, _ = builder.build()
+ base = [NeMoGymEasyInputMessage(role="user", content="q1")]
+ items = reconstruct_model_input(output, agent_step_no=2, base_input=base)
+ assert [getattr(i, "content", None) for i in items][0] == "q1"
+ assert len(items) == 3 # q1 + a1 + q2
+
+ def test_to_response_create_params_is_native(self) -> None:
+ output, base = self._episode()
+ params = to_response_create_params(output, base_input=base, model="test-model")
+ assert params.model == "test-model"
+ assert type(params).model_validate(params.model_dump(mode="json")).model == "test-model"
+
+
+class TestContractRoundTrip:
+ """Telemetry must survive NeMoGymResponse validation — the whole point of putting it
+ on the contract: every server hop revalidates, and nothing may be stripped."""
+
+ def _response(self) -> NeMoGymResponse:
+ output, generations, telemetry = _two_step_builder().build()
+ return NeMoGymResponse(
+ id="resp_x",
+ created_at=int(time()),
+ model="test-model",
+ object="response",
+ output=output,
+ tool_choice="auto",
+ tools=[],
+ parallel_tool_calls=True,
+ generations=generations,
+ agent_telemetry=telemetry,
+ )
+
+ def test_items_keep_telemetry_through_revalidation(self) -> None:
+ response = self._response()
+ revalidated = NeMoGymResponse.model_validate(response.model_dump(mode="json"))
+ assert [getattr(i, "agent_step_no", None) for i in revalidated.output] == [1, 1, 1, 1, 1, 2, 2]
+ assert revalidated.output[3].execution.duration_ms == 500.0
+ assert revalidated == response
+
+ def test_generations_and_telemetry_survive(self) -> None:
+ revalidated = NeMoGymResponse.model_validate(self._response().model_dump(mode="json"))
+ assert revalidated.generations[0].provider_usage["cache_creation_input_tokens"] == 10
+ assert revalidated.agent_telemetry.source == "unit_test"
+ assert revalidated.agent_telemetry.schema_version == "1.0"
+
+ def test_plain_payloads_stay_plain(self) -> None:
+ # A model-server response (no telemetry) must validate to the plain classes and
+ # keep generations/agent_telemetry as None.
+ plain = {
+ "id": "resp_y",
+ "created_at": 0,
+ "model": "m",
+ "object": "response",
+ "tool_choice": "auto",
+ "tools": [],
+ "parallel_tool_calls": True,
+ "output": [
+ {"type": "function_call", "call_id": "c", "name": "f", "arguments": "{}"},
+ {"type": "function_call_output", "call_id": "c", "output": "ok"},
+ ],
+ }
+ response = NeMoGymResponse.model_validate(plain)
+ assert type(response.output[0]) is NeMoGymResponseFunctionToolCall
+ assert type(response.output[1]) is NeMoGymFunctionCallOutput
+ assert response.generations is None
+ assert response.agent_telemetry is None
+
+ def test_context_boundary_survives_and_plain_user_message_stays_plain(self) -> None:
+ boundary = NeMoGymContextBoundaryMessage(role="user", content="summary", context_boundary=True)
+ plain_user = {"role": "user", "content": "hi", "type": "message"}
+ response = self._response()
+ response.output = [boundary] + response.output
+ revalidated = NeMoGymResponse.model_validate(response.model_dump(mode="json"))
+ assert isinstance(revalidated.output[0], NeMoGymContextBoundaryMessage)
+ from pydantic import TypeAdapter
+
+ from nemo_gym.openai_utils import NeMoGymResponseInputItem
+
+ assert type(TypeAdapter(NeMoGymResponseInputItem).validate_python(plain_user)) is NeMoGymEasyInputMessage