ariadne/tests/test_hermes_agent_client.py

from __future__ import annotations

from ariadne.services import hermes_agent_client as module


class FakeClock:
    def __init__(self) -> None:
        self.now = 0.0
        self.sleeps: list[float] = []

    def time(self) -> float:
        return self.now

    def sleep(self, seconds: float) -> None:
        self.sleeps.append(seconds)
        self.now += seconds


class FakeResponse:
    def __init__(self, status_code: int, payload: object = None) -> None:
        self.status_code = status_code
        self._payload = payload

    def json(self):  # type: ignore[no-untyped-def]
        if isinstance(self._payload, Exception):
            raise self._payload
        return self._payload


def _install_client(monkeypatch, posts: list, gets: list) -> dict:
    calls: dict = {"init_timeout": None, "posts": [], "gets": []}

    def next_item(queue: list, label: str):  # type: ignore[no-untyped-def]
        assert queue, f"unexpected {label} request"
        item = queue.pop(0)
        if isinstance(item, Exception):
            raise item
        return item

    class FakeClient:
        def __init__(self, *, timeout=None) -> None:  # type: ignore[no-untyped-def]
            calls["init_timeout"] = timeout

        def __enter__(self):
            return self

        def __exit__(self, *args) -> None:  # type: ignore[no-untyped-def]
            return None

        def post(self, url, headers=None, json=None):  # type: ignore[no-untyped-def]
            calls["posts"].append((url, headers, json))
            return next_item(posts, "post")

        def get(self, url, headers=None):  # type: ignore[no-untyped-def]
            calls["gets"].append((url, headers))
            return next_item(gets, "get")

    monkeypatch.setattr(module.httpx, "Client", FakeClient)
    return calls


def _install_clock(monkeypatch) -> FakeClock:
    clock = FakeClock()
    monkeypatch.setattr(module, "time", clock)
    return clock


def _config(**overrides) -> dict:  # type: ignore[no-untyped-def]
    base = {
        "base_url": "http://hermes.hermes.svc.cluster.local:8642",
        "api_key": "key-123",
        "total_timeout_seconds": 30.0,
        "poll_interval_seconds": 1.0,
        "request_timeout_seconds": 7.0,
    }
    base.update(overrides)
    return base


def _started(run_id: str = "run_" + "a" * 32) -> FakeResponse:
    return FakeResponse(202, {"run_id": run_id, "status": "started"})


def test_run_triage_completes_happy_path(monkeypatch) -> None:
    _install_clock(monkeypatch)
    calls = _install_client(
        monkeypatch,
        posts=[_started("run_abc")],
        gets=[
            FakeResponse(200, {"object": "hermes.run", "status": "running"}),
            FakeResponse(
                200,
                {
                    "object": "hermes.run",
                    "run_id": "run_abc",
                    "status": "completed",
                    "session_id": "sess-1",
                    "output": "Demo fixture failure confirmed.",
                    "usage": {"total_tokens": 10},
                    "error": None,
                },
            ),
        ],
    )

    result = module.run_triage(_config(), "triage incident inc-42")

    assert result.status == "completed"
    assert result.output == "Demo fixture failure confirmed."
    assert result.run_id == "run_abc"
    assert result.session_id == "sess-1"
    assert result.error is None
    assert result.denied_approvals == 0
    assert result.duration_seconds == 1.0
    assert calls["init_timeout"] == 7.0
    start_url, start_headers, start_body = calls["posts"][0]
    assert start_url == "http://hermes.hermes.svc.cluster.local:8642/v1/runs"
    assert start_headers == {"Authorization": "Bearer key-123"}
    assert start_body == {"input": "triage incident inc-42"}
    poll_url, poll_headers = calls["gets"][0]
    assert poll_url == "http://hermes.hermes.svc.cluster.local:8642/v1/runs/run_abc"
    assert poll_headers == {"Authorization": "Bearer key-123"}


def test_run_triage_denies_approvals_and_keeps_polling(monkeypatch) -> None:
    _install_clock(monkeypatch)
    waiting = {"status": "waiting_for_approval", "last_event": "approval_requested"}
    calls = _install_client(
        monkeypatch,
        posts=[
            _started("run_abc"),
            FakeResponse(200, {"ok": True}),
            ConnectionError("deny endpoint unreachable"),
            FakeResponse(409, {"error": {"message": "already resolved"}}),
        ],
        gets=[
            FakeResponse(200, dict(waiting)),
            FakeResponse(200, dict(waiting)),
            FakeResponse(200, dict(waiting)),
            FakeResponse(200, {"status": "completed", "output": "done", "session_id": "sess-2"}),
        ],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "completed"
    assert result.denied_approvals == 1
    deny_url, _headers, deny_body = calls["posts"][1]
    assert deny_url.endswith("/v1/runs/run_abc/approval")
    assert deny_body == {"choice": "deny"}


def test_run_triage_marks_lost_on_poll_404(monkeypatch) -> None:
    _install_clock(monkeypatch)
    _install_client(
        monkeypatch,
        posts=[_started("run_abc")],
        gets=[FakeResponse(404, {"error": {"code": "run_not_found", "message": "unknown run"}})],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "lost"
    assert result.run_id == "run_abc"
    assert result.error == "run_not_found"
    assert result.output is None


def test_run_triage_times_out_and_stops_run(monkeypatch) -> None:
    clock = _install_clock(monkeypatch)
    calls = _install_client(
        monkeypatch,
        posts=[_started("run_abc"), FakeResponse(200, {"status": "cancelling"})],
        gets=[
            FakeResponse(200, {"status": "running"}),
            FakeResponse(200, {"status": "running"}),
        ],
    )

    result = module.run_triage(_config(total_timeout_seconds=2.0), "prompt")

    assert result.status == "timeout"
    assert result.run_id == "run_abc"
    assert "total_timeout_after_2.0s" in (result.error or "")
    stop_url, stop_headers, _body = calls["posts"][-1]
    assert stop_url.endswith("/v1/runs/run_abc/stop")
    assert stop_headers == {"Authorization": "Bearer key-123"}
    assert clock.sleeps == [1.0, 1.0]


def test_run_triage_timeout_tolerates_stop_failure(monkeypatch) -> None:
    _install_clock(monkeypatch)
    _install_client(
        monkeypatch,
        posts=[_started("run_abc"), ConnectionError("stop unreachable")],
        gets=[FakeResponse(200, {"status": "queued"})],
    )

    result = module.run_triage(_config(total_timeout_seconds=1.0), "prompt")

    assert result.status == "timeout"
    assert result.run_id == "run_abc"


def test_run_triage_retries_start_once_on_connection_error(monkeypatch) -> None:
    _install_clock(monkeypatch)
    calls = _install_client(
        monkeypatch,
        posts=[ConnectionError("connection refused"), _started("run_abc")],
        gets=[FakeResponse(200, {"status": "completed", "output": "ok"})],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "completed"
    assert len(calls["posts"]) == 2


def test_run_triage_retries_start_once_on_server_error(monkeypatch) -> None:
    _install_clock(monkeypatch)
    calls = _install_client(
        monkeypatch,
        posts=[FakeResponse(503, None), _started("run_abc")],
        gets=[FakeResponse(200, {"status": "completed", "output": "ok"})],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "completed"
    assert len(calls["posts"]) == 2


def test_run_triage_start_failure_after_retry_returns_error(monkeypatch) -> None:
    _install_clock(monkeypatch)
    calls = _install_client(
        monkeypatch,
        posts=[FakeResponse(500, None), ConnectionError("still down")],
        gets=[],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "error"
    assert result.run_id is None
    assert "start_request_failed: still down" in (result.error or "")
    assert len(calls["posts"]) == 2
    assert calls["gets"] == []


def test_run_triage_does_not_retry_auth_failure(monkeypatch) -> None:
    _install_clock(monkeypatch)
    calls = _install_client(
        monkeypatch,
        posts=[FakeResponse(401, {"error": {"message": "invalid api key", "type": "invalid_request_error"}})],
        gets=[],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "error"
    assert result.error == "start_http_401: invalid api key"
    assert len(calls["posts"]) == 1


def test_run_triage_start_missing_run_id_returns_error(monkeypatch) -> None:
    _install_clock(monkeypatch)
    _install_client(monkeypatch, posts=[FakeResponse(202, {"status": "started"})], gets=[])

    result = module.run_triage(_config(), "prompt")

    assert result.status == "error"
    assert result.error == "start_missing_run_id"


def test_run_triage_tolerates_individual_poll_failures(monkeypatch) -> None:
    clock = _install_clock(monkeypatch)
    _install_client(
        monkeypatch,
        posts=[_started("run_abc")],
        gets=[
            ConnectionError("poll dropped"),
            FakeResponse(500, None),
            FakeResponse(200, ValueError("bad json body")),
            FakeResponse(200, ["not", "a", "dict"]),
            FakeResponse(200, {"status": "completed", "output": "recovered"}),
        ],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "completed"
    assert result.output == "recovered"
    assert clock.sleeps == [1.0, 1.0, 1.0, 1.0]


def test_run_triage_reports_failed_run_error(monkeypatch) -> None:
    _install_clock(monkeypatch)
    _install_client(
        monkeypatch,
        posts=[_started("run_abc")],
        gets=[FakeResponse(200, {"status": "failed", "output": None, "error": "tool exploded"})],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "failed"
    assert result.output is None
    assert result.error == "tool exploded"
    assert result.session_id is None


def test_run_triage_reports_cancelled_run(monkeypatch) -> None:
    _install_clock(monkeypatch)
    _install_client(
        monkeypatch,
        posts=[_started("run_abc")],
        gets=[FakeResponse(200, {"status": "cancelled", "output": "partial notes", "error": None})],
    )

    result = module.run_triage(_config(), "prompt")

    assert result.status == "cancelled"
    assert result.output == "partial notes"
    assert result.error is None


def test_run_triage_never_raises_on_client_setup_failure(monkeypatch) -> None:
    _install_clock(monkeypatch)

    class BoomClient:
        def __init__(self, *, timeout=None) -> None:  # type: ignore[no-untyped-def]
            raise RuntimeError("no transport available")

    monkeypatch.setattr(module.httpx, "Client", BoomClient)

    result = module.run_triage(_config(), "prompt")

    assert result.status == "error"
    assert "unexpected_client_failure: no transport available" in (result.error or "")
    assert result.denied_approvals == 0


def test_run_config_applies_defaults_and_normalization() -> None:
    defaults = module._run_config({})  # noqa: SLF001
    assert defaults.base_url == ""
    assert defaults.api_key == ""
    assert defaults.total_timeout_seconds == 420.0
    assert defaults.poll_interval_seconds == 5.0
    assert defaults.request_timeout_seconds == 15.0

    parsed = module._run_config(  # noqa: SLF001
        {
            "base_url": " http://hermes:8642/ ",
            "api_key": "abc",
            "total_timeout_seconds": "60",
            "poll_interval_seconds": 0,
            "request_timeout_seconds": "bad",
        }
    )
    assert parsed.base_url == "http://hermes:8642"
    assert parsed.total_timeout_seconds == 60.0
    assert parsed.poll_interval_seconds == 5.0
    assert parsed.request_timeout_seconds == 15.0


def test_payload_and_error_message_helpers() -> None:
    assert module._json_payload(FakeResponse(200, ["list"])) == {}  # noqa: SLF001
    assert module._json_payload(FakeResponse(200, ValueError("nope"))) == {}  # noqa: SLF001
    assert module._error_message(FakeResponse(401, {"error": "denied"})) == "denied"  # noqa: SLF001
    assert module._error_message(FakeResponse(401, {"error": {"message": "bad key"}})) == "bad key"  # noqa: SLF001
    assert module._error_message(FakeResponse(401, {})) == ""  # noqa: SLF001