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/README.md b/responses_api_agents/claude_code_agent/README.md index 573c5b394b..3fbdc600ac 100644 --- a/responses_api_agents/claude_code_agent/README.md +++ b/responses_api_agents/claude_code_agent/README.md @@ -196,5 +196,6 @@ 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` 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 0aaf459b35..dd85306962 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -22,13 +22,14 @@ import subprocess import tempfile from asyncio import Semaphore +from contextlib import suppress from pathlib import Path -from time import time -from typing import Any, Optional +from time import monotonic, time +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 +47,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 @@ -86,6 +89,9 @@ 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] = {} + compacting_sessions: set[str] = set() + compaction_attempts: list[dict[str, str]] = [] for event in raw_events: etype = event.get("type") @@ -97,6 +103,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", {}) @@ -171,12 +184,44 @@ 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 +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 returncode not in (0, None): + return "failed", f"process_exit_{returncode}" + if subtype == "success": + return "completed", None + 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) @@ -243,6 +288,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,8 +429,9 @@ async def _run_claude_code( mcp_config: Optional[str] = None, skills_path: Optional[str] = None, rollout_id: Optional[str] = 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[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. @@ -393,6 +442,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. @@ -420,28 +470,60 @@ 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, 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) - return "", model + _, 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]) - 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]) + 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 output_items, model, run_metadata finally: if claude_config_dir is not None: - shutil.rmtree(claude_config_dir, ignore_errors=True) + try: + if observation_collector is not None: + await asyncio.to_thread(observation_collector, claude_config_dir, run_metadata) + 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 +591,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, dict[str, Any]], None]] = None, ) -> NeMoGymResponse: body = body.model_copy(deep=True) if isinstance(body.input, str): @@ -518,14 +601,14 @@ 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( + output_items, model_name, run_metadata = await self._run_claude_code( user_message, system_prompt=system_prompt, mcp_config=mcp_config, skills_path=skills_path, rollout_id=rollout_id, + observation_collector=observation_collector, ) - output_items, usage = parse_stream_json(stdout) if not any( getattr(item, "type", None) == "message" and getattr(item, "role", None) == "assistant" @@ -542,8 +625,8 @@ async def _create_response( ) ) - input_tokens = usage.get("input_tokens", 0) - output_tokens = usage.get("output_tokens", 0) + input_tokens = run_metadata.get("input_tokens", 0) + output_tokens = run_metadata.get("output_tokens", 0) return NeMoGymResponse( id=f"resp_{uuid4().hex}", @@ -568,7 +651,52 @@ async def responses( request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body(), ) -> NeMoGymResponse: - return await self._create_response(body) + return await self._create_response(body, rollout_id=request.path_params.get("rollout_id")) + + async def _create_episode( + self, + 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, run_metadata: dict[str, Any]) -> None: + nonlocal observations + try: + 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"), + compaction_attempts=run_metadata.get("compaction_attempts"), + ) + 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")], + ) + observations.gaps.append(ObservationGap(code="no_sandbox_runtime")) + return AgentEpisode(response=response, observations=observations) async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> ClaudeCodeAgentVerifyResponse: async with self.sem: @@ -593,12 +721,21 @@ 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._create_episode( + 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( @@ -619,9 +756,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..5e90d5d1e7 --- /dev/null +++ b/responses_api_agents/claude_code_agent/observability.py @@ -0,0 +1,782 @@ +# 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 +import re +from collections import Counter, defaultdict +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal + +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, +) + + +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: + 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 _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" + if isinstance(result, dict): + if result.get("interrupted") is True: + return "incomplete" + value = result.get("status") + 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. + 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) 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 = _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 and isinstance(message, dict) 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=_compaction_outcome( + event, + metadata, + has_completion_marker=is_summary or is_boundary, + ), + 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(block: dict[str, Any], status: str) -> NeMoGymFunctionCallOutput: + return NeMoGymFunctionCallOutput( + call_id=block["tool_use_id"], + output=_text(block.get("content")), + status="completed" if status == "completed" else "incomplete", + ) + + +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 _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, + *, + 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, + compaction_attempts: list[dict[str, str]] | 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) + 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] = {} + 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) + 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) + 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 = attempts + + for invocation_id, entries in events_by_invocation.items(): + 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: + is_summary = _is_compaction_summary(event) + 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 ("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 + 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) + 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(block, tool_status)) + 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: + 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") + + 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]) + ): + 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: + 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, + 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) + 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=status, + duration_ms=duration_ms, + error_type=error_type, + model_calls=model_calls[invocation_id], + conversation=conversations[invocation_id], + ) + + ordered_ids = sorted(all_invocation_ids, key=lambda invocation_id: (first_seen[invocation_id], invocation_id)) + + return AgentObservationBundle( + source=SOURCE, + records=[ + *(invocations_by_id[invocation_id] for invocation_id in ordered_ids), + *tool_calls, + *compactions, + ], + 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 8209eafbb6..1625d753ee 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, ContextCompactionObservation from nemo_gym.server_utils import ServerClient from responses_api_agents.claude_code_agent.app import ( ClaudeCodeAgent, @@ -39,8 +41,10 @@ ModelServerRef, ResourcesServerRef, _extract_instruction, + _invocation_outcome, 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: @@ -73,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 @@ -260,26 +268,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", @@ -289,7 +293,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": []}, @@ -303,12 +309,90 @@ 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": []}}) - 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") + assert result.turns_used == 1 + assert result.finished_naturally is True + + +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"}], + }, + } + ) + ) + run_metadata = { + "status": "completed", + "duration_ms": 123.0, + "num_turns": 7, + "subtype": "success", + "is_error": False, + "compaction_attempts": [{"invocation_id": "session-1", "outcome": "failed"}], + } + observation_collector(tmp_path, run_metadata) + return ( + _output(_event("assistant", message={"content": [{"type": "text", "text": "done"}]})), + "model", + run_metadata, + ) + + 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.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 { + "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: @@ -320,7 +404,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"] @@ -336,7 +424,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")) + 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"] @@ -345,8 +433,9 @@ 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" def test_skills_staged_and_bare_dropped(self, tmp_path: Path) -> None: skills_dir = _write_skill_dir(tmp_path) @@ -393,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 ( @@ -416,11 +508,111 @@ 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")) + output_items, model, metadata = asyncio.run(agent._run_claude_code("hello")) - assert stdout == "" - assert killed["called"] is True + assert output_items == [] + 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() + captured: dict = {} + event_loop_thread = threading.get_ident() + + class FakeProc: + returncode = 0 + + async def communicate(self): + 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"]) + 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, run_metadata: dict) -> None: + captured["collector_thread"] = threading.get_ident() + 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), + 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 invocation.status == "completed" + assert captured["collector_thread"] != event_loop_thread + assert not captured["config_dir"].exists() class TestRolloutMCPConfig: @@ -531,10 +723,16 @@ 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 ( + _output( + _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) @@ -586,7 +784,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 ( + _output(_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) @@ -670,6 +872,21 @@ def test_empty(self) -> None: class TestParseStreamJson: + @pytest.mark.parametrize( + ("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")), + ({}, 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}}) @@ -758,15 +975,49 @@ 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"}]) _, 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 new file mode 100644 index 0000000000..e080c76199 --- /dev/null +++ b/responses_api_agents/claude_code_agent/tests/test_observability.py @@ -0,0 +1,524 @@ +# 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 +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 ( + associate_claude_code_compaction_calls, + 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, + status="timeout", + ), + ) + _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, + 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 + 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.values() 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", + ] + 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) + 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 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", + "system", + "2026-07-22T10:00:01Z", + "", + subtype="compact_boundary", + ), + _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"}), + ) + + 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.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 + + +def test_compaction_metadata_does_not_assume_success(tmp_path: Path) -> None: + _write( + tmp_path / "projects" / "work" / "session.jsonl", + _event( + "root", + "system", + "2026-07-22T10:00:00Z", + "", + compact_metadata={"tokensBefore": 100, "tokensAfter": 80}, + ), + ) + + bundle = extract_claude_code_observations(tmp_path, model_ref=MODEL_REF) + + [compaction] = _records(bundle, ContextCompactionObservation) + 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: + 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": {}}, + ), + _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( + 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", + "tool_timing_invalid", + } <= codes + 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", + _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 == [] + + +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}