From 1d49f0c917c90ef65e3180f571d1a71c8d6bd503 Mon Sep 17 00:00:00 2001
From: Alex Fournier <afournier@nvidia.com>
Date: Wed, 22 Jul 2026 10:32:59 -0700
Subject: [PATCH] fix(runtime): finalize Relay iteration summaries

Signed-off-by: Alex Fournier <afournier@nvidia.com>
---
 agent/chat_completion_helpers.py  | 11 +++++++++++
 tests/run_agent/test_run_agent.py | 15 +++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py
index 5ddaf93e5648c..8836160b6b83f 100644
--- a/agent/chat_completion_helpers.py
+++ b/agent/chat_completion_helpers.py
@@ -1912,6 +1912,7 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str:
     print(f"⚠️  Reached maximum iterations ({agent.max_iterations}). Requesting summary...")
 
     summary_api_request_id = f"iteration-summary:{uuid.uuid4()}"
+    summary_call_outcome = "failed"
 
     def _managed_summary_call(request, callback, *, retry_count: int):
         from agent import relay_llm
@@ -1929,6 +1930,7 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str:
                 "call_role": "iteration_summary",
                 "retry_count": retry_count,
             },
+            defer_logical_completion=True,
         )
 
     summary_request = (
@@ -2133,6 +2135,7 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str:
             if "<think>" in final_response:
                 final_response = re.sub(r'<think>.*?</think>\s*', '', final_response, flags=re.DOTALL).strip()
             if final_response:
+                summary_call_outcome = "success"
                 messages.append({"role": "assistant", "content": final_response})
             else:
                 final_response = "I reached the iteration limit and couldn't generate a summary."
@@ -2187,6 +2190,7 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str:
                 if "<think>" in final_response:
                     final_response = re.sub(r'<think>.*?</think>\s*', '', final_response, flags=re.DOTALL).strip()
                 if final_response:
+                    summary_call_outcome = "success"
                     messages.append({"role": "assistant", "content": final_response})
                 else:
                     final_response = "I reached the iteration limit and couldn't generate a summary."
@@ -2196,6 +2200,13 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str:
     except Exception as e:
         logger.warning(f"Failed to get summary response: {e}")
         final_response = f"I reached the maximum iterations ({agent.max_iterations}) but couldn't summarize. Error: {str(e)}"
+    finally:
+        from agent import relay_llm
+
+        relay_llm.complete_logical_call(
+            summary_api_request_id,
+            outcome=summary_call_outcome,
+        )
 
     return final_response
 
diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py
index 4bd2f6c383f3c..d710f91cb58a6 100644
--- a/tests/run_agent/test_run_agent.py
+++ b/tests/run_agent/test_run_agent.py
@@ -3865,7 +3865,10 @@ class TestHandleMaxIterations:
             relay_calls.append(kwargs)
             return callback(request)
 
-        with patch("agent.relay_llm.execute_current", side_effect=execute_current):
+        with (
+            patch("agent.relay_llm.execute_current", side_effect=execute_current),
+            patch("agent.relay_llm.complete_logical_call") as complete_logical,
+        ):
             result = agent._handle_max_iterations(
                 [{"role": "user", "content": "do stuff"}],
                 60,
@@ -3877,15 +3880,23 @@ class TestHandleMaxIterations:
             relay_calls[1]["metadata"]["api_request_id"]
         )
         assert relay_calls[0]["metadata"]["call_role"] == "iteration_summary"
+        assert all(call["defer_logical_completion"] is True for call in relay_calls)
+        complete_logical.assert_called_once_with(
+            relay_calls[0]["metadata"]["api_request_id"],
+            outcome="success",
+        )
 
     def test_api_failure_returns_error(self, agent):
         agent.client.chat.completions.create.side_effect = Exception("API down")
         agent._cached_system_prompt = "You are helpful."
         messages = [{"role": "user", "content": "do stuff"}]
-        result = agent._handle_max_iterations(messages, 60)
+        with patch("agent.relay_llm.complete_logical_call") as complete_logical:
+            result = agent._handle_max_iterations(messages, 60)
         assert isinstance(result, str)
         assert "error" in result.lower()
         assert "API down" in result
+        complete_logical.assert_called_once()
+        assert complete_logical.call_args.kwargs == {"outcome": "failed"}
 
     def test_summary_skips_reasoning_for_unsupported_openrouter_model(self, agent):
         agent.base_url = "https://openrouter.ai/api/v1"