diff --git a/src/livepeer_gateway/__init__.py b/src/livepeer_gateway/__init__.py index 9e4c284..eb9274c 100644 --- a/src/livepeer_gateway/__init__.py +++ b/src/livepeer_gateway/__init__.py @@ -37,6 +37,7 @@ LiveRunnerSessionCallback, LiveRunnerSessionEvent, LiveRunnerProxy, + LivePaymentSessionCallback, call_runner, create_proxy, create_trickle_channels, @@ -78,6 +79,7 @@ "LiveRunnerSessionEvent", "LivePaymentChallenge", "LivePaymentSession", + "LivePaymentSessionCallback", "LiveRunnerProxy", "LivepeerGatewayError", "LivepeerHTTPError", diff --git a/src/livepeer_gateway/live_runner.py b/src/livepeer_gateway/live_runner.py index 8d80b3b..95f6dcc 100644 --- a/src/livepeer_gateway/live_runner.py +++ b/src/livepeer_gateway/live_runner.py @@ -87,6 +87,7 @@ class LiveRunnerSessionEvent: LiveRunnerSessionCallback = Callable[[LiveRunnerSessionEvent], None | Awaitable[None]] +LivePaymentSessionCallback = Callable[[LivePaymentSession], None | Awaitable[None]] @dataclass(frozen=True) @@ -762,6 +763,7 @@ async def call_runner( payment_unit: str | None = ..., timeout: float = ..., max_payment_challenge_retries: int = ..., + on_payment_session: LivePaymentSessionCallback | None = ..., stream: Literal[False] = False, ) -> LiveRunnerCallResult: pass @@ -779,6 +781,7 @@ async def call_runner( payment_unit: str | None = ..., timeout: float = ..., max_payment_challenge_retries: int = ..., + on_payment_session: LivePaymentSessionCallback | None = ..., stream: Literal[True], ) -> LiveRunnerCallStream: pass @@ -795,6 +798,7 @@ async def call_runner( payment_unit: str | None = None, timeout: float = 5.0, max_payment_challenge_retries: int = 3, + on_payment_session: LivePaymentSessionCallback | None = None, stream: bool = False, ) -> LiveRunnerCallResult | LiveRunnerCallStream: """Call a runner once and return its result (or a live stream if ``stream=True``). @@ -803,6 +807,13 @@ async def call_runner( paid via the signer and retried (up to ``max_payment_challenge_retries``), one job, one upfront payment. Raises ``LivepeerHTTPError`` on non-402 errors. + ``on_payment_session`` (sync or async) is called with each payment session the + signer signs, before the paid request goes out; if it raises, the request is not + sent. The signer charges when it signs, so a reported session can cost even if the + call then fails. A runner that answers 402 again after payment gets a new session; + metered calls keep paying on the one they have. ``session.manifest_id`` identifies + the orchestrator session each payment funds. + ``application/json`` and ``+json`` types parse into ``result.data``; anything else (an image, ndjson) comes back unparsed in ``result.content`` + ``result.content_type``. @@ -841,7 +852,10 @@ async def call_runner( signer_url=signer_url or "", signer_headers=signer_headers, max_price=max_price, + on_payment_session=on_payment_session, ) + except _RunnerPaymentCallbackError as e: + raise e.error except SignerRefreshRequired as e: if attempt + 1 >= attempts: raise @@ -852,6 +866,7 @@ async def call_runner( ) challenge = None continue + _validate_runner_payment(payment) request_headers["Livepeer-Payment"] = payment.payment request_headers["Livepeer-Segment"] = payment.seg_creds or "" session_id = challenge.manifest_id @@ -995,8 +1010,24 @@ async def _get_runner_payment( signer_url: str, signer_headers: dict[str, str] | None, max_price: LiveRunnerPriceInfo | None, + on_payment_session: LivePaymentSessionCallback | None, app: str | None = None, ) -> tuple[LivePaymentSession, GetPaymentResponse]: + session: LivePaymentSession + + async def _report_payment() -> None: + if on_payment_session is None: + return + try: + result = on_payment_session(session) + if inspect.isawaitable(result): + await result + except Exception as e: + raise _RunnerPaymentCallbackError(e) from e + + callback_args: dict[str, Any] = {} + if on_payment_session is not None: + callback_args["_on_payment"] = _report_payment session = LivePaymentSession( signer_url=signer_url, signer_headers=signer_headers, @@ -1004,13 +1035,37 @@ async def _get_runner_payment( challenge=challenge, app=app, max_price=max_price.to_json() if max_price is not None else None, + **callback_args, ) - payment = await session.get_payment() + payment_request = session.get_payment() + if on_payment_session is None: + payment = await payment_request + else: + payment_task = asyncio.create_task(payment_request) + try: + payment = await asyncio.shield(payment_task) + except asyncio.CancelledError: + drain = asyncio.gather(payment_task, return_exceptions=True) + while not drain.done(): + try: + await asyncio.shield(drain) + except asyncio.CancelledError: + continue + raise + return session, payment + + +class _RunnerPaymentCallbackError(Exception): + def __init__(self, error: Exception) -> None: + super().__init__(str(error)) + self.error = error + + +def _validate_runner_payment(payment: GetPaymentResponse) -> None: if not payment.payment: raise LivepeerGatewayError("Live runner payment response missing payment") if not payment.seg_creds: raise LivepeerGatewayError("Live runner payment response missing segCreds") - return session, payment def _runner_payment_type( diff --git a/src/livepeer_gateway/remote_signer.py b/src/livepeer_gateway/remote_signer.py index 1c027bf..9e008c7 100644 --- a/src/livepeer_gateway/remote_signer.py +++ b/src/livepeer_gateway/remote_signer.py @@ -1,7 +1,9 @@ from __future__ import annotations import asyncio +import inspect import logging +from collections.abc import Awaitable, Callable from dataclasses import dataclass, replace from typing import Any, Optional @@ -121,6 +123,7 @@ def __init__( app: str | None = None, max_price: dict[str, Any] | None = None, max_refresh_retries: int = 3, + _on_payment: Callable[[], None | Awaitable[None]] | None = None, ) -> None: self._signer_url = signer_url self._signer_headers = _freeze_headers(signer_headers) @@ -130,6 +133,12 @@ def __init__( self._max_price = dict(max_price) if max_price is not None else None self._max_refresh_retries = max(0, int(max_refresh_retries)) self._state: dict[str, Any] | None = None + self._on_payment = _on_payment + + @property + def manifest_id(self) -> str: + """The manifest ID from the 402 challenge this session pays for.""" + return self._challenge.manifest_id async def get_payment(self) -> GetPaymentResponse: if not self._signer_url: @@ -218,6 +227,12 @@ async def _payment_request(self) -> GetPaymentResponse: raise PaymentError( f"GetPayment error: missing/invalid 'payment' in response (url={url})" ) + on_payment = self._on_payment + self._on_payment = None + if on_payment is not None: + result = on_payment() + if inspect.isawaitable(result): + await result seg_creds = data.get("segCreds") if seg_creds is not None and not isinstance(seg_creds, str): diff --git a/tests/test_live_payment_session.py b/tests/test_live_payment_session.py index e96af45..3a4a7c7 100644 --- a/tests/test_live_payment_session.py +++ b/tests/test_live_payment_session.py @@ -105,6 +105,7 @@ async def test_send_payment_preserves_typed_http_error(self) -> None: async def test_get_payment_sends_opaque_payment_params_and_state(self) -> None: calls: list[tuple[str, dict[str, object], dict[str, str] | None]] = [] + reported: list[str] = [] async def _post_json( url: str, @@ -133,12 +134,14 @@ async def _post_json( "currency": "wei", "unit": "seconds", }, + _on_payment=lambda: reported.append("signed"), ) first = await session.get_payment() second = await session.get_payment() assert first.payment == "payment" assert second.seg_creds == "segment" + assert reported == ["signed"] assert calls[0][0] == "https://signer.example.com/generate-live-payment" assert calls[0][1] == { "orchestrator": "opaque-payment-params", @@ -291,3 +294,45 @@ async def _post_json( assert first.address == "opaque-sender" assert first.sig == "opaque-signature" assert len(calls) == 1 + + async def test_manifest_id_comes_from_the_challenge_and_survives_refresh( + self, + ) -> None: + payment_requests = 0 + + async def _post_json( + url: str, + payload: dict[str, object], + *, + headers: dict[str, str] | None = None, + timeout: float = 5.0, + ) -> dict[str, object]: + nonlocal payment_requests + del payload, headers, timeout + if url == "https://signer.example.com/generate-live-payment": + payment_requests += 1 + if payment_requests == 2: + raise SignerRefreshRequired("refresh") + return { + "payment": "payment", + "segCreds": "segment", + "state": {"state": "one"}, + } + if url == "https://signer.example.com/sign-orchestrator-info": + return {"address": "opaque-sender", "signature": "opaque-signature"} + if url == "https://orch.example.com/refresh-payment": + return {"payment_params": "new-payment-params"} + raise AssertionError(f"unexpected POST {url}") + + session = LivePaymentSession( + "https://signer.example.com", + type="fixed", + challenge=_challenge(), + ) + assert session.manifest_id == "manifest-1" + with mock.patch("livepeer_gateway.http.post_json", side_effect=_post_json): + await session.get_payment() + await session.get_payment() + + assert session.manifest_id == "manifest-1" + assert session._challenge.payment_params == "new-payment-params" diff --git a/tests/test_live_runner.py b/tests/test_live_runner.py index 3d542fd..f4912b7 100644 --- a/tests/test_live_runner.py +++ b/tests/test_live_runner.py @@ -1,22 +1,28 @@ from __future__ import annotations import asyncio +import contextlib import json import os import sys +from collections.abc import Iterator from types import SimpleNamespace from unittest import mock import pytest +from livepeer_gateway import http as gateway_http from livepeer_gateway import live_runner from livepeer_gateway.errors import ( LivepeerGatewayError, LivepeerHTTPError, + PaymentError, SignerRefreshRequired, ) from livepeer_gateway.live_runner import ( call_runner, + LivePaymentSessionCallback, + LiveRunnerCallResult, LiveRunnerGPU, LiveRunnerInstance, LiveRunnerPriceInfo, @@ -661,6 +667,334 @@ def _json_data(data: dict[str, object]) -> tuple[bytes, str]: return json.dumps(data).encode("utf-8"), "application/json" +_RUNNER_URL = "https://service.example.com/apps/runner-1/session" + + +@contextlib.contextmanager +def _paid_runner( + replies: list[object], + signs: list[object] | None = None, +) -> Iterator[tuple[list[dict[str, str]], list[object]]]: + """Fake the runner and signer behind ``call_runner``. + + The runner answers each request with the next of ``replies``: 402 is a + payment challenge, "hang" never answers, an exception is raised, and a dict + is returned as JSON. Each payment session signs with the next of ``signs``, + raising it if it is an exception; by default every session signs. Yields + each request's headers and every session created. + """ + requests: list[dict[str, str]] = [] + sessions: list[object] = [] + pending_signs = list(signs or []) + + class _PaymentSession: + def __init__(self, signer_url: str, **kwargs: object) -> None: + del signer_url + challenge = kwargs["challenge"] + assert isinstance(challenge, LivePaymentChallenge) + self._challenge = challenge + self._on_payment = kwargs.get("_on_payment") + sessions.append(self) + + @property + def manifest_id(self) -> str: + return self._challenge.manifest_id + + async def get_payment(self) -> object: + sign = pending_signs.pop(0) if pending_signs else None + if isinstance(sign, Exception): + raise sign + payment = sign or SimpleNamespace( + payment="payment-b64", + seg_creds="seg-b64", + ) + if self._on_payment is not None: + await self._on_payment() + return payment + + async def run_payments(self) -> bool: + await asyncio.Event().wait() + return False + + async def _request_body( + url: str, + *, + headers: dict[str, str] | None = None, + **kwargs: object, + ) -> tuple[bytes, str]: + del kwargs + requests.append(dict(headers or {})) + reply = replies.pop(0) + if reply == 402: + raise LivepeerHTTPError( + 402, url, _payment_challenge_body("manifest-1"), "payment required" + ) + if reply == "hang": + await asyncio.Event().wait() + if isinstance(reply, Exception): + raise reply + assert isinstance(reply, dict) + return _json_data(reply) + + with ( + mock.patch.object(live_runner, "_request_body", _request_body), + mock.patch.object(live_runner, "LivePaymentSession", _PaymentSession), + mock.patch.object( + live_runner, + "get_signer_info", + new_callable=mock.AsyncMock, + return_value=SimpleNamespace(address="opaque-payer", sig="opaque-signature"), + ), + ): + yield requests, sessions + + +async def _paid_call( + on_payment_session: LivePaymentSessionCallback, + *, + payment_unit: str = "fixed", +) -> LiveRunnerCallResult: + return await call_runner( + _RUNNER_URL, + signer_url="https://signer.example.com", + payment_unit=payment_unit, + on_payment_session=on_payment_session, + ) + + +class TestCallRunnerPaymentSessions: + @pytest.mark.parametrize("is_async", [False, True]) + async def test_reports_the_signed_session_before_the_paid_request( + self, + is_async: bool, + ) -> None: + reported: list[tuple[object, int]] = [] + with _paid_runner([402, {"ok": True}]) as (requests, _): + + def _report(session: object) -> None: + reported.append((session, len(requests))) + + async def _report_async(session: object) -> None: + await asyncio.sleep(0) + _report(session) + + result = await _paid_call( + _report_async if is_async else _report, + payment_unit="seconds", + ) + + # Reported after the unpaid request, before the paid one. + assert result.payment_session is not None + assert reported == [(result.payment_session, 1)] + assert result.payment_session.manifest_id == "manifest-1" + assert requests[1]["Livepeer-Payment"] == "payment-b64" + + async def test_reports_every_session_of_a_fixed_price_call(self) -> None: + reported: list[object] = [] + with _paid_runner([402, 402, {"ok": True}]) as (_, sessions): + result = await _paid_call(reported.append) + + assert len(sessions) == 2 + assert reported == sessions + assert result.payment_session is None + + async def test_reports_nothing_when_the_runner_asks_no_payment(self) -> None: + reported: list[object] = [] + with _paid_runner([{"ok": True}]) as (_, sessions): + await _paid_call(reported.append) + + assert sessions == [] + assert reported == [] + + async def test_runner_error_after_payment_is_raised_unchanged(self) -> None: + reported: list[object] = [] + failure = LivepeerHTTPError(500, _RUNNER_URL, "boom", "server error") + with _paid_runner([402, failure]) as (_, sessions): + with pytest.raises(LivepeerHTTPError) as info: + await _paid_call(reported.append) + + assert info.value is failure + assert reported == sessions + + async def test_skips_a_session_the_signer_refused(self) -> None: + reported: list[object] = [] + with _paid_runner([402, 402], [None, PaymentError("signer down")]) as ( + requests, + sessions, + ): + with pytest.raises(PaymentError): + await _paid_call(reported.append) + + assert len(sessions) == 2 + assert reported == sessions[:1] + assert len(requests) == 2 + + async def test_reports_a_signed_session_whose_payment_is_unusable(self) -> None: + reported: list[object] = [] + unusable = SimpleNamespace(payment="payment-b64", seg_creds="") + with _paid_runner([402], [unusable]) as (requests, sessions): + with pytest.raises(LivepeerGatewayError, match="missing segCreds"): + await _paid_call(reported.append) + + assert reported == sessions + assert len(requests) == 1 + + @pytest.mark.parametrize( + ("signer_response", "error"), + [ + ( + {"payment": "payment-b64", "segCreds": 123, "state": {}}, + "invalid 'segCreds'", + ), + ( + {"payment": "payment-b64", "segCreds": "seg-b64"}, + "missing 'state' object", + ), + ], + ) + async def test_real_session_reports_before_signer_response_validation( + self, + signer_response: dict[str, object], + error: str, + ) -> None: + requests: list[dict[str, str]] = [] + reported: list[object] = [] + + async def _request_body( + url: str, + *, + headers: dict[str, str] | None = None, + **kwargs: object, + ) -> tuple[bytes, str]: + del kwargs + requests.append(dict(headers or {})) + raise LivepeerHTTPError( + 402, + url, + _payment_challenge_body("manifest-1"), + "payment required", + ) + + with ( + mock.patch.object(live_runner, "_request_body", _request_body), + mock.patch.object( + live_runner, + "get_signer_info", + new_callable=mock.AsyncMock, + return_value=SimpleNamespace( + address="opaque-payer", + sig="opaque-signature", + ), + ), + mock.patch.object( + gateway_http, + "post_json", + new_callable=mock.AsyncMock, + return_value=signer_response, + ), + ): + with pytest.raises(PaymentError, match=error): + await _paid_call(reported.append) + + assert len(requests) == 1 + assert len(reported) == 1 + assert reported[0].manifest_id == "manifest-1" + + async def test_failing_callback_stops_the_paid_request(self) -> None: + def _report(session: object) -> None: + del session + raise RuntimeError("ledger unavailable") + + with _paid_runner([402]) as (requests, _): + with pytest.raises(RuntimeError, match="ledger unavailable"): + await _paid_call(_report) + + assert len(requests) == 1 + + async def test_callback_refresh_error_does_not_restart_the_challenge(self) -> None: + def _report(session: object) -> None: + del session + raise SignerRefreshRequired("ledger asked to refresh") + + with _paid_runner([402, 402, 402]) as (requests, sessions): + with pytest.raises(SignerRefreshRequired, match="ledger asked to refresh"): + await _paid_call(_report) + + assert len(requests) == 1 + assert len(sessions) == 1 + + async def test_reported_sessions_outlive_a_cancelled_call(self) -> None: + reported: list[object] = [] + with _paid_runner([402, "hang"]) as (requests, sessions): + with pytest.raises(TimeoutError): + await asyncio.wait_for(_paid_call(reported.append), 0.1) + + assert len(requests) == 2 + assert reported == sessions + + async def test_cancellation_drains_a_signed_payment_response_before_returning( + self, + ) -> None: + requests: list[dict[str, str]] = [] + reported: list[object] = [] + signer_charged = asyncio.Event() + release_response = asyncio.Event() + + async def _request_body( + url: str, + *, + headers: dict[str, str] | None = None, + **kwargs: object, + ) -> tuple[bytes, str]: + del kwargs + requests.append(dict(headers or {})) + if len(requests) == 1: + raise LivepeerHTTPError( + 402, + url, + _payment_challenge_body("manifest-1"), + "payment required", + ) + pytest.fail("cancelled call sent a paid runner request") + + async def _sign_payment(*args: object, **kwargs: object) -> dict[str, object]: + del args, kwargs + signer_charged.set() + await release_response.wait() + return { + "payment": "payment-b64", + "segCreds": "seg-b64", + "state": {}, + } + + with ( + mock.patch.object(live_runner, "_request_body", _request_body), + mock.patch.object( + live_runner, + "get_signer_info", + new_callable=mock.AsyncMock, + return_value=SimpleNamespace( + address="opaque-payer", + sig="opaque-signature", + ), + ), + mock.patch.object(gateway_http, "post_json", _sign_payment), + ): + call = asyncio.create_task(_paid_call(reported.append)) + await asyncio.wait_for(signer_charged.wait(), 1.0) + call.cancel() + await asyncio.sleep(0) + assert not call.done() + + release_response.set() + with pytest.raises(asyncio.CancelledError): + await call + + assert len(requests) == 1 + assert len(reported) == 1 + assert reported[0].manifest_id == "manifest-1" + + class _FakeO2RReader: instances: list[_FakeO2RReader] = []