fix(pi): retain failed model-attempt ownership

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 ""
         )

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")

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"]

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 = [