From 3c973bb18148b7f62b3353fdecb9fa3c07a372bc Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Thu, 23 Jul 2026 15:07:55 +0200 Subject: [PATCH 01/14] Define rollout observation and correlation contract Signed-off-by: Michal Bien --- .../pages/model-server/model-call-capture.mdx | 40 ++- nemo_gym/base_responses_api_model.py | 252 ++++++++++---- nemo_gym/rollout_collection.py | 6 +- nemo_gym/rollout_observability.py | 328 ++++++++++++++++++ .../test_base_responses_api_model.py | 255 ++++++++++++-- tests/unit_tests/test_rollout_collection.py | 9 +- .../unit_tests/test_rollout_observability.py | 244 +++++++++++++ 7 files changed, 1037 insertions(+), 97 deletions(-) create mode 100644 nemo_gym/rollout_observability.py create mode 100644 tests/unit_tests/test_rollout_observability.py diff --git a/fern/versions/latest/pages/model-server/model-call-capture.mdx b/fern/versions/latest/pages/model-server/model-call-capture.mdx index 8f2a5c0ec3..fb8618a9f2 100644 --- a/fern/versions/latest/pages/model-server/model-call-capture.mdx +++ b/fern/versions/latest/pages/model-server/model-call-capture.mdx @@ -41,6 +41,8 @@ Use `/ng-rollout/` as the model-server URL prefix. SDKs append their the rollout prefix before routing. Correlation is caller-supplied. An unprefixed call is forwarded normally but is not captured. +SDK-based harnesses must configure their `model_server` field for capture; calls sent directly to an +external provider do not cross a Gym model server. Agents built on `SimpleResponsesAPIAgent` can use these helpers: `url_path_for_run(url_path, body)` prefixes a downstream call from the run request's task/rollout indices (only when @@ -77,9 +79,10 @@ totals = aggregate_model_call_metrics(store, rollout_id) ``` `ModelCallRecord` is an observability serialization model derived from captured HTTP exchanges. It -contains a unique server-generated `model_call_id`, typed `model_ref`, wall-clock `started_at` and -`completed_at`, a `call_index`, API dialect, token and cache usage, latency, error details, tool -calls, reasoning content, and the captured request and response. `started_at` is recorded immediately +contains a unique server-generated `model_call_id`, the protocol `response_id` when present, typed +`model_ref`, wall-clock `started_at` and `completed_at`, a `call_index`, API dialect, token and cache +usage, latency, error details, tool calls, reasoning content, and the captured request and response. +`started_at` is recorded immediately before invoking the downstream ASGI application; `completed_at` is recorded when that invocation returns or raises, before capture parsing and persistence. Both are UTC Unix seconds for external trace correlation; durations use the monotonic latency fields. `call_index` reflects durable append @@ -111,7 +114,32 @@ data does not mix with an earlier attempt. Before dispatch, the collector clears for that exact rollout-attempt id, including a kill-shaped attempt being redispatched. The attachment is additive: it does not replace or rewrite the existing response, reward, -`NeMoGymResponse`, token-id, or log-prob fields. Downstream consumers can choose whether to read it. +`NeMoGymResponse`, token-id, or log-prob fields. Downstream consumers can choose whether to read it, +and aggregate-metrics requests exclude it. + +## Agent observations + +Supported Agent Servers may also attach `ng_agent_observations` when observability is enabled. It +contains an unordered `records` list of typed agent invocations, tool-call intervals, explicit +context-compaction events, and sandbox observations. `gaps` reports unavailable evidence. An +invocation's `conversation` contains the ordered, normalized items exposed by that integration. + +Agent observations and model-call capture are separate evidence. Join an invocation's model-call +references by `model_call_id`, or by the exact `(model_ref, response_id)` pair when the harness sees +the protocol response ID. Do not infer ownership from timestamps, text, or list position. The full +model request and response remain in `CaptureStore`; rollout attachments intentionally omit them. + +Tool timestamps are UTC Unix seconds when the source provides wall-clock time. `duration_ms` is the +measured interval, and `timing_source` identifies executor, harness, or artifact-derived timing. + +Model-visible tool calls and results remain in `AgentInvocation.conversation` as `function_call` and +`function_call_output` items. Execution timing and outcome, when observable, live in +`ToolCallObservation` and join through `(invocation_id, tool_call_id)`. A tool observation may also +reference its enclosing `sandbox_id`; concurrent calls retain independent timing and outcome. + +Each `SandboxObservation` covers one sandbox execution. Usage fields contain measured values only; +configured limits are never reported as usage. Integrations emit only facts available at their +execution boundary or in retained artifacts and report unavailable evidence in `gaps`. ## Limitations @@ -119,3 +147,7 @@ The model-server boundary observes model HTTP requests and responses. It can rec reasoning, and tool results present in those payloads, but it does not observe the actual tool execution boundary, environment events, context compaction, semantic turns, or subagent structure. Those require instrumentation at the agent, tool, environment, or rollout layer. + +Sandbox CPU and memory observations describe the sandbox as a whole. They cannot be attributed to +individual overlapping tool calls unless each call runs in a separately measured execution scope. +Gym therefore records the shared sandbox relationship without estimating per-tool resource usage. diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 9fadeea62c..6326e06776 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -33,7 +33,6 @@ import logging import os import re -import threading import time from abc import abstractmethod from pathlib import Path @@ -66,6 +65,7 @@ synthesize_responses_sse, validate_streaming_responses_params, ) +from nemo_gym.rollout_observability import AgentObservationBundle, ObservationGap, join_model_call_observations from nemo_gym.server_utils import ( BaseRunServerInstanceConfig, BaseServer, @@ -285,7 +285,6 @@ class CaptureStore: def __init__(self, root: str | Path) -> None: self._root = Path(root) self._root.mkdir(parents=True, exist_ok=True) - self._lock = threading.Lock() @property def root(self) -> Path: @@ -294,25 +293,32 @@ def root(self) -> Path: def path_for(self, rollout_id: str) -> Path: return self._root / f"{_validate_rollout_id(rollout_id)}.capture.jsonl" + def incomplete_path_for(self, rollout_id: str) -> Path: + return self._root / f"{_validate_rollout_id(rollout_id)}.capture.incomplete" + + def mark_incomplete(self, rollout_id: str) -> None: + self.incomplete_path_for(rollout_id).touch(exist_ok=True) + + def is_incomplete(self, rollout_id: str) -> bool: + return self.incomplete_path_for(rollout_id).exists() + def record(self, rollout_id: str, exchange: dict[str, Any]) -> None: """Append one exchange and fsync (durable across a killed box). - ``flock`` serializes appends across worker processes (a model server may run with - ``num_workers > 1``, where the in-process lock can't coordinate); the in-process lock - serializes threads. This does blocking file IO + fsync, so callers run it off the event - loop (the capture middleware offloads it via ``asyncio.to_thread``). + ``flock`` serializes appends to the same rollout across worker processes and threads while + allowing independent rollouts to write concurrently. This does blocking file IO + fsync, + so callers run it off the event loop (the capture middleware uses ``asyncio.to_thread``). """ line = orjson.dumps(exchange, default=str, option=orjson.OPT_APPEND_NEWLINE) path = self.path_for(rollout_id) - with self._lock: - with path.open("ab") as handle: - fcntl.flock(handle.fileno(), fcntl.LOCK_EX) - try: - handle.write(line) - handle.flush() - os.fsync(handle.fileno()) - finally: - fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + with path.open("ab") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX) + try: + handle.write(line) + handle.flush() + os.fsync(handle.fileno()) + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) def read(self, rollout_id: str) -> list[dict[str, Any]]: path = self.path_for(rollout_id) @@ -320,19 +326,47 @@ def read(self, rollout_id: str) -> list[dict[str, Any]]: return [] exchanges: list[dict[str, Any]] = [] # Stream line-by-line; a capture can be large (token-ids / logprobs). - with self._lock: - with path.open("rb") as handle: - fcntl.flock(handle.fileno(), fcntl.LOCK_SH) - try: - for line in handle: - stripped = line.strip() - if not stripped: - continue - exchanges.append(orjson.loads(stripped)) - finally: - fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + with path.open("rb") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_SH) + try: + for line in handle: + stripped = line.strip() + if not stripped: + continue + exchanges.append(orjson.loads(stripped)) + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) return exchanges + def read_available(self, rollout_id: str) -> tuple[list[tuple[int, dict[str, Any]]], int]: + """Read valid exchanges without letting one damaged line hide the rest.""" + path = self.path_for(rollout_id) + if not path.exists(): + return [], 0 + exchanges: list[tuple[int, dict[str, Any]]] = [] + invalid_count = 0 + capture_index = 0 + with path.open("rb") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_SH) + try: + for line in handle: + stripped = line.strip() + if not stripped: + continue + try: + exchange = orjson.loads(stripped) + except orjson.JSONDecodeError: + invalid_count += 1 + else: + if isinstance(exchange, dict): + exchanges.append((capture_index, exchange)) + else: + invalid_count += 1 + capture_index += 1 + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + return exchanges, invalid_count + def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]: """Per-rollout model-call capture id from a run-request's task/rollout indices. @@ -365,6 +399,10 @@ def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) - # --- Observability records derived from captured exchanges --- +def _token_count(value: Any) -> Optional[int]: + return value if type(value) is int and value >= 0 else None + + def extract_token_stats(usage: Any) -> dict[str, Optional[int]]: """Normalize token totals across Responses, Chat Completions, and Anthropic Messages usage. @@ -378,7 +416,7 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]: cache-creation (~1.25x) differently from base input, so cost-accurate consumers should weight ``cached_tokens`` and ``cache_creation_tokens`` separately rather than summing ``tokens_in``. """ - if not usage: + if not isinstance(usage, Mapping): return { "tokens_in": None, "tokens_out": None, @@ -386,28 +424,30 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]: "tokens_total": None, "cache_creation_tokens": None, } - tokens_in = usage.get("input_tokens") + tokens_in = _token_count(usage.get("input_tokens")) if tokens_in is None: - tokens_in = usage.get("prompt_tokens") - tokens_out = usage.get("output_tokens") + tokens_in = _token_count(usage.get("prompt_tokens")) + tokens_out = _token_count(usage.get("output_tokens")) if tokens_out is None: - tokens_out = usage.get("completion_tokens") + tokens_out = _token_count(usage.get("completion_tokens")) # Anthropic-native shape: top-level cache_* keys mean input_tokens excludes cached tokens. - cache_read = usage.get("cache_read_input_tokens") - cache_creation = usage.get("cache_creation_input_tokens") + cache_read = _token_count(usage.get("cache_read_input_tokens")) + cache_creation = _token_count(usage.get("cache_creation_input_tokens")) if cache_read is not None or cache_creation is not None: # A fully-cached response can omit input_tokens; use a 0 base so the folded prompt size is # preserved rather than dropped to null. (Top-level cache_* keys are Anthropic-only, so the # OpenAI/Responses path -- nested prompt_tokens_details.cached_tokens -- never enters here.) tokens_in = (tokens_in or 0) + (cache_read or 0) + (cache_creation or 0) - tokens_total = usage.get("total_tokens") + tokens_total = _token_count(usage.get("total_tokens")) if tokens_total is None and tokens_in is not None and tokens_out is not None: tokens_total = tokens_in + tokens_out details = usage.get("output_tokens_details") or usage.get("completion_tokens_details") or {} + if not isinstance(details, Mapping): + details = {} return { "tokens_in": tokens_in, "tokens_out": tokens_out, - "tokens_reasoning": details.get("reasoning_tokens"), + "tokens_reasoning": _token_count(details.get("reasoning_tokens")), "tokens_total": tokens_total, "cache_creation_tokens": cache_creation, } @@ -415,12 +455,14 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]: def _cache_signal(usage: Any) -> tuple[Optional[bool], Optional[int]]: """Cache hit/miss + cached-token count, from usage cache fields (OpenAI / Anthropic).""" - if not usage: + if not isinstance(usage, Mapping): return None, None details = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") or {} - cached = details.get("cached_tokens") + if not isinstance(details, Mapping): + details = {} + cached = _token_count(details.get("cached_tokens")) if cached is None: - cached = usage.get("cache_read_input_tokens") # Anthropic + cached = _token_count(usage.get("cache_read_input_tokens")) # Anthropic if cached is None: return None, None return cached > 0, cached @@ -458,7 +500,7 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str, elif item.get("type") == "reasoning": for summary in item.get("summary") or []: text = summary.get("text") if isinstance(summary, dict) else None - if text: + if isinstance(text, str) and text: reasoning.append(text) return tool_calls, ("\n".join(reasoning) or None) @@ -466,19 +508,21 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str, if isinstance(choices, list): # Chat Completions for choice in choices: message = choice.get("message") if isinstance(choice, dict) else None - if not message: + if not isinstance(message, Mapping): continue for tc in message.get("tool_calls") or []: if not isinstance(tc, dict): continue fn = tc.get("function") or {} + if not isinstance(fn, Mapping): + fn = {} tool_calls.append( {"call_id": tc.get("id"), "name": fn.get("name"), "arguments": _as_arguments(fn.get("arguments"))} ) # vLLM and newer OpenAI-compatible servers emit `reasoning`; `reasoning_content` is the # older field. Accept either (reasoning_content wins when both are present). reasoning_text = message.get("reasoning_content") or message.get("reasoning") - if reasoning_text: + if isinstance(reasoning_text, str) and reasoning_text: reasoning.append(reasoning_text) return tool_calls, ("\n".join(reasoning) or None) @@ -491,7 +535,7 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str, tool_calls.append( {"call_id": block.get("id"), "name": block.get("name"), "arguments": block.get("input") or {}} ) - elif block.get("type") in ("thinking", "redacted_thinking") and block.get("thinking"): + elif block.get("type") in ("thinking", "redacted_thinking") and isinstance(block.get("thinking"), str): reasoning.append(block["thinking"]) return tool_calls, ("\n".join(reasoning) or None) @@ -503,12 +547,15 @@ class ModelCallRecord(BaseModel): # Unique server-generated identity for each persisted call. model_call_id: Optional[str] = None + response_id: Optional[str] = None # Durable append order, not a causal or semantic order for concurrent calls. call_index: int model_ref: Optional[ModelServerRef] = None + model: Optional[str] = None dialect: Optional[str] = None status_code: Optional[int] = None + finish_reason: Optional[str] = None # Wall-clock bounds around the downstream ASGI invocation, as UTC Unix timestamps. These are # for external trace correlation; durations use the monotonic latency fields below. @@ -526,6 +573,8 @@ class ModelCallRecord(BaseModel): # Model-call record. request: Optional[dict[str, Any]] = None response: Optional[dict[str, Any]] = None + request_raw: Optional[str] = None + response_raw: Optional[str] = None tool_calls: list[dict[str, Any]] = Field(default_factory=list) # Structured reasoning (not flattened into the response text). @@ -547,20 +596,46 @@ class ModelCallRecord(BaseModel): def build_model_call_record(exchange: dict[str, Any], *, call_index: int) -> ModelCallRecord: """Map one captured exchange and its transport metadata into an observability record.""" - response = exchange.get("response") or {} + raw_response = exchange.get("response") + response = raw_response if isinstance(raw_response, dict) else {} tokens = extract_token_stats(response.get("usage")) cache_hit, cached_tokens = _cache_signal(response.get("usage")) tool_calls, reasoning_content = _tool_calls_and_reasoning(response) + raw_request = exchange.get("request") + request = raw_request if isinstance(raw_request, dict) else {} + choices = response.get("choices") + first_choice = choices[0] if isinstance(choices, list) and choices and isinstance(choices[0], dict) else {} + incomplete_details = response.get("incomplete_details") + if not isinstance(incomplete_details, dict): + incomplete_details = {} + finish_reason = next( + ( + value + for value in ( + response.get("stop_reason"), + first_choice.get("finish_reason"), + incomplete_details.get("reason"), + ) + if isinstance(value, str) + ), + None, + ) + model = response.get("model") or request.get("model") return ModelCallRecord( model_call_id=exchange.get("model_call_id"), + response_id=response.get("id") if isinstance(response.get("id"), str) else None, call_index=call_index, model_ref=exchange.get("model_ref"), + model=model if isinstance(model, str) else None, dialect=exchange.get("dialect"), status_code=exchange.get("status_code"), + finish_reason=finish_reason, started_at=exchange.get("started_at"), completed_at=exchange.get("completed_at"), - request=exchange.get("request"), - response=response or None, + request=raw_request if isinstance(raw_request, dict) else None, + response=raw_response if isinstance(raw_response, dict) else None, + request_raw=exchange.get("request_raw") if isinstance(exchange.get("request_raw"), str) else None, + response_raw=exchange.get("response_raw") if isinstance(exchange.get("response_raw"), str) else None, tool_calls=tool_calls, reasoning_content=reasoning_content, cache_hit=cache_hit, @@ -579,6 +654,18 @@ def read_model_call_records(store: CaptureStore, rollout_id: str) -> list[ModelC ] +def read_available_model_call_records(store: CaptureStore, rollout_id: str) -> tuple[list[ModelCallRecord], int]: + """Read valid call records and count damaged records.""" + exchanges, invalid_count = store.read_available(rollout_id) + calls = [] + for index, exchange in exchanges: + try: + calls.append(build_model_call_record(exchange, call_index=index)) + except Exception: + invalid_count += 1 + return calls, invalid_count + + def aggregate_model_call_records(calls: list[ModelCallRecord]) -> dict[str, Any]: """Aggregate token and latency values from model-call records.""" @@ -776,11 +863,14 @@ def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, An tool_calls: dict[int, dict[str, Any]] = {} usage: Optional[dict[str, Any]] = None model: Optional[str] = None + response_id: Optional[str] = None role = "assistant" finish_reason: Optional[str] = None saw_choice = False for chunk in events: model = chunk.get("model") or model + if isinstance(chunk.get("id"), str): + response_id = chunk["id"] if chunk.get("usage"): usage = chunk["usage"] for choice in chunk.get("choices") or []: @@ -819,6 +909,8 @@ def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, An "model": model, "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], } + if response_id is not None: + result["id"] = response_id if usage: result["usage"] = usage return result @@ -864,6 +956,7 @@ def _record( error_category: Optional[str], latency_ms: float, ttft_ms: Optional[float] = None, + response_raw: Optional[str] = None, ) -> None: """Append one exchange (success or failure). Best-effort: never raises.""" request_body = None @@ -894,9 +987,15 @@ def _record( } if request_raw is not None: exchange["request_raw"] = request_raw + if response_raw is not None: + exchange["response_raw"] = response_raw store.record(rollout_id, exchange) except Exception: logger.warning("Model-call capture failed for one %s call.", dialect, exc_info=True) + try: + store.mark_incomplete(rollout_id) + except Exception: + logger.warning("Could not mark rollout %s capture as incomplete.", rollout_id, exc_info=True) class _CaptureMiddleware: @@ -1013,10 +1112,11 @@ async def _flush_deferred_response() -> None: started_at=started_at, completed_at=completed_at, response_body=None, - status_code=None, + status_code=state["status"], error_category=_classify_exception(exc), latency_ms=(time.perf_counter() - start) * 1000.0, ttft_ms=state["ttft_ms"], + response_raw=(bytes(state["body"]).decode("utf-8", errors="replace") if state["body"] else None), ) except Exception: logger.warning("Model-call capture finalization failed.", exc_info=True) @@ -1043,6 +1143,8 @@ def _parse_and_record() -> None: response_body = ( _reconstruct_streamed_response(body_bytes, dialect) if streaming else json.loads(body_bytes) ) + if not isinstance(response_body, dict): + response_body = None except Exception: response_body = None error_category = _classify_status(status) if status is not None else None @@ -1057,6 +1159,11 @@ def _parse_and_record() -> None: # doesn't silently count as a success with null tokens in reliability/cost sums. if error_category is None and body_bytes and response_body is None: error_category = "capture_parse_error" + response_raw = ( + body_bytes.decode("utf-8", errors="replace") + if body_bytes and (streaming or response_body is None) + else None + ) _record( store, dialect, @@ -1071,6 +1178,7 @@ def _parse_and_record() -> None: error_category=error_category, latency_ms=latency_ms, ttft_ms=ttft_ms, + response_raw=response_raw, ) try: @@ -1115,7 +1223,7 @@ def model_call_capture_dirs_from_config(global_config_dict: Any) -> list[Path]: def _store_for_rollout(rollout_id: str, capture_dirs: list[Path]) -> Optional[CaptureStore]: for directory in capture_dirs: store = CaptureStore(directory) - if store.path_for(rollout_id).exists(): + if store.path_for(rollout_id).exists() or store.is_incomplete(rollout_id): return store return None @@ -1135,34 +1243,60 @@ def clear_model_call_captures_for_rollouts(records: list[Any], capture_dirs: lis rollout_id = maybe_rollout_id_from_run_body(record) if rollout_id: store.path_for(rollout_id).unlink(missing_ok=True) + store.incomplete_path_for(rollout_id).unlink(missing_ok=True) -def merge_model_call_capture_into_record( - record: dict[str, Any], capture_dirs: list[Path], *, include_payloads: bool = False -) -> dict[str, Any]: +def merge_model_call_capture_into_record(record: dict[str, Any], capture_dirs: list[Path]) -> dict[str, Any]: """Attach captured model-call observability data to a rollout record in place. Keyed by the rollout id derived from the record's task/rollout/attempt indices, so the attached shape is identical for every agent harness. Adds ``ng_model_call_capture = {rollout_id, metrics, calls}`` where ``calls`` are derived observability - records. Raw request and response payloads remain in the capture store unless ``include_payloads`` - is true. No-op when no capture exists. The harness output and reward are not modified. + records. Raw request and response payloads remain in the capture store. Capture/read/join + failures are attached as ``gaps``. The harness output and reward are not modified. """ if not capture_dirs: return record rollout_id = maybe_rollout_id_from_run_body(record) if rollout_id is None: return record + gaps: list[ObservationGap] = [] store = _store_for_rollout(rollout_id, capture_dirs) if store is None: - return record - calls = read_model_call_records(store, rollout_id) - if not calls: - return record - exclude = None if include_payloads else {"request", "response"} - record["ng_model_call_capture"] = { + calls = [] + gaps.append(ObservationGap(code="model_call_capture_no_records")) + else: + try: + calls, invalid_count = read_available_model_call_records(store, rollout_id) + if invalid_count: + gaps.append( + ObservationGap( + code="model_call_capture_records_unreadable", + detail=f"count={invalid_count}", + ) + ) + if store.is_incomplete(rollout_id): + gaps.append(ObservationGap(code="model_call_capture_incomplete")) + elif not calls and not invalid_count: + gaps.append(ObservationGap(code="model_call_capture_no_records")) + except Exception: + logger.warning("Could not read model-call capture for rollout %s.", rollout_id, exc_info=True) + calls = [] + gaps.append(ObservationGap(code="model_call_capture_unreadable")) + observations = record.get("ng_agent_observations") + if observations is not None and calls: + try: + bundle = AgentObservationBundle.model_validate(observations) + record["ng_agent_observations"] = join_model_call_observations(bundle, calls).model_dump(mode="json") + except Exception: + logger.warning("Could not join agent observations for rollout %s.", rollout_id, exc_info=True) + gaps.append(ObservationGap(code="agent_observation_join_failed")) + capture = { "rollout_id": rollout_id, "metrics": aggregate_model_call_records(calls), - "calls": [call.model_dump(exclude=exclude) for call in calls], + "calls": [call.model_dump(exclude={"request", "response", "request_raw", "response_raw"}) for call in calls], } + if gaps: + capture["gaps"] = [gap.model_dump(mode="json", exclude_none=True) for gap in gaps] + record["ng_model_call_capture"] = capture return record diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index 2fb0bd15b1..17b7cee014 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -653,7 +653,11 @@ async def _fetch_agent_metrics(agent_name: str, agent_result_list: List[Dict]) - # Strip heavyweight fields before sending, but preserve response.usage stripped = [] for r in agent_result_list: - entry = {k: v for k, v in r.items() if k not in ("response", "responses_create_params")} + entry = { + k: v + for k, v in r.items() + if k not in ("response", "responses_create_params", "ng_agent_observations") + } usage = (r.get("response") or {}).get("usage") if usage: entry["response"] = {"usage": usage} diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py new file mode 100644 index 0000000000..dd722f7b2a --- /dev/null +++ b/nemo_gym/rollout_observability.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small shared contract for observations exposed by Agent integrations.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseInputItem + + +if TYPE_CHECKING: + from nemo_gym.base_responses_api_model import ModelCallRecord + + +class ObservationModel(BaseModel): + model_config = ConfigDict(extra="forbid", validate_assignment=True) + + +class ModelCallRef(ObservationModel): + """Stable identifiers an Agent integration can observe for one model call.""" + + model_call_id: Optional[str] = None + model_ref: Optional[ModelServerRef] = None + response_id: Optional[str] = None + + @model_validator(mode="after") + def validate_join_key(self) -> "ModelCallRef": + if not self.model_call_id and not (self.model_ref is not None and self.response_id): + raise ValueError("model_call_id or both model_ref and response_id are required") + return self + + +class AgentInvocation(ObservationModel): + """One root Agent or subagent conversation observed by a harness.""" + + kind: Literal["agent_invocation"] = "agent_invocation" + invocation_id: str + parent_invocation_id: Optional[str] = None + spawned_by_tool_call_id: Optional[str] = None + status: Literal["completed", "failed", "incomplete", "unknown"] = Field( + default="unknown", description="Harness-reported invocation outcome; unknown when not explicit." + ) + model_calls: list[ModelCallRef] = Field(default_factory=list) + conversation: list[NeMoGymResponseInputItem] = Field( + default_factory=list, + description="Normalized conversation items supported by this producer; gaps describe unavailable evidence.", + ) + + +class ToolCallObservation(ObservationModel): + """Timing observed for one tool call at an Agent-owned boundary.""" + + kind: Literal["tool_call"] = "tool_call" + invocation_id: str + tool_call_id: str + sandbox_id: Optional[str] = Field( + default=None, + description="Enclosing sandbox instance, shared by concurrent calls; not per-call resource attribution.", + ) + tool_name: Optional[str] = None + started_at: Optional[float] = None + completed_at: Optional[float] = None + duration_ms: Optional[float] = None + timing_source: Optional[Literal["executor", "artifact", "harness"]] = None + status: Literal["completed", "failed", "timeout", "cancelled", "incomplete", "unknown"] = "unknown" + error_type: Optional[str] = None + + +class SandboxObservation(ObservationModel): + """Outcome and lifetime resource usage reported by a sandbox-owning harness.""" + + kind: Literal["sandbox"] = "sandbox" + role: Literal["agent", "verifier", "environment"] + provider: Optional[str] = None + sandbox_id: Optional[str] = None + outcome: Literal["completed", "failed", "timeout", "oom", "sandbox_error", "cancelled", "unknown"] = "unknown" + exit_code: Optional[int] = None + wall_time_s: Optional[float] = Field(default=None, ge=0) + cpu_time_s: Optional[float] = Field( + default=None, + ge=0, + description="Cumulative CPU time for the sandbox, never an allocation or per-tool estimate.", + ) + peak_memory_mib: Optional[float] = Field( + default=None, + ge=0, + description="Measured sandbox high-water mark, never its configured memory limit.", + ) + resource_usage_source: Optional[str] = None + error_type: Optional[str] = None + + +class ContextCompactionObservation(ObservationModel): + """An explicit context-compaction event reported by the Agent harness.""" + + kind: Literal["context_compaction"] = "context_compaction" + invocation_id: str + observed_at: Optional[float] = None + trigger: Optional[str] = None + tokens_before: Optional[int] = None + tokens_after: Optional[int] = None + outcome: Literal["completed", "failed", "aborted", "unknown"] = "unknown" + summary: Optional[str] = None + first_kept_item_id: Optional[str] = None + before_model_call: Optional[ModelCallRef] = None + after_model_call: Optional[ModelCallRef] = None + + +class ObservationGap(ObservationModel): + """A fact that the selected integration could not observe or join exactly.""" + + code: str + invocation_id: Optional[str] = None + detail: Optional[str] = None + + +AgentObservationRecord = Annotated[ + AgentInvocation | ToolCallObservation | ContextCompactionObservation | SandboxObservation, + Field(discriminator="kind"), +] + + +class AgentObservationBundle(ObservationModel): + """Normalized observations returned by one Agent Server for one rollout.""" + + source: str + records: list[AgentObservationRecord] = Field( + default_factory=list, + description="Unordered typed records; list position does not imply execution order.", + ) + gaps: list[ObservationGap] = Field(default_factory=list) + + @model_validator(mode="after") + def validate_identity(self) -> "AgentObservationBundle": + invocation_ids = [record.invocation_id for record in self.records if isinstance(record, AgentInvocation)] + if len(invocation_ids) != len(set(invocation_ids)): + raise ValueError("invocation_id must be unique within an observation bundle") + return self + + +def link_tool_calls_to_sandbox(bundle: AgentObservationBundle, sandbox_id: Optional[str]) -> None: + """Link unassigned tool calls to their shared enclosing sandbox.""" + if sandbox_id is None: + return + for record in bundle.records: + if isinstance(record, ToolCallObservation) and record.sandbox_id is None: + record.sandbox_id = sandbox_id + + +def join_model_call_observations( + bundle: AgentObservationBundle, + calls: Iterable[ModelCallRecord], +) -> AgentObservationBundle: + """Resolve harness call references against captured model calls without guessing ownership.""" + + result = bundle.model_copy() + result.records = [ + record.model_copy() if isinstance(record, (AgentInvocation, ContextCompactionObservation)) else record + for record in bundle.records + ] + invocations = [record for record in result.records if isinstance(record, AgentInvocation)] + compactions = [record for record in result.records if isinstance(record, ContextCompactionObservation)] + captured = list(calls) + by_call_id: dict[str, list[ModelCallRecord]] = {} + by_response: dict[tuple[str, str, str], list[ModelCallRecord]] = {} + for call in captured: + if call.model_call_id: + by_call_id.setdefault(call.model_call_id, []).append(call) + if call.model_ref is not None and call.response_id: + key = (call.model_ref.type, call.model_ref.name, call.response_id) + by_response.setdefault(key, []).append(call) + + def matches(ref: ModelCallRef) -> list[ModelCallRecord]: + if ref.model_call_id: + candidates = by_call_id.get(ref.model_call_id, []) + return [ + call + for call in candidates + if (ref.model_ref is None or ref.model_ref == call.model_ref) + and (ref.response_id is None or ref.response_id == call.response_id) + ] + assert ref.model_ref is not None and ref.response_id is not None + return by_response.get((ref.model_ref.type, ref.model_ref.name, ref.response_id), []) + + def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: + return ModelCallRef.model_validate( + { + "model_call_id": call.model_call_id, + "model_ref": call.model_ref, + "response_id": call.response_id, + } + ) + + join_codes = { + "model_call_reference_ambiguous", + "model_call_reference_conflict", + "model_call_reference_unmatched", + } + result.gaps = [ + gap + for gap in bundle.gaps + if gap.code not in join_codes + and not ( + gap.code == "model_call_ownership_unavailable" + and gap.invocation_id is None + and (gap.detail is None or gap.detail.startswith("capture:")) + ) + ] + + claimed: set[int] = set() + join_gaps: list[ObservationGap] = [] + for invocation in invocations: + resolved: list[ModelCallRef] = [] + for ref in invocation.model_calls: + candidates = matches(ref) + detail = ref.model_call_id or ref.response_id + if len(candidates) != 1: + join_gaps.append( + ObservationGap( + code=("model_call_reference_ambiguous" if candidates else "model_call_reference_unmatched"), + invocation_id=invocation.invocation_id, + detail=detail, + ) + ) + resolved.append(ref) + continue + + call = candidates[0] + identity = id(call) + if identity in claimed: + join_gaps.append( + ObservationGap( + code="model_call_reference_conflict", + invocation_id=invocation.invocation_id, + detail=call.model_call_id or call.response_id, + ) + ) + resolved.append(ref) + continue + claimed.add(identity) + resolved.append(canonical(ref, call)) + invocation.model_calls = resolved + + for compaction in compactions: + for field_name in ("before_model_call", "after_model_call"): + ref = getattr(compaction, field_name) + if ref is None: + continue + candidates = matches(ref) + if len(candidates) == 1: + setattr(compaction, field_name, canonical(ref, candidates[0])) + else: + join_gaps.append( + ObservationGap( + code=("model_call_reference_ambiguous" if candidates else "model_call_reference_unmatched"), + invocation_id=compaction.invocation_id, + detail=f"{field_name}:{ref.model_call_id or ref.response_id}", + ) + ) + + result.gaps.extend(join_gaps) + for call in captured: + if id(call) not in claimed: + result.gaps.append( + ObservationGap( + code="model_call_ownership_unavailable", + detail=f"capture:{call.model_call_id or call.response_id or 'unknown'}:call_index={call.call_index}", + ) + ) + result.gaps = list({(gap.code, gap.invocation_id, gap.detail): gap for gap in result.gaps}.values()) + return result + + +def model_visible_tool_calls( + conversation: Iterable[NeMoGymResponseInputItem], +) -> list[tuple[str, Optional[str], str]]: + """Return model-visible call IDs, tool names, and result statuses.""" + + def field(item: Any, name: str) -> Any: + return item.get(name) if isinstance(item, dict) else getattr(item, name, None) + + items = list(conversation) + results = { + call_id: field(item, "status") + for item in items + if field(item, "type") == "function_call_output" and isinstance((call_id := field(item, "call_id")), str) + } + return [ + ( + call_id, + field(item, "name"), + ( + results[call_id] + if results.get(call_id) in {"completed", "incomplete"} + else "unknown" + if call_id in results + else "incomplete" + ), + ) + for item in items + if field(item, "type") == "function_call" and isinstance((call_id := field(item, "call_id")), str) and call_id + ] + + +@dataclass(frozen=True, slots=True) +class AgentEpisode: + """An Agent response and the observations available at its execution boundary.""" + + response: NeMoGymResponse + observations: AgentObservationBundle + + +def response_with_observations(episode: AgentEpisode) -> NeMoGymResponse: + """Carry observations through a prefixed Agent self-call.""" + return episode.response.model_copy(update={"ng_agent_observations": episode.observations.model_dump(mode="json")}) + + +def pop_response_observations(response: dict[str, Any]) -> Optional[AgentObservationBundle]: + """Remove and validate observations carried by an Agent self-call.""" + observations = response.pop("ng_agent_observations", None) + return AgentObservationBundle.model_validate(observations) if observations is not None else None diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index c1fd00304d..c401da9b9d 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -135,6 +135,7 @@ def test_build_model_call_record_from_exchange(): "latency_ms": 18.4, "request": {"input": "hi"}, "response": { + "id": "resp-1", "model": "m", "usage": { "input_tokens": 10, @@ -152,8 +153,10 @@ def test_build_model_call_record_from_exchange(): } rec = build_model_call_record(exchange, call_index=3) assert rec.model_call_id == "call-1" + assert rec.response_id == "resp-1" assert rec.call_index == 3 assert rec.model_ref is not None and rec.model_ref.name == "srv" + assert rec.model == "m" assert rec.dialect == "responses" assert rec.started_at == 100.0 and rec.completed_at == 100.02 assert (rec.tokens_in, rec.tokens_out, rec.tokens_total, rec.tokens_reasoning) == (10, 5, 15, 3) @@ -161,12 +164,19 @@ def test_build_model_call_record_from_exchange(): assert rec.reasoning_content == "thinking..." assert rec.tool_calls == [{"call_id": "c1", "name": "calc", "arguments": {"x": 1}}] assert rec.latency_total_ms == 18.4 + assert build_model_call_record({"response": {"id": 123}}, call_index=0).response_id is None + empty = build_model_call_record({"request": {}, "response": {}}, call_index=0) + assert empty.request == {} + assert empty.response == {} assert { "model_call_id", + "response_id", "call_index", "model_ref", + "model", "dialect", "status_code", + "finish_reason", "started_at", "completed_at", "tokens_in", @@ -175,6 +185,8 @@ def test_build_model_call_record_from_exchange(): "tokens_total", "request", "response", + "request_raw", + "response_raw", "tool_calls", "reasoning_content", "cache_hit", @@ -186,6 +198,36 @@ def test_build_model_call_record_from_exchange(): } <= type(rec).model_json_schema()["properties"].keys() +@pytest.mark.parametrize( + "response,expected", + [ + ({"choices": [{"finish_reason": "tool_calls"}]}, "tool_calls"), + ({"stop_reason": "end_turn"}, "end_turn"), + ({"incomplete_details": {"reason": "max_output_tokens"}}, "max_output_tokens"), + ({"status": "completed"}, None), + ({"choices": {"finish_reason": "invalid"}, "incomplete_details": []}, None), + ], +) +def test_build_model_call_record_normalizes_finish_reason(response, expected): + assert build_model_call_record({"response": response}, call_index=0).finish_reason == expected + + +def test_build_model_call_record_tolerates_malformed_nested_shapes(): + record = build_model_call_record( + { + "response": { + "usage": {"input_tokens": "invalid", "prompt_tokens_details": {"cached_tokens": []}}, + "choices": [{"message": "invalid"}], + "output": [{"type": "reasoning", "summary": [{"text": {}}]}], + } + }, + call_index=0, + ) + + assert record.tokens_total is None + assert record.tool_calls == [] + + def test_capture_is_durable_before_stream_terminal_event_is_sent(tmp_path): import asyncio @@ -232,6 +274,48 @@ async def send(message): assert durable_call_counts == [0, 1, 1] +def test_capture_retains_partial_stream_when_downstream_raises(tmp_path): + import asyncio + + from nemo_gym.base_responses_api_model import _CaptureMiddleware + + store = CaptureStore(tmp_path) + partial = b'data: {"type":"response.output_text.delta","delta":"partial"}\n\n' + + async def app(_scope, receive, send): + await receive() + await send( + {"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")]} + ) + await send({"type": "http.response.body", "body": partial, "more_body": True}) + raise RuntimeError("stream failed") + + async def receive(): + return {"type": "http.request", "body": b'{"input":"hi"}', "more_body": False} + + async def send(_message): + pass + + with pytest.raises(RuntimeError, match="stream failed"): + asyncio.run( + _CaptureMiddleware(app, store=store, model_server_name="srv")( + { + "type": "http", + "path": "/ng-rollout/partial/v1/responses", + "raw_path": b"/ng-rollout/partial/v1/responses", + "headers": [], + }, + receive, + send, + ) + ) + + [call] = read_model_call_records(store, "partial") + assert call.status_code == 200 + assert call.error_category == "exception" + assert call.response_raw == partial.decode() + + def test_stream_error_events_are_terminal(): from nemo_gym.base_responses_api_model import _consume_terminal_sse_event @@ -343,7 +427,7 @@ async def _responses(body: dict = Body()) -> dict: [record] = read_model_call_records(CaptureStore(tmp_path), "invalid-request") assert record.request is None - assert "request_raw" not in record.model_dump() + assert record.request_raw == request_bytes.decode() def test_per_rollout_url_prefix_correlates_and_is_openai_compatible(tmp_path): @@ -490,16 +574,15 @@ def _boom(_root): assert obs.make_capture_store(config) is None -def test_record_swallows_store_failure(): +def test_record_swallows_store_failure_and_marks_capture_incomplete(tmp_path, monkeypatch): from nemo_gym.base_responses_api_model import _record - class _BadStore: - def record(self, *args, **kwargs): - raise RuntimeError("disk full") + store = CaptureStore(tmp_path) + monkeypatch.setattr(store, "record", MagicMock(side_effect=RuntimeError("disk full"))) # Best-effort: a failing store must not raise out of _record. _record( - _BadStore(), + store, "chat", "srv", b"{}", @@ -512,6 +595,7 @@ def record(self, *args, **kwargs): error_category=None, latency_ms=1.0, ) + assert store.is_incomplete("r") def test_record_falls_back_to_raw_when_request_parser_raises(tmp_path, monkeypatch): @@ -555,8 +639,11 @@ async def _r(body: dict = Body()) -> PlainTextResponse: assert r.status_code == 200 and r.text == "not json" # response passed through unaltered records = CaptureStore(tmp_path).read("rnj") assert len(records) == 1 and records[0]["response"] is None # non-JSON body -> None + assert records[0]["response_raw"] == "not json" # a 2xx whose body we couldn't parse is flagged, not silently counted as a clean success assert records[0]["error_category"] == "capture_parse_error" + [record] = read_model_call_records(CaptureStore(tmp_path), "rnj") + assert record.response_raw == "not json" def test_as_arguments(): @@ -835,11 +922,13 @@ async def gen(): records = CaptureStore(tmp_path).read("3-0") assert len(records) == 1 and records[0]["response"] is not None # reassembled, not dropped + assert records[0]["response_raw"] == r.text # lossless SSE evidence is retained alongside the projection calls = read_model_call_records(CaptureStore(tmp_path), "3-0") assert len(calls) == 1 call = calls[0] assert call.model_call_id + assert call.response_id == "msg_1" assert call.model_ref is not None and call.model_ref.name == "srv" assert call.started_at is not None and call.completed_at is not None assert call.started_at <= call.completed_at @@ -850,13 +939,18 @@ async def gen(): assert call.tool_calls == [{"call_id": "t1", "name": "calc", "arguments": {"x": 1}}] assert call.latency_ttft_ms is not None assert call.error_category is None + assert call.response_raw == r.text def test_reconstruct_chat_sse(): from nemo_gym.base_responses_api_model import _reconstruct_streamed_response chunks = [ - {"model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hel"}}]}, + { + "id": "chatcmpl-1", + "model": "m", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hel"}}], + }, {"choices": [{"index": 0, "delta": {"content": "lo", "reasoning": "hmm"}}]}, # vLLM `reasoning` alias { "choices": [ @@ -882,6 +976,7 @@ def test_reconstruct_chat_sse(): raw = (b"".join(_sse("", c) for c in chunks) + b"data: [DONE]\n\n").replace(b"\n", b"\r\n") resp = _reconstruct_streamed_response(raw, "chat") msg = resp["choices"][0]["message"] + assert resp["id"] == "chatcmpl-1" assert msg["content"] == "Hello" and msg["reasoning_content"] == "hmm" assert msg["tool_calls"][0]["function"] == {"name": "f", "arguments": '{"a":1}'} assert resp["usage"]["total_tokens"] == 8 @@ -957,17 +1052,41 @@ def test_merge_capture_attaches_metrics_without_raw_payloads(tmp_path): from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record store = CaptureStore(tmp_path) - store.record( - "0-0", - _capture_exchange( - "responses", - "A", - {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, - {"output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}]}]}, - ), + exchange = _capture_exchange( + "responses", + "A", + {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, + { + "id": "resp-A", + "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}]}], + }, ) - - record = {"_ng_task_index": 0, "_ng_rollout_index": 0, "reward": 1.0, "response": {"harness": "A"}} + exchange["request_raw"] = "malformed request" + exchange["response_raw"] = "malformed response" + store.record("0-0", exchange) + + record = { + "_ng_task_index": 0, + "_ng_rollout_index": 0, + "reward": 1.0, + "response": {"harness": "A"}, + "ng_agent_observations": { + "source": "test", + "records": [ + { + "kind": "agent_invocation", + "invocation_id": "root", + "model_calls": [ + { + "model_ref": {"type": "responses_api_models", "name": "A"}, + "response_id": "resp-A", + } + ], + } + ], + "gaps": [{"code": "model_call_ownership_unavailable"}], + }, + } merge_model_call_capture_into_record(record, [tmp_path]) capture = record["ng_model_call_capture"] @@ -976,33 +1095,74 @@ def test_merge_capture_attaches_metrics_without_raw_payloads(tmp_path): assert capture["metrics"]["num_calls"] == 1 attached_call = capture["calls"][0] assert attached_call["model_call_id"] == "call-A" + assert attached_call["response_id"] == "resp-A" assert attached_call["model_ref"] == {"type": "responses_api_models", "name": "A"} + assert attached_call["model"] == "m" assert attached_call["started_at"] == 100.0 and attached_call["completed_at"] == 100.01 assert attached_call["tokens_in"] == 3 - assert "request" not in attached_call and "response" not in attached_call + assert {"request", "response", "request_raw", "response_raw"}.isdisjoint(attached_call) assert record["response"] == {"harness": "A"} and record["reward"] == 1.0 + [joined_ref] = record["ng_agent_observations"]["records"][0]["model_calls"] + assert joined_ref["model_call_id"] == "call-A" + assert record["ng_agent_observations"]["gaps"] == [] -def test_merge_capture_noop_without_capture(tmp_path): - from nemo_gym.base_responses_api_model import merge_model_call_capture_into_record +def test_merge_capture_reports_missing_capture(tmp_path): + from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record - rec = {"_ng_task_index": 9, "_ng_rollout_index": 9, "reward": 1.0} + rec = { + "_ng_task_index": 9, + "_ng_rollout_index": 9, + "reward": 1.0, + "ng_agent_observations": { + "source": "test", + "gaps": [{"code": "model_call_ownership_unavailable"}], + }, + } merge_model_call_capture_into_record(rec, [tmp_path]) # no capture file for 9-9 - assert "ng_model_call_capture" not in rec - merge_model_call_capture_into_record(rec, []) # no dirs - assert "ng_model_call_capture" not in rec + assert rec["ng_model_call_capture"]["calls"] == [] + assert [gap["code"] for gap in rec["ng_model_call_capture"]["gaps"]] == ["model_call_capture_no_records"] + assert rec["ng_agent_observations"]["gaps"] == [{"code": "model_call_ownership_unavailable"}] + CaptureStore(tmp_path).path_for("8-8").touch() + empty = {"_ng_task_index": 8, "_ng_rollout_index": 8} + merge_model_call_capture_into_record(empty, [tmp_path]) + assert [gap["code"] for gap in empty["ng_model_call_capture"]["gaps"]] == ["model_call_capture_no_records"] -def test_merge_capture_surfaces_malformed_data_only_when_active(tmp_path): + +def test_merge_capture_preserves_valid_records_around_malformed_data(tmp_path): from nemo_gym.base_responses_api_model import merge_model_call_capture_into_record store = CaptureStore(tmp_path) - store.path_for("9-9").write_bytes(b"{not-json}\n") - record = {"_ng_task_index": 9, "_ng_rollout_index": 9} + first = _capture_exchange("responses", "A", {}, {"id": "resp-A"}) + third = _capture_exchange("responses", "B", {}, {"id": "resp-B"}) + store.path_for("9-8").write_bytes(orjson.dumps(first) + b"\n{not-json}\n" + orjson.dumps(third) + b"\n") + record = {"_ng_task_index": 9, "_ng_rollout_index": 8} - merge_model_call_capture_into_record(record, []) - with pytest.raises(orjson.JSONDecodeError): - merge_model_call_capture_into_record(record, [tmp_path]) + merge_model_call_capture_into_record(record, [tmp_path]) + + assert [call["call_index"] for call in record["ng_model_call_capture"]["calls"]] == [0, 2] + assert [call["response_id"] for call in record["ng_model_call_capture"]["calls"]] == ["resp-A", "resp-B"] + assert record["ng_model_call_capture"]["gaps"] == [ + { + "code": "model_call_capture_records_unreadable", + "detail": "count=1", + } + ] + + +def test_merge_capture_reports_partial_write_loss(tmp_path): + from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record + + store = CaptureStore(tmp_path) + store.record("4-2", _capture_exchange("responses", "A", {}, {"id": "resp-A"})) + store.mark_incomplete("4-2") + record = {"_ng_task_index": 4, "_ng_rollout_index": 2} + + merge_model_call_capture_into_record(record, [tmp_path]) + + assert len(record["ng_model_call_capture"]["calls"]) == 1 + assert [gap["code"] for gap in record["ng_model_call_capture"]["gaps"]] == ["model_call_capture_incomplete"] def test_clear_model_call_captures_for_rollouts_run_scoping(tmp_path, monkeypatch): @@ -1012,11 +1172,12 @@ def test_clear_model_call_captures_for_rollouts_run_scoping(tmp_path, monkeypatc store = CaptureStore(tmp_path) store.record("0-0", {"dialect": "chat", "request": {}, "response": {}}) store.record("1-0", {"dialect": "chat", "request": {}, "response": {}}) + store.mark_incomplete("0-0") assert store.read("0-0") and store.read("1-0") # Clears only the rollout ids about to be (re)run; rows without indices are skipped, others stay. clear_model_call_captures_for_rollouts([{"_ng_task_index": 0, "_ng_rollout_index": 0}, {"no": "id"}], [tmp_path]) - assert store.read("0-0") == [] and store.read("1-0") + assert store.read("0-0") == [] and not store.is_incomplete("0-0") and store.read("1-0") clear_model_call_captures_for_rollouts([{"_ng_task_index": 1, "_ng_rollout_index": 0}], []) # no dirs -> no-op assert store.read("1-0") @@ -1126,10 +1287,40 @@ def _write(i: int) -> None: for t in threads: t.join() rows = store.read("0-0") - assert len(rows) == 20 # flock + in-process lock: no lost or corrupted appends + assert len(rows) == 20 # flock prevents lost or corrupted appends assert sorted(r["request"]["i"] for r in rows) == list(range(20)) +def test_capture_store_does_not_serialize_independent_rollouts(tmp_path, monkeypatch): + import os + import threading + + store = CaptureStore(tmp_path) + first_fsync = threading.Event() + release_first = threading.Event() + original_fsync = os.fsync + + def _fsync(fd): + if threading.current_thread().name == "first-rollout": + first_fsync.set() + assert release_first.wait(timeout=5) + original_fsync(fd) + + monkeypatch.setattr(os, "fsync", _fsync) + first = threading.Thread(name="first-rollout", target=store.record, args=("0-0", {"request": {}, "response": {}})) + second = threading.Thread(target=store.record, args=("1-0", {"request": {}, "response": {}})) + first.start() + assert first_fsync.wait(timeout=5) + second.start() + second.join(timeout=1) + try: + assert not second.is_alive() + assert store.read("1-0") == [{"request": {}, "response": {}}] + finally: + release_first.set() + first.join(timeout=5) + + def test_capture_store_read_waits_for_in_progress_append(tmp_path): import fcntl import threading diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py index ece7bacb1b..134d7d331c 100644 --- a/tests/unit_tests/test_rollout_collection.py +++ b/tests/unit_tests/test_rollout_collection.py @@ -994,7 +994,13 @@ def setup_server_client(self): {AGENT_REF_KEY_NAME: {"name": "my_agent"}, TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 1}, ] results = [ - {TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0, "reward": 1.0, "response": {"usage": {"tokens": 10}}}, + { + TASK_INDEX_KEY_NAME: 0, + ROLLOUT_INDEX_KEY_NAME: 0, + "reward": 1.0, + "response": {"usage": {"tokens": 10}}, + "ng_agent_observations": {"invocations": [{"conversation": ["large"]}]}, + }, {TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 1, "reward": 0.0, "response": {"usage": {"tokens": 12}}}, {TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 0, "reward": 1.0, "response": {"usage": {"tokens": 8}}}, {TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 1, "reward": 0.0, "response": {"usage": {"tokens": 15}}}, @@ -1023,6 +1029,7 @@ def setup_server_client(self): ) for item in sent_data: assert "responses_create_params" not in item + assert "ng_agent_observations" not in item assert "usage" in item["response"] async def test_call_aggregate_metrics_multiple_agents(self, tmp_path: Path) -> None: diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py new file mode 100644 index 0000000000..c5dadd658d --- /dev/null +++ b/tests/unit_tests/test_rollout_observability.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pydantic import ValidationError + +from nemo_gym.base_responses_api_model import ModelCallRecord +from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import NeMoGymResponse +from nemo_gym.rollout_observability import ( + AgentEpisode, + AgentInvocation, + AgentObservationBundle, + ContextCompactionObservation, + ModelCallRef, + ObservationGap, + SandboxObservation, + ToolCallObservation, + join_model_call_observations, + link_tool_calls_to_sandbox, + pop_response_observations, + response_with_observations, +) + + +@pytest.mark.parametrize( + "value", + ({}, {"response_id": "resp-1"}, {"model_ref": {"name": "policy", "type": "responses_api_models"}}), +) +def test_model_call_ref_rejects_incomplete_join_keys(value: dict) -> None: + with pytest.raises(ValidationError, match="model_call_id or both model_ref and response_id"): + ModelCallRef.model_validate(value) + + +def test_observations_round_trip_through_internal_response() -> None: + response = NeMoGymResponse.model_validate( + { + "id": "response", + "created_at": 0, + "model": "model", + "object": "response", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + } + ) + carried = response_with_observations( + AgentEpisode(response=response, observations=AgentObservationBundle(source="test")) + ).model_dump(mode="json") + + observations = pop_response_observations(carried) + + assert observations == AgentObservationBundle(source="test") + assert "ng_agent_observations" not in carried + + +def test_parallel_tool_calls_share_enclosing_sandbox_without_overwriting_existing_link() -> None: + bundle = AgentObservationBundle( + source="test", + records=[ + ToolCallObservation( + invocation_id="root", + tool_call_id="call-1", + started_at=1.0, + completed_at=3.0, + ), + ToolCallObservation( + invocation_id="root", + tool_call_id="call-2", + started_at=2.0, + completed_at=4.0, + ), + ToolCallObservation( + invocation_id="child", + tool_call_id="call-3", + sandbox_id="other-sandbox", + ), + ], + ) + + link_tool_calls_to_sandbox(bundle, "shared-sandbox") + + tools = [record for record in bundle.records if isinstance(record, ToolCallObservation)] + assert [tool.sandbox_id for tool in tools] == [ + "shared-sandbox", + "shared-sandbox", + "other-sandbox", + ] + + +def test_observation_bundle_rejects_duplicate_invocation_ids() -> None: + with pytest.raises(ValidationError, match="invocation_id must be unique"): + AgentObservationBundle( + source="test", + records=[AgentInvocation(invocation_id="root"), AgentInvocation(invocation_id="root")], + ) + + +def test_observation_models_reject_unknown_fields() -> None: + with pytest.raises(ValidationError, match="producer_extension"): + ModelCallRef.model_validate({"model_call_id": "call-1", "producer_extension": "unexpected"}) + + +def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + bundle = AgentObservationBundle( + source="test", + records=[ + AgentInvocation( + invocation_id="root", + model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")], + ) + ], + gaps=[ObservationGap(code="model_call_ownership_unavailable")], + ) + calls = [ + ModelCallRecord( + model_call_id="call-1", + response_id="resp-1", + model_ref=model_ref, + call_index=0, + ), + ModelCallRecord(model_call_id="call-2", model_ref=model_ref, call_index=1), + ] + + joined = join_model_call_observations(bundle, calls) + + [invocation] = [record for record in joined.records if isinstance(record, AgentInvocation)] + [joined_call] = invocation.model_calls + assert joined_call.model_call_id == "call-1" + assert joined_call.model_ref == model_ref + assert joined_call.response_id == "resp-1" + ownership_gaps = [gap for gap in joined.gaps if gap.code == "model_call_ownership_unavailable"] + assert [gap.detail for gap in ownership_gaps] == ["capture:call-2:call_index=1"] + + +def test_join_model_calls_does_not_guess_ambiguous_response_ids() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + bundle = AgentObservationBundle( + source="test", + records=[ + AgentInvocation( + invocation_id="root", + model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")], + ) + ], + ) + calls = [ + ModelCallRecord(model_call_id="call-1", response_id="resp-1", model_ref=model_ref, call_index=0), + ModelCallRecord(model_call_id="call-2", response_id="resp-1", model_ref=model_ref, call_index=1), + ] + + joined = join_model_call_observations(bundle, calls) + + [invocation] = [record for record in joined.records if isinstance(record, AgentInvocation)] + assert invocation.model_calls[0].model_call_id is None + assert "model_call_reference_ambiguous" in {gap.code for gap in joined.gaps} + assert [gap.detail for gap in joined.gaps if gap.code == "model_call_ownership_unavailable"] == [ + "capture:call-1:call_index=0", + "capture:call-2:call_index=1", + ] + + +def test_join_model_calls_reports_conflicting_and_unmatched_references() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + bundle = AgentObservationBundle( + source="test", + records=[ + AgentInvocation( + invocation_id="root", + model_calls=[ + ModelCallRef(model_call_id="call-1"), + ModelCallRef(model_ref=model_ref, response_id="resp-1"), + ModelCallRef(model_call_id="missing"), + ], + ) + ], + gaps=[ + ObservationGap( + code="model_call_ownership_unavailable", + invocation_id="root", + detail="producer gap", + ) + ], + ) + calls = [ + ModelCallRecord( + model_call_id="call-1", + response_id="resp-1", + model_ref=model_ref, + call_index=0, + ) + ] + + joined = join_model_call_observations(bundle, calls) + joined_again = join_model_call_observations(joined, calls) + + assert joined_again.model_dump() == joined.model_dump() + assert {gap.code for gap in joined.gaps} == { + "model_call_ownership_unavailable", + "model_call_reference_conflict", + "model_call_reference_unmatched", + } + assert any(gap.detail == "producer gap" for gap in joined.gaps) + + +def test_join_model_calls_resolves_compaction_boundaries() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + bundle = AgentObservationBundle( + source="test", + records=[ + ContextCompactionObservation( + invocation_id="root", + before_model_call=ModelCallRef(model_ref=model_ref, response_id="before"), + after_model_call=ModelCallRef(model_ref=model_ref, response_id="after"), + ) + ], + ) + calls = [ + ModelCallRecord( + model_call_id="call-before", + response_id="before", + model_ref=model_ref, + call_index=0, + ), + ModelCallRecord( + model_call_id="call-after", + response_id="after", + model_ref=model_ref, + call_index=1, + ), + ] + + joined = join_model_call_observations(bundle, calls) + [compaction] = [record for record in joined.records if isinstance(record, ContextCompactionObservation)] + + assert compaction.before_model_call.model_call_id == "call-before" + assert compaction.after_model_call.model_call_id == "call-after" + + +def test_sandbox_observation_rejects_negative_usage() -> None: + with pytest.raises(ValidationError): + SandboxObservation(role="agent", cpu_time_s=-1) From 3309fb1234e6ed5761e8af435f2155e84832cff4 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Fri, 24 Jul 2026 18:29:35 +0200 Subject: [PATCH 02/14] Propagate rollout correlation across servers Signed-off-by: Michal Bien --- nemo_gym/base_resources_server.py | 2 + nemo_gym/base_responses_api_agent.py | 16 +- nemo_gym/base_responses_api_model.py | 34 +--- nemo_gym/mcp_auto_exposure.py | 2 + nemo_gym/rollout_correlation.py | 87 ++++++++ nemo_gym/server_utils.py | 30 ++- tests/unit_tests/test_rollout_correlation.py | 198 +++++++++++++++++++ 7 files changed, 332 insertions(+), 37 deletions(-) create mode 100644 nemo_gym/rollout_correlation.py create mode 100644 tests/unit_tests/test_rollout_correlation.py diff --git a/nemo_gym/base_resources_server.py b/nemo_gym/base_resources_server.py index 1c4b6be35c..2eba9f0605 100644 --- a/nemo_gym/base_resources_server.py +++ b/nemo_gym/base_resources_server.py @@ -30,6 +30,7 @@ NeMoGymResponseCreateParamsNonStreaming, ) from nemo_gym.reward_profile import AggregateMetricsMixin, compute_aggregate_metrics +from nemo_gym.rollout_correlation import RolloutContextMiddleware from nemo_gym.server_utils import BaseRunServerInstanceConfig, BaseServer, SimpleServer @@ -110,6 +111,7 @@ def setup_webserver(self) -> FastAPI: app = FastAPI() self.setup_session_middleware(app) + app.add_middleware(RolloutContextMiddleware) app.post("/seed_session")(self.seed_session) app.post("/verify")(self.verify) diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 18fe0816ab..c6e76cd89f 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -14,6 +14,7 @@ # limitations under the License. from abc import abstractmethod from collections.abc import Mapping +from functools import wraps from typing import Any, Optional from fastapi import Body, FastAPI, Request @@ -24,7 +25,6 @@ BaseRunRequest, BaseVerifyResponse, ) -from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body from nemo_gym.config_types import ROLLOUT_PATH_PREFIX from nemo_gym.global_config import OBSERVABILITY_ENABLED_KEY_NAME, get_first_server_config_dict from nemo_gym.openai_utils import ( @@ -32,6 +32,7 @@ NeMoGymResponseCreateParamsNonStreaming, ) from nemo_gym.reward_profile import AggregateMetricsMixin, compute_aggregate_metrics +from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body, rollout_context from nemo_gym.server_utils import ( BaseRunServerInstanceConfig, BaseServer, @@ -62,7 +63,18 @@ def setup_webserver(self) -> FastAPI: # responses() recovers the rollout id from the path (see url_path_for_request) to correlate # its model calls. Same handler, so unprefixed calls are unaffected. app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/v1/responses")(self.responses) - app.post("/run")(self.run) + + run = self.run + + @wraps(run) + async def run_with_rollout_context(*args: Any, **kwargs: Any) -> BaseVerifyResponse: + body = kwargs.get("body") + if body is None: + body = next((arg for arg in args if isinstance(arg, BaseRunRequest)), None) + with rollout_context(self.rollout_id_from_run(body)): + return await run(*args, **kwargs) + + app.post("/run")(run_with_rollout_context) app.post("/aggregate_metrics")(self.aggregate_metrics) return app diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 6326e06776..733bb03e5e 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -48,11 +48,6 @@ from nemo_gym.anthropic_converter import AnthropicConverter from nemo_gym.chat_streaming import sanitize_streaming_chat_body, synthesize_chat_completion_sse from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, ModelServerRef -from nemo_gym.global_config import ( - ATTEMPT_INDEX_KEY_NAME, - ROLLOUT_INDEX_KEY_NAME, - TASK_INDEX_KEY_NAME, -) from nemo_gym.openai_utils import ( NeMoGymChatCompletion, NeMoGymChatCompletionCreateParamsNonStreaming, @@ -65,6 +60,7 @@ synthesize_responses_sse, validate_streaming_responses_params, ) +from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body from nemo_gym.rollout_observability import AgentObservationBundle, ObservationGap, join_model_call_observations from nemo_gym.server_utils import ( BaseRunServerInstanceConfig, @@ -368,34 +364,6 @@ def read_available(self, rollout_id: str) -> tuple[list[tuple[int, dict[str, Any return exchanges, invalid_count -def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]: - """Per-rollout model-call capture id from a run-request's task/rollout indices. - - Reads the canonical row keys (``_ng_task_index`` / ``_ng_rollout_index``) that - rollout_collection ships to an agent's ``/run``. When a resume re-dispatch attempt is present - (``_ng_attempt_index`` > 0), an ``-a`` suffix is appended so a retry's captured model calls - stay separable from the prior attempt; the first attempt (0) keeps the bare ``-`` - key for backward compatibility. - """ - if isinstance(body, BaseModel): - data = body.model_dump() - elif isinstance(body, Mapping): - data = body - else: - return None - task = data.get(TASK_INDEX_KEY_NAME) - rollout = data.get(ROLLOUT_INDEX_KEY_NAME) - if task is None or rollout is None: - return None - rollout_id = f"{task}-{rollout}" - attempt = data.get(ATTEMPT_INDEX_KEY_NAME) - if attempt is not None: - attempt_index = int(attempt) - if attempt_index > 0: - rollout_id = f"{rollout_id}-a{attempt_index}" - return rollout_id - - # --- Observability records derived from captured exchanges --- diff --git a/nemo_gym/mcp_auto_exposure.py b/nemo_gym/mcp_auto_exposure.py index 45de0aaaef..0e6aab4f6d 100644 --- a/nemo_gym/mcp_auto_exposure.py +++ b/nemo_gym/mcp_auto_exposure.py @@ -438,6 +438,8 @@ def harvest_tools(app: FastAPI, server: Any) -> dict[str, MCPTool]: cls = m.cls if f"{cls.__module__}.{cls.__name__}" == "starlette.middleware.sessions.SessionMiddleware": continue # Gym's SessionMiddleware — replaced by a materialized session on direct dispatch + if f"{cls.__module__}.{cls.__name__}" == "nemo_gym.rollout_correlation.RolloutContextMiddleware": + continue # Correlation prefixes are handled before resource routes; direct MCP dispatch has no prefix. dispatch = m.kwargs.get("dispatch") if dispatch is not None and getattr(dispatch, "__module__", None) in _GYM_MIDDLEWARE_MODULES: continue # Gym's add_session_id / exception middleware diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py new file mode 100644 index 0000000000..e41a70b240 --- /dev/null +++ b/nemo_gym/rollout_correlation.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import re +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, Optional + +from pydantic import BaseModel + +from nemo_gym.config_types import ROLLOUT_PATH_PREFIX +from nemo_gym.global_config import ( + ATTEMPT_INDEX_KEY_NAME, + ROLLOUT_INDEX_KEY_NAME, + TASK_INDEX_KEY_NAME, +) + + +_ROLLOUT_ID: ContextVar[Optional[str]] = ContextVar("nemo_gym_rollout_id", default=None) + + +def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]: + """Build the capture key stamped by rollout collection.""" + if isinstance(body, BaseModel): + data = body.model_dump() + elif isinstance(body, Mapping): + data = body + else: + return None + + task = data.get(TASK_INDEX_KEY_NAME) + rollout = data.get(ROLLOUT_INDEX_KEY_NAME) + if task is None or rollout is None: + return None + + rollout_id = f"{task}-{rollout}" + attempt = data.get(ATTEMPT_INDEX_KEY_NAME) + if attempt is not None and int(attempt) > 0: + rollout_id = f"{rollout_id}-a{int(attempt)}" + return rollout_id + + +def current_rollout_id() -> Optional[str]: + return _ROLLOUT_ID.get() + + +@contextmanager +def rollout_context(rollout_id: Optional[str]) -> Iterator[None]: + token = _ROLLOUT_ID.set(rollout_id) + try: + yield + finally: + _ROLLOUT_ID.reset(token) + + +class RolloutContextMiddleware: + """Strip a rollout prefix and expose it to downstream Gym calls for this request.""" + + _PREFIX = re.compile( + rf"^/{re.escape(ROLLOUT_PATH_PREFIX)}/(?P[A-Za-z0-9][A-Za-z0-9._-]*)(?P/.*)$" + ) + + def __init__(self, app: Any) -> None: + self._app = app + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + match = self._PREFIX.match(scope.get("path", "")) if scope.get("type") == "http" else None + if match is None: + await self._app(scope, receive, send) + return + + path = match.group("rest") + scope = {**scope, "path": path, "raw_path": path.encode()} + with rollout_context(match.group("rollout_id")): + await self._app(scope, receive, send) diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index 03070f237f..d05f3100cd 100644 --- a/nemo_gym/server_utils.py +++ b/nemo_gym/server_utils.py @@ -64,6 +64,7 @@ DRY_RUN_KEY_NAME, HEAD_SERVER_KEY_NAME, NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME, + OBSERVABILITY_ENABLED_KEY_NAME, RAY_HEAD_NODE_ADDRESS_KEY_NAME, GlobalConfigDictParser, GlobalConfigDictParserConfig, @@ -71,6 +72,7 @@ get_global_config_dict, ) from nemo_gym.profiling import Profiler +from nemo_gym.rollout_correlation import current_rollout_id, maybe_rollout_id_from_run_body _GLOBAL_AIOHTTP_CLIENT: Union[None, ClientSession] = None @@ -330,10 +332,34 @@ async def request( server_config_dict = get_first_server_config_dict(self.global_config_dict, server_name) base_url = self._build_server_base_url(server_config_dict) + json_obj = kwargs.get("json") if "json" in kwargs: - json_obj = kwargs["json"] if isinstance(json_obj, BaseModel): - kwargs["json"] = json_obj.model_dump(exclude_unset=True) + json_obj = json_obj.model_dump(exclude_unset=True) + kwargs["json"] = json_obj + + observability_enabled = self.global_config_dict.get(OBSERVABILITY_ENABLED_KEY_NAME, False) + server_entry = self.global_config_dict.get(server_name) + rollout_id = current_rollout_id() + if ( + observability_enabled + and server_entry is not None + and "resources_servers" in server_entry + and url_path == "/verify" + ): + rollout_id = rollout_id or maybe_rollout_id_from_run_body(json_obj) + if rollout_id is not None: + url_path = f"{rollout_path_prefix(rollout_id)}{url_path}" + + if ( + rollout_id is not None + and observability_enabled + and server_entry is not None + and "responses_api_models" in server_entry + and url_path.partition("?")[0] in {"/v1/responses", "/v1/chat/completions", "/v1/messages"} + and not url_path.startswith(f"/{ROLLOUT_PATH_PREFIX}/") + ): + url_path = f"{rollout_path_prefix(rollout_id)}{url_path}" return await request(method=method, url=f"{base_url}{url_path}", _internal=True, **kwargs) diff --git a/tests/unit_tests/test_rollout_correlation.py b/tests/unit_tests/test_rollout_correlation.py new file mode 100644 index 0000000000..25f18ba62f --- /dev/null +++ b/tests/unit_tests/test_rollout_correlation.py @@ -0,0 +1,198 @@ +# 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 urllib.parse import urlsplit + +import orjson +import pytest +from fastapi import Body, FastAPI +from fastapi.responses import JSONResponse +from omegaconf import OmegaConf +from pydantic import ConfigDict +from starlette.testclient import TestClient + +import nemo_gym.server_utils +from nemo_gym.base_resources_server import ( + BaseResourcesServerConfig, + BaseRunRequest, + BaseVerifyRequest, + BaseVerifyResponse, + SimpleResourcesServer, +) +from nemo_gym.base_responses_api_agent import BaseResponsesAPIAgentConfig, SimpleResponsesAPIAgent +from nemo_gym.base_responses_api_model import ( + CaptureStore, + ModelCallCaptureConfig, + install_model_call_capture, + merge_model_call_capture_into_record, +) +from nemo_gym.config_types import BaseServerConfig +from nemo_gym.server_utils import ServerClient, get_response_json + + +def _model_response(model: str, text: str = "") -> dict: + return { + "id": f"resp-{model}", + "created_at": 0.0, + "model": model, + "object": "response", + "output": [ + { + "id": f"msg-{model}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + } + + +def _model_app(capture_dir, name: str) -> FastAPI: + app = FastAPI() + + @app.post("/v1/responses") + async def responses(body: dict = Body()) -> JSONResponse: + return JSONResponse(_model_response(name, "[[A=B]]" if name == "judge" else "answer")) + + install_model_call_capture( + app, + ModelCallCaptureConfig(observability_enabled=True, model_call_capture_dir=capture_dir), + model_server_name=name, + ) + return app + + +class _JudgeResourcesServer(SimpleResourcesServer): + async def verify(self, body: BaseVerifyRequest) -> BaseVerifyResponse: + judge = await self.server_client.post( + server_name="judge", + url_path="/v1/responses", + json={"input": "grade"}, + ) + await get_response_json(judge) + return BaseVerifyResponse(**body.model_dump(), reward=1.0) + + +class _AgentRunRequest(BaseRunRequest): + model_config = ConfigDict(extra="allow") + + +class _Agent(SimpleResponsesAPIAgent): + async def responses(self, body): + raise NotImplementedError + + async def run(self, body: _AgentRunRequest) -> BaseVerifyResponse: + policy = await self.server_client.post( + server_name="policy", + url_path="/v1/responses", + json=body.responses_create_params, + ) + response = orjson.loads(await policy.read()) + verify = await self.server_client.post( + server_name="resources", + url_path="/verify", + json={ + "responses_create_params": body.responses_create_params.model_dump(), + "response": response, + }, + ) + return BaseVerifyResponse.model_validate(orjson.loads(await verify.read())) + + +class _Response: + def __init__(self, response) -> None: + self.status = response.status_code + self.ok = response.is_success + self.cookies = response.cookies + self._content = response.content + + async def read(self) -> bytes: + return self._content + + +@pytest.mark.asyncio +async def test_verify_correlates_policy_and_judge_calls_and_preserves_raw_capture(tmp_path, monkeypatch) -> None: + capture_dir = tmp_path / "captures" + config = OmegaConf.create( + { + "observability_enabled": True, + "policy": {"responses_api_models": {"model": {"host": "policy.test", "port": 80}}}, + "judge": {"responses_api_models": {"model": {"host": "judge.test", "port": 80}}}, + "resources": {"resources_servers": {"judge": {"host": "resources.test", "port": 80}}}, + "agent": {"responses_api_agents": {"agent": {"host": "agent.test", "port": 80}}}, + } + ) + server_client = ServerClient( + head_server_config=BaseServerConfig(host="head.test", port=80), + global_config_dict=config, + ) + resources = _JudgeResourcesServer( + config=BaseResourcesServerConfig( + host="resources.test", + port=80, + entrypoint="app.py", + name="resources", + ), + server_client=server_client, + ) + agent = _Agent( + config=BaseResponsesAPIAgentConfig( + host="agent.test", + port=80, + entrypoint="app.py", + name="agent", + ), + server_client=server_client, + ) + clients = { + "policy.test": TestClient(_model_app(capture_dir, "policy")), + "judge.test": TestClient(_model_app(capture_dir, "judge")), + "resources.test": TestClient(resources.setup_webserver()), + "agent.test": TestClient(agent.setup_webserver()), + } + + async def dispatch(method: str, url: str, **kwargs): + parsed = urlsplit(url) + response = clients[parsed.hostname].request(method, parsed.path, json=kwargs.get("json")) + return _Response(response) + + monkeypatch.setattr(nemo_gym.server_utils, "request", dispatch) + + verify = await server_client.post( + server_name="agent", + url_path="/run", + json={ + "_ng_task_index": 4, + "_ng_rollout_index": 2, + "responses_create_params": {"input": "solve"}, + }, + ) + assert orjson.loads(await verify.read())["reward"] == 1.0 + + store = CaptureStore(capture_dir) + capture_path = store.path_for("4-2") + assert capture_path.is_file() + exchanges = store.read("4-2") + assert [exchange["model_ref"]["name"] for exchange in exchanges] == ["policy", "judge"] + assert all(exchange.get("request") is not None or exchange.get("request_raw") for exchange in exchanges) + assert all(exchange.get("response") is not None or exchange.get("response_raw") for exchange in exchanges) + + rollout = {"_ng_task_index": 4, "_ng_rollout_index": 2} + merge_model_call_capture_into_record(rollout, [capture_dir]) + assert capture_path.is_file() + assert len(capture_path.read_bytes()) > 0 From dc136eaaf61beaf09280ed5a1649f07146238691 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Mon, 27 Jul 2026 17:38:53 +0200 Subject: [PATCH 03/14] Tighten rollout observation foundation Signed-off-by: Michal Bien --- nemo_gym/rollout_collection.py | 8 +- nemo_gym/rollout_correlation.py | 15 ++-- nemo_gym/rollout_observability.py | 61 ++------------- nemo_gym/server_utils.py | 12 +-- tests/unit_tests/test_rollout_collection.py | 2 + tests/unit_tests/test_rollout_correlation.py | 40 +++++++++- .../unit_tests/test_rollout_observability.py | 74 +++---------------- 7 files changed, 79 insertions(+), 133 deletions(-) diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index 17b7cee014..d503687980 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -656,7 +656,13 @@ async def _fetch_agent_metrics(agent_name: str, agent_result_list: List[Dict]) - entry = { k: v for k, v in r.items() - if k not in ("response", "responses_create_params", "ng_agent_observations") + if k + not in ( + "response", + "responses_create_params", + "ng_agent_observations", + "ng_model_call_capture", + ) } usage = (r.get("response") or {}).get("usage") if usage: diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py index e41a70b240..c1d58a2694 100644 --- a/nemo_gym/rollout_correlation.py +++ b/nemo_gym/rollout_correlation.py @@ -33,20 +33,19 @@ def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]: """Build the capture key stamped by rollout collection.""" - if isinstance(body, BaseModel): - data = body.model_dump() - elif isinstance(body, Mapping): - data = body - else: + if not isinstance(body, (BaseModel, Mapping)): return None - task = data.get(TASK_INDEX_KEY_NAME) - rollout = data.get(ROLLOUT_INDEX_KEY_NAME) + def field(key: str) -> Any: + return body.get(key) if isinstance(body, Mapping) else getattr(body, key, None) + + task = field(TASK_INDEX_KEY_NAME) + rollout = field(ROLLOUT_INDEX_KEY_NAME) if task is None or rollout is None: return None rollout_id = f"{task}-{rollout}" - attempt = data.get(ATTEMPT_INDEX_KEY_NAME) + attempt = field(ATTEMPT_INDEX_KEY_NAME) if attempt is not None and int(attempt) > 0: rollout_id = f"{rollout_id}-a{int(attempt)}" return rollout_id diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py index dd722f7b2a..2ea664de12 100644 --- a/nemo_gym/rollout_observability.py +++ b/nemo_gym/rollout_observability.py @@ -7,7 +7,7 @@ from collections.abc import Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional +from typing import TYPE_CHECKING, Annotated, Literal, Optional from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -67,11 +67,17 @@ class ToolCallObservation(ObservationModel): tool_name: Optional[str] = None started_at: Optional[float] = None completed_at: Optional[float] = None - duration_ms: Optional[float] = None + duration_ms: Optional[float] = Field(default=None, ge=0) timing_source: Optional[Literal["executor", "artifact", "harness"]] = None status: Literal["completed", "failed", "timeout", "cancelled", "incomplete", "unknown"] = "unknown" error_type: Optional[str] = None + @model_validator(mode="after") + def validate_timing(self) -> "ToolCallObservation": + if self.started_at is not None and self.completed_at is not None and self.completed_at < self.started_at: + raise ValueError("completed_at must not precede started_at") + return self + class SandboxObservation(ObservationModel): """Outcome and lifetime resource usage reported by a sandbox-owning harness.""" @@ -145,15 +151,6 @@ def validate_identity(self) -> "AgentObservationBundle": return self -def link_tool_calls_to_sandbox(bundle: AgentObservationBundle, sandbox_id: Optional[str]) -> None: - """Link unassigned tool calls to their shared enclosing sandbox.""" - if sandbox_id is None: - return - for record in bundle.records: - if isinstance(record, ToolCallObservation) and record.sandbox_id is None: - record.sandbox_id = sandbox_id - - def join_model_call_observations( bundle: AgentObservationBundle, calls: Iterable[ModelCallRecord], @@ -278,51 +275,9 @@ def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: return result -def model_visible_tool_calls( - conversation: Iterable[NeMoGymResponseInputItem], -) -> list[tuple[str, Optional[str], str]]: - """Return model-visible call IDs, tool names, and result statuses.""" - - def field(item: Any, name: str) -> Any: - return item.get(name) if isinstance(item, dict) else getattr(item, name, None) - - items = list(conversation) - results = { - call_id: field(item, "status") - for item in items - if field(item, "type") == "function_call_output" and isinstance((call_id := field(item, "call_id")), str) - } - return [ - ( - call_id, - field(item, "name"), - ( - results[call_id] - if results.get(call_id) in {"completed", "incomplete"} - else "unknown" - if call_id in results - else "incomplete" - ), - ) - for item in items - if field(item, "type") == "function_call" and isinstance((call_id := field(item, "call_id")), str) and call_id - ] - - @dataclass(frozen=True, slots=True) class AgentEpisode: """An Agent response and the observations available at its execution boundary.""" response: NeMoGymResponse observations: AgentObservationBundle - - -def response_with_observations(episode: AgentEpisode) -> NeMoGymResponse: - """Carry observations through a prefixed Agent self-call.""" - return episode.response.model_copy(update={"ng_agent_observations": episode.observations.model_dump(mode="json")}) - - -def pop_response_observations(response: dict[str, Any]) -> Optional[AgentObservationBundle]: - """Remove and validate observations carried by an Agent self-call.""" - observations = response.pop("ng_agent_observations", None) - return AgentObservationBundle.model_validate(observations) if observations is not None else None diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index d05f3100cd..b676c8f49a 100644 --- a/nemo_gym/server_utils.py +++ b/nemo_gym/server_utils.py @@ -341,14 +341,10 @@ async def request( observability_enabled = self.global_config_dict.get(OBSERVABILITY_ENABLED_KEY_NAME, False) server_entry = self.global_config_dict.get(server_name) rollout_id = current_rollout_id() - if ( - observability_enabled - and server_entry is not None - and "resources_servers" in server_entry - and url_path == "/verify" - ): - rollout_id = rollout_id or maybe_rollout_id_from_run_body(json_obj) - if rollout_id is not None: + if observability_enabled and server_entry is not None and "resources_servers" in server_entry: + if url_path == "/verify": + rollout_id = rollout_id or maybe_rollout_id_from_run_body(json_obj) + if rollout_id is not None and not url_path.startswith(f"/{ROLLOUT_PATH_PREFIX}/"): url_path = f"{rollout_path_prefix(rollout_id)}{url_path}" if ( diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py index 134d7d331c..d3145eb403 100644 --- a/tests/unit_tests/test_rollout_collection.py +++ b/tests/unit_tests/test_rollout_collection.py @@ -1000,6 +1000,7 @@ def setup_server_client(self): "reward": 1.0, "response": {"usage": {"tokens": 10}}, "ng_agent_observations": {"invocations": [{"conversation": ["large"]}]}, + "ng_model_call_capture": {"calls": [{"request": "large"}]}, }, {TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 1, "reward": 0.0, "response": {"usage": {"tokens": 12}}}, {TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 0, "reward": 1.0, "response": {"usage": {"tokens": 8}}}, @@ -1030,6 +1031,7 @@ def setup_server_client(self): for item in sent_data: assert "responses_create_params" not in item assert "ng_agent_observations" not in item + assert "ng_model_call_capture" not in item assert "usage" in item["response"] async def test_call_aggregate_metrics_multiple_agents(self, tmp_path: Path) -> None: diff --git a/tests/unit_tests/test_rollout_correlation.py b/tests/unit_tests/test_rollout_correlation.py index 25f18ba62f..3954ece3fc 100644 --- a/tests/unit_tests/test_rollout_correlation.py +++ b/tests/unit_tests/test_rollout_correlation.py @@ -38,6 +38,7 @@ merge_model_call_capture_into_record, ) from nemo_gym.config_types import BaseServerConfig +from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body from nemo_gym.server_utils import ServerClient, get_response_json @@ -78,6 +79,19 @@ async def responses(body: dict = Body()) -> JSONResponse: class _JudgeResourcesServer(SimpleResourcesServer): + def setup_webserver(self) -> FastAPI: + app = super().setup_webserver() + app.post("/tool")(self.tool) + return app + + async def tool(self, body: dict = Body()) -> dict: + tool_model = await self.server_client.post( + server_name="tool_model", + url_path="/v1/responses", + json={"input": body["input"]}, + ) + return await get_response_json(tool_model) + async def verify(self, body: BaseVerifyRequest) -> BaseVerifyResponse: judge = await self.server_client.post( server_name="judge", @@ -103,6 +117,12 @@ async def run(self, body: _AgentRunRequest) -> BaseVerifyResponse: json=body.responses_create_params, ) response = orjson.loads(await policy.read()) + tool = await self.server_client.post( + server_name="resources", + url_path="/tool", + json={"input": "lookup"}, + ) + await get_response_json(tool) verify = await self.server_client.post( server_name="resources", url_path="/verify", @@ -132,6 +152,7 @@ async def test_verify_correlates_policy_and_judge_calls_and_preserves_raw_captur { "observability_enabled": True, "policy": {"responses_api_models": {"model": {"host": "policy.test", "port": 80}}}, + "tool_model": {"responses_api_models": {"model": {"host": "tool-model.test", "port": 80}}}, "judge": {"responses_api_models": {"model": {"host": "judge.test", "port": 80}}}, "resources": {"resources_servers": {"judge": {"host": "resources.test", "port": 80}}}, "agent": {"responses_api_agents": {"agent": {"host": "agent.test", "port": 80}}}, @@ -161,6 +182,7 @@ async def test_verify_correlates_policy_and_judge_calls_and_preserves_raw_captur ) clients = { "policy.test": TestClient(_model_app(capture_dir, "policy")), + "tool-model.test": TestClient(_model_app(capture_dir, "tool_model")), "judge.test": TestClient(_model_app(capture_dir, "judge")), "resources.test": TestClient(resources.setup_webserver()), "agent.test": TestClient(agent.setup_webserver()), @@ -188,7 +210,7 @@ async def dispatch(method: str, url: str, **kwargs): capture_path = store.path_for("4-2") assert capture_path.is_file() exchanges = store.read("4-2") - assert [exchange["model_ref"]["name"] for exchange in exchanges] == ["policy", "judge"] + assert [exchange["model_ref"]["name"] for exchange in exchanges] == ["policy", "tool_model", "judge"] assert all(exchange.get("request") is not None or exchange.get("request_raw") for exchange in exchanges) assert all(exchange.get("response") is not None or exchange.get("response_raw") for exchange in exchanges) @@ -196,3 +218,19 @@ async def dispatch(method: str, url: str, **kwargs): merge_model_call_capture_into_record(rollout, [capture_dir]) assert capture_path.is_file() assert len(capture_path.read_bytes()) > 0 + + +def test_rollout_id_does_not_serialize_run_body() -> None: + class UndumpableRunRequest(_AgentRunRequest): + def model_dump(self, *args, **kwargs): + raise AssertionError("run body must not be serialized") + + body = UndumpableRunRequest.model_validate( + { + "_ng_task_index": 4, + "_ng_rollout_index": 2, + "responses_create_params": {"input": "solve"}, + } + ) + + assert maybe_rollout_id_from_run_body(body) == "4-2" diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py index c5dadd658d..daf42e5b5a 100644 --- a/tests/unit_tests/test_rollout_observability.py +++ b/tests/unit_tests/test_rollout_observability.py @@ -6,9 +6,7 @@ from nemo_gym.base_responses_api_model import ModelCallRecord from nemo_gym.config_types import ModelServerRef -from nemo_gym.openai_utils import NeMoGymResponse from nemo_gym.rollout_observability import ( - AgentEpisode, AgentInvocation, AgentObservationBundle, ContextCompactionObservation, @@ -17,9 +15,6 @@ SandboxObservation, ToolCallObservation, join_model_call_observations, - link_tool_calls_to_sandbox, - pop_response_observations, - response_with_observations, ) @@ -32,63 +27,6 @@ def test_model_call_ref_rejects_incomplete_join_keys(value: dict) -> None: ModelCallRef.model_validate(value) -def test_observations_round_trip_through_internal_response() -> None: - response = NeMoGymResponse.model_validate( - { - "id": "response", - "created_at": 0, - "model": "model", - "object": "response", - "output": [], - "parallel_tool_calls": True, - "tool_choice": "auto", - "tools": [], - } - ) - carried = response_with_observations( - AgentEpisode(response=response, observations=AgentObservationBundle(source="test")) - ).model_dump(mode="json") - - observations = pop_response_observations(carried) - - assert observations == AgentObservationBundle(source="test") - assert "ng_agent_observations" not in carried - - -def test_parallel_tool_calls_share_enclosing_sandbox_without_overwriting_existing_link() -> None: - bundle = AgentObservationBundle( - source="test", - records=[ - ToolCallObservation( - invocation_id="root", - tool_call_id="call-1", - started_at=1.0, - completed_at=3.0, - ), - ToolCallObservation( - invocation_id="root", - tool_call_id="call-2", - started_at=2.0, - completed_at=4.0, - ), - ToolCallObservation( - invocation_id="child", - tool_call_id="call-3", - sandbox_id="other-sandbox", - ), - ], - ) - - link_tool_calls_to_sandbox(bundle, "shared-sandbox") - - tools = [record for record in bundle.records if isinstance(record, ToolCallObservation)] - assert [tool.sandbox_id for tool in tools] == [ - "shared-sandbox", - "shared-sandbox", - "other-sandbox", - ] - - def test_observation_bundle_rejects_duplicate_invocation_ids() -> None: with pytest.raises(ValidationError, match="invocation_id must be unique"): AgentObservationBundle( @@ -102,6 +40,18 @@ def test_observation_models_reject_unknown_fields() -> None: ModelCallRef.model_validate({"model_call_id": "call-1", "producer_extension": "unexpected"}) +@pytest.mark.parametrize( + "timing", + ( + {"duration_ms": -1}, + {"started_at": 2.0, "completed_at": 1.0}, + ), +) +def test_tool_call_observation_rejects_invalid_timing(timing: dict) -> None: + with pytest.raises(ValidationError): + ToolCallObservation(invocation_id="root", tool_call_id="call-1", **timing) + + def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() -> None: model_ref = ModelServerRef(name="policy", type="responses_api_models") bundle = AgentObservationBundle( From 671177cdc2047169931942ed9c307be9899a849b Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Mon, 27 Jul 2026 19:06:15 +0200 Subject: [PATCH 04/14] Record agent invocation outcomes Signed-off-by: Michal Bien --- nemo_gym/rollout_observability.py | 2 ++ tests/unit_tests/test_rollout_observability.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py index 2ea664de12..162359e059 100644 --- a/nemo_gym/rollout_observability.py +++ b/nemo_gym/rollout_observability.py @@ -47,6 +47,8 @@ class AgentInvocation(ObservationModel): status: Literal["completed", "failed", "incomplete", "unknown"] = Field( default="unknown", description="Harness-reported invocation outcome; unknown when not explicit." ) + duration_ms: Optional[float] = Field(default=None, ge=0) + error_type: Optional[str] = None model_calls: list[ModelCallRef] = Field(default_factory=list) conversation: list[NeMoGymResponseInputItem] = Field( default_factory=list, diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py index daf42e5b5a..2a09972cf5 100644 --- a/tests/unit_tests/test_rollout_observability.py +++ b/tests/unit_tests/test_rollout_observability.py @@ -52,6 +52,11 @@ def test_tool_call_observation_rejects_invalid_timing(timing: dict) -> None: ToolCallObservation(invocation_id="root", tool_call_id="call-1", **timing) +def test_agent_invocation_rejects_negative_duration() -> None: + with pytest.raises(ValidationError): + AgentInvocation(invocation_id="root", duration_ms=-1) + + def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() -> None: model_ref = ModelServerRef(name="policy", type="responses_api_models") bundle = AgentObservationBundle( From ffc1eaf56ad5647c5ea9ca08d29efd9712de90c7 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Mon, 27 Jul 2026 19:34:30 +0200 Subject: [PATCH 05/14] Preserve upstream failure evidence Signed-off-by: Michal Bien --- nemo_gym/base_responses_api_model.py | 29 +++++++++++++++++-- .../test_base_responses_api_model.py | 23 +++++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 733bb03e5e..57abfe76d9 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -755,6 +755,24 @@ def _classify_exception(exc: BaseException) -> str: return "exception" +def _exception_http_details(exc: BaseException) -> tuple[Optional[int], bytes]: + response = getattr(exc, "response", None) + status = getattr(exc, "status", None) + if not isinstance(status, int): + status = getattr(exc, "status_code", None) + if not isinstance(status, int) and response is not None: + status = getattr(response, "status_code", None) + + body = getattr(exc, "response_content", None) + if body is None and response is not None: + body = getattr(response, "content", None) + if body is None: + body = getattr(response, "text", None) + if isinstance(body, str): + body = body.encode() + return (status if isinstance(status, int) else None, bytes(body) if isinstance(body, (bytes, bytearray)) else b"") + + # --- SSE reconstruction: rebuild a final response object from a streamed body --- def _parse_sse_events(raw: bytes) -> list[dict[str, Any]]: """Parse an SSE byte stream into its JSON ``data:`` payloads (best-effort; non-JSON skipped).""" @@ -1067,6 +1085,11 @@ async def _flush_deferred_response() -> None: await self._app(scope, _receive, _send) except Exception as exc: completed_at = time.time() + exception_status, exception_body = _exception_http_details(exc) + upstream_status = state["status"] or exception_status + upstream_body = bytes(state["body"]) or exception_body + error_category = _classify_status(upstream_status) if isinstance(upstream_status, int) else None + error_category = error_category or _classify_exception(exc) # Offload the blocking write+fsync so it never stalls the event loop. try: await asyncio.to_thread( @@ -1080,11 +1103,11 @@ async def _flush_deferred_response() -> None: started_at=started_at, completed_at=completed_at, response_body=None, - status_code=state["status"], - error_category=_classify_exception(exc), + status_code=upstream_status, + error_category=error_category, latency_ms=(time.perf_counter() - start) * 1000.0, ttft_ms=state["ttft_ms"], - response_raw=(bytes(state["body"]).decode("utf-8", errors="replace") if state["body"] else None), + response_raw=upstream_body.decode("utf-8", errors="replace") if upstream_body else None, ) except Exception: logger.warning("Model-call capture finalization failed.", exc_info=True) diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index c401da9b9d..0e1775ff3f 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -402,6 +402,29 @@ async def _boom(body: dict = Body()) -> dict: assert calls[0].latency_ttft_ms is None # nothing streamed before the raise +def test_raised_upstream_error_preserves_status_and_body(tmp_path): + class UpstreamError(RuntimeError): + response = SimpleNamespace(status_code=403, content=b'{"error":{"message":"denied"}}') + + app = FastAPI() + + @app.post("/v1/responses") + async def _boom(body: dict = Body()) -> dict: + raise UpstreamError + + _install_capture(app, tmp_path) + response = TestClient(app, raise_server_exceptions=False).post( + "/ng-rollout/r-upstream/v1/responses", + json={"input": "x"}, + ) + + assert response.status_code == 500 + [exchange] = CaptureStore(tmp_path).read("r-upstream") + assert exchange["status_code"] == 403 + assert exchange["error_category"] == "auth" + assert json.loads(exchange["response_raw"]) == {"error": {"message": "denied"}} + + @pytest.mark.parametrize("request_bytes", [b"{not-json", b"[]"]) def test_invalid_request_body_does_not_drop_capture(tmp_path, request_bytes): app = FastAPI() From 662c2b0b20a7bd3f34f7488f9947abb042dba98f Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Tue, 28 Jul 2026 10:25:38 +0200 Subject: [PATCH 06/14] Harden rollout observation contract Signed-off-by: Michal Bien --- nemo_gym/base_responses_api_model.py | 14 +++++++----- nemo_gym/rollout_observability.py | 22 ++++++++++++++++--- resources_servers/gymnasium/base.py | 2 ++ resources_servers/gymnasium/tests/test_app.py | 19 ++++++++++++++++ .../test_base_responses_api_model.py | 20 ++++++++++++++++- .../unit_tests/test_rollout_observability.py | 15 +++++++++++++ 6 files changed, 83 insertions(+), 9 deletions(-) diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 57abfe76d9..b9c7e19b9b 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -1237,14 +1237,17 @@ def clear_model_call_captures_for_rollouts(records: list[Any], capture_dirs: lis store.incomplete_path_for(rollout_id).unlink(missing_ok=True) -def merge_model_call_capture_into_record(record: dict[str, Any], capture_dirs: list[Path]) -> dict[str, Any]: +def merge_model_call_capture_into_record( + record: dict[str, Any], capture_dirs: list[Path], *, include_payloads: bool = False +) -> dict[str, Any]: """Attach captured model-call observability data to a rollout record in place. Keyed by the rollout id derived from the record's task/rollout/attempt indices, so the attached shape is identical for every agent harness. Adds ``ng_model_call_capture = {rollout_id, metrics, calls}`` where ``calls`` are derived observability - records. Raw request and response payloads remain in the capture store. Capture/read/join - failures are attached as ``gaps``. The harness output and reward are not modified. + records. Raw request and response payloads remain in the capture store and are omitted from the + attachment unless ``include_payloads`` is true. Capture/read/join failures are attached as + ``gaps``. The harness output and reward are not modified. """ if not capture_dirs: return record @@ -1275,17 +1278,18 @@ def merge_model_call_capture_into_record(record: dict[str, Any], capture_dirs: l calls = [] gaps.append(ObservationGap(code="model_call_capture_unreadable")) observations = record.get("ng_agent_observations") - if observations is not None and calls: + if observations is not None: try: bundle = AgentObservationBundle.model_validate(observations) record["ng_agent_observations"] = join_model_call_observations(bundle, calls).model_dump(mode="json") except Exception: logger.warning("Could not join agent observations for rollout %s.", rollout_id, exc_info=True) gaps.append(ObservationGap(code="agent_observation_join_failed")) + exclude = None if include_payloads else {"request", "response", "request_raw", "response_raw"} capture = { "rollout_id": rollout_id, "metrics": aggregate_model_call_records(calls), - "calls": [call.model_dump(exclude={"request", "response", "request_raw", "response_raw"}) for call in calls], + "calls": [call.model_dump(exclude=exclude) for call in calls], } if gaps: capture["gaps"] = [gap.model_dump(mode="json", exclude_none=True) for gap in gaps] diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py index 162359e059..db770cea6d 100644 --- a/nemo_gym/rollout_observability.py +++ b/nemo_gym/rollout_observability.py @@ -147,9 +147,24 @@ class AgentObservationBundle(ObservationModel): @model_validator(mode="after") def validate_identity(self) -> "AgentObservationBundle": - invocation_ids = [record.invocation_id for record in self.records if isinstance(record, AgentInvocation)] - if len(invocation_ids) != len(set(invocation_ids)): + invocation_records = [record for record in self.records if isinstance(record, AgentInvocation)] + invocations = {record.invocation_id: record for record in invocation_records} + if len(invocation_records) != len(invocations): raise ValueError("invocation_id must be unique within an observation bundle") + + resolved: set[str] = set() + for invocation_id in invocations: + chain: set[str] = set() + current = invocation_id + while current in invocations and current not in resolved: + if current in chain: + raise ValueError("parent_invocation_id must not form a cycle") + chain.add(current) + parent = invocations[current].parent_invocation_id + if parent is None: + break + current = parent + resolved.update(chain) return self @@ -207,7 +222,8 @@ def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: for gap in bundle.gaps if gap.code not in join_codes and not ( - gap.code == "model_call_ownership_unavailable" + captured + and gap.code == "model_call_ownership_unavailable" and gap.invocation_id is None and (gap.detail is None or gap.detail.startswith("capture:")) ) diff --git a/resources_servers/gymnasium/base.py b/resources_servers/gymnasium/base.py index 211e30d03f..9c9044c9a4 100644 --- a/resources_servers/gymnasium/base.py +++ b/resources_servers/gymnasium/base.py @@ -26,6 +26,7 @@ NeMoGymResponseCreateParamsNonStreaming, NeMoGymResponseFunctionToolCall, ) +from nemo_gym.rollout_correlation import RolloutContextMiddleware from nemo_gym.server_utils import SESSION_ID_KEY @@ -79,6 +80,7 @@ class GymnasiumServer(SimpleResourcesServer): def setup_webserver(self) -> FastAPI: app = FastAPI() self.setup_session_middleware(app) + app.add_middleware(RolloutContextMiddleware) app.post("/reset")(self._reset_endpoint) app.post("/step")(self._step_endpoint) app.post("/aggregate_metrics")(self.aggregate_metrics) diff --git a/resources_servers/gymnasium/tests/test_app.py b/resources_servers/gymnasium/tests/test_app.py index a6faac8e00..87bd3570a3 100644 --- a/resources_servers/gymnasium/tests/test_app.py +++ b/resources_servers/gymnasium/tests/test_app.py @@ -16,6 +16,7 @@ from unittest.mock import MagicMock import pytest +from fastapi.testclient import TestClient from nemo_gym.base_resources_server import BaseResourcesServerConfig from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseOutputMessage, NeMoGymResponseOutputText @@ -88,6 +89,24 @@ def test_routes_registered(self): routes = {r.path for r in env.setup_webserver().routes} assert {"/reset", "/step", "/aggregate_metrics"}.issubset(routes) + def test_rollout_prefixed_reset_and_step(self): + client = TestClient(_make_env(_TerminatingEnv).setup_webserver()) + reset = client.post( + "/ng-rollout/4-2/reset", + json={"responses_create_params": {"input": []}}, + ) + assert reset.status_code == 200 + + step = client.post( + "/ng-rollout/4-2/step", + json={ + "responses_create_params": {"input": []}, + "response": _make_response("x").model_dump(mode="json"), + }, + ) + assert step.status_code == 200 + assert step.json()["terminated"] is True + def test_verify_raises(self): env = _make_env(_TerminatingEnv) with pytest.raises(NotImplementedError): diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index 0e1775ff3f..cef7abeb53 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -1129,6 +1129,14 @@ def test_merge_capture_attaches_metrics_without_raw_payloads(tmp_path): assert joined_ref["model_call_id"] == "call-A" assert record["ng_agent_observations"]["gaps"] == [] + with_payloads = {"_ng_task_index": 0, "_ng_rollout_index": 0} + merge_model_call_capture_into_record(with_payloads, [tmp_path], include_payloads=True) + attached_call = with_payloads["ng_model_call_capture"]["calls"][0] + assert attached_call["request"] == exchange["request"] + assert attached_call["response"] == exchange["response"] + assert attached_call["request_raw"] == "malformed request" + assert attached_call["response_raw"] == "malformed response" + def test_merge_capture_reports_missing_capture(tmp_path): from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record @@ -1139,13 +1147,23 @@ def test_merge_capture_reports_missing_capture(tmp_path): "reward": 1.0, "ng_agent_observations": { "source": "test", + "records": [ + { + "kind": "agent_invocation", + "invocation_id": "root", + "model_calls": [{"model_call_id": "missing"}], + } + ], "gaps": [{"code": "model_call_ownership_unavailable"}], }, } merge_model_call_capture_into_record(rec, [tmp_path]) # no capture file for 9-9 assert rec["ng_model_call_capture"]["calls"] == [] assert [gap["code"] for gap in rec["ng_model_call_capture"]["gaps"]] == ["model_call_capture_no_records"] - assert rec["ng_agent_observations"]["gaps"] == [{"code": "model_call_ownership_unavailable"}] + assert {gap["code"] for gap in rec["ng_agent_observations"]["gaps"]} == { + "model_call_ownership_unavailable", + "model_call_reference_unmatched", + } CaptureStore(tmp_path).path_for("8-8").touch() empty = {"_ng_task_index": 8, "_ng_rollout_index": 8} diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py index 2a09972cf5..2cba07d563 100644 --- a/tests/unit_tests/test_rollout_observability.py +++ b/tests/unit_tests/test_rollout_observability.py @@ -35,6 +35,21 @@ def test_observation_bundle_rejects_duplicate_invocation_ids() -> None: ) +@pytest.mark.parametrize( + "records", + ( + [AgentInvocation(invocation_id="root", parent_invocation_id="root")], + [ + AgentInvocation(invocation_id="a", parent_invocation_id="b"), + AgentInvocation(invocation_id="b", parent_invocation_id="a"), + ], + ), +) +def test_observation_bundle_rejects_parent_cycles(records: list[AgentInvocation]) -> None: + with pytest.raises(ValidationError, match="parent_invocation_id must not form a cycle"): + AgentObservationBundle(source="test", records=records) + + def test_observation_models_reject_unknown_fields() -> None: with pytest.raises(ValidationError, match="producer_extension"): ModelCallRef.model_validate({"model_call_id": "call-1", "producer_extension": "unexpected"}) From 698e1e4140a8de7f2879f57ff30391b1405d77ca Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Tue, 28 Jul 2026 18:10:04 +0200 Subject: [PATCH 07/14] Address rollout observation review feedback Signed-off-by: Michal Bien --- .../pages/model-server/model-call-capture.mdx | 8 +- nemo_gym/base_responses_api_model.py | 20 ++-- nemo_gym/rollout_observability.py | 99 +++++++++++++++---- .../test_base_responses_api_model.py | 21 +++- .../unit_tests/test_rollout_observability.py | 95 ++++++++++++++++-- 5 files changed, 210 insertions(+), 33 deletions(-) diff --git a/fern/versions/latest/pages/model-server/model-call-capture.mdx b/fern/versions/latest/pages/model-server/model-call-capture.mdx index fb8618a9f2..261bbd3da5 100644 --- a/fern/versions/latest/pages/model-server/model-call-capture.mdx +++ b/fern/versions/latest/pages/model-server/model-call-capture.mdx @@ -41,8 +41,8 @@ Use `/ng-rollout/` as the model-server URL prefix. SDKs append their the rollout prefix before routing. Correlation is caller-supplied. An unprefixed call is forwarded normally but is not captured. -SDK-based harnesses must configure their `model_server` field for capture; calls sent directly to an -external provider do not cross a Gym model server. +SDK-based harnesses must route calls through a Gym Model Server; calls sent directly to an external +provider are not captured. Agents built on `SimpleResponsesAPIAgent` can use these helpers: `url_path_for_run(url_path, body)` prefixes a downstream call from the run request's task/rollout indices (only when @@ -129,6 +129,10 @@ references by `model_call_id`, or by the exact `(model_ref, response_id)` pair w the protocol response ID. Do not infer ownership from timestamps, text, or list position. The full model request and response remain in `CaptureStore`; rollout attachments intentionally omit them. +Compaction records distinguish the calls immediately before and after the context change from +`model_calls` used to perform the compaction. Compaction calls are exact references to calls owned by +the enclosing invocation; opaque integrations leave them empty rather than infer them. + Tool timestamps are UTC Unix seconds when the source provides wall-clock time. `duration_ms` is the measured interval, and `timing_source` identifies executor, harness, or artifact-derived timing. diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index b9c7e19b9b..a8dc1aab4e 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -756,18 +756,24 @@ def _classify_exception(exc: BaseException) -> str: def _exception_http_details(exc: BaseException) -> tuple[Optional[int], bytes]: - response = getattr(exc, "response", None) - status = getattr(exc, "status", None) + def read_attr(value: Any, name: str) -> Any: + try: + return getattr(value, name, None) + except Exception: + return None + + response = read_attr(exc, "response") + status = read_attr(exc, "status") if not isinstance(status, int): - status = getattr(exc, "status_code", None) + status = read_attr(exc, "status_code") if not isinstance(status, int) and response is not None: - status = getattr(response, "status_code", None) + status = read_attr(response, "status_code") - body = getattr(exc, "response_content", None) + body = read_attr(exc, "response_content") if body is None and response is not None: - body = getattr(response, "content", None) + body = read_attr(response, "content") if body is None: - body = getattr(response, "text", None) + body = read_attr(response, "text") if isinstance(body, str): body = body.encode() return (status if isinstance(status, int) else None, bytes(body) if isinstance(body, (bytes, bytearray)) else b"") diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py index db770cea6d..4ca3828a55 100644 --- a/nemo_gym/rollout_observability.py +++ b/nemo_gym/rollout_observability.py @@ -112,13 +112,31 @@ class ContextCompactionObservation(ObservationModel): invocation_id: str observed_at: Optional[float] = None trigger: Optional[str] = None - tokens_before: Optional[int] = None - tokens_after: Optional[int] = None + tokens_before: Optional[int] = Field( + default=None, + ge=0, + description="Producer-reported token count before compaction; accounting may differ from tokens_after.", + ) + tokens_after: Optional[int] = Field( + default=None, + ge=0, + description="Producer-reported token count after compaction; accounting may differ from tokens_before.", + ) outcome: Literal["completed", "failed", "aborted", "unknown"] = "unknown" summary: Optional[str] = None first_kept_item_id: Optional[str] = None - before_model_call: Optional[ModelCallRef] = None - after_model_call: Optional[ModelCallRef] = None + before_model_call: Optional[ModelCallRef] = Field( + default=None, + description="Last invocation model call observed before compaction.", + ) + model_calls: list[ModelCallRef] = Field( + default_factory=list, + description=("Invocation-owned model calls used for compaction, in producer-observed order; never inferred."), + ) + after_model_call: Optional[ModelCallRef] = Field( + default=None, + description="First invocation model call observed after compaction.", + ) class ObservationGap(ObservationModel): @@ -147,16 +165,18 @@ class AgentObservationBundle(ObservationModel): @model_validator(mode="after") def validate_identity(self) -> "AgentObservationBundle": + """Require unique invocation IDs and reject cycles in the observed parent graph.""" invocation_records = [record for record in self.records if isinstance(record, AgentInvocation)] invocations = {record.invocation_id: record for record in invocation_records} if len(invocation_records) != len(invocations): raise ValueError("invocation_id must be unique within an observation bundle") - resolved: set[str] = set() + # Missing parents are valid when an opaque producer can observe only part of the invocation tree. + checked: set[str] = set() for invocation_id in invocations: chain: set[str] = set() current = invocation_id - while current in invocations and current not in resolved: + while current in invocations and current not in checked: if current in chain: raise ValueError("parent_invocation_id must not form a cycle") chain.add(current) @@ -164,7 +184,7 @@ def validate_identity(self) -> "AgentObservationBundle": if parent is None: break current = parent - resolved.update(chain) + checked.update(chain) return self @@ -203,7 +223,7 @@ def matches(ref: ModelCallRef) -> list[ModelCallRecord]: assert ref.model_ref is not None and ref.response_id is not None return by_response.get((ref.model_ref.type, ref.model_ref.name, ref.response_id), []) - def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: + def canonical(call: ModelCallRecord) -> ModelCallRef: return ModelCallRef.model_validate( { "model_call_id": call.model_call_id, @@ -222,14 +242,14 @@ def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: for gap in bundle.gaps if gap.code not in join_codes and not ( - captured - and gap.code == "model_call_ownership_unavailable" + gap.code == "model_call_ownership_unavailable" and gap.invocation_id is None - and (gap.detail is None or gap.detail.startswith("capture:")) + and gap.detail is not None + and gap.detail.startswith("capture:") ) ] - claimed: set[int] = set() + owner_by_call: dict[int, str] = {} join_gaps: list[ObservationGap] = [] for invocation in invocations: resolved: list[ModelCallRef] = [] @@ -249,7 +269,7 @@ def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: call = candidates[0] identity = id(call) - if identity in claimed: + if identity in owner_by_call: join_gaps.append( ObservationGap( code="model_call_reference_conflict", @@ -259,18 +279,63 @@ def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: ) resolved.append(ref) continue - claimed.add(identity) - resolved.append(canonical(ref, call)) + owner_by_call[identity] = invocation.invocation_id + resolved.append(canonical(call)) invocation.model_calls = resolved for compaction in compactions: + resolved: list[ModelCallRef] = [] + for ref in compaction.model_calls: + candidates = matches(ref) + detail = ref.model_call_id or ref.response_id + if len(candidates) != 1: + join_gaps.append( + ObservationGap( + code=("model_call_reference_ambiguous" if candidates else "model_call_reference_unmatched"), + invocation_id=compaction.invocation_id, + detail=f"model_calls:{detail}", + ) + ) + resolved.append(ref) + continue + + call = candidates[0] + owner = owner_by_call.get(id(call)) + if owner != compaction.invocation_id: + join_gaps.append( + ObservationGap( + code=( + "model_call_reference_conflict" + if owner is not None + else "model_call_ownership_unavailable" + ), + invocation_id=compaction.invocation_id, + detail=f"model_calls:{call.model_call_id or call.response_id}", + ) + ) + resolved.append(ref) + continue + resolved.append(canonical(call)) + compaction.model_calls = resolved + for field_name in ("before_model_call", "after_model_call"): ref = getattr(compaction, field_name) if ref is None: continue candidates = matches(ref) if len(candidates) == 1: - setattr(compaction, field_name, canonical(ref, candidates[0])) + call = candidates[0] + owner = owner_by_call.get(id(call)) + if owner is not None and owner != compaction.invocation_id: + join_gaps.append( + ObservationGap( + code="model_call_reference_conflict", + invocation_id=compaction.invocation_id, + detail=f"{field_name}:{call.model_call_id or call.response_id}", + ) + ) + continue + setattr(compaction, field_name, canonical(call)) else: join_gaps.append( ObservationGap( @@ -282,7 +347,7 @@ def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: result.gaps.extend(join_gaps) for call in captured: - if id(call) not in claimed: + if id(call) not in owner_by_call: result.gaps.append( ObservationGap( code="model_call_ownership_unavailable", diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index cef7abeb53..dead40c4f3 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -560,6 +560,25 @@ class _ReadTimeout(Exception): assert _classify_exception(ValueError("x")) == "exception" +def test_exception_http_details_tolerates_lazy_response_failure(): + from nemo_gym.base_responses_api_model import _exception_http_details + + class _Response: + status_code = 503 + + @property + def content(self): + raise RuntimeError("body unavailable") + + @property + def text(self): + raise RuntimeError("body unavailable") + + error = RuntimeError("upstream") + error.response = _Response() + assert _exception_http_details(error) == (503, b"") + + # --- capture-store config + init failure --- def test_model_call_capture_keys_are_reserved_global_config(): assert {"observability_enabled", "model_call_capture_dir"} <= set(NEMO_GYM_RESERVED_TOP_LEVEL_KEYS) @@ -1107,7 +1126,7 @@ def test_merge_capture_attaches_metrics_without_raw_payloads(tmp_path): ], } ], - "gaps": [{"code": "model_call_ownership_unavailable"}], + "gaps": [{"code": "model_call_ownership_unavailable", "detail": "capture:stale"}], }, } merge_model_call_capture_into_record(record, [tmp_path]) diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py index 2cba07d563..b3e4b295ec 100644 --- a/tests/unit_tests/test_rollout_observability.py +++ b/tests/unit_tests/test_rollout_observability.py @@ -50,6 +50,14 @@ def test_observation_bundle_rejects_parent_cycles(records: list[AgentInvocation] AgentObservationBundle(source="test", records=records) +def test_observation_bundle_allows_missing_parent() -> None: + bundle = AgentObservationBundle( + source="test", + records=[AgentInvocation(invocation_id="child", parent_invocation_id="unobserved")], + ) + assert bundle.records[0].parent_invocation_id == "unobserved" + + def test_observation_models_reject_unknown_fields() -> None: with pytest.raises(ValidationError, match="producer_extension"): ModelCallRef.model_validate({"model_call_id": "call-1", "producer_extension": "unexpected"}) @@ -72,6 +80,28 @@ def test_agent_invocation_rejects_negative_duration() -> None: AgentInvocation(invocation_id="root", duration_ms=-1) +@pytest.mark.parametrize( + "values", + ( + {"tokens_before": -1}, + {"tokens_after": -1}, + ), +) +def test_context_compaction_rejects_invalid_token_counts(values: dict) -> None: + with pytest.raises(ValidationError): + ContextCompactionObservation(invocation_id="root", **values) + + +def test_context_compaction_preserves_producer_reported_token_counts() -> None: + observation = ContextCompactionObservation( + invocation_id="root", + tokens_before=10, + tokens_after=11, + outcome="completed", + ) + assert observation.tokens_after == 11 + + def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() -> None: model_ref = ModelServerRef(name="policy", type="responses_api_models") bundle = AgentObservationBundle( @@ -82,7 +112,10 @@ def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")], ) ], - gaps=[ObservationGap(code="model_call_ownership_unavailable")], + gaps=[ + ObservationGap(code="model_call_ownership_unavailable", detail="direct_provider"), + ObservationGap(code="model_call_ownership_unavailable", detail="capture:stale"), + ], ) calls = [ ModelCallRecord( @@ -102,7 +135,7 @@ def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() assert joined_call.model_ref == model_ref assert joined_call.response_id == "resp-1" ownership_gaps = [gap for gap in joined.gaps if gap.code == "model_call_ownership_unavailable"] - assert [gap.detail for gap in ownership_gaps] == ["capture:call-2:call_index=1"] + assert [gap.detail for gap in ownership_gaps] == ["direct_provider", "capture:call-2:call_index=1"] def test_join_model_calls_does_not_guess_ambiguous_response_ids() -> None: @@ -177,14 +210,23 @@ def test_join_model_calls_reports_conflicting_and_unmatched_references() -> None def test_join_model_calls_resolves_compaction_boundaries() -> None: model_ref = ModelServerRef(name="policy", type="responses_api_models") + references = { + response_id: ModelCallRef(model_ref=model_ref, response_id=response_id) + for response_id in ("before", "compaction", "after") + } bundle = AgentObservationBundle( source="test", records=[ + AgentInvocation( + invocation_id="root", + model_calls=list(references.values()), + ), ContextCompactionObservation( invocation_id="root", - before_model_call=ModelCallRef(model_ref=model_ref, response_id="before"), - after_model_call=ModelCallRef(model_ref=model_ref, response_id="after"), - ) + before_model_call=references["before"], + model_calls=[references["compaction"]], + after_model_call=references["after"], + ), ], ) calls = [ @@ -194,11 +236,17 @@ def test_join_model_calls_resolves_compaction_boundaries() -> None: model_ref=model_ref, call_index=0, ), + ModelCallRecord( + model_call_id="call-compaction", + response_id="compaction", + model_ref=model_ref, + call_index=1, + ), ModelCallRecord( model_call_id="call-after", response_id="after", model_ref=model_ref, - call_index=1, + call_index=2, ), ] @@ -206,7 +254,42 @@ def test_join_model_calls_resolves_compaction_boundaries() -> None: [compaction] = [record for record in joined.records if isinstance(record, ContextCompactionObservation)] assert compaction.before_model_call.model_call_id == "call-before" + assert compaction.model_calls[0].model_call_id == "call-compaction" assert compaction.after_model_call.model_call_id == "call-after" + assert joined.gaps == [] + + +def test_join_model_calls_rejects_cross_invocation_compaction_ownership() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + owned = ModelCallRef(model_ref=model_ref, response_id="owned") + unowned = ModelCallRef(model_ref=model_ref, response_id="unowned") + bundle = AgentObservationBundle( + source="test", + records=[ + AgentInvocation(invocation_id="root"), + AgentInvocation(invocation_id="other", model_calls=[owned]), + ContextCompactionObservation( + invocation_id="root", + before_model_call=owned, + model_calls=[owned, unowned], + ), + ], + ) + calls = [ + ModelCallRecord(model_call_id="call-owned", response_id="owned", model_ref=model_ref, call_index=0), + ModelCallRecord(model_call_id="call-unowned", response_id="unowned", model_ref=model_ref, call_index=1), + ] + + joined = join_model_call_observations(bundle, calls) + compaction = next(record for record in joined.records if isinstance(record, ContextCompactionObservation)) + + assert compaction.before_model_call.model_call_id is None + assert all(reference.model_call_id is None for reference in compaction.model_calls) + assert {(gap.code, gap.detail) for gap in joined.gaps if gap.invocation_id == "root"} == { + ("model_call_reference_conflict", "before_model_call:call-owned"), + ("model_call_reference_conflict", "model_calls:call-owned"), + ("model_call_ownership_unavailable", "model_calls:call-unowned"), + } def test_sandbox_observation_rejects_negative_usage() -> None: From fd2f7d5f5896da9aaa3601786c5ba6e004c89ec9 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Thu, 23 Jul 2026 15:08:35 +0200 Subject: [PATCH 08/14] Add Claude Code and AnyTerminal rollout observations Signed-off-by: Michal Bien --- nemo_gym/anthropic_converter.py | 2 +- responses_api_agents/anyterminal_agent/app.py | 157 +++++- .../anyterminal_agent/tests/test_app.py | 315 ++++++++++- responses_api_agents/claude_code_agent/app.py | 95 +++- .../claude_code_agent/observability.py | 498 ++++++++++++++++++ .../claude_code_agent/tests/test_app.py | 125 ++++- .../tests/test_observability.py | 291 ++++++++++ tests/unit_tests/test_anthropic_converter.py | 1 + 8 files changed, 1441 insertions(+), 43 deletions(-) create mode 100644 responses_api_agents/claude_code_agent/observability.py create mode 100644 responses_api_agents/claude_code_agent/tests/test_observability.py diff --git a/nemo_gym/anthropic_converter.py b/nemo_gym/anthropic_converter.py index ad3e7c6d82..1a30a18be2 100644 --- a/nemo_gym/anthropic_converter.py +++ b/nemo_gym/anthropic_converter.py @@ -372,7 +372,7 @@ def flush_message() -> None: arguments=json.dumps(block.get("input", {})), call_id=block["id"], name=block["name"], - id=block["id"], + id=f"fc_{uuid4().hex}", status="completed", type="function_call", ) diff --git a/responses_api_agents/anyterminal_agent/app.py b/responses_api_agents/anyterminal_agent/app.py index 716c61dfa8..d6fe1c595c 100644 --- a/responses_api_agents/anyterminal_agent/app.py +++ b/responses_api_agents/anyterminal_agent/app.py @@ -35,6 +35,12 @@ from nemo_gym.config_types import ModelServerRef from nemo_gym.global_config import get_first_server_config_dict from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming +from nemo_gym.rollout_observability import ( + AgentObservationBundle, + ObservationGap, + SandboxObservation, + link_tool_calls_to_sandbox, +) from nemo_gym.sandbox import AsyncSandbox, SandboxSpec from nemo_gym.sandbox.providers.apptainer import ApptainerProvider from nemo_gym.sandbox.providers.docker import DockerCreateConfig, DockerProvider @@ -109,6 +115,12 @@ class TerminalBenchMetrics(BaseModel): agent_run_time: Optional[float] = None eval_run_time: Optional[float] = None total_run_time: Optional[float] = None + sandbox_provider: Optional[str] = None + sandbox_id: Optional[str] = None + sandbox_wall_time_s: Optional[float] = None + sandbox_cpu_time_s: Optional[float] = None + sandbox_peak_memory_mib: Optional[float] = None + sandbox_resource_usage_source: Optional[str] = None def update_metrics(metrics_fpath: Path, update_dict: Dict[str, Any]) -> None: @@ -125,7 +137,9 @@ def redact(value: Any) -> Any: return { key: ( "***" - if any(secret in key.lower() for secret in ("api_key", "secret", "password", "token")) + if any(secret in key.lower() for secret in ("api_key", "secret", "password")) + or key.lower() == "token" + or key.lower().endswith("_token") else redact(item) ) for key, item in value.items() @@ -139,6 +153,49 @@ def redact(value: Any) -> Any: return json.dumps(redact(d), indent=indent) +def _load_agent_observations(path: Path) -> tuple[str, str]: + """Return canonical observation JSON and an error code without raising.""" + try: + raw = path.read_text() + except FileNotFoundError: + return "", "" + except OSError: + return "", "observation_read_failed" + + try: + bundle = AgentObservationBundle.model_validate_json(raw) + except Exception: + return "", "observation_parse_failed" + return bundle.model_dump_json(), "" + + +def _observation_gap(source: str, code: str) -> AgentObservationBundle: + return AgentObservationBundle( + source=source, + gaps=[ObservationGap(code=code)], + ) + + +def _terminal_sandbox_observations(metrics: TerminalBenchMetrics) -> list[SandboxObservation]: + if metrics.sandbox_id is None: + return [] + outcome = "sandbox_error" if metrics.sandbox_failed else "timeout" + if not metrics.sandbox_failed and not metrics.agent_timed_out and not metrics.container_timed_out: + outcome = "completed" + return [ + SandboxObservation( + role="environment", + provider=metrics.sandbox_provider, + sandbox_id=metrics.sandbox_id, + outcome=outcome, + wall_time_s=metrics.sandbox_wall_time_s, + cpu_time_s=metrics.sandbox_cpu_time_s, + peak_memory_mib=metrics.sandbox_peak_memory_mib, + resource_usage_source=metrics.sandbox_resource_usage_source, + ) + ] + + ### Agent runner template # Injected into the task container; imports any agent class and calls responses(). @@ -157,6 +214,8 @@ def redact(value: Any) -> Any: INSTRUCTION = Path("/trajectories_mount/instruction.txt").read_text() AGENT_KWARGS = json.loads(os.environ.get("NGTB_AGENT_KWARGS", "{{}}")) SAMPLING = json.loads(os.environ.get("NGTB_SAMPLING", "{{}}")) +OBSERVABILITY = os.environ.get("NGTB_OBSERVABILITY") == "1" +MODEL_SERVER_REF = json.loads(os.environ.get("NGTB_MODEL_SERVER_REF", "null")) from nemo_gym.openai_utils import NeMoGymResponseCreateParamsNonStreaming, NeMoGymEasyInputMessage from nemo_gym.config_types import ModelServerRef, ResourcesServerRef @@ -168,7 +227,7 @@ def redact(value: Any) -> Any: _cfg_sampling = {{k: v for k, v in SAMPLING.items() if k in {agent_cfg_class}.model_fields}} -_model_server = ModelServerRef(name="policy_model", type="responses_api_models") if MODEL_URL else None +_model_server = ModelServerRef.model_validate(MODEL_SERVER_REF) if MODEL_URL and MODEL_SERVER_REF else None config = {agent_cfg_class}( host="0.0.0.0", port=0, @@ -194,8 +253,20 @@ def redact(value: Any) -> Any: model=MODEL_NAME, **SAMPLING, ) -response = asyncio.run(agent.responses(request=None, body=body)) +observed = getattr(agent, "responses_with_observations", None) +if OBSERVABILITY and callable(observed): + episode = asyncio.run(observed(request=None, body=body)) + response = episode.response +else: + response = asyncio.run(agent.responses(request=None, body=body)) Path("/trajectories_mount/response.json").write_text(response.model_dump_json()) +if OBSERVABILITY and callable(observed): + try: + Path("/trajectories_mount/agent_observations.json").write_text( + episode.observations.model_dump_json() + ) + except Exception as exc: + print(f"failed to persist agent observations: {{type(exc).__name__}}", file=sys.stderr, flush=True) print(f"agent finished: {{len(response.output)}} output items", flush=True) """ @@ -312,6 +383,7 @@ class AnyTerminalInstanceConfig(AnyTerminalAgentConfig, AnyTerminalServerConfig) metrics_fpath: Path container: str ray_queue_timestamp: float + observability_enabled: bool = Field(default=False, exclude_if=lambda value: not value) agent_command_str: Optional[str] = None @property @@ -325,6 +397,10 @@ def instance_id(self) -> str: class AnyTerminalVerifyResponse(TerminalBenchMetrics, BaseVerifyResponse): instance_config: Dict[str, Any] + ng_agent_observations: AgentObservationBundle | None = Field( + default=None, + exclude_if=lambda value: value is None, + ) ### Sandbox provider selection @@ -456,6 +532,10 @@ async def _stage_remote_tests(self, sandbox: AsyncSandbox, cfg: AnyTerminalInsta async def _collect_remote_outputs(self, sandbox: AsyncSandbox, cfg: AnyTerminalInstanceConfig) -> None: for remote, local in ( ("/trajectories_mount/response.json", cfg.persistent_dir / "response.json"), + ( + "/trajectories_mount/agent_observations.json", + cfg.persistent_dir / "agent_observations.json", + ), ("/logs/verifier/reward.txt", cfg.verifier_dir / "reward.txt"), ("/logs/verifier/test-stdout.txt", cfg.verifier_dir / "test-stdout.txt"), ): @@ -477,6 +557,10 @@ def _agent_env(self, cfg: AnyTerminalInstanceConfig) -> Dict[str, str]: } if cfg.model_server_url: env["NGTB_MODEL_URL"] = cfg.model_server_url + if cfg.model_server is not None: + env["NGTB_MODEL_SERVER_REF"] = cfg.model_server.model_dump_json() + if cfg.observability_enabled: + env["NGTB_OBSERVABILITY"] = "1" return env async def _run_agent(self, sandbox: AsyncSandbox, cfg: AnyTerminalInstanceConfig) -> tuple[float, bool]: @@ -530,8 +614,12 @@ async def process_single_datapoint(self) -> bool: agent_timed_out = container_timed_out = False sandbox_failed = False agent_run_time = eval_run_time = None + sandbox_id = None + sandbox_usage = None + sandbox_provider = next(iter(cfg.sandbox_provider), None) try: await sandbox.start() + sandbox_id = getattr(sandbox, "sandbox_id", None) if not self._uses_bind_mounts(cfg): await self._stage_remote_runtime(sandbox, cfg) agent_run_time, agent_timed_out = await self._run_agent(sandbox, cfg) @@ -545,6 +633,11 @@ async def process_single_datapoint(self) -> bool: sandbox_failed = True print(f"[{cfg.task_name}] sandbox run failed: {e}", flush=True) finally: + try: + if sandbox_id is not None: + sandbox_usage = await sandbox.resource_usage() + except Exception as e: + print(f"[{cfg.task_name}] sandbox resource usage unavailable: {e}", flush=True) try: await sandbox.stop() except Exception as e: @@ -572,6 +665,12 @@ async def process_single_datapoint(self) -> bool: agent_run_time=agent_run_time, eval_run_time=eval_run_time, total_run_time=total_run_time, + sandbox_provider=sandbox_provider, + sandbox_id=sandbox_id, + sandbox_wall_time_s=None if sandbox_usage is None else sandbox_usage.wall_time_s, + sandbox_cpu_time_s=None if sandbox_usage is None else sandbox_usage.cpu_time_s, + sandbox_peak_memory_mib=None if sandbox_usage is None else sandbox_usage.peak_memory_mib, + sandbox_resource_usage_source=None if sandbox_usage is None else sandbox_usage.source, ) update_metrics(cfg.metrics_fpath, metrics.model_dump()) return resolved @@ -713,6 +812,7 @@ def _setup_params( self.config.container_formatter, task_name, problem_info.get("docker_image", "ubuntu:22.04") ), ray_queue_timestamp=time.time(), + observability_enabled=rollout_id is not None, ) params.metrics_fpath.write_text("{}") @@ -763,6 +863,20 @@ async def _inner_responses(self, params: AnyTerminalInstanceConfig) -> NeMoGymRe except (json.JSONDecodeError, ValueError) as e: print(f"[{params.task_name}] response.json unreadable ({e}), treating as empty response", flush=True) + metadata = { + "input": json.dumps(params.body.model_dump(mode="json").get("input") or []), + "metrics": params.metrics_fpath.read_text(), + "instance_config": _safe_config_json(params), + } + if params.observability_enabled: + observations_json, observations_error = _load_agent_observations( + params.persistent_dir / "agent_observations.json" + ) + metadata.update( + agent_observations=observations_json, + agent_observations_error=observations_error, + ) + return NeMoGymResponse( id=f"anyterminal-{params.instance_id}", created_at=int(time.time()), @@ -772,11 +886,7 @@ async def _inner_responses(self, params: AnyTerminalInstanceConfig) -> NeMoGymRe parallel_tool_calls=params.body.parallel_tool_calls, tool_choice=params.body.tool_choice, tools=tools, - metadata={ - "input": json.dumps(params.body.model_dump(mode="json").get("input") or []), - "metrics": params.metrics_fpath.read_text(), - "instance_config": _safe_config_json(params), - }, + metadata=metadata, ) async def run(self, body: AnyTerminalRunRequest) -> AnyTerminalVerifyResponse: @@ -787,8 +897,9 @@ async def run(self, body: AnyTerminalRunRequest) -> AnyTerminalVerifyResponse: meta, response.metadata = response.metadata, None metrics = TerminalBenchMetrics.model_validate_json(meta["metrics"]) + instance_config = AnyTerminalInstanceConfig.model_validate_json(meta["instance_config"]) - return AnyTerminalVerifyResponse( + result = dict( responses_create_params=body.responses_create_params.model_dump() | { "input": json.loads(meta["input"]), @@ -798,8 +909,34 @@ async def run(self, body: AnyTerminalRunRequest) -> AnyTerminalVerifyResponse: response=response, reward=1.0 if metrics.resolved else 0.0, **metrics.model_dump(), - instance_config=AnyTerminalInstanceConfig.model_validate_json(meta["instance_config"]).model_dump(), + instance_config=instance_config.model_dump(), ) + observations_json = meta.get("agent_observations") + source = self.config.agent_server_module.split(".")[-2].removesuffix("_agent") + observations: Optional[AgentObservationBundle] = None + if observations_json: + try: + observations = AgentObservationBundle.model_validate_json(observations_json) + except Exception: + observations = _observation_gap(source, "observation_parse_failed") + elif meta.get("agent_observations_error"): + observations = _observation_gap(source, meta["agent_observations_error"]) + elif instance_config.observability_enabled: + observations = _observation_gap(source, "agent_observations_unavailable") + if observations is not None: + sandbox_observations = _terminal_sandbox_observations(metrics) + observations.records.extend(sandbox_observations) + if sandbox_observations: + sandbox_observation = sandbox_observations[0] + link_tool_calls_to_sandbox(observations, sandbox_observation.sandbox_id) + if sandbox_observation.cpu_time_s is None: + observations.gaps.append(ObservationGap(code="sandbox_cpu_time_unavailable")) + if sandbox_observation.peak_memory_mib is None: + observations.gaps.append(ObservationGap(code="sandbox_memory_usage_unavailable")) + else: + observations.gaps.append(ObservationGap(code="sandbox_observation_unavailable")) + result["ng_agent_observations"] = observations.model_dump(mode="json") + return AnyTerminalVerifyResponse.model_validate(result) if __name__ == "__main__": diff --git a/responses_api_agents/anyterminal_agent/tests/test_app.py b/responses_api_agents/anyterminal_agent/tests/test_app.py index 1bdeb5a06f..0bb27f4e28 100644 --- a/responses_api_agents/anyterminal_agent/tests/test_app.py +++ b/responses_api_agents/anyterminal_agent/tests/test_app.py @@ -20,15 +20,24 @@ constructing the agent. """ +import asyncio import json from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock, PropertyMock, patch +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest from nemo_gym import PARENT_DIR -from nemo_gym.openai_utils import NeMoGymResponseCreateParamsNonStreaming +from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming +from nemo_gym.rollout_observability import ( + AgentInvocation, + AgentObservationBundle, + ModelCallRef, + SandboxObservation, + ToolCallObservation, +) from nemo_gym.sandbox.providers.apptainer import ApptainerProvider from nemo_gym.sandbox.providers.apptainer import provider as apptainer_provider from nemo_gym.sandbox.providers.docker import DockerProvider @@ -37,13 +46,17 @@ AnyTerminalAgent, AnyTerminalAgentConfig, AnyTerminalInstanceConfig, + AnyTerminalRunRequest, GymAgentHarnessProcessor, RunTerminalAgent, + TerminalBenchMetrics, _build_provider, _format_container, _instruction_from_input, + _load_agent_observations, _read_task_meta, _safe_config_json, + _terminal_sandbox_observations, update_metrics, ) @@ -63,6 +76,31 @@ def _config(**overrides) -> AnyTerminalAgentConfig: return AnyTerminalAgentConfig(**base) +def test_terminal_sandbox_observations_only_map_available_metrics() -> None: + observations = _terminal_sandbox_observations( + TerminalBenchMetrics( + agent_timed_out=True, + container_timed_out=True, + agent_run_time=30.0, + eval_run_time=4.0, + total_run_time=35.0, + sandbox_provider="docker", + sandbox_id="sandbox-1", + sandbox_wall_time_s=35.0, + sandbox_cpu_time_s=6.0, + sandbox_peak_memory_mib=512.0, + sandbox_resource_usage_source="docker_container_cgroup_v2", + ) + ) + + assert [(item.role, item.provider, item.sandbox_id, item.outcome, item.wall_time_s) for item in observations] == [ + ("environment", "docker", "sandbox-1", "timeout", 35.0), + ] + assert observations[0].cpu_time_s == 6.0 + assert observations[0].peak_memory_mib == 512.0 + assert observations[0].resource_usage_source == "docker_container_cgroup_v2" + + class TestRunnerTemplate: def _render(self) -> str: return _RUNNER_TEMPLATE.format( @@ -92,6 +130,21 @@ def test_sampling_is_forwarded(self) -> None: assert "**SAMPLING," in rendered assert "HermesAgentConfig.model_fields" in rendered + def test_observed_episode_is_written_when_enabled(self) -> None: + rendered = self._render() + assert "responses_with_observations" in rendered + assert "/trajectories_mount/agent_observations.json" in rendered + assert 'os.environ.get("NGTB_OBSERVABILITY") == "1"' in rendered + + def test_uses_configured_model_ref_and_writes_response_first(self) -> None: + rendered = self._render() + assert "ModelServerRef.model_validate(MODEL_SERVER_REF)" in rendered + assert 'ModelServerRef(name="policy_model"' not in rendered + assert rendered.index('Path("/trajectories_mount/response.json").write_text') < rendered.index( + 'Path("/trajectories_mount/agent_observations.json").write_text' + ) + assert "failed to persist agent observations" in rendered + class TestAgentKey: def test_key_from_module(self) -> None: @@ -318,6 +371,39 @@ def test_indent_produces_multiline(self, tmp_path: Path) -> None: cfg = _make_instance_config(tmp_path) assert "\n" in _safe_config_json(cfg, indent=2) + def test_disabled_observability_does_not_change_serialized_config(self, tmp_path: Path) -> None: + serialized = _safe_config_json(_make_instance_config(tmp_path)) + + assert "observability_enabled" not in json.loads(serialized) + assert AnyTerminalInstanceConfig.model_validate_json(serialized).observability_enabled is False + + +class TestLoadAgentObservations: + def test_validates_and_canonicalizes_sidecar(self, tmp_path: Path) -> None: + path = tmp_path / "observations.json" + bundle = AgentObservationBundle(source="claude_code", records=[AgentInvocation(invocation_id="root")]) + path.write_text(bundle.model_dump_json()) + + raw, error = _load_agent_observations(path) + + assert error == "" + assert AgentObservationBundle.model_validate_json(raw) == bundle + + @pytest.mark.parametrize( + ("contents", "expected_error"), + [("not json", "observation_parse_failed"), ("{}", "observation_parse_failed")], + ) + def test_rejects_invalid_sidecar(self, tmp_path: Path, contents: str, expected_error: str) -> None: + path = tmp_path / "observations.json" + path.write_text(contents) + assert _load_agent_observations(path) == ("", expected_error) + + def test_read_failure_is_nonfatal(self, tmp_path: Path) -> None: + path = tmp_path / "observations.json" + path.write_text("{}") + with patch.object(Path, "read_text", side_effect=PermissionError): + assert _load_agent_observations(path) == ("", "observation_read_failed") + # ── AnyTerminalInstanceConfig properties ────────────────────────────────────────── @@ -500,6 +586,33 @@ def test_no_model_url_when_empty(self, tmp_path: Path) -> None: env = RunTerminalAgent(config=cfg)._agent_env(cfg) assert "NGTB_MODEL_URL" not in env + def test_model_ref_included_when_set(self, tmp_path: Path) -> None: + cfg = _make_instance_config( + tmp_path, + model_server=ModelServerRef(type="responses_api_models", name="custom_policy"), + model_server_url="http://model:8000/ng-rollout/2-1", + ) + env = RunTerminalAgent(config=cfg)._agent_env(cfg) + assert env["NGTB_MODEL_URL"] == "http://model:8000/ng-rollout/2-1" + assert json.loads(env["NGTB_MODEL_SERVER_REF"]) == { + "type": "responses_api_models", + "name": "custom_policy", + } + + def test_no_model_ref_when_unconfigured(self, tmp_path: Path) -> None: + cfg = _make_instance_config(tmp_path, model_server=None, model_server_url="") + env = RunTerminalAgent(config=cfg)._agent_env(cfg) + assert "NGTB_MODEL_SERVER_REF" not in env + + def test_observability_env_is_opt_in(self, tmp_path: Path) -> None: + disabled_cfg = _make_instance_config(tmp_path / "off") + disabled = RunTerminalAgent(config=disabled_cfg)._agent_env(disabled_cfg) + enabled_cfg = _make_instance_config(tmp_path / "on", observability_enabled=True) + enabled = RunTerminalAgent(config=enabled_cfg)._agent_env(enabled_cfg) + + assert "NGTB_OBSERVABILITY" not in disabled + assert enabled["NGTB_OBSERVABILITY"] == "1" + def test_sampling_and_kwargs_forwarded(self, tmp_path: Path) -> None: cfg = _make_instance_config( tmp_path, @@ -511,6 +624,189 @@ def test_sampling_and_kwargs_forwarded(self, tmp_path: Path) -> None: assert json.loads(env["NGTB_AGENT_KWARGS"]) == {"model": "my-model"} +class TestObservationRoundTrip: + @staticmethod + def _agent() -> AnyTerminalAgent: + server_client = MagicMock() + server_client.global_config_dict = {"observability_enabled": True} + with patch.object(AnyTerminalAgent, "model_post_init"): + agent = AnyTerminalAgent.model_construct( + config=_config(model_server={"type": "responses_api_models", "name": "custom_policy"}), + server_client=server_client, + ) + agent._sem = asyncio.Semaphore(1) + return agent + + @staticmethod + def _response( + instance: AnyTerminalInstanceConfig, + observations: str, + error: str = "", + metrics: dict | None = None, + ) -> NeMoGymResponse: + return NeMoGymResponse( + id="anyterminal-test", + created_at=1, + model="test-model", + object="response", + output=[], + parallel_tool_calls=True, + tool_choice="auto", + tools=[], + metadata={ + "input": "[]", + "metrics": json.dumps(metrics or {"resolved": True}), + "instance_config": _safe_config_json(instance), + "agent_observations": observations, + "agent_observations_error": error, + }, + ) + + @staticmethod + def _body() -> AnyTerminalRunRequest: + return AnyTerminalRunRequest.model_validate( + { + "responses_create_params": {"input": "solve"}, + "_ng_task_index": 1, + "_ng_rollout_index": 2, + } + ) + + @pytest.mark.asyncio + async def test_preserves_nondefault_model_ref(self, tmp_path: Path) -> None: + model_ref = ModelServerRef(type="responses_api_models", name="custom_policy") + bundle = AgentObservationBundle( + source="claude_code", + records=[ + AgentInvocation( + invocation_id="root", + model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")], + ) + ], + ) + instance = _make_instance_config( + tmp_path, + model_server=model_ref, + observability_enabled=True, + ) + agent = self._agent() + + with patch.object( + agent, "_responses", AsyncMock(return_value=self._response(instance, bundle.model_dump_json())) + ): + result = await agent.run(self._body()) + + emitted = result.ng_agent_observations + assert result.reward == 1.0 + assert emitted is not None + invocation = next(record for record in emitted.records if isinstance(record, AgentInvocation)) + assert invocation.model_calls[0].model_ref == model_ref + assert invocation.model_calls[0].response_id == "resp-1" + + @pytest.mark.asyncio + async def test_appends_terminal_sandbox_metrics_to_agent_observations(self, tmp_path: Path) -> None: + bundle = AgentObservationBundle( + source="hermes", + records=[ + AgentInvocation(invocation_id="root"), + ToolCallObservation( + invocation_id="root", + tool_call_id="call-1", + started_at=1.0, + completed_at=3.0, + ), + ToolCallObservation( + invocation_id="root", + tool_call_id="call-2", + started_at=2.0, + completed_at=4.0, + ), + ], + ) + instance = _make_instance_config(tmp_path, observability_enabled=True) + response = self._response( + instance, + bundle.model_dump_json(), + metrics={ + "resolved": False, + "agent_run_time": 12.0, + "total_run_time": 15.0, + "sandbox_provider": "docker", + "sandbox_id": "sandbox-1", + "sandbox_wall_time_s": 15.0, + }, + ) + agent = self._agent() + + with patch.object(agent, "_responses", AsyncMock(return_value=response)): + result = await agent.run(self._body()) + + assert result.ng_agent_observations is not None + sandboxes = [ + record for record in result.ng_agent_observations.records if isinstance(record, SandboxObservation) + ] + tools = [record for record in result.ng_agent_observations.records if isinstance(record, ToolCallObservation)] + assert [(item.role, item.wall_time_s) for item in sandboxes] == [ + ("environment", 15.0), + ] + assert sandboxes[0].provider == "docker" + assert sandboxes[0].sandbox_id == "sandbox-1" + assert {tool.sandbox_id for tool in tools} == {"sandbox-1"} + assert [gap.code for gap in result.ng_agent_observations.gaps] == [ + "sandbox_cpu_time_unavailable", + "sandbox_memory_usage_unavailable", + ] + + @pytest.mark.parametrize( + ("observations", "error", "gap_code"), + [ + ("[]", "", "observation_parse_failed"), + ("", "observation_read_failed", "observation_read_failed"), + ], + ) + @pytest.mark.asyncio + async def test_sidecar_failure_keeps_reward_and_emits_typed_gap( + self, + tmp_path: Path, + observations: str, + error: str, + gap_code: str, + ) -> None: + instance = _make_instance_config(tmp_path, observability_enabled=True) + agent = self._agent() + + with patch.object( + agent, + "_responses", + AsyncMock(return_value=self._response(instance, observations, error)), + ): + result = await agent.run(self._body()) + + assert result.reward == 1.0 + assert result.ng_agent_observations is not None + assert result.ng_agent_observations.records == [] + assert [gap.code for gap in result.ng_agent_observations.gaps] == [ + gap_code, + "sandbox_observation_unavailable", + ] + + @pytest.mark.asyncio + async def test_disabled_observability_does_not_change_response(self, tmp_path: Path) -> None: + instance = _make_instance_config(tmp_path) + agent = self._agent() + + with patch.object(agent, "_responses", AsyncMock(return_value=self._response(instance, ""))): + result = await agent.run(self._body()) + + assert result.reward == 1.0 + assert result.ng_agent_observations is None + assert "ng_agent_observations" not in result.model_dump(mode="json") + assert "observability_enabled" not in result.instance_config + + +# ── RunTerminalAgent.process_single_datapoint ──────────────────────────────────── + + class TestProcessSingleDatapoint: @pytest.fixture(autouse=True) def _no_real_provider(self): @@ -524,8 +820,17 @@ async def test_resolved_when_reward_positive(self, tmp_path: Path) -> None: (cfg.verifier_dir / "reward.txt").write_text("1.0") sandbox = SimpleNamespace( + sandbox_id="sandbox-1", start=AsyncMock(), exec=AsyncMock(return_value=_sandbox_result()), + resource_usage=AsyncMock( + return_value=SimpleNamespace( + wall_time_s=12.0, + cpu_time_s=3.0, + peak_memory_mib=256.0, + source="docker_container_cgroup_v2", + ) + ), stop=AsyncMock(), ) with patch("responses_api_agents.anyterminal_agent.app.AsyncSandbox", return_value=sandbox): @@ -533,7 +838,11 @@ async def test_resolved_when_reward_positive(self, tmp_path: Path) -> None: result = await RunTerminalAgent(config=cfg).process_single_datapoint() assert result is True - assert json.loads(cfg.metrics_fpath.read_text())["resolved"] is True + metrics = json.loads(cfg.metrics_fpath.read_text()) + assert metrics["resolved"] is True + assert metrics["sandbox_id"] == "sandbox-1" + assert metrics["sandbox_cpu_time_s"] == 3.0 + assert metrics["sandbox_peak_memory_mib"] == 256.0 sandbox.stop.assert_awaited_once() async def test_unresolved_when_no_reward_file(self, tmp_path: Path) -> None: diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index 0aaf459b35..f7c7aa723b 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -24,11 +24,11 @@ from asyncio import Semaphore from pathlib import Path from time import time -from typing import Any, Optional +from typing import Any, Callable, Optional from uuid import uuid4 from fastapi import Request -from pydantic import ConfigDict, PrivateAttr +from pydantic import ConfigDict, Field, PrivateAttr from nemo_gym.base_resources_server import NEMO_GYM_MCP_METADATA_KEY, BaseRunRequest, BaseVerifyResponse from nemo_gym.base_responses_api_agent import BaseResponsesAPIAgentConfig, Body, SimpleResponsesAPIAgent @@ -46,8 +46,10 @@ NeMoGymResponseOutputTokensDetails, NeMoGymResponseUsage, ) +from nemo_gym.rollout_observability import AgentEpisode, AgentObservationBundle, ObservationGap from nemo_gym.server_utils import apply_rollout_prefix, get_response_json, raise_for_status from nemo_gym.skills import stage_skills +from responses_api_agents.claude_code_agent.observability import extract_claude_code_observations from responses_api_agents.claude_code_agent.setup_claude_code import ensure_claude_code @@ -243,6 +245,9 @@ class ClaudeCodeAgentVerifyResponse(BaseVerifyResponse): model_config = ConfigDict(extra="allow") turns_used: int = 0 finished_naturally: bool = False + ng_agent_observations: Optional[AgentObservationBundle] = Field( + default=None, exclude_if=lambda value: value is None + ) class ClaudeCodeAgent(SimpleResponsesAPIAgent): @@ -381,6 +386,7 @@ async def _run_claude_code( mcp_config: Optional[str] = None, skills_path: Optional[str] = None, rollout_id: Optional[str] = None, + observation_collector: Optional[Callable[[Path], None]] = None, ) -> tuple[str, str]: """Run claude -p --output-format=stream-json and return (stdout, model_name). @@ -441,7 +447,14 @@ async def _run_claude_code( return stdout.decode(errors="replace"), model finally: if claude_config_dir is not None: - shutil.rmtree(claude_config_dir, ignore_errors=True) + try: + if observation_collector is not None: + try: + await asyncio.to_thread(observation_collector, claude_config_dir) + except Exception: + LOG.exception("failed to collect Claude Code observations") + finally: + shutil.rmtree(claude_config_dir, ignore_errors=True) def _resources_server_base_url(self) -> str: cfg = get_first_server_config_dict( @@ -509,6 +522,7 @@ async def _create_response( mcp_config: Optional[str] = None, skills_path: Optional[str] = None, rollout_id: Optional[str] = None, + observation_collector: Optional[Callable[[Path], None]] = None, ) -> NeMoGymResponse: body = body.model_copy(deep=True) if isinstance(body.input, str): @@ -524,6 +538,7 @@ async def _create_response( mcp_config=mcp_config, skills_path=skills_path, rollout_id=rollout_id, + observation_collector=observation_collector, ) output_items, usage = parse_stream_json(stdout) @@ -568,7 +583,46 @@ async def responses( request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body(), ) -> NeMoGymResponse: - return await self._create_response(body) + rollout_id = request.path_params.get("rollout_id") if request is not None else None + return await self._create_response(body, rollout_id=rollout_id) + + async def responses_with_observations( + self, + request: Optional[Request], + body: NeMoGymResponseCreateParamsNonStreaming, + *, + mcp_config: Optional[str] = None, + skills_path: Optional[str] = None, + rollout_id: Optional[str] = None, + ) -> AgentEpisode: + observations: Optional[AgentObservationBundle] = None + + def collect(config_dir: Path) -> None: + nonlocal observations + try: + observations = extract_claude_code_observations(config_dir, model_ref=self.config.model_server) + if self.config.model_server is None: + observations.gaps.append(ObservationGap(code="model_call_ownership_unavailable")) + except Exception: + LOG.exception("failed to extract Claude Code observations") + observations = AgentObservationBundle( + source="claude_code", + gaps=[ObservationGap(code="observation_parse_failed")], + ) + + response = await self._create_response( + body, + mcp_config=mcp_config, + skills_path=skills_path, + rollout_id=rollout_id, + observation_collector=collect, + ) + if observations is None: + observations = AgentObservationBundle( + source="claude_code", + gaps=[ObservationGap(code="agent_transcript_unavailable")], + ) + return AgentEpisode(response=response, observations=observations) async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> ClaudeCodeAgentVerifyResponse: async with self.sem: @@ -593,18 +647,30 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude with tempfile.TemporaryDirectory(prefix="nemo_gym_claude_mcp_") as mcp_config_dir: mcp_config = self._write_rollout_mcp_config(seed_resp_json, Path(mcp_config_dir)) - agent_resp = await self._create_response( - body.responses_create_params, - mcp_config=mcp_config, - skills_path=skills_path, - rollout_id=rollout_id, - ) + if rollout_id is not None: + episode = await self.responses_with_observations( + request, + body.responses_create_params, + mcp_config=mcp_config, + skills_path=skills_path, + rollout_id=rollout_id, + ) + agent_resp, observations = episode.response, episode.observations + else: + agent_resp = await self._create_response( + body.responses_create_params, + mcp_config=mcp_config, + skills_path=skills_path, + ) + observations = None agent_resp_json = agent_resp.model_dump(mode="json") verify_resp = await self.server_client.post( server_name=self.config.resources_server.name, url_path="/verify", - json=body.model_dump() | {"response": agent_resp_json}, + json=body.model_dump() + | {"response": agent_resp_json} + | ({"rollout_id": rollout_id} if rollout_id is not None else {}), cookies=cookies, ) await raise_for_status(verify_resp) @@ -619,9 +685,10 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude last = gym_resp.output[-1] if gym_resp.output else None naturally = getattr(last, "type", None) == "message" and getattr(last, "role", None) == "assistant" - return ClaudeCodeAgentVerifyResponse.model_validate( - verify_json | {"turns_used": turns, "finished_naturally": naturally} - ) + result = verify_json | {"turns_used": turns, "finished_naturally": naturally} + if observations is not None: + result["ng_agent_observations"] = observations.model_dump(mode="json") + return ClaudeCodeAgentVerifyResponse.model_validate(result) if __name__ == "__main__": diff --git a/responses_api_agents/claude_code_agent/observability.py b/responses_api_agents/claude_code_agent/observability.py new file mode 100644 index 0000000000..65110f17d2 --- /dev/null +++ b/responses_api_agents/claude_code_agent/observability.py @@ -0,0 +1,498 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Read Claude Code's per-run transcripts into Gym observability records.""" + +from __future__ import annotations + +import json +import math +from collections import defaultdict +from datetime import datetime +from pathlib import Path +from typing import Any + +from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import ( + NeMoGymEasyInputMessage, + NeMoGymFunctionCallOutput, + NeMoGymResponseFunctionToolCall, + NeMoGymResponseOutputMessage, + NeMoGymResponseOutputText, + NeMoGymResponseReasoningItem, + NeMoGymSummary, +) +from nemo_gym.rollout_observability import ( + AgentInvocation, + AgentObservationBundle, + ContextCompactionObservation, + ModelCallRef, + ObservationGap, + ToolCallObservation, +) + + +SOURCE = "claude_code" + + +def _gap(code: str, *, invocation_id: str | None = None, detail: str | None = None) -> ObservationGap: + return ObservationGap(code=code, invocation_id=invocation_id, detail=detail) + + +def _timestamp(value: Any) -> float | None: + try: + result = ( + float(value) + if isinstance(value, (int, float)) and not isinstance(value, bool) + else datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() + ) + except (AttributeError, TypeError, ValueError): + return None + return result if math.isfinite(result) else None + + +def _text(value: Any) -> str: + if isinstance(value, str): + return value + if isinstance(value, list): + if all(isinstance(item, dict) and item.get("type") == "text" for item in value): + return "".join(str(item.get("text") or "") for item in value) + return json.dumps(value, ensure_ascii=False, sort_keys=True) + if value is None: + return "" + return json.dumps(value, ensure_ascii=False, sort_keys=True) + + +def _status(block: dict[str, Any], result: Any) -> str: + if block.get("is_error") is True: + return "failed" + if isinstance(result, dict): + if result.get("interrupted") is True: + return "incomplete" + value = result.get("status") + if value in {"completed", "failed", "timeout", "incomplete"}: + return value + # A tool_result block is an explicit terminal observation even when Claude Code + # does not attach a separate status object. + return "completed" + + +def _metadata(event: dict[str, Any]) -> dict[str, Any]: + message = event.get("message") + for owner in (event, message if isinstance(message, dict) else {}): + for key in ("compactMetadata", "compact_metadata"): + if isinstance(metadata := owner.get(key), dict): + return metadata + return {} + + +def _integer(metadata: dict[str, Any], *keys: str) -> int | None: + for key in keys: + value = metadata.get(key) + if isinstance(value, int) and not isinstance(value, bool): + return value + return None + + +def _compaction(event: dict[str, Any], invocation_id: str) -> ContextCompactionObservation | None: + message = event.get("message") + is_summary = isinstance(message, dict) and message.get("isCompactSummary") is True + is_boundary = event.get("type") == "system" and event.get("subtype") == "compact_boundary" + metadata = _metadata(event) + if not is_summary and not is_boundary and not metadata: + return None + + trigger = metadata.get("trigger") + summary = _text(message.get("content")) if is_summary else None + return ContextCompactionObservation( + invocation_id=invocation_id, + observed_at=_timestamp(event.get("timestamp")), + trigger=trigger if isinstance(trigger, str) else None, + tokens_before=_integer(metadata, "tokensBefore", "preTokens"), + tokens_after=_integer(metadata, "tokensAfter", "postTokens"), + outcome="completed", + summary=summary or None, + ) + + +def _message_id(event: dict[str, Any], block_index: int, kind: str) -> str | None: + event_id = event.get("uuid") + if isinstance(event_id, str) and event_id: + return f"{event_id}:{kind}:{block_index}" + message = event.get("message") + response_id = message.get("id") if isinstance(message, dict) else None + if isinstance(response_id, str) and response_id: + return f"{response_id}:{kind}:{block_index}" + return None + + +def _message(item_id: str, text: str) -> NeMoGymResponseOutputMessage: + return NeMoGymResponseOutputMessage(id=item_id, content=[NeMoGymResponseOutputText(text=text, annotations=[])]) + + +def _reasoning(item_id: str, block: dict[str, Any]) -> NeMoGymResponseReasoningItem: + signature = block.get("signature") + return NeMoGymResponseReasoningItem( + id=item_id, + summary=[NeMoGymSummary(text=block["thinking"], type="summary_text")], + encrypted_content=signature if isinstance(signature, str) else None, + ) + + +def _tool_call(tool_call_id: str, block: dict[str, Any]) -> NeMoGymResponseFunctionToolCall: + return NeMoGymResponseFunctionToolCall( + arguments=json.dumps(block.get("input", {}), ensure_ascii=False, sort_keys=True), + call_id=tool_call_id, + name=block.get("name") if isinstance(block.get("name"), str) else "", + id=tool_call_id, + status="completed", + ) + + +def _tool_result(event: dict[str, Any], block: dict[str, Any]) -> NeMoGymFunctionCallOutput: + event_id = event.get("uuid") + return NeMoGymFunctionCallOutput( + call_id=block["tool_use_id"], + output=_text(block.get("content")), + id=event_id if isinstance(event_id, str) else None, + status="completed", + ) + + +def _read_events(config_dir: Path, gaps: list[ObservationGap]) -> list[tuple[int, dict[str, Any]]]: + if not config_dir.is_dir(): + gaps.append(_gap("transcript_dir_missing")) + return [] + + transcript_dir = config_dir / "projects" + if not transcript_dir.is_dir(): + gaps.append(_gap("transcript_dir_missing", detail="projects")) + return [] + + try: + # Claude Code stores session and subagent transcripts below ``projects``. + # Other JSONL files in CLAUDE_CONFIG_DIR may belong to staged skills or + # unrelated CLI state and must not be interpreted as rollout evidence. + paths = sorted(transcript_dir.rglob("*.jsonl")) + except OSError: + gaps.append(_gap("transcript_dir_unreadable")) + return [] + + events: list[tuple[int, dict[str, Any]]] = [] + for path in paths: + try: + lines = path.open(encoding="utf-8", errors="replace") + except OSError: + gaps.append(_gap("transcript_unreadable", detail=path.name)) + continue + with lines: + for line_number, line in enumerate(lines, start=1): + if not line.strip(): + continue + try: + event = json.loads(line) + except (json.JSONDecodeError, UnicodeError): + gaps.append(_gap("malformed_transcript_line", detail=f"{path.name}:{line_number}")) + continue + if not isinstance(event, dict): + gaps.append(_gap("invalid_transcript_record", detail=f"{path.name}:{line_number}")) + continue + if not isinstance(event.get("sessionId"), str): + continue + events.append((len(events), event)) + return events + + +def extract_claude_code_observations( + config_dir: Path, + *, + model_ref: ModelServerRef | None = None, +) -> AgentObservationBundle: + """Extract exact relationships available in one ``CLAUDE_CONFIG_DIR``. + + Transcript IDs and timestamps are used directly. Missing or ambiguous evidence + is reported as a gap; the extractor never joins calls by text or proximity. + """ + + gaps: list[ObservationGap] = [] + raw_events = _read_events(Path(config_dir), gaps) + if not raw_events: + gaps.append(_gap("agent_transcript_unavailable")) + return AgentObservationBundle(source=SOURCE, gaps=gaps) + + events_by_invocation: dict[str, list[tuple[int, dict[str, Any]]]] = defaultdict(list) + first_seen: dict[str, int] = {} + agent_invocations: set[str] = set() + + for ordinal, event in raw_events: + agent_id = event.get("agentId") + invocation_id = agent_id if isinstance(agent_id, str) and agent_id else event["sessionId"] + events_by_invocation[invocation_id].append((ordinal, event)) + first_seen.setdefault(invocation_id, ordinal) + if isinstance(agent_id, str) and agent_id: + agent_invocations.add(invocation_id) + + starts: dict[tuple[str, str], list[tuple[float | None, str]]] = defaultdict(list) + finishes: dict[tuple[str, str], list[tuple[float | None, str]]] = defaultdict(list) + parents: dict[str, tuple[str, str, str, int]] = {} + ambiguous_parents: set[str] = set() + conversations: dict[str, list[Any]] = defaultdict(list) + model_calls: dict[str, list[ModelCallRef]] = defaultdict(list) + compactions: list[ContextCompactionObservation] = [] + + for invocation_id, entries in events_by_invocation.items(): + entries.sort( + key=lambda pair: ( + _timestamp(pair[1].get("timestamp")) is None, + _timestamp(pair[1].get("timestamp")) or 0, + pair[0], + ) + ) + items = conversations[invocation_id] + refs = model_calls[invocation_id] + + def add_gap(code: str, detail: str | None = None) -> None: + gaps.append(_gap(code, invocation_id=invocation_id, detail=detail)) + + seen_response_ids: set[str] = set() + last_model_call: ModelCallRef | None = None + pending_compactions: list[ContextCompactionObservation] = [] + previous_compaction: tuple[int, bool, ContextCompactionObservation] | None = None + for entry_index, (ordinal, event) in enumerate(entries): + compaction = _compaction(event, invocation_id) + if compaction is not None: + message = event.get("message") + is_summary = isinstance(message, dict) and message.get("isCompactSummary") is True + if ( + previous_compaction is not None + and previous_compaction[0] + 1 == entry_index + and previous_compaction[1] != is_summary + ): + prior = previous_compaction[2] + for field in ("trigger", "tokens_before", "tokens_after", "summary"): + if getattr(prior, field) is None: + setattr(prior, field, getattr(compaction, field)) + compaction = prior + else: + compaction.before_model_call = last_model_call + compactions.append(compaction) + pending_compactions.append(compaction) + if last_model_call is None: + add_gap("compaction_before_model_call_unavailable") + previous_compaction = (entry_index, is_summary, compaction) + if compaction.observed_at is None: + add_gap("compaction_timestamp_missing") + else: + previous_compaction = None + + message = event.get("message") + if not isinstance(message, dict): + continue + role = message.get("role") or event.get("type") + content = message.get("content") + + if role == "assistant": + response_id = message.get("id") + model_call = None + if not isinstance(response_id, str) or not response_id: + add_gap("model_response_id_missing") + elif model_ref is not None and response_id not in seen_response_ids: + model_call = ModelCallRef(model_ref=model_ref, response_id=response_id) + refs.append(model_call) + seen_response_ids.add(response_id) + if pending_compactions: + for pending in pending_compactions: + pending.after_model_call = model_call + if model_call is None: + add_gap("compaction_after_model_call_unavailable") + pending_compactions.clear() + if model_call is not None: + last_model_call = model_call + + if isinstance(content, list): + blocks = content + elif isinstance(content, str): + blocks = [{"type": "text", "text": content}] + else: + add_gap("unsupported_assistant_content_block", type(content).__name__) + blocks = [] + for block_index, block in enumerate(blocks): + if not isinstance(block, dict): + add_gap("invalid_assistant_content") + continue + block_type = block.get("type") + item_id = _message_id(event, block_index, str(block_type or "content")) + if block_type == "text": + text = block.get("text") + if not isinstance(text, str) or not text: + continue + if item_id is None: + add_gap("assistant_item_id_missing") + continue + items.append(_message(item_id, text)) + elif block_type == "thinking": + thinking = block.get("thinking") + if not isinstance(thinking, str) or not thinking: + continue + if item_id is None: + add_gap("reasoning_item_id_missing") + continue + items.append(_reasoning(item_id, block)) + elif block_type == "tool_use": + tool_call_id = block.get("id") + if not isinstance(tool_call_id, str) or not tool_call_id: + add_gap("tool_call_id_missing") + continue + tool_name = block.get("name") if isinstance(block.get("name"), str) else "" + items.append(_tool_call(tool_call_id, block)) + starts[(invocation_id, tool_call_id)].append((_timestamp(event.get("timestamp")), tool_name)) + else: + add_gap( + "unsupported_assistant_content_block", + block_type if isinstance(block_type, str) else None, + ) + + elif role in {"user", "system", "developer"}: + if isinstance(content, str): + if content: + items.append(NeMoGymEasyInputMessage(role=role, content=content)) + continue + if not isinstance(content, list): + if content is not None: + add_gap("unsupported_user_content_block", type(content).__name__) + continue + + tool_results: list[dict[str, Any]] = [] + result_metadata = event.get("toolUseResult") + for block in content: + if not isinstance(block, dict): + add_gap("invalid_user_content") + continue + if block.get("type") == "tool_result": + tool_results.append(block) + tool_call_id = block.get("tool_use_id") + if not isinstance(tool_call_id, str) or not tool_call_id: + add_gap("tool_result_id_missing") + continue + tool_status = _status(block, result_metadata) + items.append(_tool_result(event, block)) + finishes[(invocation_id, tool_call_id)].append( + (_timestamp(event.get("timestamp")), tool_status) + ) + elif block.get("type") == "text": + if isinstance(block.get("text"), str): + items.append(NeMoGymEasyInputMessage(role=role, content=block["text"])) + else: + add_gap("unsupported_user_content_block", "text") + else: + block_type = block.get("type") + add_gap( + "unsupported_user_content_block", + block_type if isinstance(block_type, str) else None, + ) + + child_id = result_metadata.get("agentId") if isinstance(result_metadata, dict) else None + if isinstance(child_id, str) and child_id: + if len(tool_results) == 1 and isinstance(tool_results[0].get("tool_use_id"), str): + parent = ( + invocation_id, + tool_results[0]["tool_use_id"], + _status(tool_results[0], result_metadata), + ordinal, + ) + if child_id in parents and parents[child_id][:2] != parent[:2]: + parents.pop(child_id) + ambiguous_parents.add(child_id) + gaps.append(_gap("conflicting_subagent_parent", invocation_id=child_id)) + elif child_id not in ambiguous_parents: + parents.setdefault(child_id, parent) + else: + add_gap("ambiguous_subagent_relation") + + for _ in pending_compactions: + add_gap("compaction_after_model_call_unavailable") + + tool_calls: list[ToolCallObservation] = [] + for invocation_id, tool_call_id in sorted( + set(starts) | set(finishes), key=lambda key: (first_seen.get(key[0], math.inf), key[1]) + ): + call_starts = starts.get((invocation_id, tool_call_id), []) + call_finishes = finishes.get((invocation_id, tool_call_id), []) + + def add_tool_gap(code: str) -> None: + gaps.append(_gap(code, invocation_id=invocation_id, detail=tool_call_id)) + + if len(call_starts) > 1 or len(call_finishes) > 1: + add_tool_gap("ambiguous_tool_artifact") + continue + started_at, tool_name = call_starts[0] if call_starts else (None, "") + completed_at, tool_status = call_finishes[0] if call_finishes else (None, "incomplete") + if not call_starts: + add_tool_gap("tool_start_missing") + if not call_finishes: + add_tool_gap("tool_result_missing") + if call_starts and started_at is None: + add_tool_gap("tool_start_timestamp_missing") + if call_finishes and completed_at is None: + add_tool_gap("tool_result_timestamp_missing") + duration_ms = None + if started_at is not None and completed_at is not None and completed_at >= started_at: + duration_ms = (completed_at - started_at) * 1000 + tool_calls.append( + ToolCallObservation( + invocation_id=invocation_id, + tool_call_id=tool_call_id, + tool_name=tool_name or None, + started_at=started_at, + completed_at=completed_at, + duration_ms=duration_ms, + timing_source="artifact" if started_at is not None or completed_at is not None else None, + status=tool_status, + ) + ) + + parent_by_invocation = {child_id: parent[:3] for child_id, parent in parents.items()} + for child_id, parent in parents.items(): + if child_id not in events_by_invocation: + first_seen[child_id] = parent[3] + gaps.append(_gap("subagent_transcript_missing", invocation_id=child_id)) + + for invocation_id in agent_invocations - set(parent_by_invocation): + gaps.append(_gap("subagent_parent_unavailable", invocation_id=invocation_id)) + + all_invocation_ids = set(events_by_invocation) | set(parent_by_invocation) + invocations_by_id: dict[str, AgentInvocation] = {} + for invocation_id in all_invocation_ids: + parent = parent_by_invocation.get(invocation_id) + invocations_by_id[invocation_id] = AgentInvocation( + invocation_id=invocation_id, + parent_invocation_id=parent[0] if parent else None, + spawned_by_tool_call_id=parent[1] if parent else None, + status=parent[2] if parent else "unknown", + model_calls=model_calls[invocation_id], + conversation=conversations[invocation_id], + ) + + def order_key(invocation_id: str) -> tuple[tuple[float, str], ...]: + path: list[tuple[float, str]] = [] + seen: set[str] = set() + while invocation_id not in seen: + seen.add(invocation_id) + path.append((first_seen.get(invocation_id, math.inf), invocation_id)) + parent = parent_by_invocation.get(invocation_id) + if parent is None: + break + invocation_id = parent[0] + return tuple(reversed(path)) + + ordered_ids = sorted(all_invocation_ids, key=order_key) + + return AgentObservationBundle( + source=SOURCE, + records=[ + *(invocations_by_id[invocation_id] for invocation_id in ordered_ids), + *tool_calls, + *compactions, + ], + gaps=gaps, + ) 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 8209eafbb6..faaac204a8 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -15,6 +15,7 @@ import asyncio import json +import threading from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -31,6 +32,7 @@ NeMoGymResponseFunctionToolCall, NeMoGymResponseOutputMessage, ) +from nemo_gym.rollout_observability import AgentInvocation from nemo_gym.server_utils import ServerClient from responses_api_agents.claude_code_agent.app import ( ClaudeCodeAgent, @@ -41,6 +43,7 @@ _extract_instruction, parse_stream_json, ) +from responses_api_agents.claude_code_agent.observability import extract_claude_code_observations def _write_skill_dir(root: Path, name: str = "cot_enhanced") -> Path: @@ -260,26 +263,22 @@ def _gym_response(text: str = "done") -> dict: } -class TestRunForwardsSkillsPath: - """run() reads skills_ref off the request's model_extra (extra='allow') and forwards its path - directly to _create_response/_run_claude_code.""" +def _seed_and_verify_post(): + async def _post(server_name, url_path, json=None, cookies=None, **kw): + if url_path == "/verify": + return _FakeHttpResp( + {"responses_create_params": {"input": []}, "response": _gym_response(), "reward": 1.0} + ) + return _FakeHttpResp({}) - def _seed_and_verify_post(self): - async def _post(server_name, url_path, json=None, cookies=None, **kw): - if url_path == "/verify": - return _FakeHttpResp( - {"responses_create_params": {"input": []}, "response": _gym_response(), "reward": 1.0} - ) - return _FakeHttpResp({}) + return AsyncMock(side_effect=_post) - return AsyncMock(side_effect=_post) +class TestRunForwardsSkillsPath: def _run(self, agent: ClaudeCodeAgent, body: ClaudeCodeAgentRunRequest, run_claude_code: AsyncMock): - agent.server_client.post = self._seed_and_verify_post() + agent.server_client.post = _seed_and_verify_post() req = MagicMock() req.cookies = {} - # Stub the CLI invocation; _create_response still runs for real, so we exercise the full - # run() -> _create_response -> _run_claude_code argument threading. with patch.object( ClaudeCodeAgent, "_run_claude_code", @@ -306,9 +305,57 @@ def test_no_skills_ref_forwards_none(self) -> None: 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) + result = self._run(agent, body, run_claude_code) assert run_claude_code.call_args.kwargs["skills_path"] is None + assert "ng_agent_observations" not in result.model_dump(mode="json") + + +class TestObservability: + def test_run_returns_observations_when_enabled(self, tmp_path: Path) -> None: + agent = _make_agent(model_server=ModelServerRef(type="responses_api_models", name="policy")) + agent.server_client.global_config_dict = {"observability_enabled": True} + agent.server_client.post = _seed_and_verify_post() + + async def run_claude_code(*args, observation_collector=None, **kwargs): + transcript = tmp_path / "projects" / "session.jsonl" + transcript.parent.mkdir(parents=True) + transcript.write_text( + json.dumps( + { + "type": "assistant", + "sessionId": "session-1", + "timestamp": "2026-07-22T10:00:00Z", + "uuid": "event-1", + "message": { + "role": "assistant", + "id": "msg-1", + "content": [{"type": "text", "text": "done"}], + }, + } + ) + ) + observation_collector(tmp_path) + return _event("assistant", message={"content": [{"type": "text", "text": "done"}]}), "model" + + request = MagicMock() + request.cookies = {} + body = ClaudeCodeAgentRunRequest.model_validate( + { + "responses_create_params": {"input": "solve"}, + "_ng_task_index": 1, + "_ng_rollout_index": 2, + } + ) + with patch.object(ClaudeCodeAgent, "_run_claude_code", run_claude_code): + result = asyncio.run(agent.run(request, body)) + + observations = result.ng_agent_observations + assert observations is not None + invocation = next(record for record in observations.records if isinstance(record, AgentInvocation)) + assert invocation.invocation_id == "session-1" + assert invocation.model_calls[0].response_id == "msg-1" + assert agent.server_client.post.await_args_list[-1].kwargs["json"]["rollout_id"] == "1-2" class TestRunClaudeCode: @@ -422,6 +469,54 @@ async def fake_wait_for(coro, timeout): assert killed["called"] is True assert model == "claude-sonnet-4-6" + def test_collects_observations_before_cleanup(self, tmp_path: Path) -> None: + agent = _make_agent() + captured: dict = {} + event_loop_thread = threading.get_ident() + + class FakeProc: + returncode = 0 + + async def communicate(self): + return b'{"type":"result","usage":{}}\n', b"" + + async def fake_exec(*cmd, **kwargs): + config_dir = Path(kwargs["env"]["CLAUDE_CONFIG_DIR"]) + transcript = config_dir / "projects" / "run.jsonl" + transcript.parent.mkdir(parents=True) + transcript.write_text( + json.dumps( + { + "type": "assistant", + "sessionId": "session-1", + "timestamp": "2026-07-22T10:00:00Z", + "uuid": "event-1", + "message": { + "role": "assistant", + "id": "msg-1", + "content": [{"type": "text", "text": "done"}], + }, + } + ) + ) + captured["config_dir"] = config_dir + return FakeProc() + + def collect(config_dir: Path) -> None: + captured["collector_thread"] = threading.get_ident() + captured["observations"] = extract_claude_code_observations(config_dir) + + 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), + ): + asyncio.run(agent._run_claude_code("hello", observation_collector=collect)) + + invocation = next(record for record in captured["observations"].records if isinstance(record, AgentInvocation)) + assert invocation.invocation_id == "session-1" + assert captured["collector_thread"] != event_loop_thread + assert not captured["config_dir"].exists() + class TestRolloutMCPConfig: def test_no_metadata_preserves_static_config(self, tmp_path: Path) -> None: diff --git a/responses_api_agents/claude_code_agent/tests/test_observability.py b/responses_api_agents/claude_code_agent/tests/test_observability.py new file mode 100644 index 0000000000..a7b704aee9 --- /dev/null +++ b/responses_api_agents/claude_code_agent/tests/test_observability.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path +from typing import TypeVar + +import pytest + +from nemo_gym.config_types import ModelServerRef +from nemo_gym.rollout_observability import ( + AgentInvocation, + AgentObservationBundle, + ContextCompactionObservation, + ToolCallObservation, +) +from responses_api_agents.claude_code_agent.observability import extract_claude_code_observations + + +MODEL_REF = ModelServerRef(type="responses_api_models", name="policy") +T = TypeVar("T") + + +def _records(bundle: AgentObservationBundle, record_type: type[T]) -> list[T]: + return [record for record in bundle.records if isinstance(record, record_type)] + + +def _write(path: Path, *events: dict | str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(event if isinstance(event, str) else json.dumps(event) for event in events)) + + +def _event( + session: str, + role: str, + timestamp: str, + content: str | list[dict], + *, + agent: str | None = None, + message_id: str | None = None, + message_extra: dict | None = None, + **extra: object, +) -> dict: + message = {"role": role, "content": content, **(message_extra or {})} + if message_id: + message["id"] = message_id + event = { + "type": role, + "sessionId": session, + "timestamp": timestamp, + "message": message, + **extra, + } + if agent: + event["agentId"] = agent + return event + + +def _assistant(session: str, timestamp: str, message_id: str, *content: dict, agent: str | None = None) -> dict: + return _event( + session, + "assistant", + timestamp, + list(content), + agent=agent, + message_id=message_id, + uuid=f"{message_id}-event", + ) + + +def _tool_result( + session: str, + timestamp: str, + tool_call_id: str, + *, + agent: str | None = None, + child_id: str | None = None, + status: str = "completed", + is_error: bool = False, +) -> dict: + content = [{"type": "tool_result", "tool_use_id": tool_call_id, "content": "result", "is_error": is_error}] + event = _event(session, "user", timestamp, content, agent=agent, uuid=f"{tool_call_id}-result") + if child_id is not None: + event["toolUseResult"] = {"agentId": child_id, "status": status} + return event + + +def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path) -> None: + session = "session-root" + child = "agent-child" + grandchild = "agent-grandchild" + root = tmp_path / "projects" / "work" / f"{session}.jsonl" + subagents = root.parent / session / "subagents" + + _write( + root, + _event(session, "user", "2026-07-22T10:00:00Z", "solve", uuid="root-user"), + _assistant( + session, + "2026-07-22T10:00:01Z", + "msg-root", + {"type": "thinking", "thinking": "plan", "signature": "sig"}, + ), + _assistant( + session, + "2026-07-22T10:00:02Z", + "msg-root", + {"type": "tool_use", "id": "tool-fast", "name": "Read", "input": {"path": "a"}}, + {"type": "tool_use", "id": "tool-child", "name": "Agent", "input": {"prompt": "delegate"}}, + ), + _tool_result(session, "2026-07-22T10:00:03Z", "tool-fast"), + _tool_result( + session, + "2026-07-22T10:00:05Z", + "tool-child", + child_id=child, + ), + ) + _write( + subagents / f"{child}.jsonl", + _event(session, "user", "2026-07-22T10:00:02.100Z", "child task", agent=child, uuid="child-user"), + _assistant( + session, + "2026-07-22T10:00:03Z", + "msg-child", + {"type": "tool_use", "id": "tool-grandchild", "name": "Agent", "input": {}}, + agent=child, + ), + _tool_result( + session, + "2026-07-22T10:00:04Z", + "tool-grandchild", + agent=child, + child_id=grandchild, + ), + ) + _write( + subagents / f"{grandchild}.jsonl", + _assistant( + session, + "2026-07-22T10:00:03.100Z", + "msg-grandchild", + {"type": "text", "text": "done"}, + agent=grandchild, + ), + ) + + bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) + + invocations = _records(bundle, AgentInvocation) + assert [invocation.invocation_id for invocation in invocations] == [session, child, grandchild] + root_invocation, child_invocation, grandchild_invocation = invocations + assert child_invocation.parent_invocation_id == session + assert child_invocation.spawned_by_tool_call_id == "tool-child" + assert grandchild_invocation.parent_invocation_id == child + assert grandchild_invocation.spawned_by_tool_call_id == "tool-grandchild" + assert [reference.response_id for reference in root_invocation.model_calls] == ["msg-root"] + assert [reference.response_id for reference in child_invocation.model_calls] == ["msg-child"] + assert [reference.response_id for reference in grandchild_invocation.model_calls] == ["msg-grandchild"] + assert all(reference.model_ref == MODEL_REF for invocation in invocations for reference in invocation.model_calls) + assert [item.type for item in root_invocation.conversation] == [ + "message", + "reasoning", + "function_call", + "function_call", + "function_call_output", + "function_call_output", + ] + + timings = {tool.tool_call_id: tool for tool in _records(bundle, ToolCallObservation)} + assert timings["tool-fast"].duration_ms == pytest.approx(1000) + assert timings["tool-child"].duration_ms == pytest.approx(3000) + assert timings["tool-grandchild"].duration_ms == pytest.approx(1000) + assert all(tool.timing_source == "artifact" for tool in timings.values()) + assert bundle.gaps == [] + + +def test_extracts_explicit_compaction_markers(tmp_path: Path) -> None: + _write( + tmp_path / "projects" / "work" / "session.jsonl", + _assistant("root", "2026-07-22T09:59:59Z", "msg-before", {"type": "text", "text": "before"}), + _event( + "root", + "user", + "2026-07-22T10:00:00Z", + "summary", + message_extra={ + "isCompactSummary": True, + "compactMetadata": {"tokensBefore": 1000, "tokensAfter": 200, "trigger": "auto"}, + }, + ), + _event( + "root", + "system", + "2026-07-22T10:00:01Z", + "", + subtype="compact_boundary", + compact_metadata={"preTokens": 900, "postTokens": 180}, + ), + _assistant("root", "2026-07-22T10:00:02Z", "msg-after", {"type": "text", "text": "after"}), + ) + + bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) + + [compaction] = _records(bundle, ContextCompactionObservation) + assert compaction.trigger == "auto" + assert compaction.tokens_before == 1000 + assert compaction.tokens_after == 200 + assert compaction.summary == "summary" + assert compaction.outcome == "completed" + assert compaction.before_model_call.response_id == "msg-before" + assert compaction.after_model_call.response_id == "msg-after" + + +def test_malformed_and_incomplete_artifacts_produce_sanitized_gaps(tmp_path: Path) -> None: + sentinel = "redacted-payload-line" + _write( + tmp_path / "projects" / "work" / "root.jsonl", + f'{{"private":"{sentinel}"', + _assistant( + "root", + "bad-timestamp", + "msg-root", + {"type": "tool_use", "id": "pending", "name": "Bash", "input": {}}, + ), + _tool_result("root", "2026-07-22T10:00:03Z", "orphan"), + ) + _write( + tmp_path / "projects" / "work" / "subagents" / "agent-orphan.jsonl", + _assistant( + "root", + "2026-07-22T10:00:01Z", + "msg-orphan", + {"type": "text", "text": "answer"}, + agent="agent-orphan", + ), + ) + + bundle = extract_claude_code_observations(tmp_path) + codes = {gap.code for gap in bundle.gaps} + + assert { + "malformed_transcript_line", + "subagent_parent_unavailable", + "tool_result_missing", + "tool_start_timestamp_missing", + "tool_start_missing", + } <= codes + assert all(not invocation.model_calls for invocation in _records(bundle, AgentInvocation)) + assert sentinel not in bundle.model_dump_json() + + +def test_ignores_non_transcript_jsonl_and_reports_no_usable_transcript(tmp_path: Path) -> None: + _write( + tmp_path / "skills" / "fixture.jsonl", + _assistant("unrelated", "2026-07-22T10:00:00Z", "msg-unrelated", {"type": "text", "text": "x"}), + ) + (tmp_path / "projects").mkdir() + + bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) + + assert _records(bundle, AgentInvocation) == [] + assert "agent_transcript_unavailable" in {gap.code for gap in bundle.gaps} + + +def test_reports_missing_response_id_and_unsupported_content_blocks(tmp_path: Path) -> None: + _write( + tmp_path / "projects" / "work" / "session.jsonl", + _event( + "root", + "assistant", + "2026-07-22T10:00:00Z", + [{"type": "image", "source": "omitted"}], + uuid="assistant-event", + ), + _event( + "root", + "user", + "2026-07-22T10:00:01Z", + [{"type": "image", "source": "omitted"}], + uuid="user-event", + ), + ) + + bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) + codes = {gap.code for gap in bundle.gaps} + + assert "model_response_id_missing" in codes + assert "unsupported_assistant_content_block" in codes + assert "unsupported_user_content_block" in codes + assert _records(bundle, AgentInvocation)[0].model_calls == [] diff --git a/tests/unit_tests/test_anthropic_converter.py b/tests/unit_tests/test_anthropic_converter.py index 6fe997705a..ca9607205c 100644 --- a/tests/unit_tests/test_anthropic_converter.py +++ b/tests/unit_tests/test_anthropic_converter.py @@ -145,6 +145,7 @@ def test_assistant_tool_use_becomes_function_call(self) -> None: assert params.input[0].content == "calling" fc = params.input[1] assert fc.type == "function_call" + assert fc.id.startswith("fc_") assert fc.call_id == "toolu_1" assert fc.name == "lookup" assert json.loads(fc.arguments) == {"city": "Paris"} From 8d13d5cfab550e86e6a1a660d43e2e29992c2e03 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Fri, 24 Jul 2026 18:33:29 +0200 Subject: [PATCH 09/14] Report Claude and AnyTerminal observation gaps Signed-off-by: Michal Bien --- responses_api_agents/anyterminal_agent/app.py | 1 + responses_api_agents/anyterminal_agent/tests/test_app.py | 2 ++ responses_api_agents/claude_code_agent/app.py | 1 + responses_api_agents/claude_code_agent/observability.py | 2 ++ responses_api_agents/claude_code_agent/tests/test_app.py | 1 + .../claude_code_agent/tests/test_observability.py | 2 +- 6 files changed, 8 insertions(+), 1 deletion(-) diff --git a/responses_api_agents/anyterminal_agent/app.py b/responses_api_agents/anyterminal_agent/app.py index d6fe1c595c..74e8292fe4 100644 --- a/responses_api_agents/anyterminal_agent/app.py +++ b/responses_api_agents/anyterminal_agent/app.py @@ -927,6 +927,7 @@ async def run(self, body: AnyTerminalRunRequest) -> AnyTerminalVerifyResponse: sandbox_observations = _terminal_sandbox_observations(metrics) observations.records.extend(sandbox_observations) if sandbox_observations: + observations.gaps = [gap for gap in observations.gaps if gap.code != "no_sandbox_runtime"] sandbox_observation = sandbox_observations[0] link_tool_calls_to_sandbox(observations, sandbox_observation.sandbox_id) if sandbox_observation.cpu_time_s is None: diff --git a/responses_api_agents/anyterminal_agent/tests/test_app.py b/responses_api_agents/anyterminal_agent/tests/test_app.py index 0bb27f4e28..8fde330fa9 100644 --- a/responses_api_agents/anyterminal_agent/tests/test_app.py +++ b/responses_api_agents/anyterminal_agent/tests/test_app.py @@ -35,6 +35,7 @@ AgentInvocation, AgentObservationBundle, ModelCallRef, + ObservationGap, SandboxObservation, ToolCallObservation, ) @@ -722,6 +723,7 @@ async def test_appends_terminal_sandbox_metrics_to_agent_observations(self, tmp_ completed_at=4.0, ), ], + gaps=[ObservationGap(code="no_sandbox_runtime")], ) instance = _make_instance_config(tmp_path, observability_enabled=True) response = self._response( diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index f7c7aa723b..0a05ddcc8d 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -622,6 +622,7 @@ def collect(config_dir: Path) -> None: source="claude_code", gaps=[ObservationGap(code="agent_transcript_unavailable")], ) + observations.gaps.append(ObservationGap(code="no_sandbox_runtime")) return AgentEpisode(response=response, observations=observations) async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> ClaudeCodeAgentVerifyResponse: diff --git a/responses_api_agents/claude_code_agent/observability.py b/responses_api_agents/claude_code_agent/observability.py index 65110f17d2..7fa679dc49 100644 --- a/responses_api_agents/claude_code_agent/observability.py +++ b/responses_api_agents/claude_code_agent/observability.py @@ -464,6 +464,8 @@ def add_tool_gap(code: str) -> None: invocations_by_id: dict[str, AgentInvocation] = {} for invocation_id in all_invocation_ids: parent = parent_by_invocation.get(invocation_id) + if parent is None: + gaps.append(_gap("invocation_outcome_unavailable", invocation_id=invocation_id)) invocations_by_id[invocation_id] = AgentInvocation( invocation_id=invocation_id, parent_invocation_id=parent[0] if parent else None, 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 faaac204a8..6b4e5a032d 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -355,6 +355,7 @@ async def run_claude_code(*args, observation_collector=None, **kwargs): invocation = next(record for record in observations.records if isinstance(record, AgentInvocation)) assert invocation.invocation_id == "session-1" assert invocation.model_calls[0].response_id == "msg-1" + assert "no_sandbox_runtime" in {gap.code for gap in observations.gaps} assert agent.server_client.post.await_args_list[-1].kwargs["json"]["rollout_id"] == "1-2" diff --git a/responses_api_agents/claude_code_agent/tests/test_observability.py b/responses_api_agents/claude_code_agent/tests/test_observability.py index a7b704aee9..b5e61260db 100644 --- a/responses_api_agents/claude_code_agent/tests/test_observability.py +++ b/responses_api_agents/claude_code_agent/tests/test_observability.py @@ -172,7 +172,7 @@ def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path assert timings["tool-child"].duration_ms == pytest.approx(3000) assert timings["tool-grandchild"].duration_ms == pytest.approx(1000) assert all(tool.timing_source == "artifact" for tool in timings.values()) - assert bundle.gaps == [] + assert [(gap.code, gap.invocation_id) for gap in bundle.gaps] == [("invocation_outcome_unavailable", session)] def test_extracts_explicit_compaction_markers(tmp_path: Path) -> None: From 08d8e3aee25924a1b1bb7464919e04aed8dad0f0 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Mon, 27 Jul 2026 18:55:02 +0200 Subject: [PATCH 10/14] Focus Claude Code rollout observations Signed-off-by: Michal Bien --- nemo_gym/anthropic_converter.py | 2 +- responses_api_agents/anyterminal_agent/app.py | 158 +-------- .../anyterminal_agent/tests/test_app.py | 317 +----------------- responses_api_agents/claude_code_agent/app.py | 20 +- .../claude_code_agent/observability.py | 32 +- .../claude_code_agent/tests/test_app.py | 5 +- .../tests/test_observability.py | 23 +- tests/unit_tests/test_anthropic_converter.py | 1 - 8 files changed, 57 insertions(+), 501 deletions(-) diff --git a/nemo_gym/anthropic_converter.py b/nemo_gym/anthropic_converter.py index 1a30a18be2..ad3e7c6d82 100644 --- a/nemo_gym/anthropic_converter.py +++ b/nemo_gym/anthropic_converter.py @@ -372,7 +372,7 @@ def flush_message() -> None: arguments=json.dumps(block.get("input", {})), call_id=block["id"], name=block["name"], - id=f"fc_{uuid4().hex}", + id=block["id"], status="completed", type="function_call", ) diff --git a/responses_api_agents/anyterminal_agent/app.py b/responses_api_agents/anyterminal_agent/app.py index 74e8292fe4..716c61dfa8 100644 --- a/responses_api_agents/anyterminal_agent/app.py +++ b/responses_api_agents/anyterminal_agent/app.py @@ -35,12 +35,6 @@ from nemo_gym.config_types import ModelServerRef from nemo_gym.global_config import get_first_server_config_dict from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming -from nemo_gym.rollout_observability import ( - AgentObservationBundle, - ObservationGap, - SandboxObservation, - link_tool_calls_to_sandbox, -) from nemo_gym.sandbox import AsyncSandbox, SandboxSpec from nemo_gym.sandbox.providers.apptainer import ApptainerProvider from nemo_gym.sandbox.providers.docker import DockerCreateConfig, DockerProvider @@ -115,12 +109,6 @@ class TerminalBenchMetrics(BaseModel): agent_run_time: Optional[float] = None eval_run_time: Optional[float] = None total_run_time: Optional[float] = None - sandbox_provider: Optional[str] = None - sandbox_id: Optional[str] = None - sandbox_wall_time_s: Optional[float] = None - sandbox_cpu_time_s: Optional[float] = None - sandbox_peak_memory_mib: Optional[float] = None - sandbox_resource_usage_source: Optional[str] = None def update_metrics(metrics_fpath: Path, update_dict: Dict[str, Any]) -> None: @@ -137,9 +125,7 @@ def redact(value: Any) -> Any: return { key: ( "***" - if any(secret in key.lower() for secret in ("api_key", "secret", "password")) - or key.lower() == "token" - or key.lower().endswith("_token") + if any(secret in key.lower() for secret in ("api_key", "secret", "password", "token")) else redact(item) ) for key, item in value.items() @@ -153,49 +139,6 @@ def redact(value: Any) -> Any: return json.dumps(redact(d), indent=indent) -def _load_agent_observations(path: Path) -> tuple[str, str]: - """Return canonical observation JSON and an error code without raising.""" - try: - raw = path.read_text() - except FileNotFoundError: - return "", "" - except OSError: - return "", "observation_read_failed" - - try: - bundle = AgentObservationBundle.model_validate_json(raw) - except Exception: - return "", "observation_parse_failed" - return bundle.model_dump_json(), "" - - -def _observation_gap(source: str, code: str) -> AgentObservationBundle: - return AgentObservationBundle( - source=source, - gaps=[ObservationGap(code=code)], - ) - - -def _terminal_sandbox_observations(metrics: TerminalBenchMetrics) -> list[SandboxObservation]: - if metrics.sandbox_id is None: - return [] - outcome = "sandbox_error" if metrics.sandbox_failed else "timeout" - if not metrics.sandbox_failed and not metrics.agent_timed_out and not metrics.container_timed_out: - outcome = "completed" - return [ - SandboxObservation( - role="environment", - provider=metrics.sandbox_provider, - sandbox_id=metrics.sandbox_id, - outcome=outcome, - wall_time_s=metrics.sandbox_wall_time_s, - cpu_time_s=metrics.sandbox_cpu_time_s, - peak_memory_mib=metrics.sandbox_peak_memory_mib, - resource_usage_source=metrics.sandbox_resource_usage_source, - ) - ] - - ### Agent runner template # Injected into the task container; imports any agent class and calls responses(). @@ -214,8 +157,6 @@ def _terminal_sandbox_observations(metrics: TerminalBenchMetrics) -> list[Sandbo INSTRUCTION = Path("/trajectories_mount/instruction.txt").read_text() AGENT_KWARGS = json.loads(os.environ.get("NGTB_AGENT_KWARGS", "{{}}")) SAMPLING = json.loads(os.environ.get("NGTB_SAMPLING", "{{}}")) -OBSERVABILITY = os.environ.get("NGTB_OBSERVABILITY") == "1" -MODEL_SERVER_REF = json.loads(os.environ.get("NGTB_MODEL_SERVER_REF", "null")) from nemo_gym.openai_utils import NeMoGymResponseCreateParamsNonStreaming, NeMoGymEasyInputMessage from nemo_gym.config_types import ModelServerRef, ResourcesServerRef @@ -227,7 +168,7 @@ def _terminal_sandbox_observations(metrics: TerminalBenchMetrics) -> list[Sandbo _cfg_sampling = {{k: v for k, v in SAMPLING.items() if k in {agent_cfg_class}.model_fields}} -_model_server = ModelServerRef.model_validate(MODEL_SERVER_REF) if MODEL_URL and MODEL_SERVER_REF else None +_model_server = ModelServerRef(name="policy_model", type="responses_api_models") if MODEL_URL else None config = {agent_cfg_class}( host="0.0.0.0", port=0, @@ -253,20 +194,8 @@ def _terminal_sandbox_observations(metrics: TerminalBenchMetrics) -> list[Sandbo model=MODEL_NAME, **SAMPLING, ) -observed = getattr(agent, "responses_with_observations", None) -if OBSERVABILITY and callable(observed): - episode = asyncio.run(observed(request=None, body=body)) - response = episode.response -else: - response = asyncio.run(agent.responses(request=None, body=body)) +response = asyncio.run(agent.responses(request=None, body=body)) Path("/trajectories_mount/response.json").write_text(response.model_dump_json()) -if OBSERVABILITY and callable(observed): - try: - Path("/trajectories_mount/agent_observations.json").write_text( - episode.observations.model_dump_json() - ) - except Exception as exc: - print(f"failed to persist agent observations: {{type(exc).__name__}}", file=sys.stderr, flush=True) print(f"agent finished: {{len(response.output)}} output items", flush=True) """ @@ -383,7 +312,6 @@ class AnyTerminalInstanceConfig(AnyTerminalAgentConfig, AnyTerminalServerConfig) metrics_fpath: Path container: str ray_queue_timestamp: float - observability_enabled: bool = Field(default=False, exclude_if=lambda value: not value) agent_command_str: Optional[str] = None @property @@ -397,10 +325,6 @@ def instance_id(self) -> str: class AnyTerminalVerifyResponse(TerminalBenchMetrics, BaseVerifyResponse): instance_config: Dict[str, Any] - ng_agent_observations: AgentObservationBundle | None = Field( - default=None, - exclude_if=lambda value: value is None, - ) ### Sandbox provider selection @@ -532,10 +456,6 @@ async def _stage_remote_tests(self, sandbox: AsyncSandbox, cfg: AnyTerminalInsta async def _collect_remote_outputs(self, sandbox: AsyncSandbox, cfg: AnyTerminalInstanceConfig) -> None: for remote, local in ( ("/trajectories_mount/response.json", cfg.persistent_dir / "response.json"), - ( - "/trajectories_mount/agent_observations.json", - cfg.persistent_dir / "agent_observations.json", - ), ("/logs/verifier/reward.txt", cfg.verifier_dir / "reward.txt"), ("/logs/verifier/test-stdout.txt", cfg.verifier_dir / "test-stdout.txt"), ): @@ -557,10 +477,6 @@ def _agent_env(self, cfg: AnyTerminalInstanceConfig) -> Dict[str, str]: } if cfg.model_server_url: env["NGTB_MODEL_URL"] = cfg.model_server_url - if cfg.model_server is not None: - env["NGTB_MODEL_SERVER_REF"] = cfg.model_server.model_dump_json() - if cfg.observability_enabled: - env["NGTB_OBSERVABILITY"] = "1" return env async def _run_agent(self, sandbox: AsyncSandbox, cfg: AnyTerminalInstanceConfig) -> tuple[float, bool]: @@ -614,12 +530,8 @@ async def process_single_datapoint(self) -> bool: agent_timed_out = container_timed_out = False sandbox_failed = False agent_run_time = eval_run_time = None - sandbox_id = None - sandbox_usage = None - sandbox_provider = next(iter(cfg.sandbox_provider), None) try: await sandbox.start() - sandbox_id = getattr(sandbox, "sandbox_id", None) if not self._uses_bind_mounts(cfg): await self._stage_remote_runtime(sandbox, cfg) agent_run_time, agent_timed_out = await self._run_agent(sandbox, cfg) @@ -633,11 +545,6 @@ async def process_single_datapoint(self) -> bool: sandbox_failed = True print(f"[{cfg.task_name}] sandbox run failed: {e}", flush=True) finally: - try: - if sandbox_id is not None: - sandbox_usage = await sandbox.resource_usage() - except Exception as e: - print(f"[{cfg.task_name}] sandbox resource usage unavailable: {e}", flush=True) try: await sandbox.stop() except Exception as e: @@ -665,12 +572,6 @@ async def process_single_datapoint(self) -> bool: agent_run_time=agent_run_time, eval_run_time=eval_run_time, total_run_time=total_run_time, - sandbox_provider=sandbox_provider, - sandbox_id=sandbox_id, - sandbox_wall_time_s=None if sandbox_usage is None else sandbox_usage.wall_time_s, - sandbox_cpu_time_s=None if sandbox_usage is None else sandbox_usage.cpu_time_s, - sandbox_peak_memory_mib=None if sandbox_usage is None else sandbox_usage.peak_memory_mib, - sandbox_resource_usage_source=None if sandbox_usage is None else sandbox_usage.source, ) update_metrics(cfg.metrics_fpath, metrics.model_dump()) return resolved @@ -812,7 +713,6 @@ def _setup_params( self.config.container_formatter, task_name, problem_info.get("docker_image", "ubuntu:22.04") ), ray_queue_timestamp=time.time(), - observability_enabled=rollout_id is not None, ) params.metrics_fpath.write_text("{}") @@ -863,20 +763,6 @@ async def _inner_responses(self, params: AnyTerminalInstanceConfig) -> NeMoGymRe except (json.JSONDecodeError, ValueError) as e: print(f"[{params.task_name}] response.json unreadable ({e}), treating as empty response", flush=True) - metadata = { - "input": json.dumps(params.body.model_dump(mode="json").get("input") or []), - "metrics": params.metrics_fpath.read_text(), - "instance_config": _safe_config_json(params), - } - if params.observability_enabled: - observations_json, observations_error = _load_agent_observations( - params.persistent_dir / "agent_observations.json" - ) - metadata.update( - agent_observations=observations_json, - agent_observations_error=observations_error, - ) - return NeMoGymResponse( id=f"anyterminal-{params.instance_id}", created_at=int(time.time()), @@ -886,7 +772,11 @@ async def _inner_responses(self, params: AnyTerminalInstanceConfig) -> NeMoGymRe parallel_tool_calls=params.body.parallel_tool_calls, tool_choice=params.body.tool_choice, tools=tools, - metadata=metadata, + metadata={ + "input": json.dumps(params.body.model_dump(mode="json").get("input") or []), + "metrics": params.metrics_fpath.read_text(), + "instance_config": _safe_config_json(params), + }, ) async def run(self, body: AnyTerminalRunRequest) -> AnyTerminalVerifyResponse: @@ -897,9 +787,8 @@ async def run(self, body: AnyTerminalRunRequest) -> AnyTerminalVerifyResponse: meta, response.metadata = response.metadata, None metrics = TerminalBenchMetrics.model_validate_json(meta["metrics"]) - instance_config = AnyTerminalInstanceConfig.model_validate_json(meta["instance_config"]) - result = dict( + return AnyTerminalVerifyResponse( responses_create_params=body.responses_create_params.model_dump() | { "input": json.loads(meta["input"]), @@ -909,35 +798,8 @@ async def run(self, body: AnyTerminalRunRequest) -> AnyTerminalVerifyResponse: response=response, reward=1.0 if metrics.resolved else 0.0, **metrics.model_dump(), - instance_config=instance_config.model_dump(), + instance_config=AnyTerminalInstanceConfig.model_validate_json(meta["instance_config"]).model_dump(), ) - observations_json = meta.get("agent_observations") - source = self.config.agent_server_module.split(".")[-2].removesuffix("_agent") - observations: Optional[AgentObservationBundle] = None - if observations_json: - try: - observations = AgentObservationBundle.model_validate_json(observations_json) - except Exception: - observations = _observation_gap(source, "observation_parse_failed") - elif meta.get("agent_observations_error"): - observations = _observation_gap(source, meta["agent_observations_error"]) - elif instance_config.observability_enabled: - observations = _observation_gap(source, "agent_observations_unavailable") - if observations is not None: - sandbox_observations = _terminal_sandbox_observations(metrics) - observations.records.extend(sandbox_observations) - if sandbox_observations: - observations.gaps = [gap for gap in observations.gaps if gap.code != "no_sandbox_runtime"] - sandbox_observation = sandbox_observations[0] - link_tool_calls_to_sandbox(observations, sandbox_observation.sandbox_id) - if sandbox_observation.cpu_time_s is None: - observations.gaps.append(ObservationGap(code="sandbox_cpu_time_unavailable")) - if sandbox_observation.peak_memory_mib is None: - observations.gaps.append(ObservationGap(code="sandbox_memory_usage_unavailable")) - else: - observations.gaps.append(ObservationGap(code="sandbox_observation_unavailable")) - result["ng_agent_observations"] = observations.model_dump(mode="json") - return AnyTerminalVerifyResponse.model_validate(result) if __name__ == "__main__": diff --git a/responses_api_agents/anyterminal_agent/tests/test_app.py b/responses_api_agents/anyterminal_agent/tests/test_app.py index 8fde330fa9..1bdeb5a06f 100644 --- a/responses_api_agents/anyterminal_agent/tests/test_app.py +++ b/responses_api_agents/anyterminal_agent/tests/test_app.py @@ -20,25 +20,15 @@ constructing the agent. """ -import asyncio import json from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch +from unittest.mock import AsyncMock, PropertyMock, patch import pytest from nemo_gym import PARENT_DIR -from nemo_gym.config_types import ModelServerRef -from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming -from nemo_gym.rollout_observability import ( - AgentInvocation, - AgentObservationBundle, - ModelCallRef, - ObservationGap, - SandboxObservation, - ToolCallObservation, -) +from nemo_gym.openai_utils import NeMoGymResponseCreateParamsNonStreaming from nemo_gym.sandbox.providers.apptainer import ApptainerProvider from nemo_gym.sandbox.providers.apptainer import provider as apptainer_provider from nemo_gym.sandbox.providers.docker import DockerProvider @@ -47,17 +37,13 @@ AnyTerminalAgent, AnyTerminalAgentConfig, AnyTerminalInstanceConfig, - AnyTerminalRunRequest, GymAgentHarnessProcessor, RunTerminalAgent, - TerminalBenchMetrics, _build_provider, _format_container, _instruction_from_input, - _load_agent_observations, _read_task_meta, _safe_config_json, - _terminal_sandbox_observations, update_metrics, ) @@ -77,31 +63,6 @@ def _config(**overrides) -> AnyTerminalAgentConfig: return AnyTerminalAgentConfig(**base) -def test_terminal_sandbox_observations_only_map_available_metrics() -> None: - observations = _terminal_sandbox_observations( - TerminalBenchMetrics( - agent_timed_out=True, - container_timed_out=True, - agent_run_time=30.0, - eval_run_time=4.0, - total_run_time=35.0, - sandbox_provider="docker", - sandbox_id="sandbox-1", - sandbox_wall_time_s=35.0, - sandbox_cpu_time_s=6.0, - sandbox_peak_memory_mib=512.0, - sandbox_resource_usage_source="docker_container_cgroup_v2", - ) - ) - - assert [(item.role, item.provider, item.sandbox_id, item.outcome, item.wall_time_s) for item in observations] == [ - ("environment", "docker", "sandbox-1", "timeout", 35.0), - ] - assert observations[0].cpu_time_s == 6.0 - assert observations[0].peak_memory_mib == 512.0 - assert observations[0].resource_usage_source == "docker_container_cgroup_v2" - - class TestRunnerTemplate: def _render(self) -> str: return _RUNNER_TEMPLATE.format( @@ -131,21 +92,6 @@ def test_sampling_is_forwarded(self) -> None: assert "**SAMPLING," in rendered assert "HermesAgentConfig.model_fields" in rendered - def test_observed_episode_is_written_when_enabled(self) -> None: - rendered = self._render() - assert "responses_with_observations" in rendered - assert "/trajectories_mount/agent_observations.json" in rendered - assert 'os.environ.get("NGTB_OBSERVABILITY") == "1"' in rendered - - def test_uses_configured_model_ref_and_writes_response_first(self) -> None: - rendered = self._render() - assert "ModelServerRef.model_validate(MODEL_SERVER_REF)" in rendered - assert 'ModelServerRef(name="policy_model"' not in rendered - assert rendered.index('Path("/trajectories_mount/response.json").write_text') < rendered.index( - 'Path("/trajectories_mount/agent_observations.json").write_text' - ) - assert "failed to persist agent observations" in rendered - class TestAgentKey: def test_key_from_module(self) -> None: @@ -372,39 +318,6 @@ def test_indent_produces_multiline(self, tmp_path: Path) -> None: cfg = _make_instance_config(tmp_path) assert "\n" in _safe_config_json(cfg, indent=2) - def test_disabled_observability_does_not_change_serialized_config(self, tmp_path: Path) -> None: - serialized = _safe_config_json(_make_instance_config(tmp_path)) - - assert "observability_enabled" not in json.loads(serialized) - assert AnyTerminalInstanceConfig.model_validate_json(serialized).observability_enabled is False - - -class TestLoadAgentObservations: - def test_validates_and_canonicalizes_sidecar(self, tmp_path: Path) -> None: - path = tmp_path / "observations.json" - bundle = AgentObservationBundle(source="claude_code", records=[AgentInvocation(invocation_id="root")]) - path.write_text(bundle.model_dump_json()) - - raw, error = _load_agent_observations(path) - - assert error == "" - assert AgentObservationBundle.model_validate_json(raw) == bundle - - @pytest.mark.parametrize( - ("contents", "expected_error"), - [("not json", "observation_parse_failed"), ("{}", "observation_parse_failed")], - ) - def test_rejects_invalid_sidecar(self, tmp_path: Path, contents: str, expected_error: str) -> None: - path = tmp_path / "observations.json" - path.write_text(contents) - assert _load_agent_observations(path) == ("", expected_error) - - def test_read_failure_is_nonfatal(self, tmp_path: Path) -> None: - path = tmp_path / "observations.json" - path.write_text("{}") - with patch.object(Path, "read_text", side_effect=PermissionError): - assert _load_agent_observations(path) == ("", "observation_read_failed") - # ── AnyTerminalInstanceConfig properties ────────────────────────────────────────── @@ -587,33 +500,6 @@ def test_no_model_url_when_empty(self, tmp_path: Path) -> None: env = RunTerminalAgent(config=cfg)._agent_env(cfg) assert "NGTB_MODEL_URL" not in env - def test_model_ref_included_when_set(self, tmp_path: Path) -> None: - cfg = _make_instance_config( - tmp_path, - model_server=ModelServerRef(type="responses_api_models", name="custom_policy"), - model_server_url="http://model:8000/ng-rollout/2-1", - ) - env = RunTerminalAgent(config=cfg)._agent_env(cfg) - assert env["NGTB_MODEL_URL"] == "http://model:8000/ng-rollout/2-1" - assert json.loads(env["NGTB_MODEL_SERVER_REF"]) == { - "type": "responses_api_models", - "name": "custom_policy", - } - - def test_no_model_ref_when_unconfigured(self, tmp_path: Path) -> None: - cfg = _make_instance_config(tmp_path, model_server=None, model_server_url="") - env = RunTerminalAgent(config=cfg)._agent_env(cfg) - assert "NGTB_MODEL_SERVER_REF" not in env - - def test_observability_env_is_opt_in(self, tmp_path: Path) -> None: - disabled_cfg = _make_instance_config(tmp_path / "off") - disabled = RunTerminalAgent(config=disabled_cfg)._agent_env(disabled_cfg) - enabled_cfg = _make_instance_config(tmp_path / "on", observability_enabled=True) - enabled = RunTerminalAgent(config=enabled_cfg)._agent_env(enabled_cfg) - - assert "NGTB_OBSERVABILITY" not in disabled - assert enabled["NGTB_OBSERVABILITY"] == "1" - def test_sampling_and_kwargs_forwarded(self, tmp_path: Path) -> None: cfg = _make_instance_config( tmp_path, @@ -625,190 +511,6 @@ def test_sampling_and_kwargs_forwarded(self, tmp_path: Path) -> None: assert json.loads(env["NGTB_AGENT_KWARGS"]) == {"model": "my-model"} -class TestObservationRoundTrip: - @staticmethod - def _agent() -> AnyTerminalAgent: - server_client = MagicMock() - server_client.global_config_dict = {"observability_enabled": True} - with patch.object(AnyTerminalAgent, "model_post_init"): - agent = AnyTerminalAgent.model_construct( - config=_config(model_server={"type": "responses_api_models", "name": "custom_policy"}), - server_client=server_client, - ) - agent._sem = asyncio.Semaphore(1) - return agent - - @staticmethod - def _response( - instance: AnyTerminalInstanceConfig, - observations: str, - error: str = "", - metrics: dict | None = None, - ) -> NeMoGymResponse: - return NeMoGymResponse( - id="anyterminal-test", - created_at=1, - model="test-model", - object="response", - output=[], - parallel_tool_calls=True, - tool_choice="auto", - tools=[], - metadata={ - "input": "[]", - "metrics": json.dumps(metrics or {"resolved": True}), - "instance_config": _safe_config_json(instance), - "agent_observations": observations, - "agent_observations_error": error, - }, - ) - - @staticmethod - def _body() -> AnyTerminalRunRequest: - return AnyTerminalRunRequest.model_validate( - { - "responses_create_params": {"input": "solve"}, - "_ng_task_index": 1, - "_ng_rollout_index": 2, - } - ) - - @pytest.mark.asyncio - async def test_preserves_nondefault_model_ref(self, tmp_path: Path) -> None: - model_ref = ModelServerRef(type="responses_api_models", name="custom_policy") - bundle = AgentObservationBundle( - source="claude_code", - records=[ - AgentInvocation( - invocation_id="root", - model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")], - ) - ], - ) - instance = _make_instance_config( - tmp_path, - model_server=model_ref, - observability_enabled=True, - ) - agent = self._agent() - - with patch.object( - agent, "_responses", AsyncMock(return_value=self._response(instance, bundle.model_dump_json())) - ): - result = await agent.run(self._body()) - - emitted = result.ng_agent_observations - assert result.reward == 1.0 - assert emitted is not None - invocation = next(record for record in emitted.records if isinstance(record, AgentInvocation)) - assert invocation.model_calls[0].model_ref == model_ref - assert invocation.model_calls[0].response_id == "resp-1" - - @pytest.mark.asyncio - async def test_appends_terminal_sandbox_metrics_to_agent_observations(self, tmp_path: Path) -> None: - bundle = AgentObservationBundle( - source="hermes", - records=[ - AgentInvocation(invocation_id="root"), - ToolCallObservation( - invocation_id="root", - tool_call_id="call-1", - started_at=1.0, - completed_at=3.0, - ), - ToolCallObservation( - invocation_id="root", - tool_call_id="call-2", - started_at=2.0, - completed_at=4.0, - ), - ], - gaps=[ObservationGap(code="no_sandbox_runtime")], - ) - instance = _make_instance_config(tmp_path, observability_enabled=True) - response = self._response( - instance, - bundle.model_dump_json(), - metrics={ - "resolved": False, - "agent_run_time": 12.0, - "total_run_time": 15.0, - "sandbox_provider": "docker", - "sandbox_id": "sandbox-1", - "sandbox_wall_time_s": 15.0, - }, - ) - agent = self._agent() - - with patch.object(agent, "_responses", AsyncMock(return_value=response)): - result = await agent.run(self._body()) - - assert result.ng_agent_observations is not None - sandboxes = [ - record for record in result.ng_agent_observations.records if isinstance(record, SandboxObservation) - ] - tools = [record for record in result.ng_agent_observations.records if isinstance(record, ToolCallObservation)] - assert [(item.role, item.wall_time_s) for item in sandboxes] == [ - ("environment", 15.0), - ] - assert sandboxes[0].provider == "docker" - assert sandboxes[0].sandbox_id == "sandbox-1" - assert {tool.sandbox_id for tool in tools} == {"sandbox-1"} - assert [gap.code for gap in result.ng_agent_observations.gaps] == [ - "sandbox_cpu_time_unavailable", - "sandbox_memory_usage_unavailable", - ] - - @pytest.mark.parametrize( - ("observations", "error", "gap_code"), - [ - ("[]", "", "observation_parse_failed"), - ("", "observation_read_failed", "observation_read_failed"), - ], - ) - @pytest.mark.asyncio - async def test_sidecar_failure_keeps_reward_and_emits_typed_gap( - self, - tmp_path: Path, - observations: str, - error: str, - gap_code: str, - ) -> None: - instance = _make_instance_config(tmp_path, observability_enabled=True) - agent = self._agent() - - with patch.object( - agent, - "_responses", - AsyncMock(return_value=self._response(instance, observations, error)), - ): - result = await agent.run(self._body()) - - assert result.reward == 1.0 - assert result.ng_agent_observations is not None - assert result.ng_agent_observations.records == [] - assert [gap.code for gap in result.ng_agent_observations.gaps] == [ - gap_code, - "sandbox_observation_unavailable", - ] - - @pytest.mark.asyncio - async def test_disabled_observability_does_not_change_response(self, tmp_path: Path) -> None: - instance = _make_instance_config(tmp_path) - agent = self._agent() - - with patch.object(agent, "_responses", AsyncMock(return_value=self._response(instance, ""))): - result = await agent.run(self._body()) - - assert result.reward == 1.0 - assert result.ng_agent_observations is None - assert "ng_agent_observations" not in result.model_dump(mode="json") - assert "observability_enabled" not in result.instance_config - - -# ── RunTerminalAgent.process_single_datapoint ──────────────────────────────────── - - class TestProcessSingleDatapoint: @pytest.fixture(autouse=True) def _no_real_provider(self): @@ -822,17 +524,8 @@ async def test_resolved_when_reward_positive(self, tmp_path: Path) -> None: (cfg.verifier_dir / "reward.txt").write_text("1.0") sandbox = SimpleNamespace( - sandbox_id="sandbox-1", start=AsyncMock(), exec=AsyncMock(return_value=_sandbox_result()), - resource_usage=AsyncMock( - return_value=SimpleNamespace( - wall_time_s=12.0, - cpu_time_s=3.0, - peak_memory_mib=256.0, - source="docker_container_cgroup_v2", - ) - ), stop=AsyncMock(), ) with patch("responses_api_agents.anyterminal_agent.app.AsyncSandbox", return_value=sandbox): @@ -840,11 +533,7 @@ async def test_resolved_when_reward_positive(self, tmp_path: Path) -> None: result = await RunTerminalAgent(config=cfg).process_single_datapoint() assert result is True - metrics = json.loads(cfg.metrics_fpath.read_text()) - assert metrics["resolved"] is True - assert metrics["sandbox_id"] == "sandbox-1" - assert metrics["sandbox_cpu_time_s"] == 3.0 - assert metrics["sandbox_peak_memory_mib"] == 256.0 + assert json.loads(cfg.metrics_fpath.read_text())["resolved"] is True sandbox.stop.assert_awaited_once() async def test_unresolved_when_no_reward_file(self, tmp_path: Path) -> None: diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index 0a05ddcc8d..a3af7f7ee9 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -449,10 +449,9 @@ async def _run_claude_code( if claude_config_dir is not None: try: if observation_collector is not None: - try: - await asyncio.to_thread(observation_collector, claude_config_dir) - except Exception: - LOG.exception("failed to collect Claude Code observations") + await asyncio.to_thread(observation_collector, claude_config_dir) + except Exception: + LOG.exception("failed to collect Claude Code observations") finally: shutil.rmtree(claude_config_dir, ignore_errors=True) @@ -583,12 +582,10 @@ async def responses( request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body(), ) -> NeMoGymResponse: - rollout_id = request.path_params.get("rollout_id") if request is not None else None - return await self._create_response(body, rollout_id=rollout_id) + return await self._create_response(body, rollout_id=request.path_params.get("rollout_id")) - async def responses_with_observations( + async def _create_episode( self, - request: Optional[Request], body: NeMoGymResponseCreateParamsNonStreaming, *, mcp_config: Optional[str] = None, @@ -649,8 +646,7 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude with tempfile.TemporaryDirectory(prefix="nemo_gym_claude_mcp_") as mcp_config_dir: mcp_config = self._write_rollout_mcp_config(seed_resp_json, Path(mcp_config_dir)) if rollout_id is not None: - episode = await self.responses_with_observations( - request, + episode = await self._create_episode( body.responses_create_params, mcp_config=mcp_config, skills_path=skills_path, @@ -669,9 +665,7 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude verify_resp = await self.server_client.post( server_name=self.config.resources_server.name, url_path="/verify", - json=body.model_dump() - | {"response": agent_resp_json} - | ({"rollout_id": rollout_id} if rollout_id is not None else {}), + json=body.model_dump() | {"response": agent_resp_json}, cookies=cookies, ) await raise_for_status(verify_resp) diff --git a/responses_api_agents/claude_code_agent/observability.py b/responses_api_agents/claude_code_agent/observability.py index 7fa679dc49..8d77c84184 100644 --- a/responses_api_agents/claude_code_agent/observability.py +++ b/responses_api_agents/claude_code_agent/observability.py @@ -70,7 +70,7 @@ def _status(block: dict[str, Any], result: Any) -> str: if result.get("interrupted") is True: return "incomplete" value = result.get("status") - if value in {"completed", "failed", "timeout", "incomplete"}: + if value in {"completed", "failed", "timeout", "cancelled", "incomplete"}: return value # A tool_result block is an explicit terminal observation even when Claude Code # does not attach a separate status object. @@ -436,8 +436,12 @@ def add_tool_gap(code: str) -> None: if call_finishes and completed_at is None: add_tool_gap("tool_result_timestamp_missing") duration_ms = None - if started_at is not None and completed_at is not None and completed_at >= started_at: - duration_ms = (completed_at - started_at) * 1000 + if started_at is not None and completed_at is not None: + if completed_at >= started_at: + duration_ms = (completed_at - started_at) * 1000 + else: + add_tool_gap("tool_timing_invalid") + completed_at = None tool_calls.append( ToolCallObservation( invocation_id=invocation_id, @@ -470,24 +474,18 @@ def add_tool_gap(code: str) -> None: invocation_id=invocation_id, parent_invocation_id=parent[0] if parent else None, spawned_by_tool_call_id=parent[1] if parent else None, - status=parent[2] if parent else "unknown", + status=( + "incomplete" + if parent and parent[2] in {"timeout", "cancelled"} + else parent[2] + if parent + else "unknown" + ), model_calls=model_calls[invocation_id], conversation=conversations[invocation_id], ) - def order_key(invocation_id: str) -> tuple[tuple[float, str], ...]: - path: list[tuple[float, str]] = [] - seen: set[str] = set() - while invocation_id not in seen: - seen.add(invocation_id) - path.append((first_seen.get(invocation_id, math.inf), invocation_id)) - parent = parent_by_invocation.get(invocation_id) - if parent is None: - break - invocation_id = parent[0] - return tuple(reversed(path)) - - ordered_ids = sorted(all_invocation_ids, key=order_key) + ordered_ids = sorted(all_invocation_ids, key=lambda invocation_id: (first_seen[invocation_id], invocation_id)) return AgentObservationBundle( source=SOURCE, 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 6b4e5a032d..efe0922a9f 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -356,7 +356,6 @@ async def run_claude_code(*args, observation_collector=None, **kwargs): assert invocation.invocation_id == "session-1" assert invocation.model_calls[0].response_id == "msg-1" assert "no_sandbox_runtime" in {gap.code for gap in observations.gaps} - assert agent.server_client.post.await_args_list[-1].kwargs["json"]["rollout_id"] == "1-2" class TestRunClaudeCode: @@ -543,7 +542,7 @@ def test_writes_rollout_mcp_config_with_session_header(self, tmp_path: Path) -> "mcp": { "server_name": "example_mcp_weather", "url_path": "/mcp", - "headers": {"X-NeMo-Gym-Session-Token": "secret-token"}, + "headers": {"X-NeMo-Gym-Session-Token": "mcp-session-value"}, } }, tmp_path, @@ -554,7 +553,7 @@ def test_writes_rollout_mcp_config_with_session_header(self, tmp_path: Path) -> server = config["mcpServers"]["example_mcp_weather"] assert server["type"] == "http" assert server["url"] == "http://127.0.0.1:8123/mcp" - assert server["headers"]["X-NeMo-Gym-Session-Token"] == "secret-token" + assert server["headers"]["X-NeMo-Gym-Session-Token"] == "mcp-session-value" def test_merges_static_mcp_config_when_metadata_present(self, tmp_path: Path) -> None: static_config = tmp_path / "static_mcp.json" diff --git a/responses_api_agents/claude_code_agent/tests/test_observability.py b/responses_api_agents/claude_code_agent/tests/test_observability.py index b5e61260db..11f343a6fe 100644 --- a/responses_api_agents/claude_code_agent/tests/test_observability.py +++ b/responses_api_agents/claude_code_agent/tests/test_observability.py @@ -132,6 +132,7 @@ def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path "tool-grandchild", agent=child, child_id=grandchild, + status="timeout", ), ) _write( @@ -147,17 +148,22 @@ def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) - invocations = _records(bundle, AgentInvocation) - assert [invocation.invocation_id for invocation in invocations] == [session, child, grandchild] - root_invocation, child_invocation, grandchild_invocation = invocations + invocations = {invocation.invocation_id: invocation for invocation in _records(bundle, AgentInvocation)} + assert set(invocations) == {session, child, grandchild} + root_invocation = invocations[session] + child_invocation = invocations[child] + grandchild_invocation = invocations[grandchild] assert child_invocation.parent_invocation_id == session assert child_invocation.spawned_by_tool_call_id == "tool-child" assert grandchild_invocation.parent_invocation_id == child assert grandchild_invocation.spawned_by_tool_call_id == "tool-grandchild" + assert grandchild_invocation.status == "incomplete" assert [reference.response_id for reference in root_invocation.model_calls] == ["msg-root"] assert [reference.response_id for reference in child_invocation.model_calls] == ["msg-child"] assert [reference.response_id for reference in grandchild_invocation.model_calls] == ["msg-grandchild"] - assert all(reference.model_ref == MODEL_REF for invocation in invocations for reference in invocation.model_calls) + assert all( + reference.model_ref == MODEL_REF for invocation in invocations.values() for reference in invocation.model_calls + ) assert [item.type for item in root_invocation.conversation] == [ "message", "reasoning", @@ -171,6 +177,7 @@ def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path assert timings["tool-fast"].duration_ms == pytest.approx(1000) assert timings["tool-child"].duration_ms == pytest.approx(3000) assert timings["tool-grandchild"].duration_ms == pytest.approx(1000) + assert timings["tool-grandchild"].status == "timeout" assert all(tool.timing_source == "artifact" for tool in timings.values()) assert [(gap.code, gap.invocation_id) for gap in bundle.gaps] == [("invocation_outcome_unavailable", session)] @@ -223,6 +230,13 @@ def test_malformed_and_incomplete_artifacts_produce_sanitized_gaps(tmp_path: Pat "msg-root", {"type": "tool_use", "id": "pending", "name": "Bash", "input": {}}, ), + _assistant( + "root", + "2026-07-22T10:00:05Z", + "msg-later", + {"type": "tool_use", "id": "reversed", "name": "Bash", "input": {}}, + ), + _tool_result("root", "2026-07-22T09:59:59Z", "reversed"), _tool_result("root", "2026-07-22T10:00:03Z", "orphan"), ) _write( @@ -245,6 +259,7 @@ def test_malformed_and_incomplete_artifacts_produce_sanitized_gaps(tmp_path: Pat "tool_result_missing", "tool_start_timestamp_missing", "tool_start_missing", + "tool_timing_invalid", } <= codes assert all(not invocation.model_calls for invocation in _records(bundle, AgentInvocation)) assert sentinel not in bundle.model_dump_json() diff --git a/tests/unit_tests/test_anthropic_converter.py b/tests/unit_tests/test_anthropic_converter.py index ca9607205c..6fe997705a 100644 --- a/tests/unit_tests/test_anthropic_converter.py +++ b/tests/unit_tests/test_anthropic_converter.py @@ -145,7 +145,6 @@ def test_assistant_tool_use_becomes_function_call(self) -> None: assert params.input[0].content == "calling" fc = params.input[1] assert fc.type == "function_call" - assert fc.id.startswith("fc_") assert fc.call_id == "toolu_1" assert fc.name == "lookup" assert json.loads(fc.arguments) == {"city": "Paris"} From 46539539a59eab5cf3cb83da3bd538b7e1b3d10f Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Mon, 27 Jul 2026 19:20:35 +0200 Subject: [PATCH 11/14] Preserve Claude Code run outcomes Signed-off-by: Michal Bien --- responses_api_agents/claude_code_agent/app.py | 143 +++++++++++++----- .../claude_code_agent/observability.py | 35 +++-- .../claude_code_agent/tests/test_app.py | 104 ++++++++++--- .../tests/test_observability.py | 13 +- 4 files changed, 223 insertions(+), 72 deletions(-) diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index a3af7f7ee9..659e7cd6b2 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -23,7 +23,7 @@ import tempfile from asyncio import Semaphore from pathlib import Path -from time import time +from time import monotonic, time from typing import Any, Callable, Optional from uuid import uuid4 @@ -88,6 +88,7 @@ def parse_stream_json(stdout: str) -> tuple[list[Any], dict]: total_input = 0 total_output = 0 num_turns: Optional[int] = None + result_metadata: dict[str, Any] = {} for event in raw_events: etype = event.get("type") @@ -99,6 +100,13 @@ def parse_stream_json(stdout: str) -> tuple[list[Any], dict]: # Claude Code's authoritative turn counter (what --max-turns bounds). if event.get("num_turns") is not None: num_turns = int(event["num_turns"]) + if isinstance(event.get("subtype"), str): + result_metadata["subtype"] = event["subtype"] + if isinstance(event.get("is_error"), bool): + result_metadata["is_error"] = event["is_error"] + duration_ms = event.get("duration_ms") + if isinstance(duration_ms, (int, float)) and not isinstance(duration_ms, bool) and duration_ms >= 0: + result_metadata["duration_ms"] = float(duration_ms) elif etype == "assistant": message = event.get("message", {}) @@ -176,9 +184,23 @@ def parse_stream_json(stdout: str) -> tuple[list[Any], dict]: metadata: dict = {"input_tokens": total_input, "output_tokens": total_output} if num_turns is not None: metadata["num_turns"] = num_turns + metadata.update(result_metadata) return output_items, metadata +def _invocation_outcome(metadata: dict[str, Any], returncode: int | None) -> tuple[str, str | None]: + subtype = metadata.get("subtype") + if subtype == "error_max_turns": + return "incomplete", subtype + if metadata.get("is_error") is True or (isinstance(subtype, str) and subtype.startswith("error_")): + return "failed", subtype if isinstance(subtype, str) else "agent_error" + if subtype == "success": + return "completed", None + if returncode not in (0, None): + return "failed", f"process_exit_{returncode}" + return "incomplete", "result_missing" + + def _extract_instruction(body_input) -> tuple[str, Optional[str]]: """Return (user_message, system_message) from a responses body input list.""" items = list(body_input) @@ -386,9 +408,9 @@ async def _run_claude_code( mcp_config: Optional[str] = None, skills_path: Optional[str] = None, rollout_id: Optional[str] = None, - observation_collector: Optional[Callable[[Path], None]] = None, - ) -> tuple[str, str]: - """Run claude -p --output-format=stream-json and return (stdout, model_name). + observation_collector: Optional[Callable[[Path, dict[str, Any]], None]] = None, + ) -> tuple[str, str, dict[str, Any]]: + """Run claude -p --output-format=stream-json and return stdout, model name, and run metadata. When ``rollout_id`` is set and a model server is configured, the per-rollout capture prefix is applied to ANTHROPIC_BASE_URL so the CLI's streaming /v1/messages calls correlate to this rollout. @@ -399,6 +421,7 @@ async def _run_claude_code( api_key = self.config.anthropic_api_key claude_config_dir = None + run_metadata: dict[str, Any] = {"status": "unknown"} try: # Inside the try so a bad skills.path (raising in stage_skills) still cleans up the # partially-created config dir in the finally rather than leaking it per failing request. @@ -426,6 +449,7 @@ async def _run_claude_code( skills_active=bool(skills_path), ) + process_started_at = monotonic() proc = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, @@ -438,18 +462,30 @@ async def _run_claude_code( proc.kill() await proc.communicate() LOG.warning("claude-code timed out after %ds", self.config.timeout) - return "", model + run_metadata = { + "status": "incomplete", + "error_type": "timeout", + "duration_ms": (monotonic() - process_started_at) * 1000, + } + return "", model, run_metadata 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 + stdout_text = stdout.decode(errors="replace") + LOG.debug("claude-code stdout (%d chars): %s", len(stdout), stdout_text[:2000]) + _, run_metadata = parse_stream_json(stdout_text) + run_metadata.setdefault("duration_ms", (monotonic() - process_started_at) * 1000) + status, error_type = _invocation_outcome(run_metadata, proc.returncode) + run_metadata["status"] = status + if error_type is not None: + run_metadata["error_type"] = error_type + return stdout_text, model, run_metadata finally: if claude_config_dir is not None: try: if observation_collector is not None: - await asyncio.to_thread(observation_collector, claude_config_dir) + await asyncio.to_thread(observation_collector, claude_config_dir, run_metadata) except Exception: LOG.exception("failed to collect Claude Code observations") finally: @@ -521,8 +557,8 @@ async def _create_response( mcp_config: Optional[str] = None, skills_path: Optional[str] = None, rollout_id: Optional[str] = None, - observation_collector: Optional[Callable[[Path], None]] = None, - ) -> NeMoGymResponse: + observation_collector: Optional[Callable[[Path, dict[str, Any]], None]] = None, + ) -> tuple[NeMoGymResponse, dict[str, Any]]: body = body.model_copy(deep=True) if isinstance(body.input, str): body.input = [NeMoGymEasyInputMessage(role="user", content=body.input)] @@ -531,7 +567,7 @@ 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, run_metadata = await self._run_claude_code( user_message, system_prompt=system_prompt, mcp_config=mcp_config, @@ -540,6 +576,15 @@ async def _create_response( observation_collector=observation_collector, ) output_items, usage = parse_stream_json(stdout) + run_metadata = usage | run_metadata + run_metadata.setdefault( + "num_turns", + sum( + 1 + for item in output_items + if getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" + ), + ) if not any( getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" @@ -559,22 +604,25 @@ async def _create_response( input_tokens = usage.get("input_tokens", 0) output_tokens = usage.get("output_tokens", 0) - return NeMoGymResponse( - id=f"resp_{uuid4().hex}", - created_at=int(time()), - model=model_name, - object="response", - output=output_items, - tool_choice=body.tool_choice, - tools=body.tools, - parallel_tool_calls=body.parallel_tool_calls, - usage=NeMoGymResponseUsage( - input_tokens=input_tokens, - input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=0), - output_tokens=output_tokens, - output_tokens_details=NeMoGymResponseOutputTokensDetails(reasoning_tokens=0), - total_tokens=input_tokens + output_tokens, + return ( + NeMoGymResponse( + id=f"resp_{uuid4().hex}", + created_at=int(time()), + model=model_name, + object="response", + output=output_items, + tool_choice=body.tool_choice, + tools=body.tools, + parallel_tool_calls=body.parallel_tool_calls, + usage=NeMoGymResponseUsage( + input_tokens=input_tokens, + input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=0), + output_tokens=output_tokens, + output_tokens_details=NeMoGymResponseOutputTokensDetails(reasoning_tokens=0), + total_tokens=input_tokens + output_tokens, + ), ), + run_metadata, ) async def responses( @@ -582,7 +630,8 @@ async def responses( request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body(), ) -> NeMoGymResponse: - return await self._create_response(body, rollout_id=request.path_params.get("rollout_id")) + response, _ = await self._create_response(body, rollout_id=request.path_params.get("rollout_id")) + return response async def _create_episode( self, @@ -591,13 +640,19 @@ async def _create_episode( mcp_config: Optional[str] = None, skills_path: Optional[str] = None, rollout_id: Optional[str] = None, - ) -> AgentEpisode: + ) -> tuple[AgentEpisode, dict[str, Any]]: observations: Optional[AgentObservationBundle] = None - def collect(config_dir: Path) -> None: + def collect(config_dir: Path, run_metadata: dict[str, Any]) -> None: nonlocal observations try: - observations = extract_claude_code_observations(config_dir, model_ref=self.config.model_server) + observations = extract_claude_code_observations( + config_dir, + model_ref=self.config.model_server, + root_status=run_metadata["status"], + root_duration_ms=run_metadata.get("duration_ms"), + root_error_type=run_metadata.get("error_type"), + ) if self.config.model_server is None: observations.gaps.append(ObservationGap(code="model_call_ownership_unavailable")) except Exception: @@ -607,7 +662,7 @@ def collect(config_dir: Path) -> None: gaps=[ObservationGap(code="observation_parse_failed")], ) - response = await self._create_response( + response, run_metadata = await self._create_response( body, mcp_config=mcp_config, skills_path=skills_path, @@ -620,7 +675,7 @@ def collect(config_dir: Path) -> None: gaps=[ObservationGap(code="agent_transcript_unavailable")], ) observations.gaps.append(ObservationGap(code="no_sandbox_runtime")) - return AgentEpisode(response=response, observations=observations) + return AgentEpisode(response=response, observations=observations), run_metadata async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> ClaudeCodeAgentVerifyResponse: async with self.sem: @@ -646,7 +701,7 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude with tempfile.TemporaryDirectory(prefix="nemo_gym_claude_mcp_") as mcp_config_dir: mcp_config = self._write_rollout_mcp_config(seed_resp_json, Path(mcp_config_dir)) if rollout_id is not None: - episode = await self._create_episode( + episode, run_metadata = await self._create_episode( body.responses_create_params, mcp_config=mcp_config, skills_path=skills_path, @@ -654,7 +709,7 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude ) agent_resp, observations = episode.response, episode.observations else: - agent_resp = await self._create_response( + agent_resp, run_metadata = await self._create_response( body.responses_create_params, mcp_config=mcp_config, skills_path=skills_path, @@ -672,13 +727,21 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude verify_json = await get_response_json(verify_resp) gym_resp = NeMoGymResponse.model_validate(agent_resp_json) - turns = sum( - 1 - for item in gym_resp.output - if getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" - ) + turns = run_metadata.get("num_turns") + if not isinstance(turns, int): + turns = sum( + 1 + for item in gym_resp.output + if getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" + ) last = gym_resp.output[-1] if gym_resp.output else None - naturally = getattr(last, "type", None) == "message" and getattr(last, "role", None) == "assistant" + subtype = run_metadata.get("subtype") + if run_metadata.get("status") != "completed": + naturally = False + elif isinstance(subtype, str): + naturally = subtype == "success" and run_metadata.get("is_error") is not True + else: + naturally = getattr(last, "type", None) == "message" and getattr(last, "role", None) == "assistant" result = verify_json | {"turns_used": turns, "finished_naturally": naturally} if observations is not None: diff --git a/responses_api_agents/claude_code_agent/observability.py b/responses_api_agents/claude_code_agent/observability.py index 8d77c84184..91a84adaca 100644 --- a/responses_api_agents/claude_code_agent/observability.py +++ b/responses_api_agents/claude_code_agent/observability.py @@ -10,7 +10,7 @@ from collections import defaultdict from datetime import datetime from pathlib import Path -from typing import Any +from typing import Any, Literal from nemo_gym.config_types import ModelServerRef from nemo_gym.openai_utils import ( @@ -149,12 +149,10 @@ def _tool_call(tool_call_id: str, block: dict[str, Any]) -> NeMoGymResponseFunct ) -def _tool_result(event: dict[str, Any], block: dict[str, Any]) -> NeMoGymFunctionCallOutput: - event_id = event.get("uuid") +def _tool_result(block: dict[str, Any]) -> NeMoGymFunctionCallOutput: return NeMoGymFunctionCallOutput( call_id=block["tool_use_id"], output=_text(block.get("content")), - id=event_id if isinstance(event_id, str) else None, status="completed", ) @@ -207,6 +205,9 @@ def extract_claude_code_observations( config_dir: Path, *, model_ref: ModelServerRef | None = None, + root_status: Literal["completed", "failed", "incomplete", "unknown"] = "unknown", + root_duration_ms: float | None = None, + root_error_type: str | None = None, ) -> AgentObservationBundle: """Extract exact relationships available in one ``CLAUDE_CONFIG_DIR``. @@ -375,7 +376,7 @@ def add_gap(code: str, detail: str | None = None) -> None: add_gap("tool_result_id_missing") continue tool_status = _status(block, result_metadata) - items.append(_tool_result(event, block)) + items.append(_tool_result(block)) finishes[(invocation_id, tool_call_id)].append( (_timestamp(event.get("timestamp")), tool_status) ) @@ -468,19 +469,27 @@ def add_tool_gap(code: str) -> None: invocations_by_id: dict[str, AgentInvocation] = {} for invocation_id in all_invocation_ids: parent = parent_by_invocation.get(invocation_id) - if parent is None: + is_root = parent is None and invocation_id not in agent_invocations + if parent is None and (not is_root or root_status == "unknown"): gaps.append(_gap("invocation_outcome_unavailable", invocation_id=invocation_id)) + status = "unknown" + duration_ms = None + error_type = None + if is_root: + status = root_status + duration_ms = root_duration_ms + error_type = root_error_type + elif parent is not None: + status = "incomplete" if parent[2] in {"timeout", "cancelled"} else parent[2] + if parent[2] in {"failed", "timeout", "cancelled"}: + error_type = parent[2] invocations_by_id[invocation_id] = AgentInvocation( invocation_id=invocation_id, parent_invocation_id=parent[0] if parent else None, spawned_by_tool_call_id=parent[1] if parent else None, - status=( - "incomplete" - if parent and parent[2] in {"timeout", "cancelled"} - else parent[2] - if parent - else "unknown" - ), + status=status, + duration_ms=duration_ms, + error_type=error_type, model_calls=model_calls[invocation_id], conversation=conversations[invocation_id], ) 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 efe0922a9f..6ad2ed8886 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -41,6 +41,7 @@ ModelServerRef, ResourcesServerRef, _extract_instruction, + _invocation_outcome, parse_stream_json, ) from responses_api_agents.claude_code_agent.observability import extract_claude_code_observations @@ -288,7 +289,9 @@ 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", {"status": "incomplete", "error_type": "result_missing"}) + ) body = ClaudeCodeAgentRunRequest.model_validate( { "responses_create_params": {"input": []}, @@ -302,13 +305,17 @@ 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", {"status": "incomplete", "error_type": "result_missing"}) + ) body = ClaudeCodeAgentRunRequest.model_validate({"responses_create_params": {"input": []}}) result = self._run(agent, body, run_claude_code) assert run_claude_code.call_args.kwargs["skills_path"] is None assert "ng_agent_observations" not in result.model_dump(mode="json") + assert result.turns_used == 0 + assert result.finished_naturally is False class TestObservability: @@ -335,8 +342,19 @@ async def run_claude_code(*args, observation_collector=None, **kwargs): } ) ) - observation_collector(tmp_path) - return _event("assistant", message={"content": [{"type": "text", "text": "done"}]}), "model" + run_metadata = { + "status": "completed", + "duration_ms": 123.0, + "num_turns": 7, + "subtype": "success", + "is_error": False, + } + observation_collector(tmp_path, run_metadata) + return ( + _event("assistant", message={"content": [{"type": "text", "text": "done"}]}), + "model", + run_metadata, + ) request = MagicMock() request.cookies = {} @@ -354,7 +372,11 @@ async def run_claude_code(*args, observation_collector=None, **kwargs): assert observations is not None invocation = next(record for record in observations.records if isinstance(record, AgentInvocation)) assert invocation.invocation_id == "session-1" + assert invocation.status == "completed" + assert invocation.duration_ms == 123 assert invocation.model_calls[0].response_id == "msg-1" + assert result.turns_used == 7 + assert result.finished_naturally is True assert "no_sandbox_runtime" in {gap.code for gap in observations.gaps} @@ -367,7 +389,11 @@ class FakeProc: returncode = 0 async def communicate(self): - return b'{"type":"result","usage":{"input_tokens":3,"output_tokens":4}}\n', b"" + return ( + b'{"type":"result","subtype":"success","is_error":false,' + b'"usage":{"input_tokens":3,"output_tokens":4}}\n', + b"", + ) async def fake_exec(*cmd, **kwargs): env = kwargs["env"] @@ -383,7 +409,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, metadata = asyncio.run(agent._run_claude_code("hello", system_prompt="be terse")) assert "claude" in captured["cmd"][0] assert "--mcp-config" in captured["cmd"] @@ -394,6 +420,7 @@ async def fake_exec(*cmd, **kwargs): assert not Path(captured["config_dir"]).exists() assert "result" in stdout assert model == "claude-sonnet-4-6" + assert metadata["status"] == "completed" def test_skills_staged_and_bare_dropped(self, tmp_path: Path) -> None: skills_dir = _write_skill_dir(tmp_path) @@ -463,11 +490,14 @@ 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, metadata = asyncio.run(agent._run_claude_code("hello")) assert stdout == "" assert killed["called"] is True assert model == "claude-sonnet-4-6" + assert metadata["status"] == "incomplete" + assert metadata["error_type"] == "timeout" + assert metadata["duration_ms"] >= 0 def test_collects_observations_before_cleanup(self, tmp_path: Path) -> None: agent = _make_agent() @@ -478,7 +508,7 @@ class FakeProc: returncode = 0 async def communicate(self): - return b'{"type":"result","usage":{}}\n', b"" + return b'{"type":"result","subtype":"success","is_error":false,"usage":{}}\n', b"" async def fake_exec(*cmd, **kwargs): config_dir = Path(kwargs["env"]["CLAUDE_CONFIG_DIR"]) @@ -502,9 +532,12 @@ async def fake_exec(*cmd, **kwargs): captured["config_dir"] = config_dir return FakeProc() - def collect(config_dir: Path) -> None: + def collect(config_dir: Path, run_metadata: dict) -> None: captured["collector_thread"] = threading.get_ident() - captured["observations"] = extract_claude_code_observations(config_dir) + captured["observations"] = extract_claude_code_observations( + config_dir, + root_status=run_metadata["status"], + ) with ( patch("responses_api_agents.claude_code_agent.app.Path.home", return_value=tmp_path), @@ -514,6 +547,7 @@ def collect(config_dir: Path) -> None: invocation = next(record for record in captured["observations"].records if isinstance(record, AgentInvocation)) assert invocation.invocation_id == "session-1" + assert invocation.status == "completed" assert captured["collector_thread"] != event_loop_thread assert not captured["config_dir"].exists() @@ -626,10 +660,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", + {"status": "completed"}, + ) agent.server_client.post.side_effect = fake_post object.__setattr__(agent, "_run_claude_code", fake_run_claude_code) @@ -681,7 +719,11 @@ 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", + {"status": "completed"}, + ) agent.server_client.post.side_effect = fake_post object.__setattr__(agent, "_run_claude_code", fake_run_claude_code) @@ -765,6 +807,20 @@ def test_empty(self) -> None: class TestParseStreamJson: + @pytest.mark.parametrize( + ("metadata", "returncode", "expected"), + [ + ({"subtype": "success"}, 0, ("completed", None)), + ({"subtype": "error_max_turns", "is_error": True}, 0, ("incomplete", "error_max_turns")), + ({"subtype": "error_during_execution", "is_error": True}, 0, ("failed", "error_during_execution")), + ({"subtype": "error_max_turns", "is_error": True}, 7, ("incomplete", "error_max_turns")), + ({}, 7, ("failed", "process_exit_7")), + ({}, 0, ("incomplete", "result_missing")), + ], + ) + def test_invocation_outcome(self, metadata: dict, returncode: int, expected: tuple[str, str | None]) -> None: + assert _invocation_outcome(metadata, returncode) == expected + def _assistant(self, content: list) -> str: return _event("assistant", message={"content": content, "usage": {"input_tokens": 10, "output_tokens": 5}}) @@ -853,9 +909,23 @@ def test_result_event_accumulates_usage(self) -> None: assert usage["output_tokens"] == 50 def test_result_event_exposes_num_turns(self) -> None: - result = _event("result", num_turns=9, usage={"input_tokens": 1, "output_tokens": 1}) + result = _event( + "result", + num_turns=9, + subtype="success", + is_error=False, + duration_ms=1234, + usage={"input_tokens": 1, "output_tokens": 1}, + ) _, usage = parse_stream_json(result) - assert usage["num_turns"] == 9 + assert usage == { + "input_tokens": 1, + "output_tokens": 1, + "num_turns": 9, + "subtype": "success", + "is_error": False, + "duration_ms": 1234, + } def test_num_turns_absent_when_no_result_event(self) -> None: assistant = self._assistant([{"type": "text", "text": "hi"}]) diff --git a/responses_api_agents/claude_code_agent/tests/test_observability.py b/responses_api_agents/claude_code_agent/tests/test_observability.py index 11f343a6fe..f364a724fa 100644 --- a/responses_api_agents/claude_code_agent/tests/test_observability.py +++ b/responses_api_agents/claude_code_agent/tests/test_observability.py @@ -8,6 +8,7 @@ import pytest from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import NeMoGymFunctionCallOutput from nemo_gym.rollout_observability import ( AgentInvocation, AgentObservationBundle, @@ -146,13 +147,20 @@ def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path ), ) - bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) + bundle = extract_claude_code_observations( + tmp_path, + model_ref=MODEL_REF, + root_status="completed", + root_duration_ms=5000, + ) invocations = {invocation.invocation_id: invocation for invocation in _records(bundle, AgentInvocation)} assert set(invocations) == {session, child, grandchild} root_invocation = invocations[session] child_invocation = invocations[child] grandchild_invocation = invocations[grandchild] + assert root_invocation.status == "completed" + assert root_invocation.duration_ms == 5000 assert child_invocation.parent_invocation_id == session assert child_invocation.spawned_by_tool_call_id == "tool-child" assert grandchild_invocation.parent_invocation_id == child @@ -172,6 +180,7 @@ def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path "function_call_output", "function_call_output", ] + assert all(item.id is None for item in root_invocation.conversation if isinstance(item, NeMoGymFunctionCallOutput)) timings = {tool.tool_call_id: tool for tool in _records(bundle, ToolCallObservation)} assert timings["tool-fast"].duration_ms == pytest.approx(1000) @@ -179,7 +188,7 @@ def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path assert timings["tool-grandchild"].duration_ms == pytest.approx(1000) assert timings["tool-grandchild"].status == "timeout" assert all(tool.timing_source == "artifact" for tool in timings.values()) - assert [(gap.code, gap.invocation_id) for gap in bundle.gaps] == [("invocation_outcome_unavailable", session)] + assert bundle.gaps == [] def test_extracts_explicit_compaction_markers(tmp_path: Path) -> None: From 9f68db2d43c0cb2a558ef2e539ccba5232b26d19 Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Tue, 28 Jul 2026 10:25:49 +0200 Subject: [PATCH 12/14] Harden Claude Code observation extraction Signed-off-by: Michal Bien --- .../claude_code_agent/README.md | 6 ++-- .../claude_code_agent/observability.py | 33 ++++++++++++++----- .../tests/test_observability.py | 31 ++++++++++++++++- 3 files changed, 59 insertions(+), 11 deletions(-) diff --git a/responses_api_agents/claude_code_agent/README.md b/responses_api_agents/claude_code_agent/README.md index 573c5b394b..02adce1db2 100644 --- a/responses_api_agents/claude_code_agent/README.md +++ b/responses_api_agents/claude_code_agent/README.md @@ -196,5 +196,7 @@ The skills path is resolved like `input_jsonl_fpath` (relative paths check the w ## Limitations - Eval only for now. Token IDs and logprobs are not wired up yet. -- Does not go through Gym's model server. Token counts come from Claude Code's own usage reporting. -- `turns_used` counts assistant messages right now, not tool calls. +- With `model_server`, model calls go through Gym and can be captured. Direct Anthropic or + `anthropic_base_url` runs bypass Gym capture. +- `turns_used` uses Claude Code's terminal `num_turns` when present and otherwise counts assistant + messages. diff --git a/responses_api_agents/claude_code_agent/observability.py b/responses_api_agents/claude_code_agent/observability.py index 91a84adaca..6873584097 100644 --- a/responses_api_agents/claude_code_agent/observability.py +++ b/responses_api_agents/claude_code_agent/observability.py @@ -201,6 +201,21 @@ def _read_events(config_dir: Path, gaps: list[ObservationGap]) -> list[tuple[int return events +def _would_create_parent_cycle( + child_id: str, + parent_id: str, + parents: dict[str, tuple[str, str, str, int]], +) -> bool: + seen = {child_id} + current = parent_id + while current in parents: + if current in seen: + return True + seen.add(current) + current = parents[current][0] + return current in seen + + def extract_claude_code_observations( config_dir: Path, *, @@ -242,13 +257,6 @@ def extract_claude_code_observations( compactions: list[ContextCompactionObservation] = [] for invocation_id, entries in events_by_invocation.items(): - entries.sort( - key=lambda pair: ( - _timestamp(pair[1].get("timestamp")) is None, - _timestamp(pair[1].get("timestamp")) or 0, - pair[0], - ) - ) items = conversations[invocation_id] refs = model_calls[invocation_id] @@ -406,7 +414,16 @@ def add_gap(code: str, detail: str | None = None) -> None: ambiguous_parents.add(child_id) gaps.append(_gap("conflicting_subagent_parent", invocation_id=child_id)) elif child_id not in ambiguous_parents: - parents.setdefault(child_id, parent) + if _would_create_parent_cycle(child_id, invocation_id, parents): + gaps.append( + _gap( + "cyclic_subagent_parent", + invocation_id=child_id, + detail=invocation_id, + ) + ) + else: + parents.setdefault(child_id, parent) else: add_gap("ambiguous_subagent_relation") diff --git a/responses_api_agents/claude_code_agent/tests/test_observability.py b/responses_api_agents/claude_code_agent/tests/test_observability.py index f364a724fa..4dad6c0ad8 100644 --- a/responses_api_agents/claude_code_agent/tests/test_observability.py +++ b/responses_api_agents/claude_code_agent/tests/test_observability.py @@ -270,10 +270,39 @@ def test_malformed_and_incomplete_artifacts_produce_sanitized_gaps(tmp_path: Pat "tool_start_missing", "tool_timing_invalid", } <= codes - assert all(not invocation.model_calls for invocation in _records(bundle, AgentInvocation)) + invocations = {invocation.invocation_id: invocation for invocation in _records(bundle, AgentInvocation)} + assert all(not invocation.model_calls for invocation in invocations.values()) + assert [item.type for item in invocations["root"].conversation] == [ + "function_call", + "function_call", + "function_call_output", + "function_call_output", + ] assert sentinel not in bundle.model_dump_json() +def test_rejects_cyclic_subagent_parents(tmp_path: Path) -> None: + _write( + tmp_path / "projects" / "work" / "root.jsonl", + _tool_result("root", "2026-07-22T10:00:00Z", "self", child_id="root"), + _tool_result("root", "2026-07-22T10:00:01Z", "spawn", child_id="agent-a"), + _tool_result( + "root", + "2026-07-22T10:00:02Z", + "back-edge", + agent="agent-a", + child_id="root", + ), + ) + + bundle = extract_claude_code_observations(tmp_path) + invocations = {invocation.invocation_id: invocation for invocation in _records(bundle, AgentInvocation)} + + assert invocations["root"].parent_invocation_id is None + assert invocations["agent-a"].parent_invocation_id == "root" + assert [gap.code for gap in bundle.gaps].count("cyclic_subagent_parent") == 2 + + def test_ignores_non_transcript_jsonl_and_reports_no_usable_transcript(tmp_path: Path) -> None: _write( tmp_path / "skills" / "fixture.jsonl", From 9af9174872c51af15dd6b6a353a9211d8dae630b Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Tue, 28 Jul 2026 18:10:22 +0200 Subject: [PATCH 13/14] Address Claude Code observation review feedback Signed-off-by: Michal Bien --- .../claude_code_agent/README.md | 3 +- responses_api_agents/claude_code_agent/app.py | 100 +++++++----------- .../claude_code_agent/observability.py | 58 ++++++++-- .../claude_code_agent/tests/test_app.py | 39 ++++--- .../tests/test_observability.py | 64 ++++++++++- 5 files changed, 176 insertions(+), 88 deletions(-) diff --git a/responses_api_agents/claude_code_agent/README.md b/responses_api_agents/claude_code_agent/README.md index 02adce1db2..3fbdc600ac 100644 --- a/responses_api_agents/claude_code_agent/README.md +++ b/responses_api_agents/claude_code_agent/README.md @@ -198,5 +198,4 @@ The skills path is resolved like `input_jsonl_fpath` (relative paths check the w - Eval only for now. Token IDs and logprobs are not wired up yet. - With `model_server`, model calls go through Gym and can be captured. Direct Anthropic or `anthropic_base_url` runs bypass Gym capture. -- `turns_used` uses Claude Code's terminal `num_turns` when present and otherwise counts assistant - messages. +- `turns_used` counts assistant messages, not tool calls. diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index 659e7cd6b2..ed634c9085 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -194,10 +194,10 @@ def _invocation_outcome(metadata: dict[str, Any], returncode: int | None) -> tup return "incomplete", subtype if metadata.get("is_error") is True or (isinstance(subtype, str) and subtype.startswith("error_")): return "failed", subtype if isinstance(subtype, str) else "agent_error" - if subtype == "success": - return "completed", None if returncode not in (0, None): return "failed", f"process_exit_{returncode}" + if subtype == "success": + return "completed", None return "incomplete", "result_missing" @@ -409,8 +409,8 @@ async def _run_claude_code( skills_path: Optional[str] = None, rollout_id: Optional[str] = None, observation_collector: Optional[Callable[[Path, dict[str, Any]], None]] = None, - ) -> tuple[str, str, dict[str, Any]]: - """Run claude -p --output-format=stream-json and return stdout, model name, and run metadata. + ) -> tuple[list[Any], str, dict[str, Any]]: + """Run Claude Code and return parsed output, model name, and run metadata. When ``rollout_id`` is set and a model server is configured, the per-rollout capture prefix is applied to ANTHROPIC_BASE_URL so the CLI's streaming /v1/messages calls correlate to this rollout. @@ -467,20 +467,20 @@ async def _run_claude_code( "error_type": "timeout", "duration_ms": (monotonic() - process_started_at) * 1000, } - return "", model, run_metadata + return [], model, run_metadata if proc.returncode not in (0, None): LOG.warning("claude-code exited %d: %s", proc.returncode, stderr.decode(errors="replace")[:500]) stdout_text = stdout.decode(errors="replace") LOG.debug("claude-code stdout (%d chars): %s", len(stdout), stdout_text[:2000]) - _, run_metadata = parse_stream_json(stdout_text) + output_items, run_metadata = parse_stream_json(stdout_text) run_metadata.setdefault("duration_ms", (monotonic() - process_started_at) * 1000) status, error_type = _invocation_outcome(run_metadata, proc.returncode) run_metadata["status"] = status if error_type is not None: run_metadata["error_type"] = error_type - return stdout_text, model, run_metadata + return output_items, model, run_metadata finally: if claude_config_dir is not None: try: @@ -558,7 +558,7 @@ async def _create_response( skills_path: Optional[str] = None, rollout_id: Optional[str] = None, observation_collector: Optional[Callable[[Path, dict[str, Any]], None]] = None, - ) -> tuple[NeMoGymResponse, dict[str, Any]]: + ) -> NeMoGymResponse: body = body.model_copy(deep=True) if isinstance(body.input, str): body.input = [NeMoGymEasyInputMessage(role="user", content=body.input)] @@ -567,7 +567,7 @@ 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, run_metadata = await self._run_claude_code( + output_items, model_name, run_metadata = await self._run_claude_code( user_message, system_prompt=system_prompt, mcp_config=mcp_config, @@ -575,16 +575,6 @@ async def _create_response( rollout_id=rollout_id, observation_collector=observation_collector, ) - output_items, usage = parse_stream_json(stdout) - run_metadata = usage | run_metadata - run_metadata.setdefault( - "num_turns", - sum( - 1 - for item in output_items - if getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" - ), - ) if not any( getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" @@ -601,28 +591,25 @@ async def _create_response( ) ) - input_tokens = usage.get("input_tokens", 0) - output_tokens = usage.get("output_tokens", 0) - - return ( - NeMoGymResponse( - id=f"resp_{uuid4().hex}", - created_at=int(time()), - model=model_name, - object="response", - output=output_items, - tool_choice=body.tool_choice, - tools=body.tools, - parallel_tool_calls=body.parallel_tool_calls, - usage=NeMoGymResponseUsage( - input_tokens=input_tokens, - input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=0), - output_tokens=output_tokens, - output_tokens_details=NeMoGymResponseOutputTokensDetails(reasoning_tokens=0), - total_tokens=input_tokens + output_tokens, - ), + input_tokens = run_metadata.get("input_tokens", 0) + output_tokens = run_metadata.get("output_tokens", 0) + + return NeMoGymResponse( + id=f"resp_{uuid4().hex}", + created_at=int(time()), + model=model_name, + object="response", + output=output_items, + tool_choice=body.tool_choice, + tools=body.tools, + parallel_tool_calls=body.parallel_tool_calls, + usage=NeMoGymResponseUsage( + input_tokens=input_tokens, + input_tokens_details=NeMoGymResponseInputTokensDetails(cached_tokens=0), + output_tokens=output_tokens, + output_tokens_details=NeMoGymResponseOutputTokensDetails(reasoning_tokens=0), + total_tokens=input_tokens + output_tokens, ), - run_metadata, ) async def responses( @@ -630,8 +617,7 @@ async def responses( request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body(), ) -> NeMoGymResponse: - response, _ = await self._create_response(body, rollout_id=request.path_params.get("rollout_id")) - return response + return await self._create_response(body, rollout_id=request.path_params.get("rollout_id")) async def _create_episode( self, @@ -640,7 +626,7 @@ async def _create_episode( mcp_config: Optional[str] = None, skills_path: Optional[str] = None, rollout_id: Optional[str] = None, - ) -> tuple[AgentEpisode, dict[str, Any]]: + ) -> AgentEpisode: observations: Optional[AgentObservationBundle] = None def collect(config_dir: Path, run_metadata: dict[str, Any]) -> None: @@ -662,7 +648,7 @@ def collect(config_dir: Path, run_metadata: dict[str, Any]) -> None: gaps=[ObservationGap(code="observation_parse_failed")], ) - response, run_metadata = await self._create_response( + response = await self._create_response( body, mcp_config=mcp_config, skills_path=skills_path, @@ -675,7 +661,7 @@ def collect(config_dir: Path, run_metadata: dict[str, Any]) -> None: gaps=[ObservationGap(code="agent_transcript_unavailable")], ) observations.gaps.append(ObservationGap(code="no_sandbox_runtime")) - return AgentEpisode(response=response, observations=observations), run_metadata + return AgentEpisode(response=response, observations=observations) async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> ClaudeCodeAgentVerifyResponse: async with self.sem: @@ -701,7 +687,7 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude with tempfile.TemporaryDirectory(prefix="nemo_gym_claude_mcp_") as mcp_config_dir: mcp_config = self._write_rollout_mcp_config(seed_resp_json, Path(mcp_config_dir)) if rollout_id is not None: - episode, run_metadata = await self._create_episode( + episode = await self._create_episode( body.responses_create_params, mcp_config=mcp_config, skills_path=skills_path, @@ -709,7 +695,7 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude ) agent_resp, observations = episode.response, episode.observations else: - agent_resp, run_metadata = await self._create_response( + agent_resp = await self._create_response( body.responses_create_params, mcp_config=mcp_config, skills_path=skills_path, @@ -727,21 +713,13 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude verify_json = await get_response_json(verify_resp) gym_resp = NeMoGymResponse.model_validate(agent_resp_json) - turns = run_metadata.get("num_turns") - if not isinstance(turns, int): - turns = sum( - 1 - for item in gym_resp.output - if getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" - ) + turns = sum( + 1 + for item in gym_resp.output + if getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" + ) last = gym_resp.output[-1] if gym_resp.output else None - subtype = run_metadata.get("subtype") - if run_metadata.get("status") != "completed": - naturally = False - elif isinstance(subtype, str): - naturally = subtype == "success" and run_metadata.get("is_error") is not True - else: - naturally = getattr(last, "type", None) == "message" and getattr(last, "role", None) == "assistant" + naturally = getattr(last, "type", None) == "message" and getattr(last, "role", None) == "assistant" result = verify_json | {"turns_used": turns, "finished_naturally": naturally} if observations is not None: diff --git a/responses_api_agents/claude_code_agent/observability.py b/responses_api_agents/claude_code_agent/observability.py index 6873584097..3a59d0cf3d 100644 --- a/responses_api_agents/claude_code_agent/observability.py +++ b/responses_api_agents/claude_code_agent/observability.py @@ -89,11 +89,37 @@ def _metadata(event: dict[str, Any]) -> dict[str, Any]: def _integer(metadata: dict[str, Any], *keys: str) -> int | None: for key in keys: value = metadata.get(key) - if isinstance(value, int) and not isinstance(value, bool): + if isinstance(value, int) and not isinstance(value, bool) and value >= 0: return value return None +def _compaction_outcome( + event: dict[str, Any], + metadata: dict[str, Any], + *, + has_completion_marker: bool, +) -> Literal["completed", "failed", "aborted", "unknown"]: + value = metadata.get("outcome") + if not isinstance(value, str): + value = metadata.get("status") + normalized = value.lower() if isinstance(value, str) else None + if normalized in {"failed", "failure", "error"}: + return "failed" + if normalized in {"aborted", "cancelled", "canceled", "interrupted"}: + return "aborted" + message = event.get("message") + if ( + metadata.get("is_error") is True + or event.get("is_error") is True + or (isinstance(message, dict) and message.get("is_error") is True) + ): + return "failed" + if normalized in {"completed", "complete", "success", "succeeded"}: + return "completed" + return "completed" if has_completion_marker else "unknown" + + def _compaction(event: dict[str, Any], invocation_id: str) -> ContextCompactionObservation | None: message = event.get("message") is_summary = isinstance(message, dict) and message.get("isCompactSummary") is True @@ -110,7 +136,11 @@ def _compaction(event: dict[str, Any], invocation_id: str) -> ContextCompactionO trigger=trigger if isinstance(trigger, str) else None, tokens_before=_integer(metadata, "tokensBefore", "preTokens"), tokens_after=_integer(metadata, "tokensAfter", "postTokens"), - outcome="completed", + outcome=_compaction_outcome( + event, + metadata, + has_completion_marker=is_summary or is_boundary, + ), summary=summary or None, ) @@ -149,11 +179,11 @@ def _tool_call(tool_call_id: str, block: dict[str, Any]) -> NeMoGymResponseFunct ) -def _tool_result(block: dict[str, Any]) -> NeMoGymFunctionCallOutput: +def _tool_result(block: dict[str, Any], status: str) -> NeMoGymFunctionCallOutput: return NeMoGymFunctionCallOutput( call_id=block["tool_use_id"], output=_text(block.get("content")), - status="completed", + status="completed" if status == "completed" else "incomplete", ) @@ -278,9 +308,11 @@ def add_gap(code: str, detail: str | None = None) -> None: and previous_compaction[1] != is_summary ): prior = previous_compaction[2] - for field in ("trigger", "tokens_before", "tokens_after", "summary"): + for field in ("observed_at", "trigger", "tokens_before", "tokens_after", "summary"): if getattr(prior, field) is None: setattr(prior, field, getattr(compaction, field)) + if prior.outcome == "unknown" or compaction.outcome in {"failed", "aborted"}: + prior.outcome = compaction.outcome compaction = prior else: compaction.before_model_call = last_model_call @@ -289,8 +321,6 @@ def add_gap(code: str, detail: str | None = None) -> None: if last_model_call is None: add_gap("compaction_before_model_call_unavailable") previous_compaction = (entry_index, is_summary, compaction) - if compaction.observed_at is None: - add_gap("compaction_timestamp_missing") else: previous_compaction = None @@ -384,7 +414,7 @@ def add_gap(code: str, detail: str | None = None) -> None: add_gap("tool_result_id_missing") continue tool_status = _status(block, result_metadata) - items.append(_tool_result(block)) + items.append(_tool_result(block, tool_status)) finishes[(invocation_id, tool_call_id)].append( (_timestamp(event.get("timestamp")), tool_status) ) @@ -430,6 +460,18 @@ def add_gap(code: str, detail: str | None = None) -> None: for _ in pending_compactions: add_gap("compaction_after_model_call_unavailable") + for compaction in compactions: + gaps.append( + _gap( + "compaction_model_call_reference_unavailable", + invocation_id=compaction.invocation_id, + ) + ) + if compaction.outcome == "unknown": + gaps.append(_gap("compaction_outcome_unavailable", invocation_id=compaction.invocation_id)) + if compaction.observed_at is None: + gaps.append(_gap("compaction_timestamp_missing", invocation_id=compaction.invocation_id)) + tool_calls: list[ToolCallObservation] = [] for invocation_id, tool_call_id in sorted( set(starts) | set(finishes), key=lambda key: (first_seen.get(key[0], math.inf), key[1]) 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 6ad2ed8886..437b2f1c85 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -77,6 +77,10 @@ def _event(type_: str, **kwargs) -> str: return json.dumps({"type": type_, **kwargs}) +def _output(*events: str) -> list: + return parse_stream_json("\n".join(events))[0] + + class FakeAioHTTPResponse: ok = True @@ -290,7 +294,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", {"status": "incomplete", "error_type": "result_missing"}) + return_value=([], "claude-sonnet-4-6", {"status": "incomplete", "error_type": "result_missing"}) ) body = ClaudeCodeAgentRunRequest.model_validate( { @@ -306,7 +310,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", {"status": "incomplete", "error_type": "result_missing"}) + return_value=([], "claude-sonnet-4-6", {"status": "incomplete", "error_type": "result_missing"}) ) body = ClaudeCodeAgentRunRequest.model_validate({"responses_create_params": {"input": []}}) @@ -314,8 +318,8 @@ def test_no_skills_ref_forwards_none(self) -> None: assert run_claude_code.call_args.kwargs["skills_path"] is None assert "ng_agent_observations" not in result.model_dump(mode="json") - assert result.turns_used == 0 - assert result.finished_naturally is False + assert result.turns_used == 1 + assert result.finished_naturally is True class TestObservability: @@ -351,7 +355,7 @@ async def run_claude_code(*args, observation_collector=None, **kwargs): } observation_collector(tmp_path, run_metadata) return ( - _event("assistant", message={"content": [{"type": "text", "text": "done"}]}), + _output(_event("assistant", message={"content": [{"type": "text", "text": "done"}]})), "model", run_metadata, ) @@ -375,7 +379,7 @@ async def run_claude_code(*args, observation_collector=None, **kwargs): assert invocation.status == "completed" assert invocation.duration_ms == 123 assert invocation.model_calls[0].response_id == "msg-1" - assert result.turns_used == 7 + assert result.turns_used == 1 assert result.finished_naturally is True assert "no_sandbox_runtime" in {gap.code for gap in observations.gaps} @@ -409,7 +413,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, metadata = asyncio.run(agent._run_claude_code("hello", system_prompt="be terse")) + output_items, model, metadata = asyncio.run(agent._run_claude_code("hello", system_prompt="be terse")) assert "claude" in captured["cmd"][0] assert "--mcp-config" in captured["cmd"] @@ -418,7 +422,7 @@ async def fake_exec(*cmd, **kwargs): assert captured["dir_exists_during_run"] is True # config dir is removed after the run (no leakage between rollouts) assert not Path(captured["config_dir"]).exists() - assert "result" in stdout + assert output_items == [] assert model == "claude-sonnet-4-6" assert metadata["status"] == "completed" @@ -490,9 +494,9 @@ 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, metadata = asyncio.run(agent._run_claude_code("hello")) + output_items, model, metadata = asyncio.run(agent._run_claude_code("hello")) - assert stdout == "" + assert output_items == [] assert killed["called"] is True assert model == "claude-sonnet-4-6" assert metadata["status"] == "incomplete" @@ -576,7 +580,7 @@ def test_writes_rollout_mcp_config_with_session_header(self, tmp_path: Path) -> "mcp": { "server_name": "example_mcp_weather", "url_path": "/mcp", - "headers": {"X-NeMo-Gym-Session-Token": "mcp-session-value"}, + "headers": {"X-NeMo-Gym-Session-Token": "secret-token"}, } }, tmp_path, @@ -587,7 +591,7 @@ def test_writes_rollout_mcp_config_with_session_header(self, tmp_path: Path) -> server = config["mcpServers"]["example_mcp_weather"] assert server["type"] == "http" assert server["url"] == "http://127.0.0.1:8123/mcp" - assert server["headers"]["X-NeMo-Gym-Session-Token"] == "mcp-session-value" + assert server["headers"]["X-NeMo-Gym-Session-Token"] == "secret-token" def test_merges_static_mcp_config_when_metadata_present(self, tmp_path: Path) -> None: static_config = tmp_path / "static_mcp.json" @@ -661,9 +665,11 @@ async def fake_run_claude_code(instruction, system_prompt=None, mcp_config=None, 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."}]}, + _output( + _event( + "assistant", + message={"content": [{"type": "text", "text": "The weather in Paris is sunny and 72 F."}]}, + ) ), "claude-sonnet-4-6", {"status": "completed"}, @@ -720,7 +726,7 @@ async def fake_run_claude_code(instruction, system_prompt=None, mcp_config=None, "headers" ]["X-NeMo-Gym-Session-Token"] return ( - _event("assistant", message={"content": [{"type": "text", "text": "ok"}]}), + _output(_event("assistant", message={"content": [{"type": "text", "text": "ok"}]})), "claude-sonnet-4-6", {"status": "completed"}, ) @@ -811,6 +817,7 @@ class TestParseStreamJson: ("metadata", "returncode", "expected"), [ ({"subtype": "success"}, 0, ("completed", None)), + ({"subtype": "success"}, 7, ("failed", "process_exit_7")), ({"subtype": "error_max_turns", "is_error": True}, 0, ("incomplete", "error_max_turns")), ({"subtype": "error_during_execution", "is_error": True}, 0, ("failed", "error_during_execution")), ({"subtype": "error_max_turns", "is_error": True}, 7, ("incomplete", "error_max_turns")), diff --git a/responses_api_agents/claude_code_agent/tests/test_observability.py b/responses_api_agents/claude_code_agent/tests/test_observability.py index 4dad6c0ad8..20ac04b84b 100644 --- a/responses_api_agents/claude_code_agent/tests/test_observability.py +++ b/responses_api_agents/claude_code_agent/tests/test_observability.py @@ -181,6 +181,10 @@ def test_extracts_nested_tree_model_refs_and_parallel_tool_timing(tmp_path: Path "function_call_output", ] assert all(item.id is None for item in root_invocation.conversation if isinstance(item, NeMoGymFunctionCallOutput)) + [grandchild_result] = [ + item for item in child_invocation.conversation if isinstance(item, NeMoGymFunctionCallOutput) + ] + assert grandchild_result.status == "incomplete" timings = {tool.tool_call_id: tool for tool in _records(bundle, ToolCallObservation)} assert timings["tool-fast"].duration_ms == pytest.approx(1000) @@ -198,7 +202,7 @@ def test_extracts_explicit_compaction_markers(tmp_path: Path) -> None: _event( "root", "user", - "2026-07-22T10:00:00Z", + "bad-timestamp", "summary", message_extra={ "isCompactSummary": True, @@ -224,8 +228,66 @@ def test_extracts_explicit_compaction_markers(tmp_path: Path) -> None: assert compaction.tokens_after == 200 assert compaction.summary == "summary" assert compaction.outcome == "completed" + assert compaction.observed_at == pytest.approx(1784714401) assert compaction.before_model_call.response_id == "msg-before" assert compaction.after_model_call.response_id == "msg-after" + assert compaction.model_calls == [] + codes = {gap.code for gap in bundle.gaps} + assert "compaction_model_call_reference_unavailable" in codes + assert "compaction_timestamp_missing" not in codes + + +@pytest.mark.parametrize( + ("metadata", "expected_outcome", "expected_gap"), + [ + ({"outcome": None, "status": "failed", "tokensBefore": 100, "tokensAfter": 100}, "failed", None), + ({"tokensBefore": 100, "tokensAfter": 80}, "unknown", "compaction_outcome_unavailable"), + ], +) +def test_compaction_metadata_does_not_assume_success( + tmp_path: Path, + metadata: dict, + expected_outcome: str, + expected_gap: str | None, +) -> None: + _write( + tmp_path / "projects" / "work" / "session.jsonl", + _event( + "root", + "system", + "2026-07-22T10:00:00Z", + "", + compact_metadata=metadata, + ), + ) + + bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) + + [compaction] = _records(bundle, ContextCompactionObservation) + assert compaction.outcome == expected_outcome + codes = {gap.code for gap in bundle.gaps} + if expected_gap is None: + assert "compaction_outcome_unavailable" not in codes + else: + assert expected_gap in codes + + +def test_compaction_summary_preserves_explicit_failure(tmp_path: Path) -> None: + _write( + tmp_path / "projects" / "work" / "session.jsonl", + _event( + "root", + "user", + "2026-07-22T10:00:00Z", + "summary", + message_extra={"isCompactSummary": True, "is_error": True}, + ), + ) + + bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) + + [compaction] = _records(bundle, ContextCompactionObservation) + assert compaction.outcome == "failed" def test_malformed_and_incomplete_artifacts_produce_sanitized_gaps(tmp_path: Path) -> None: From 4562b9bebff614715bf24c0a976fe62b2e285bdf Mon Sep 17 00:00:00 2001 From: Michal Bien Date: Tue, 28 Jul 2026 20:22:22 +0200 Subject: [PATCH 14/14] Complete Claude Code compaction observations Signed-off-by: Michal Bien --- .../pages/model-server/model-call-capture.mdx | 9 +- nemo_gym/base_responses_api_model.py | 17 +- nemo_gym/rollout_observability.py | 5 +- responses_api_agents/claude_code_agent/app.py | 51 +++- .../claude_code_agent/observability.py | 232 +++++++++++++++++- .../claude_code_agent/tests/test_app.py | 93 ++++++- .../tests/test_observability.py | 218 ++++++++++++---- 7 files changed, 547 insertions(+), 78 deletions(-) diff --git a/fern/versions/latest/pages/model-server/model-call-capture.mdx b/fern/versions/latest/pages/model-server/model-call-capture.mdx index 261bbd3da5..96ff0b7184 100644 --- a/fern/versions/latest/pages/model-server/model-call-capture.mdx +++ b/fern/versions/latest/pages/model-server/model-call-capture.mdx @@ -126,12 +126,15 @@ invocation's `conversation` contains the ordered, normalized items exposed by th Agent observations and model-call capture are separate evidence. Join an invocation's model-call references by `model_call_id`, or by the exact `(model_ref, response_id)` pair when the harness sees -the protocol response ID. Do not infer ownership from timestamps, text, or list position. The full -model request and response remain in `CaptureStore`; rollout attachments intentionally omit them. +the protocol response ID. An integration may resolve an otherwise hidden call only through a +producer-specific, unique exact match against its retained artifact and the raw capture. Ambiguous +matches remain unowned; timestamps or list position alone are never sufficient. The full model +request and response remain in `CaptureStore`; rollout attachments intentionally omit them. Compaction records distinguish the calls immediately before and after the context change from `model_calls` used to perform the compaction. Compaction calls are exact references to calls owned by -the enclosing invocation; opaque integrations leave them empty rather than infer them. +the enclosing invocation. Integrations without an explicit identifier or a unique exact match leave +them empty and report the gap. Tool timestamps are UTC Unix seconds when the source provides wall-clock time. `duration_ms` is the measured interval, and `timing_source` identifies executor, harness, or artifact-derived timing. diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index a8dc1aab4e..99426db471 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -1287,7 +1287,22 @@ def merge_model_call_capture_into_record( if observations is not None: try: bundle = AgentObservationBundle.model_validate(observations) - record["ng_agent_observations"] = join_model_call_observations(bundle, calls).model_dump(mode="json") + if bundle.source == "claude_code": + try: + from responses_api_agents.claude_code_agent.observability import ( + associate_claude_code_compaction_calls, + ) + + bundle = associate_claude_code_compaction_calls(bundle, calls) + except Exception: + logger.warning( + "Could not associate Claude Code compaction calls for rollout %s.", + rollout_id, + exc_info=True, + ) + bundle.gaps.append(ObservationGap(code="compaction_model_call_join_failed")) + bundle = join_model_call_observations(bundle, calls) + record["ng_agent_observations"] = bundle.model_dump(mode="json") except Exception: logger.warning("Could not join agent observations for rollout %s.", rollout_id, exc_info=True) gaps.append(ObservationGap(code="agent_observation_join_failed")) diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py index 4ca3828a55..e36040b3d0 100644 --- a/nemo_gym/rollout_observability.py +++ b/nemo_gym/rollout_observability.py @@ -131,7 +131,10 @@ class ContextCompactionObservation(ObservationModel): ) model_calls: list[ModelCallRef] = Field( default_factory=list, - description=("Invocation-owned model calls used for compaction, in producer-observed order; never inferred."), + description=( + "Invocation-owned model calls used for compaction, joined by explicit identifiers or a unique " + "producer-specific exact match." + ), ) after_model_call: Optional[ModelCallRef] = Field( default=None, diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index ed634c9085..dd85306962 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -22,6 +22,7 @@ import subprocess import tempfile from asyncio import Semaphore +from contextlib import suppress from pathlib import Path from time import monotonic, time from typing import Any, Callable, Optional @@ -89,6 +90,8 @@ def parse_stream_json(stdout: str) -> tuple[list[Any], dict]: total_output = 0 num_turns: Optional[int] = None result_metadata: dict[str, Any] = {} + compacting_sessions: set[str] = set() + compaction_attempts: list[dict[str, str]] = [] for event in raw_events: etype = event.get("type") @@ -181,9 +184,27 @@ def parse_stream_json(stdout: str) -> tuple[list[Any], dict]: ) ) + elif etype == "system" and event.get("subtype") == "status": + session_id = event.get("session_id") + if not isinstance(session_id, str) or not session_id: + continue + if event.get("status") == "compacting": + compacting_sessions.add(session_id) + continue + compact_result = event.get("compact_result") + if compact_result in {"failed", "success"}: + if compact_result == "failed": + compaction_attempts.append({"invocation_id": session_id, "outcome": "failed"}) + compacting_sessions.discard(session_id) + + compaction_attempts.extend( + {"invocation_id": session_id, "outcome": "unknown"} for session_id in compacting_sessions + ) metadata: dict = {"input_tokens": total_input, "output_tokens": total_output} if num_turns is not None: metadata["num_turns"] = num_turns + if compaction_attempts: + metadata["compaction_attempts"] = compaction_attempts metadata.update(result_metadata) return output_items, metadata @@ -456,18 +477,31 @@ async def _run_claude_code( stderr=asyncio.subprocess.PIPE, env=env, ) + communication = asyncio.create_task(proc.communicate()) try: - stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=self.config.timeout) + stdout, stderr = await asyncio.wait_for( + asyncio.shield(communication), + timeout=self.config.timeout, + ) except asyncio.TimeoutError: - proc.kill() - await proc.communicate() + if proc.returncode is None: + with suppress(ProcessLookupError): + proc.kill() + stdout, _ = await communication LOG.warning("claude-code timed out after %ds", self.config.timeout) - run_metadata = { - "status": "incomplete", - "error_type": "timeout", - "duration_ms": (monotonic() - process_started_at) * 1000, - } + _, run_metadata = parse_stream_json(stdout.decode(errors="replace")) + run_metadata.update( + status="incomplete", + error_type="timeout", + duration_ms=(monotonic() - process_started_at) * 1000, + ) return [], model, run_metadata + except asyncio.CancelledError: + if proc.returncode is None: + with suppress(ProcessLookupError): + proc.kill() + await asyncio.gather(communication, return_exceptions=True) + raise if proc.returncode not in (0, None): LOG.warning("claude-code exited %d: %s", proc.returncode, stderr.decode(errors="replace")[:500]) @@ -638,6 +672,7 @@ def collect(config_dir: Path, run_metadata: dict[str, Any]) -> None: root_status=run_metadata["status"], root_duration_ms=run_metadata.get("duration_ms"), root_error_type=run_metadata.get("error_type"), + compaction_attempts=run_metadata.get("compaction_attempts"), ) if self.config.model_server is None: observations.gaps.append(ObservationGap(code="model_call_ownership_unavailable")) diff --git a/responses_api_agents/claude_code_agent/observability.py b/responses_api_agents/claude_code_agent/observability.py index 3a59d0cf3d..5e90d5d1e7 100644 --- a/responses_api_agents/claude_code_agent/observability.py +++ b/responses_api_agents/claude_code_agent/observability.py @@ -7,10 +7,11 @@ import json import math -from collections import defaultdict +import re +from collections import Counter, defaultdict from datetime import datetime from pathlib import Path -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal from nemo_gym.config_types import ModelServerRef from nemo_gym.openai_utils import ( @@ -32,7 +33,34 @@ ) +if TYPE_CHECKING: + from nemo_gym.base_responses_api_model import ModelCallRecord + + SOURCE = "claude_code" +_COMPACTION_PROMPT_MARKERS = ( + "Your task is to create a detailed summary of this conversation.", + "Your task is to create a detailed summary of the conversation so far", + "Your task is to create a detailed summary of the RECENT portion of the conversation", +) +_COMPACTION_SUMMARY_PREFIX = ( + "This session is being continued from a previous conversation that ran out of context. " + "The summary below covers the earlier portion of the conversation.\n\n" +) +_COMPACTION_SUFFIX_RE = re.compile( + r"(?:\n\nIf you need specific details from before compaction " + r"\(like exact code snippets, error messages, or content you generated\), " + r"read the full transcript at: [^\n]+)?" + r"(?:\n\nRecent messages are preserved verbatim\.)?" + r"(?:\n\nYour REPL VM state has been cleared as part of this compaction\. " + r"Variables defined in REPL calls before this point are no longer accessible " + r"— redefine any you still need\.)?" + r"(?:\n\nContinue the conversation from where it left off without asking the user " + r"any further questions\. Resume directly — do not acknowledge the summary, do not recap " + r'what was happening, do not preface with "I\'ll continue" or similar\. ' + r"Pick up the last task as if the break never happened\.)?" + r"$" +) def _gap(code: str, *, invocation_id: str | None = None, detail: str | None = None) -> ObservationGap: @@ -63,6 +91,53 @@ def _text(value: Any) -> str: return json.dumps(value, ensure_ascii=False, sort_keys=True) +def _is_compaction_request(request: Any) -> bool: + if not isinstance(request, dict): + return False + messages = request.get("messages") + if not isinstance(messages, list) or not messages: + return False + final_message = messages[-1] + return ( + isinstance(final_message, dict) + and final_message.get("role") == "user" + and any(marker in _text(final_message.get("content")) for marker in _COMPACTION_PROMPT_MARKERS) + ) + + +def _messages_text(response: Any) -> str | None: + if not isinstance(response, dict) or not isinstance(response.get("content"), list): + return None + parts = [ + block.get("text") + for block in response["content"] + if isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str) + ] + return "".join(parts) or None + + +def _normalize_compaction_output(text: str) -> str: + normalized = re.sub(r"[\s\S]*?", "", text, count=1) + summary = re.search(r"([\s\S]*?)", normalized) + if summary is not None: + normalized = ( + normalized[: summary.start()] + f"Summary:\n{summary.group(1).strip()}" + normalized[summary.end() :] + ) + return re.sub(r"\n\n+", "\n\n", normalized).strip() + + +def _matches_compaction_summary(summary: str, model_output: str) -> bool: + expected = _COMPACTION_SUMMARY_PREFIX + _normalize_compaction_output(model_output) + return summary.startswith(expected) and _COMPACTION_SUFFIX_RE.fullmatch(summary[len(expected) :]) is not None + + +def _is_compaction_summary(event: dict[str, Any]) -> bool: + message = event.get("message") + return event.get("isCompactSummary") is True or ( + isinstance(message, dict) and message.get("isCompactSummary") is True + ) + + def _status(block: dict[str, Any], result: Any) -> str: if block.get("is_error") is True: return "failed" @@ -122,14 +197,14 @@ def _compaction_outcome( def _compaction(event: dict[str, Any], invocation_id: str) -> ContextCompactionObservation | None: message = event.get("message") - is_summary = isinstance(message, dict) and message.get("isCompactSummary") is True + is_summary = _is_compaction_summary(event) is_boundary = event.get("type") == "system" and event.get("subtype") == "compact_boundary" metadata = _metadata(event) if not is_summary and not is_boundary and not metadata: return None trigger = metadata.get("trigger") - summary = _text(message.get("content")) if is_summary else None + summary = _text(message.get("content")) if is_summary and isinstance(message, dict) else None return ContextCompactionObservation( invocation_id=invocation_id, observed_at=_timestamp(event.get("timestamp")), @@ -253,6 +328,7 @@ def extract_claude_code_observations( root_status: Literal["completed", "failed", "incomplete", "unknown"] = "unknown", root_duration_ms: float | None = None, root_error_type: str | None = None, + compaction_attempts: list[dict[str, str]] | None = None, ) -> AgentObservationBundle: """Extract exact relationships available in one ``CLAUDE_CONFIG_DIR``. @@ -262,9 +338,22 @@ def extract_claude_code_observations( gaps: list[ObservationGap] = [] raw_events = _read_events(Path(config_dir), gaps) - if not raw_events: + attempts = [ + ContextCompactionObservation( + invocation_id=attempt["invocation_id"], + outcome=attempt["outcome"], + ) + for attempt in compaction_attempts or [] + if isinstance(attempt, dict) + and isinstance(attempt.get("invocation_id"), str) + and attempt.get("invocation_id") + and attempt.get("outcome") in {"failed", "aborted", "unknown"} + ] + if not raw_events and not attempts: gaps.append(_gap("agent_transcript_unavailable")) return AgentObservationBundle(source=SOURCE, gaps=gaps) + if not raw_events: + gaps.append(_gap("agent_transcript_unavailable")) events_by_invocation: dict[str, list[tuple[int, dict[str, Any]]]] = defaultdict(list) first_seen: dict[str, int] = {} @@ -277,6 +366,15 @@ def extract_claude_code_observations( first_seen.setdefault(invocation_id, ordinal) if isinstance(agent_id, str) and agent_id: agent_invocations.add(invocation_id) + for index, attempt in enumerate(attempts, start=len(raw_events)): + events_by_invocation.setdefault(attempt.invocation_id, []) + first_seen.setdefault(attempt.invocation_id, index) + gaps.extend( + ( + _gap("compaction_before_model_call_unavailable", invocation_id=attempt.invocation_id), + _gap("compaction_after_model_call_unavailable", invocation_id=attempt.invocation_id), + ) + ) starts: dict[tuple[str, str], list[tuple[float | None, str]]] = defaultdict(list) finishes: dict[tuple[str, str], list[tuple[float | None, str]]] = defaultdict(list) @@ -284,7 +382,7 @@ def extract_claude_code_observations( ambiguous_parents: set[str] = set() conversations: dict[str, list[Any]] = defaultdict(list) model_calls: dict[str, list[ModelCallRef]] = defaultdict(list) - compactions: list[ContextCompactionObservation] = [] + compactions = attempts for invocation_id, entries in events_by_invocation.items(): items = conversations[invocation_id] @@ -300,8 +398,7 @@ def add_gap(code: str, detail: str | None = None) -> None: for entry_index, (ordinal, event) in enumerate(entries): compaction = _compaction(event, invocation_id) if compaction is not None: - message = event.get("message") - is_summary = isinstance(message, dict) and message.get("isCompactSummary") is True + is_summary = _is_compaction_summary(event) if ( previous_compaction is not None and previous_compaction[0] + 1 == entry_index @@ -564,3 +661,122 @@ def add_tool_gap(code: str) -> None: ], gaps=gaps, ) + + +def associate_claude_code_compaction_calls( + bundle: AgentObservationBundle, + calls: list[ModelCallRecord], +) -> AgentObservationBundle: + """Associate hidden summary calls only when Claude's persisted summary matches exactly.""" + if bundle.source != SOURCE: + return bundle + + result = bundle.model_copy() + result.records = [ + record.model_copy(update={"model_calls": list(record.model_calls)}) + if isinstance(record, (AgentInvocation, ContextCompactionObservation)) + else record + for record in bundle.records + ] + result.gaps = list(bundle.gaps) + invocations = {record.invocation_id: record for record in result.records if isinstance(record, AgentInvocation)} + compactions = [record for record in result.records if isinstance(record, ContextCompactionObservation)] + + def ref_matches_call(reference: ModelCallRef, call: ModelCallRecord) -> bool: + if reference.model_call_id: + return reference.model_call_id == call.model_call_id + return reference.model_ref == call.model_ref and reference.response_id == call.response_id + + def resolve_call_index(reference: ModelCallRef | None) -> int | None: + if reference is None: + return None + matches = [call.call_index for call in calls if ref_matches_call(reference, call)] + return matches[0] if len(matches) == 1 else None + + owned = { + index + for invocation in invocations.values() + for reference in invocation.model_calls + for index, call in enumerate(calls) + if ref_matches_call(reference, call) + } + candidates: list[list[int]] = [] + for compaction in compactions: + before_index = resolve_call_index(compaction.before_model_call) + after_index = resolve_call_index(compaction.after_model_call) + model_refs = [ + reference.model_ref + for reference in (compaction.before_model_call, compaction.after_model_call) + if reference is not None and reference.model_ref is not None + ] + candidates.append( + [ + index + for index, call in enumerate(calls) + if index not in owned + and compaction.outcome == "completed" + and not compaction.model_calls + and compaction.summary is not None + and call.dialect == "messages" + and call.status_code is not None + and 200 <= call.status_code < 300 + and call.error_category is None + and (call.model_call_id is not None or (call.model_ref is not None and call.response_id is not None)) + and (not model_refs or call.model_ref in model_refs) + and _is_compaction_request(call.request) + and (response_text := _messages_text(call.response)) is not None + and _matches_compaction_summary(compaction.summary, response_text) + and ( + compaction.before_model_call is None + or (before_index is not None and call.call_index > before_index) + ) + and ( + compaction.after_model_call is None or (after_index is not None and call.call_index < after_index) + ) + ] + ) + + candidate_counts = Counter(index for compaction_candidates in candidates for index in compaction_candidates) + for compaction, compaction_candidates in zip(compactions, candidates): + ambiguous = len(compaction_candidates) > 1 or ( + len(compaction_candidates) == 1 and candidate_counts[compaction_candidates[0]] > 1 + ) + if ambiguous: + result.gaps.append( + _gap( + "compaction_model_call_match_ambiguous", + invocation_id=compaction.invocation_id, + detail=f"candidate_count={len(compaction_candidates)}", + ) + ) + if ( + len(compaction_candidates) != 1 + or candidate_counts[compaction_candidates[0]] != 1 + or compaction.invocation_id not in invocations + ): + continue + call = calls[compaction_candidates[0]] + reference = ModelCallRef( + model_call_id=call.model_call_id, + model_ref=call.model_ref, + response_id=call.response_id, + ) + compaction.model_calls = [reference] + + invocation = invocations[compaction.invocation_id] + insertion_index = len(invocation.model_calls) + for index, existing in enumerate(invocation.model_calls): + if compaction.after_model_call is not None and existing == compaction.after_model_call: + insertion_index = index + break + if compaction.before_model_call is not None and existing == compaction.before_model_call: + insertion_index = index + 1 + invocation.model_calls.insert(insertion_index, reference) + + unresolved = {compaction.invocation_id for compaction in compactions if not compaction.model_calls} + result.gaps = [ + gap + for gap in result.gaps + if gap.code != "compaction_model_call_reference_unavailable" or gap.invocation_id in unresolved + ] + return result 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 437b2f1c85..1625d753ee 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -32,7 +32,7 @@ NeMoGymResponseFunctionToolCall, NeMoGymResponseOutputMessage, ) -from nemo_gym.rollout_observability import AgentInvocation +from nemo_gym.rollout_observability import AgentInvocation, ContextCompactionObservation from nemo_gym.server_utils import ServerClient from responses_api_agents.claude_code_agent.app import ( ClaudeCodeAgent, @@ -352,6 +352,7 @@ async def run_claude_code(*args, observation_collector=None, **kwargs): "num_turns": 7, "subtype": "success", "is_error": False, + "compaction_attempts": [{"invocation_id": "session-1", "outcome": "failed"}], } observation_collector(tmp_path, run_metadata) return ( @@ -379,9 +380,19 @@ async def run_claude_code(*args, observation_collector=None, **kwargs): assert invocation.status == "completed" assert invocation.duration_ms == 123 assert invocation.model_calls[0].response_id == "msg-1" + compaction = next( + record for record in observations.records if isinstance(record, ContextCompactionObservation) + ) + assert compaction.invocation_id == "session-1" + assert compaction.outcome == "failed" assert result.turns_used == 1 assert result.finished_naturally is True - assert "no_sandbox_runtime" in {gap.code for gap in observations.gaps} + assert { + "compaction_before_model_call_unavailable", + "compaction_after_model_call_unavailable", + "compaction_model_call_reference_unavailable", + "no_sandbox_runtime", + } <= {gap.code for gap in observations.gaps} class TestRunClaudeCode: @@ -471,22 +482,25 @@ def test_bad_skills_path_does_not_leak_config_dir(self, tmp_path: Path) -> None: def test_timeout_returns_empty(self, tmp_path: Path) -> None: agent = _make_agent(timeout=1) - killed = {"called": False} + state = {"killed": False, "communicate_calls": 0} class SlowProc: returncode = None def kill(self): - killed["called"] = True + state["killed"] = True async def communicate(self): - return b"", b"" + state["communicate_calls"] += 1 + return ( + _event("system", subtype="status", status="compacting", session_id="session-1").encode(), + b"", + ) async def fake_exec(*cmd, **kwargs): return SlowProc() async def fake_wait_for(coro, timeout): - coro.close() # avoid un-awaited coroutine warning raise asyncio.TimeoutError with ( @@ -497,11 +511,56 @@ async def fake_wait_for(coro, timeout): output_items, model, metadata = asyncio.run(agent._run_claude_code("hello")) assert output_items == [] - assert killed["called"] is True + assert state == {"killed": True, "communicate_calls": 1} assert model == "claude-sonnet-4-6" assert metadata["status"] == "incomplete" assert metadata["error_type"] == "timeout" assert metadata["duration_ms"] >= 0 + assert metadata["compaction_attempts"] == [{"invocation_id": "session-1", "outcome": "unknown"}] + + def test_cancellation_stops_process_before_observation_cleanup(self, tmp_path: Path) -> None: + agent = _make_agent() + state: list[str] = [] + + async def run() -> None: + communicating = asyncio.Event() + stopped = asyncio.Event() + + class SlowProc: + returncode = None + + def kill(self): + state.append("kill") + self.returncode = -9 + stopped.set() + + async def communicate(self): + state.append("communicate") + communicating.set() + await stopped.wait() + state.append("stopped") + return b"", b"" + + async def fake_exec(*cmd, **kwargs): + return SlowProc() + + def collect(config_dir: Path, metadata: dict) -> None: + assert config_dir.exists() + state.append("collect") + + 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), + ): + task = asyncio.create_task(agent._run_claude_code("hello", observation_collector=collect)) + await communicating.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + + assert state == ["communicate", "kill", "stopped", "collect"] def test_collects_observations_before_cleanup(self, tmp_path: Path) -> None: agent = _make_agent() @@ -939,6 +998,26 @@ def test_num_turns_absent_when_no_result_event(self) -> None: _, usage = parse_stream_json(assistant) assert "num_turns" not in usage + def test_compaction_status_events_report_failed_and_unknown_attempts(self) -> None: + events = [ + _event("system", subtype="status", status="compacting", session_id="failed"), + _event("system", subtype="status", status="requesting", session_id="failed"), + _event("system", subtype="status", compact_result="failed", session_id="failed"), + _event("system", subtype="status", compact_result="failed", session_id="failed-without-opener"), + _event("system", subtype="status", status="compacting", session_id="success"), + _event("system", subtype="status", status="requesting", session_id="success"), + _event("system", subtype="status", compact_result="success", session_id="success"), + _event("system", subtype="status", status="compacting", session_id="open"), + ] + + _, metadata = parse_stream_json("\n".join(events)) + + assert metadata["compaction_attempts"] == [ + {"invocation_id": "failed", "outcome": "failed"}, + {"invocation_id": "failed-without-opener", "outcome": "failed"}, + {"invocation_id": "open", "outcome": "unknown"}, + ] + class TestConfigYaml: def test_module_parses(self) -> None: diff --git a/responses_api_agents/claude_code_agent/tests/test_observability.py b/responses_api_agents/claude_code_agent/tests/test_observability.py index 20ac04b84b..e080c76199 100644 --- a/responses_api_agents/claude_code_agent/tests/test_observability.py +++ b/responses_api_agents/claude_code_agent/tests/test_observability.py @@ -4,18 +4,29 @@ import json from pathlib import Path from typing import TypeVar +from unittest.mock import patch import pytest +from nemo_gym.base_responses_api_model import ( + CaptureStore, + ModelCallRecord, + merge_model_call_capture_into_record, +) from nemo_gym.config_types import ModelServerRef from nemo_gym.openai_utils import NeMoGymFunctionCallOutput from nemo_gym.rollout_observability import ( AgentInvocation, AgentObservationBundle, ContextCompactionObservation, + ModelCallRef, + ObservationGap, ToolCallObservation, ) -from responses_api_agents.claude_code_agent.observability import extract_claude_code_observations +from responses_api_agents.claude_code_agent.observability import ( + associate_claude_code_compaction_calls, + extract_claude_code_observations, +) MODEL_REF = ModelServerRef(type="responses_api_models", name="policy") @@ -199,23 +210,20 @@ def test_extracts_explicit_compaction_markers(tmp_path: Path) -> None: _write( tmp_path / "projects" / "work" / "session.jsonl", _assistant("root", "2026-07-22T09:59:59Z", "msg-before", {"type": "text", "text": "before"}), - _event( - "root", - "user", - "bad-timestamp", - "summary", - message_extra={ - "isCompactSummary": True, - "compactMetadata": {"tokensBefore": 1000, "tokensAfter": 200, "trigger": "auto"}, - }, - ), _event( "root", "system", "2026-07-22T10:00:01Z", "", subtype="compact_boundary", - compact_metadata={"preTokens": 900, "postTokens": 180}, + ), + _event( + "root", + "user", + "bad-timestamp", + "summary", + isCompactSummary=True, + compactMetadata={"tokensBefore": 1000, "tokensAfter": 200, "trigger": "auto"}, ), _assistant("root", "2026-07-22T10:00:02Z", "msg-after", {"type": "text", "text": "after"}), ) @@ -237,19 +245,7 @@ def test_extracts_explicit_compaction_markers(tmp_path: Path) -> None: assert "compaction_timestamp_missing" not in codes -@pytest.mark.parametrize( - ("metadata", "expected_outcome", "expected_gap"), - [ - ({"outcome": None, "status": "failed", "tokensBefore": 100, "tokensAfter": 100}, "failed", None), - ({"tokensBefore": 100, "tokensAfter": 80}, "unknown", "compaction_outcome_unavailable"), - ], -) -def test_compaction_metadata_does_not_assume_success( - tmp_path: Path, - metadata: dict, - expected_outcome: str, - expected_gap: str | None, -) -> None: +def test_compaction_metadata_does_not_assume_success(tmp_path: Path) -> None: _write( tmp_path / "projects" / "work" / "session.jsonl", _event( @@ -257,37 +253,15 @@ def test_compaction_metadata_does_not_assume_success( "system", "2026-07-22T10:00:00Z", "", - compact_metadata=metadata, - ), - ) - - bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) - - [compaction] = _records(bundle, ContextCompactionObservation) - assert compaction.outcome == expected_outcome - codes = {gap.code for gap in bundle.gaps} - if expected_gap is None: - assert "compaction_outcome_unavailable" not in codes - else: - assert expected_gap in codes - - -def test_compaction_summary_preserves_explicit_failure(tmp_path: Path) -> None: - _write( - tmp_path / "projects" / "work" / "session.jsonl", - _event( - "root", - "user", - "2026-07-22T10:00:00Z", - "summary", - message_extra={"isCompactSummary": True, "is_error": True}, + compact_metadata={"tokensBefore": 100, "tokensAfter": 80}, ), ) bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) [compaction] = _records(bundle, ContextCompactionObservation) - assert compaction.outcome == "failed" + assert compaction.outcome == "unknown" + assert "compaction_outcome_unavailable" in {gap.code for gap in bundle.gaps} def test_malformed_and_incomplete_artifacts_produce_sanitized_gaps(tmp_path: Path) -> None: @@ -404,3 +378,147 @@ def test_reports_missing_response_id_and_unsupported_content_blocks(tmp_path: Pa assert "unsupported_assistant_content_block" in codes assert "unsupported_user_content_block" in codes assert _records(bundle, AgentInvocation)[0].model_calls == [] + + +def _compaction_bundle() -> AgentObservationBundle: + before = ModelCallRef(model_ref=MODEL_REF, response_id="msg-before") + after = ModelCallRef(model_ref=MODEL_REF, response_id="msg-after") + return AgentObservationBundle( + source="claude_code", + records=[ + AgentInvocation(invocation_id="root", model_calls=[before, after]), + ContextCompactionObservation( + invocation_id="root", + outcome="completed", + summary=( + "This session is being continued from a previous conversation that ran out of context. " + "The summary below covers the earlier portion of the conversation.\n\n" + "Summary:\nKeep this.\n\nRecent messages are preserved verbatim." + ), + before_model_call=before, + after_model_call=after, + ), + ], + gaps=[ObservationGap(code="compaction_model_call_reference_unavailable", invocation_id="root")], + ) + + +def _captured_call(call_id: str, response_id: str, *, compact: bool = False) -> dict: + return { + "model_call_id": call_id, + "response_id": response_id, + "dialect": "messages", + "status_code": 200, + "model_ref": MODEL_REF.model_dump(mode="json"), + "request": { + "messages": [ + { + "role": "user", + "content": ( + "Your task is to create a detailed summary of this conversation." if compact else "continue" + ), + } + ] + }, + "response": { + "id": response_id, + "content": [ + { + "type": "text", + "text": ("privateKeep this." if compact else "ok"), + } + ], + }, + } + + +def _rollout_record(tmp_path: Path, *calls: dict) -> dict: + store = CaptureStore(tmp_path) + for call in calls: + store.record("0-0", call) + return { + "_ng_task_index": 0, + "_ng_rollout_index": 0, + "ng_agent_observations": _compaction_bundle().model_dump(mode="json"), + } + + +def test_merge_correlates_unique_compaction_call_into_rollout(tmp_path: Path) -> None: + record = _rollout_record( + tmp_path, + _captured_call("call-before", "msg-before"), + _captured_call("call-compact", "msg-compact", compact=True), + _captured_call("call-after", "msg-after"), + ) + + merge_model_call_capture_into_record(record, [tmp_path]) + + observations = AgentObservationBundle.model_validate(record["ng_agent_observations"]) + [invocation] = _records(observations, AgentInvocation) + [compaction] = _records(observations, ContextCompactionObservation) + assert [reference.model_call_id for reference in invocation.model_calls] == [ + "call-before", + "call-compact", + "call-after", + ] + assert [reference.model_call_id for reference in compaction.model_calls] == ["call-compact"] + assert "compaction_model_call_reference_unavailable" not in {gap.code for gap in observations.gaps} + + +def test_compaction_resolver_failure_preserves_generic_model_call_join(tmp_path: Path) -> None: + record = _rollout_record( + tmp_path, + _captured_call("call-before", "msg-before"), + _captured_call("call-after", "msg-after"), + ) + + with patch( + "responses_api_agents.claude_code_agent.observability.associate_claude_code_compaction_calls", + side_effect=RuntimeError, + ): + merge_model_call_capture_into_record(record, [tmp_path]) + + observations = AgentObservationBundle.model_validate(record["ng_agent_observations"]) + [invocation] = _records(observations, AgentInvocation) + assert [reference.model_call_id for reference in invocation.model_calls] == ["call-before", "call-after"] + assert "compaction_model_call_join_failed" in {gap.code for gap in observations.gaps} + + +def test_compaction_call_correlation_rejects_ambiguous_matches() -> None: + calls = [ + ModelCallRecord.model_validate(_captured_call("call-before", "msg-before") | {"call_index": 0}), + ModelCallRecord.model_validate(_captured_call("call-1", "msg-1", compact=True) | {"call_index": 1}), + ModelCallRecord.model_validate(_captured_call("call-2", "msg-2", compact=True) | {"call_index": 2}), + ModelCallRecord.model_validate(_captured_call("call-after", "msg-after") | {"call_index": 3}), + ] + + associated = associate_claude_code_compaction_calls(_compaction_bundle(), calls) + + [compaction] = _records(associated, ContextCompactionObservation) + assert compaction.model_calls == [] + assert "compaction_model_call_reference_unavailable" in {gap.code for gap in associated.gaps} + assert "compaction_model_call_match_ambiguous" in {gap.code for gap in associated.gaps} + + +@pytest.mark.parametrize("case", ["marker_in_history", "failed_call", "outside_boundaries"]) +def test_compaction_call_correlation_rejects_inexact_matches(case: str) -> None: + calls = [ + ModelCallRecord.model_validate(_captured_call("call-before", "msg-before") | {"call_index": 0}), + ModelCallRecord.model_validate( + _captured_call("call-compact", "msg-compact", compact=True) | {"call_index": 1} + ), + ModelCallRecord.model_validate(_captured_call("call-after", "msg-after") | {"call_index": 2}), + ] + compact_call = calls[1] + if case == "marker_in_history": + compact_call.request["messages"].append({"role": "user", "content": "continue"}) + elif case == "failed_call": + compact_call.status_code = 500 + else: + compact_call.call_index = 3 + + associated = associate_claude_code_compaction_calls(_compaction_bundle(), calls) + + [compaction] = _records(associated, ContextCompactionObservation) + assert compaction.model_calls == [] + assert "compaction_model_call_reference_unavailable" in {gap.code for gap in associated.gaps}