fix(relay): build the chat-completions Relay response from collector-observed chunks

Relay invokes its finalizer as soon as the provider stream ends, concurrently with
Hermes' consumer thread, so any finalizer that reads the consumer loop's closures
races the last chunk. #103104 moved `usage` onto the collector; the same race still
truncated the final tool-call `arguments` delta, dropped the last content delta, and
downgraded `finish_reason` to "stop" in Relay's annotated response (reproduced ~50%
of runs). Every other Relay integration (Anthropic, Bedrock, Codex) already
rebuilds from collector-observed chunks; this makes chat_completions match via
`_RelayChatAccumulator` and removes the single-field nonlocal stopgap.

Test: the ordering harness is shared, and a second invariant test pins the
tool-call/finish_reason case (fails on the stopgap, passes here). Hermes' own
returned response was never affected.

agent/chat_completion_helpers.py

@@ -2922,31 +2922,16 @@ class _StreamingCall:
         tool_calls = _ToolCallAccumulator()
         tool_calls_acc = tool_calls.acc
         finish_reason = model_name = usage_obj = None
-        relay_observed_usage = None
         role = "assistant"
         _diag = self._new_diag()
         self._writer_token = self._attempt_request_client = self._attempt_stream_response = None
+        from agent.chat_completion_helpers_relay import RelayChatAccumulator
+        relay_response = RelayChatAccumulator()
 
         def _open_stream(next_api_kwargs: dict[str, Any]):
             timeout = _httpx.Timeout(connect=conn_cap, read=read_timeout, write=base_timeout, pool=conn_cap)
             return self._open_chat_stream({**next_api_kwargs, "stream": True, "timeout": timeout})
 
-        def _relay_final_response() -> dict[str, Any]:
-            tool_calls.materialize()
-            message = {"role": role, "content": "".join(content_parts) or None,
-                "reasoning_content": "".join(reasoning_parts) or None,
-                "tool_calls": [tool_calls_acc[i] for i in sorted(tool_calls_acc)] or None}
-            return {"model": model_name,
-                "usage": relay_observed_usage if relay_observed_usage is not None else usage_obj,
-                "choices": [{"message": message, "finish_reason": finish_reason or "stop"}]}
-
-        def _capture_relay_usage(chunk: Any) -> None:
-            # Relay orders this collector before its finalizer; the consumer may not
-            # have copied a terminal usage-only frame into ``usage_obj`` yet.
-            nonlocal relay_observed_usage
-            if isinstance(chunk, dict) and chunk.get("usage") is not None:
-                relay_observed_usage = chunk["usage"]
-
         def _flush_pending_stream_text():
             pending_parts = list(pending_text_parts)
             pending_text_parts.clear()
@@ -2955,8 +2940,8 @@ class _StreamingCall:
 
         from agent import relay_llm
         stream = self._set_managed_stream(relay_llm.stream(self.api_kwargs, _open_stream,
-            **_relay_stream_identity(self.agent, "provider"), finalizer=_relay_final_response,
+            **_relay_stream_identity(self.agent, "provider"), finalizer=relay_response.finalize,
-            on_stream_created=self._chat_stream_created, on_chunk=_capture_relay_usage,
+            on_stream_created=self._chat_stream_created, on_chunk=relay_response.observe,
             accept_chunk=lambda chunk: self._accept_chat_chunk(stream_attempt_id, chunk),
             completed_response_predicate=lambda value: hasattr(value, "choices"),
             metadata=_relay_stream_metadata(self.agent, "chat_completions"), defer_logical_completion=True))

agent/chat_completion_helpers_relay.py

@@ -0,0 +1,74 @@
+"""Relay-side accumulator for the chat_completions streaming wire.
+
+Relay invokes its collector for every post-intercept chunk and then its finalizer as soon
+as the provider stream ends — concurrently with Hermes' consumer thread, which may not have
+read the last chunk yet. The finalizer therefore builds Relay's recorded response from
+collector-observed state only, never from the consumer loop's closures. Sibling of
+``relay_llm.AnthropicStreamAccumulator``; Bedrock and Codex follow the same contract.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+from agent.chat_completion_helpers import _ToolCallAccumulator
+from agent.message_content import flatten_message_text
+from agent.reasoning_summaries import separate_glued_reasoning_blocks
+
+
+def _tool_call_delta_view(tc_delta: Any) -> Any:
+    """Attribute view of a JSON tool-call delta for ``_ToolCallAccumulator.feed`` (written
+    against SDK objects). Only ``function`` is wrapped: ``feed`` passes ``extra_content``
+    (a dict) straight through ``_dump_if_model``, so a recursive view would corrupt it."""
+    if not isinstance(tc_delta, dict):
+        return tc_delta
+    function = tc_delta.get("function")
+    return SimpleNamespace(**{**tc_delta,
+        "function": SimpleNamespace(**function) if isinstance(function, dict) else function})
+
+
+class RelayChatAccumulator:
+    """Rebuild a chat.completion from Relay's post-intercept chunk dicts."""
+
+    def __init__(self) -> None:
+        self._content: list[str] = []
+        self._reasoning: list[str] = []
+        self._tool_calls = _ToolCallAccumulator()
+        self._model = self._usage = self._finish_reason = None
+        self._role = "assistant"
+
+    def observe(self, chunk: Any) -> None:
+        if not isinstance(chunk, dict):
+            return
+        self._model = chunk.get("model") or self._model
+        if chunk.get("usage"):
+            self._usage = chunk["usage"]
+        choices = chunk.get("choices") or []
+        choice = choices[0] if choices else None  # Hermes never requests n>1
+        if not isinstance(choice, dict):
+            return
+        self._finish_reason = choice.get("finish_reason") or self._finish_reason
+        delta = choice.get("delta")
+        if not isinstance(delta, dict):
+            return
+        if delta.get("role"):
+            self._role = delta["role"]
+        text = flatten_message_text(delta.get("content"), sep="")
+        if text:
+            self._content.append(text)
+        reasoning = delta.get("reasoning_content") or delta.get("reasoning")
+        if reasoning:
+            self._reasoning.append(separate_glued_reasoning_blocks(
+                self._reasoning[-1] if self._reasoning else "", reasoning))
+        for tc_delta in delta.get("tool_calls") or []:
+            self._tool_calls.feed(_tool_call_delta_view(tc_delta))
+
+    def finalize(self) -> dict[str, Any]:
+        acc = self._tool_calls.materialize()
+        message = {"role": self._role, "content": "".join(self._content) or None,
+            "reasoning_content": "".join(self._reasoning) or None,
+            "tool_calls": [acc[i] for i in sorted(acc)] or None}
+        # "stop" also covers Nous Portal ``lastOne`` usage frames, which carry no finish_reason.
+        return {"model": self._model, "usage": self._usage,
+            "choices": [{"message": message, "finish_reason": self._finish_reason or "stop"}]}

tests/e2e/test_relay_native_openai_stream.py

@@ -1,4 +1,10 @@
-"""Native OpenAI SDK streaming through Relay's managed execution path."""
+"""Native OpenAI SDK streaming through Relay's managed execution path.
+
+Relay runs its finalizer as soon as the provider stream ends — concurrently with Hermes'
+consumer thread, which may not have processed the last chunk yet. Each test forces that
+ordering deterministically (finalizer runs BEFORE the consumer sees a chosen chunk) and
+asserts Relay's LLM end event still records the full response.
+"""
 
 from __future__ import annotations
 
@@ -6,9 +12,17 @@ import threading
 
 import pytest
 
+_CHUNK_PREFIX = b'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model",'
 
-def test_openai_stream_usage_reaches_relay_parent_event(tmp_path, monkeypatch):
+
-    """A trailing usage-only chunk is retained on Relay's parent LLM event."""
+def _sse(*chunk_bodies: bytes) -> bytes:
+    return b"".join(_CHUNK_PREFIX + body + b"}\n\n" for body in chunk_bodies) + b"data: [DONE]\n\n"
+
+
+def _stream_through_relay(tmp_path, monkeypatch, response_body: bytes, *, finalize_before):
+    """Stream ``response_body`` through Relay; Relay's finalizer is forced to complete before
+    the consumer thread processes the first chunk matching ``finalize_before(chunk)``.
+    Returns ``(hermes_result, relay_llm_end_event)``."""
     httpx = pytest.importorskip("httpx")
     nemo_relay = pytest.importorskip("nemo_relay")
     openai = pytest.importorskip("openai")
@@ -18,55 +32,26 @@ def test_openai_stream_usage_reaches_relay_parent_event(tmp_path, monkeypatch):
 
     monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes-home"))
     monkeypatch.setenv("HERMES_STREAM_RETRIES", "0")
-    response_body = b"""data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model","choices":[{"index":0,"delta":{"role":"assistant","content":"done"},"finish_reason":null}]}
-
-data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
-
-data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model","choices":[],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110}}
-
-data: [DONE]
-
-"""
 
     def respond(request):
-        return httpx.Response(
+        return httpx.Response(200, headers={"content-type": "text/event-stream"},
-            200,
+                              content=response_body, request=request)
-            headers={"content-type": "text/event-stream"},
-            content=response_body,
-            request=request,
-        )
 
-    client = openai.OpenAI(
+    client = openai.OpenAI(api_key="test-key", base_url="https://example.com/v1",
-        api_key="test-key",
+                           http_client=httpx.Client(transport=httpx.MockTransport(respond)))
-        base_url="https://example.com/v1",
-        http_client=httpx.Client(transport=httpx.MockTransport(respond)),
-    )
     relay_runtime._reset_for_tests()
-    agent = AIAgent(
+    agent = AIAgent(api_key="test-key", base_url="https://example.com/v1", provider="test-provider",
-        api_key="test-key",
+                    model="test/model", quiet_mode=True, skip_context_files=True, skip_memory=True)
-        base_url="https://example.com/v1",
-        provider="test-provider",
-        model="test/model",
-        quiet_mode=True,
-        skip_context_files=True,
-        skip_memory=True,
-    )
     agent.api_mode = "chat_completions"
     agent.session_id = "openai-relay-session"
     agent._interrupt_requested = False
     agent._create_request_openai_client = lambda *args, **kwargs: client
     lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation(
-        profile_key=relay_runtime.current_profile_key(),
+        profile_key=relay_runtime.current_profile_key(), session_id=agent.session_id, platform="cli")
-        session_id=agent.session_id,
-        platform="cli",
-    )
     turn = relay_runtime.SESSION_COORDINATOR.begin_turn(
-        lease,
+        lease, turn_id="openai-relay-turn", task_id="openai-relay-task")
-        turn_id="openai-relay-turn",
-        task_id="openai-relay-task",
-    )
     consumer = "test.openai_relay"
-    subscriber_name = "test.openai_stream_usage"
+    subscriber_name = "test.openai_stream"
     events = []
     relay_finalizer_started = threading.Event()
     allow_relay_finalizer = threading.Event()
@@ -75,43 +60,30 @@ data: [DONE]
 
     def run_synchronized_relay_finalizer(managed_stream, attempt):
         relay_finalizer_started.set()
-        assert allow_relay_finalizer.wait(5), (
+        assert allow_relay_finalizer.wait(5), "consumer did not release Relay's finalizer"
-            "consumer did not release Relay's finalizer"
-        )
         try:
             return run_relay_finalizer(managed_stream, attempt)
         finally:
             relay_finalizer_finished.set()
 
-    monkeypatch.setattr(
+    monkeypatch.setattr(relay_llm.ManagedLlmStream, "_relay_finalizer", run_synchronized_relay_finalizer)
-        relay_llm.ManagedLlmStream,
-        "_relay_finalizer",
-        run_synchronized_relay_finalizer,
-    )
 
-    parse_choiceless_chunk = chat_completion_helpers._StreamingCall._choiceless_chunk
+    count_chunk = chat_completion_helpers._StreamingCall._count_chunk
 
-    def parse_usage_after_relay_finalizes(chunk, finish_reason):
+    def count_chunk_after_relay_finalizes(self, diag, chunk):
-        if not chunk.choices and getattr(chunk, "usage", None) is not None:
+        # ``_count_chunk`` is the first thing the consumer does with every chunk.
-            # Force Relay to finalize before the consumer copies the usage frame.
+        if finalize_before(chunk):
             assert relay_finalizer_started.wait(5), "Relay's finalizer did not start"
             allow_relay_finalizer.set()
             assert relay_finalizer_finished.wait(5), "Relay's finalizer did not finish"
-        return parse_choiceless_chunk(chunk, finish_reason)
+        return count_chunk(self, diag, chunk)
 
-    monkeypatch.setattr(
+    monkeypatch.setattr(chat_completion_helpers._StreamingCall, "_count_chunk", count_chunk_after_relay_finalizes)
-        chat_completion_helpers._StreamingCall,
-        "_choiceless_chunk",
-        staticmethod(parse_usage_after_relay_finalizes),
-    )
     lease.host.retain_managed_execution(consumer)
     lease.host.relay.subscribers.register(subscriber_name, events.append)
-
     try:
         result = agent._interruptible_streaming_api_call({
-            "model": "test/model",
+            "model": "test/model", "messages": [{"role": "user", "content": "hi"}]})
-            "messages": [{"role": "user", "content": "hi"}],
-        })
         lease.host.relay.subscribers.flush()
     finally:
         lease.host.relay.subscribers.deregister(subscriber_name)
@@ -121,23 +93,52 @@ data: [DONE]
         relay_runtime._reset_for_tests()
         client.close()
 
-    assert result.usage is not None
-    assert result.usage.prompt_tokens == 100
-    assert result.usage.completion_tokens == 10
-    assert result.usage.total_tokens == 110
     llm_end_events = [
-        event
+        event for event in events
-        for event in events
+        if isinstance(event, nemo_relay.ScopeEvent) and event.name == "openai.chat_completions"
-        if isinstance(event, nemo_relay.ScopeEvent)
+        and event.category == "llm" and event.scope_category == "end"
-        and event.name == "openai.chat_completions"
-        and event.category == "llm"
-        and event.scope_category == "end"
     ]
     assert len(llm_end_events) == 1
-    llm_end = llm_end_events[0]
+    assert llm_end_events[0].annotated_response is not None
-    assert llm_end.annotated_response is not None
+    return result, llm_end_events[0]
+
+
+def test_openai_stream_usage_reaches_relay_parent_event(tmp_path, monkeypatch):
+    """A trailing usage-only chunk is retained on Relay's parent LLM event."""
+    body = _sse(
+        b'"choices":[{"index":0,"delta":{"role":"assistant","content":"done"},"finish_reason":null}]',
+        b'"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]',
+        b'"choices":[],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110}',
+    )
+    result, llm_end = _stream_through_relay(
+        tmp_path, monkeypatch, body,
+        finalize_before=lambda chunk: not chunk.choices and getattr(chunk, "usage", None) is not None)
+
+    assert result.usage is not None
+    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (100, 10, 110)
     assert llm_end.annotated_response.usage == {
-        "prompt_tokens": 100,
+        "prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110}
-        "completion_tokens": 10,
+    assert llm_end.annotated_response.message == "done"
-        "total_tokens": 110,
+
-    }
+
+def test_openai_stream_final_tool_call_delta_reaches_relay_parent_event(tmp_path, monkeypatch):
+    """The last chunk's tool-call arguments and finish_reason are retained on Relay's parent
+    LLM event — the same finalizer-before-consumer race as the usage frame, without one."""
+    body = _sse(
+        b'"choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,'
+        b'"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\\"path\\": "}}]},'
+        b'"finish_reason":null}]',
+        b'"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\\"/tmp/x\\"}"}}]},'
+        b'"finish_reason":"tool_calls"}]',
+    )
+    result, llm_end = _stream_through_relay(
+        tmp_path, monkeypatch, body,
+        finalize_before=lambda chunk: bool(chunk.choices) and chunk.choices[0].finish_reason == "tool_calls")
+
+    hermes_call = result.choices[0].message.tool_calls[0]
+    assert (hermes_call.function.name, hermes_call.function.arguments) == ("read_file", '{"path": "/tmp/x"}')
+    assert result.choices[0].finish_reason == "tool_calls"
+    assert llm_end.annotated_response.message is None
+    (relay_call,) = llm_end.annotated_response.tool_calls
+    assert (relay_call["name"], relay_call["arguments"]) == ("read_file", {"path": "/tmp/x"})
+    assert llm_end.annotated_response.finish_reason == "tool_use"