hermes-agent/tests/e2e/test_relay_native_openai_stream.py

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

import threading

import pytest

_CHUNK_PREFIX = b'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test/model",'


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

    from agent import chat_completion_helpers, relay_llm, relay_runtime
    from run_agent import AIAgent

    monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes-home"))
    monkeypatch.setenv("HERMES_STREAM_RETRIES", "0")

    def respond(request):
        return httpx.Response(200, headers={"content-type": "text/event-stream"},
                              content=response_body, request=request)

    client = openai.OpenAI(api_key="test-key", base_url="https://example.com/v1",
                           http_client=httpx.Client(transport=httpx.MockTransport(respond)))
    relay_runtime._reset_for_tests()
    agent = AIAgent(api_key="test-key", 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(), session_id=agent.session_id, platform="cli")
    turn = relay_runtime.SESSION_COORDINATOR.begin_turn(
        lease, turn_id="openai-relay-turn", task_id="openai-relay-task")
    consumer = "test.openai_relay"
    subscriber_name = "test.openai_stream"
    events = []
    relay_finalizer_started = threading.Event()
    allow_relay_finalizer = threading.Event()
    relay_finalizer_finished = threading.Event()
    run_relay_finalizer = relay_llm.ManagedLlmStream._relay_finalizer

    def run_synchronized_relay_finalizer(managed_stream, attempt):
        relay_finalizer_started.set()
        assert allow_relay_finalizer.wait(30), "consumer did not release Relay's finalizer"
        try:
            return run_relay_finalizer(managed_stream, attempt)
        finally:
            relay_finalizer_finished.set()

    monkeypatch.setattr(relay_llm.ManagedLlmStream, "_relay_finalizer", run_synchronized_relay_finalizer)

    count_chunk = chat_completion_helpers._StreamingCall._count_chunk

    def count_chunk_after_relay_finalizes(self, diag, chunk):
        # Relay's producer is pumped by the consumer thread's OWN event loop (``ManagedLlmStream.
        # __next__`` -> ``run_until_complete``), so the finalizer can only start while that loop runs.
        # Blocking the consumer thread here and waiting for it therefore deadlocked whenever the loop
        # had not reached EOF yet (~1 run in 6). Instead: release the finalizer and PUMP THE STREAM'S
        # LOOP until it has completed, then hand the chunk to the consumer. That is a real
        # happens-before (finalizer done -> consumer sees chunk) on every schedule, not a retry lottery.
        if finalize_before(chunk):
            import asyncio

            stream = self.managed_stream_holder["stream"]
            allow_relay_finalizer.set()
            # A schedule probe may hold the provider generator at this chunk until the consumer has
            # taken it (adverse consumer-first ordering); release it so the pump below can reach EOF.
            gate = getattr(stream, "_review_gate", None)
            if gate is not None:
                gate.set()

            async def finalizer_done():
                return await asyncio.to_thread(relay_finalizer_finished.wait, 30)

            assert stream._loop.run_until_complete(finalizer_done()), "Relay's finalizer did not finish"
            assert relay_finalizer_started.is_set()
        return count_chunk(self, diag, chunk)

    monkeypatch.setattr(chat_completion_helpers._StreamingCall, "_count_chunk", count_chunk_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", "messages": [{"role": "user", "content": "hi"}]})
        lease.host.relay.subscribers.flush()
    finally:
        lease.host.relay.subscribers.deregister(subscriber_name)
        lease.host.release_managed_execution(consumer)
        relay_runtime.SESSION_COORDINATOR.end_turn(turn, outcome="success")
        relay_runtime.SESSION_COORDINATOR.release_conversation(lease)
        relay_runtime._reset_for_tests()
        client.close()

    llm_end_events = [
        event for event in events
        if isinstance(event, nemo_relay.ScopeEvent) and event.name == "openai.chat_completions"
        and event.category == "llm" and event.scope_category == "end"
    ]
    assert len(llm_end_events) == 1
    assert llm_end_events[0].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, "completion_tokens": 10, "total_tokens": 110}
    assert llm_end.annotated_response.message == "done"


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"