diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index 16a63f1ea95..064705a5379 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -499,10 +499,10 @@ def _turns_from_model_calls( turn_counts: Counter = Counter() tool_counts: Counter = Counter() for call in model_calls: - if call.response is None: - # A call that returned nothing is not a model decision. - continue metadata = call.response_metadata + if call.response is None or (metadata.status_code is not None and metadata.status_code >= 400): + # HTTP error payloads are retained attempts, not model decisions. + continue invocation_id = invocation_by_call_id.get(call.model_call_id or "") or invocation_by_response_id.get( metadata.response_id or "" ) diff --git a/responses_api_agents/pi_agent/app.py b/responses_api_agents/pi_agent/app.py index a6e2232f092..bc4a49acd28 100644 --- a/responses_api_agents/pi_agent/app.py +++ b/responses_api_agents/pi_agent/app.py @@ -178,6 +178,7 @@ def _build_pi_observations( conversation: list[Any], *, transcript_available: bool = True, + capture_correlated: bool = False, ) -> AgentObservationBundle: def gap(code: str, detail: Optional[str] = None) -> ObservationGap: return ObservationGap(code=code, invocation_id=invocation_id, detail=detail) @@ -316,7 +317,7 @@ def gap(code: str, detail: Optional[str] = None) -> ObservationGap: if outcome == "unknown": gaps.append(gap("compaction_outcome_unavailable")) compaction_start = None - if not model_calls or model_call_join_missing: + if not capture_correlated and (not model_calls or model_call_join_missing): gaps.append(gap("model_call_ownership_unavailable")) if invocation_status == "unknown": gaps.append(gap("invocation_outcome_unavailable")) @@ -520,6 +521,10 @@ def _build_models_config(self, rollout_id: Optional[str] = None) -> dict[str, An } ], } + if rollout_id is not None: + # Response IDs do not exist for HTTP errors; correlate every retry + # with the same invocation used by _build_pi_observations. + providers["nemo"]["headers"] = {"x-session-id": rollout_id} return config async def _run_pi( @@ -685,6 +690,7 @@ async def _create_episode( self.config.model_server, [*conversation_input, *observed_output_items], transcript_available=bool(observed_output_items), + capture_correlated=rollout_id is not None and self.config.model_server is not None, ) except Exception: LOG.exception("failed to build Pi observations") diff --git a/responses_api_agents/pi_agent/tests/test_app.py b/responses_api_agents/pi_agent/tests/test_app.py index 6391186cc88..cb6d474e33c 100644 --- a/responses_api_agents/pi_agent/tests/test_app.py +++ b/responses_api_agents/pi_agent/tests/test_app.py @@ -245,6 +245,7 @@ def test_builds_pi_provider_config(self) -> None: provider = config["providers"]["nemo"] assert agent._effective_model() == "nemo/Qwen3.6-35B-A3B" assert provider["baseUrl"] == "http://model/ng-rollout/1-2/v1" + assert provider["headers"] == {"x-session-id": "1-2"} assert provider["models"][0]["id"] == "Qwen3.6-35B-A3B" assert provider["models"][0]["maxTokens"] == 131072 assert agent.config.models_config == models_config @@ -618,3 +619,22 @@ async def test_mcp_setup_exit_is_request_failure_and_cleans_workspace(tmp_path, with pytest.raises(RuntimeError, match="Required Gym MCP"): await agent._run_pi("task", None, collect_observations=collect_observations) assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("correlated", [False, True]) +def test_failed_model_attempt_uses_transport_ownership_when_correlated(correlated: bool) -> None: + bundle = _build_pi_observations( + [(1.0, {"type": "message_end", "message": {"role": "assistant", "stopReason": "error"}})], + "rollout", + ModelServerRef(type="responses_api_models", name="policy"), + [], + capture_correlated=correlated, + ) + assert any(g.code == "model_call_ownership_unavailable" for g in bundle.gaps) is not correlated + assert _records(bundle, AgentInvocation)[0].invocation_id == "rollout" + + +def test_uncorrelated_pi_provider_has_no_invocation_header() -> None: + agent = _make_agent(model_server=ModelServerRef(type="responses_api_models", name="policy")) + with patch.object(PiAgent, "resolve_model_base_url", return_value="http://model/v1"): + assert "headers" not in agent._build_models_config()["providers"]["nemo"] diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py index 47915214a4d..bde1a9a96e4 100644 --- a/tests/unit_tests/test_rollout_collection.py +++ b/tests/unit_tests/test_rollout_collection.py @@ -8201,6 +8201,19 @@ def test_a_call_that_returned_nothing_is_not_a_turn(self) -> None: assert [(turn.model_calls[0].model_call_id, turn.turn_no) for turn in turns] == [("answered", 1)] + @pytest.mark.parametrize("status", [400, 429, 500]) + def test_http_error_payload_is_retained_without_becoming_a_turn(self, status: int) -> None: + failed = self._call("failed", {"error": {"message": "try again"}}, started_at=1.0) + failed.response_metadata.status_code = status + answered = self._call("answered", {"output": []}, started_at=2.0) + invocation = self._invocation("root", ["failed", "answered"]) + turns = nemo_gym.rollout_collection._turns_from_model_calls( + "task", "rollout", [invocation], [failed, answered], None + ) + assert [(turn.model_calls[0].model_call_id, turn.turn_no) for turn in turns] == [("answered", 1)] + assert [ref.model_call_id for ref in invocation.model_calls] == ["failed", "answered"] + assert failed.response == {"error": {"message": "try again"}} + @pytest.mark.parametrize("invocation_count", [1, 2]) def test_a_call_no_invocation_references_is_skipped(self, invocation_count: int) -> None: calls = [