hermes-agent/tests/agent/test_auxiliary_relay.py

from types import SimpleNamespace

import pytest

pytest.importorskip("nemo_relay")

from agent import auxiliary_client, relay_llm, relay_runtime


@pytest.fixture()
def relay_turn(tmp_path, monkeypatch):
    monkeypatch.setenv("HERMES_HOME", str(tmp_path / "profile"))
    relay_runtime._reset_for_tests()
    lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation(
        profile_key=relay_runtime.current_profile_key(),
        session_id="session-1",
        platform="cli",
    )
    turn = relay_runtime.SESSION_COORDINATOR.begin_turn(
        lease,
        turn_id="turn-1",
        task_id="task-1",
    )
    try:
        yield lease.host.relay, turn
    finally:
        relay_runtime.SESSION_COORDINATOR.end_turn(turn, outcome="success")
        relay_runtime.SESSION_COORDINATOR.release_conversation(lease)
        relay_runtime._reset_for_tests()


def test_auxiliary_retries_share_logical_relay_identity(monkeypatch):
    attempts = []
    logical_completions = []
    responses = iter([
        SimpleNamespace(choices=[]),
        SimpleNamespace(
            choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))]
        ),
    ])
    client = SimpleNamespace(
        chat=SimpleNamespace(
            completions=SimpleNamespace(
                create=lambda **_kwargs: next(responses),
            )
        )
    )

    def execute_current(request, callback, **kwargs):
        attempts.append(kwargs)
        return callback(request)

    monkeypatch.setattr(relay_llm, "execute_current", execute_current)
    monkeypatch.setattr(
        relay_llm,
        "complete_logical_call",
        lambda request_id, *, outcome: logical_completions.append(
            (request_id, outcome)
        ),
    )

    @auxiliary_client._relay_auxiliary_call
    def run(task):
        auxiliary_client._set_relay_auxiliary_route(
            "openrouter",
            "test-model",
            "chat_completions",
        )
        with pytest.raises(RuntimeError, match="invalid response"):
            auxiliary_client._validate_llm_response(
                auxiliary_client._relay_sync_completion(
                    client,
                    {"model": "test-model", "messages": []},
                ),
                task,
            )
        return auxiliary_client._validate_llm_response(
            auxiliary_client._relay_sync_completion(
                client,
                {"model": "test-model", "messages": []},
            ),
            task,
        )

    result = run("compression")

    assert result.choices[0].message.content == "ok"
    assert attempts[0]["metadata"]["api_request_id"] == (
        attempts[1]["metadata"]["api_request_id"]
    )
    assert [attempt["metadata"]["retry_count"] for attempt in attempts] == [0, 1]
    assert attempts[0]["metadata"]["call_role"] == "auxiliary:compression"
    assert all(attempt["defer_logical_completion"] is True for attempt in attempts)
    assert logical_completions == [
        (attempts[0]["metadata"]["api_request_id"], "success")
    ]


@pytest.mark.asyncio
async def test_async_auxiliary_attempt_uses_inherited_relay_adapter(monkeypatch):
    captured = {}
    logical_completions = []

    async def create(**kwargs):
        return SimpleNamespace(
            request=kwargs,
            choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))],
        )

    client = SimpleNamespace(
        chat=SimpleNamespace(completions=SimpleNamespace(create=create))
    )

    async def execute_current_async(request, callback, **kwargs):
        captured.update(kwargs)
        return await callback(request)

    monkeypatch.setattr(
        relay_llm,
        "execute_current_async",
        execute_current_async,
    )
    monkeypatch.setattr(
        relay_llm,
        "complete_logical_call",
        lambda request_id, *, outcome: logical_completions.append(
            (request_id, outcome)
        ),
    )

    @auxiliary_client._relay_auxiliary_call_async
    async def run(task):
        auxiliary_client._set_relay_auxiliary_route(
            "anthropic",
            "claude-test",
            "chat_completions",
        )
        return auxiliary_client._validate_llm_response(
            await auxiliary_client._relay_async_completion(
                client,
                {"model": "claude-test", "messages": []},
            ),
            task,
        )

    result = await run("title_generation")

    assert result.request["model"] == "claude-test"
    assert captured["name"] == "anthropic"
    assert captured["metadata"]["call_role"] == "auxiliary:title_generation"
    assert captured["defer_logical_completion"] is True
    assert logical_completions == [
        (captured["metadata"]["api_request_id"], "success")
    ]


def test_terminal_auxiliary_failure_stays_failed_when_caller_catches_it(
    relay_turn, monkeypatch
):
    _relay, turn = relay_turn
    consumer = "test.terminal-auxiliary-failure"
    turn.lease.host.retain_managed_execution(consumer)
    outcomes = []
    original_pop = turn.lease.host.relay.scope.pop

    def record_pop(*args, **kwargs):
        outcomes.append((kwargs.get("output") or {}).get("outcome"))
        return original_pop(*args, **kwargs)

    monkeypatch.setattr(turn.lease.host.relay.scope, "pop", record_pop)
    client = SimpleNamespace(
        chat=SimpleNamespace(
            completions=SimpleNamespace(
                create=lambda **_kwargs: SimpleNamespace(choices=[]),
            )
        )
    )

    @auxiliary_client._relay_auxiliary_call
    def run(task):
        auxiliary_client._set_relay_auxiliary_route(
            "openrouter",
            "test-model",
            "chat_completions",
        )
        with pytest.raises(RuntimeError, match="invalid response"):
            auxiliary_client._validate_llm_response(
                auxiliary_client._relay_sync_completion(
                    client,
                    {"model": "test-model", "messages": []},
                ),
                task,
            )
        assert len(turn.logical_llm_calls) == 1
        return auxiliary_client._validate_llm_response(
            auxiliary_client._relay_sync_completion(
                client,
                {"model": "test-model", "messages": []},
            ),
            task,
        )

    try:
        with pytest.raises(RuntimeError, match="invalid response"):
            run("compression")

        assert outcomes == ["failed"]
        assert turn.logical_llm_calls == {}

        relay_runtime.SESSION_COORDINATOR.end_turn(turn, outcome="success")

        assert outcomes == ["failed", "success"]
    finally:
        turn.lease.host.release_managed_execution(consumer)


@pytest.mark.asyncio
async def test_async_terminal_auxiliary_failure_closes_logical_call(relay_turn):
    _relay, turn = relay_turn
    consumer = "test.async-terminal-auxiliary-failure"
    turn.lease.host.retain_managed_execution(consumer)

    async def create(**_kwargs):
        return SimpleNamespace(choices=[])

    client = SimpleNamespace(
        chat=SimpleNamespace(completions=SimpleNamespace(create=create))
    )

    @auxiliary_client._relay_auxiliary_call_async
    async def run(task):
        auxiliary_client._set_relay_auxiliary_route(
            "anthropic",
            "claude-test",
            "chat_completions",
        )
        with pytest.raises(RuntimeError, match="invalid response"):
            auxiliary_client._validate_llm_response(
                await auxiliary_client._relay_async_completion(
                    client,
                    {"model": "claude-test", "messages": []},
                ),
                task,
            )
        assert len(turn.logical_llm_calls) == 1
        return auxiliary_client._validate_llm_response(
            await auxiliary_client._relay_async_completion(
                client,
                {"model": "claude-test", "messages": []},
            ),
            task,
        )

    try:
        with pytest.raises(RuntimeError, match="invalid response"):
            await run("title_generation")

        assert turn.logical_llm_calls == {}
    finally:
        turn.lease.host.release_managed_execution(consumer)


def test_auxiliary_stream_uses_streaming_relay_primitive(monkeypatch):
    captured = {}
    raw_stream = iter([{"delta": "one"}, {"delta": "two"}])
    client = SimpleNamespace(
        chat=SimpleNamespace(
            completions=SimpleNamespace(create=lambda **_kwargs: raw_stream)
        )
    )

    def stream_current(request, stream_factory, **kwargs):
        captured.update(kwargs)
        return stream_factory(request)

    monkeypatch.setattr(relay_llm, "stream_current", stream_current)

    @auxiliary_client._relay_auxiliary_call
    def run(task):
        auxiliary_client._set_relay_auxiliary_route(
            "openrouter",
            "moa-model",
            "chat_completions",
        )
        return auxiliary_client._relay_sync_stream(
            client,
            {"model": "moa-model", "messages": [], "stream": True},
        )

    assert list(run("moa")) == [{"delta": "one"}, {"delta": "two"}]
    assert captured["metadata"]["call_role"] == "auxiliary:moa"


def test_partial_auxiliary_stream_failure_closes_before_recovery(
    relay_turn, monkeypatch
):
    _relay, turn = relay_turn
    consumer = "test.partial-auxiliary-stream-failure"
    turn.lease.host.retain_managed_execution(consumer)
    outcomes = []
    original_pop = turn.lease.host.relay.scope.pop

    def record_pop(*args, **kwargs):
        outcomes.append((kwargs.get("output") or {}).get("outcome"))
        return original_pop(*args, **kwargs)

    monkeypatch.setattr(turn.lease.host.relay.scope, "pop", record_pop)

    class ProviderError(Exception):
        pass

    provider_error = ProviderError("stream failed")
    partial_chunk = SimpleNamespace(
        model="test-model",
        choices=[
            SimpleNamespace(
                delta=SimpleNamespace(content="partial", tool_calls=None),
                finish_reason=None,
            )
        ],
        usage=None,
    )

    def partial_stream():
        yield partial_chunk
        raise provider_error

    stream_client = SimpleNamespace(
        chat=SimpleNamespace(
            completions=SimpleNamespace(
                create=lambda **_kwargs: partial_stream(),
            )
        )
    )
    recovery_client = SimpleNamespace(
        chat=SimpleNamespace(
            completions=SimpleNamespace(
                create=lambda **_kwargs: SimpleNamespace(
                    choices=[
                        SimpleNamespace(message=SimpleNamespace(content="recovered"))
                    ]
                ),
            )
        )
    )

    @auxiliary_client._relay_auxiliary_call
    def start_stream(task):
        auxiliary_client._set_relay_auxiliary_route(
            "openrouter",
            "test-model",
            "chat_completions",
        )
        return auxiliary_client._relay_sync_stream(
            stream_client,
            {"model": "test-model", "messages": [], "stream": True},
        )

    @auxiliary_client._relay_auxiliary_call
    def recover(task):
        auxiliary_client._set_relay_auxiliary_route(
            "openrouter",
            "test-model",
            "chat_completions",
        )
        return auxiliary_client._validate_llm_response(
            auxiliary_client._relay_sync_completion(
                recovery_client,
                {"model": "test-model", "messages": []},
            ),
            task,
        )

    try:
        stream = start_stream("moa")
        assert next(stream) is partial_chunk

        with pytest.raises(ProviderError) as caught:
            next(stream)

        assert caught.value is provider_error
        assert outcomes == ["failed"]
        assert turn.logical_llm_calls == {}

        result = recover("moa")

        assert result.choices[0].message.content == "recovered"
        assert outcomes == ["failed", "success"]
        assert turn.logical_llm_calls == {}
    finally:
        turn.lease.host.release_managed_execution(consumer)


def test_auxiliary_attempt_uses_real_relay_request_intercepts(relay_turn):
    relay, turn = relay_turn
    consumer = "test.auxiliary-request-intercept"
    turn.lease.host.retain_managed_execution(consumer)
    captured_requests = []
    client = SimpleNamespace(
        chat=SimpleNamespace(
            completions=SimpleNamespace(
                create=lambda **kwargs: captured_requests.append(kwargs)
                or SimpleNamespace(
                    choices=[
                        SimpleNamespace(message=SimpleNamespace(content="ok"))
                    ]
                ),
            )
        )
    )

    def rewrite_request(_name, request, annotated):
        annotated.params = {**(annotated.params or {}), "temperature": 0.25}
        return relay.LLMRequestInterceptOutcome(request, annotated)

    relay.intercepts.register_llm_request(
        "hermes-auxiliary-request",
        1,
        False,
        rewrite_request,
    )
    try:
        @auxiliary_client._relay_auxiliary_call
        def run(task):
            auxiliary_client._set_relay_auxiliary_route(
                "openrouter",
                "test-model",
                "chat_completions",
            )
            return auxiliary_client._validate_llm_response(
                auxiliary_client._relay_sync_completion(
                    client,
                    {"model": "test-model", "messages": []},
                ),
                task,
            )

        result = run("compression")
    finally:
        relay.intercepts.deregister_llm_request("hermes-auxiliary-request")
        turn.lease.host.release_managed_execution(consumer)

    assert result.choices[0].message.content == "ok"
    assert captured_requests[0]["temperature"] == 0.25
    assert turn.logical_llm_calls == {}