From 64b67e73282ee947c7794168b42bc59ad7c49d0f Mon Sep 17 00:00:00 2001 From: Josh Allmann Date: Mon, 19 Oct 2026 20:21:37 -0700 Subject: [PATCH 1/2] Add payment session callback for runner calls Support sync and async on_payment_session callbacks in call_runner and expose LivePaymentSession.manifest_id for payment attribution. Invoke the callback for each new session. Validate segCreds when parsing signer responses and remove redundant caller checks. Cover callback ordering, failures, and cancellation. --- CHANGELOG.md | 5 + src/livepeer_gateway/live_runner.py | 56 ++++----- src/livepeer_gateway/remote_signer.py | 11 +- tests/test_live_payment_session.py | 1 + tests/test_live_runner.py | 171 +++++++++++++++++++++++++- 5 files changed, 203 insertions(+), 41 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a067f9f..bb0299c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/). ## [Unreleased] +### Added + +- `call_runner(on_payment_session=...)` callback for each new payment session +- `LivePaymentSession.manifest_id` for payment attribution + ### Changed - `discover_orchestrators` is now asynchronous and must be awaited, including diff --git a/src/livepeer_gateway/live_runner.py b/src/livepeer_gateway/live_runner.py index 8d80b3b..7b06e69 100644 --- a/src/livepeer_gateway/live_runner.py +++ b/src/livepeer_gateway/live_runner.py @@ -30,7 +30,6 @@ from .errors import LivepeerGatewayError, LivepeerHTTPError, SignerRefreshRequired from .http import _post_empty, _request_body, open_stream, post_json, request_json from .remote_signer import ( - GetPaymentResponse, LivePaymentChallenge, LivePaymentSession, _freeze_headers, @@ -762,6 +761,7 @@ async def call_runner( payment_unit: str | None = ..., timeout: float = ..., max_payment_challenge_retries: int = ..., + on_payment_session: Callable[[LivePaymentSession], None | Awaitable[None]] | None = ..., stream: Literal[False] = False, ) -> LiveRunnerCallResult: pass @@ -779,6 +779,7 @@ async def call_runner( payment_unit: str | None = ..., timeout: float = ..., max_payment_challenge_retries: int = ..., + on_payment_session: Callable[[LivePaymentSession], None | Awaitable[None]] | None = ..., stream: Literal[True], ) -> LiveRunnerCallStream: pass @@ -795,6 +796,7 @@ async def call_runner( payment_unit: str | None = None, timeout: float = 5.0, max_payment_challenge_retries: int = 3, + on_payment_session: Callable[[LivePaymentSession], None | Awaitable[None]] | None = None, stream: bool = False, ) -> LiveRunnerCallResult | LiveRunnerCallStream: """Call a runner once and return its result (or a live stream if ``stream=True``). @@ -803,6 +805,10 @@ 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. + After a new payment session is obtained, `on_payment_session` is called. + Metered renewals do not invoke it. The callback may be sync or + async. Its exceptions propagate and abort the call. + ``application/json`` and ``+json`` types parse into ``result.data``; anything else (an image, ndjson) comes back unparsed in ``result.content`` + ``result.content_type``. @@ -833,15 +839,16 @@ async def call_runner( request_headers[_LIVE_RUNNER_PAYER_ADDRESS_HEADER] = payer_address # Pending challenge means payment is needed. if challenge is not None: + payment_session = LivePaymentSession( + signer_url=signer_url or "", + signer_headers=signer_headers, + type=payment_type, + challenge=challenge, + app=runner.app if runner is not None else None, + max_price=max_price.to_json() if max_price is not None else None, + ) try: - payment_session, payment = await _get_runner_payment( - challenge, - payment_type=payment_type, - app=runner.app if runner is not None else None, - signer_url=signer_url or "", - signer_headers=signer_headers, - max_price=max_price, - ) + payment = await payment_session.get_payment() except SignerRefreshRequired as e: if attempt + 1 >= attempts: raise @@ -852,8 +859,12 @@ async def call_runner( ) challenge = None continue + if on_payment_session is not None: + result = on_payment_session(payment_session) + if inspect.isawaitable(result): + await result request_headers["Livepeer-Payment"] = payment.payment - request_headers["Livepeer-Segment"] = payment.seg_creds or "" + request_headers["Livepeer-Segment"] = payment.seg_creds session_id = challenge.manifest_id # Metered pricing bills for as long as the work runs. @@ -988,31 +999,6 @@ async def _stop_funding(task: asyncio.Task[None] | None) -> None: await task -async def _get_runner_payment( - challenge: LivePaymentChallenge, - *, - payment_type: str, - signer_url: str, - signer_headers: dict[str, str] | None, - max_price: LiveRunnerPriceInfo | None, - app: str | None = None, -) -> tuple[LivePaymentSession, GetPaymentResponse]: - session = LivePaymentSession( - signer_url=signer_url, - signer_headers=signer_headers, - type=payment_type, - challenge=challenge, - app=app, - max_price=max_price.to_json() if max_price is not None else None, - ) - payment = await session.get_payment() - 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( runner: LiveRunnerInstance | None, payment_unit: str | None = None, diff --git a/src/livepeer_gateway/remote_signer.py b/src/livepeer_gateway/remote_signer.py index 1c027bf..6196aea 100644 --- a/src/livepeer_gateway/remote_signer.py +++ b/src/livepeer_gateway/remote_signer.py @@ -131,6 +131,11 @@ def __init__( self._max_refresh_retries = max(0, int(max_refresh_retries)) self._state: dict[str, Any] | None = None + @property + def manifest_id(self) -> str: + """The manifest ID of the payment challenge for this session.""" + return self._challenge.manifest_id + async def get_payment(self) -> GetPaymentResponse: if not self._signer_url: return GetPaymentResponse(payment="", seg_creds=None) @@ -161,10 +166,6 @@ async def send_payment(self) -> None: from .http import _post_empty payment = await self.get_payment() - if not payment.seg_creds: - # An empty segment header fails the orchestrator's sig check and - # comes back 403, which reads as a dead session, not a bad signer. - raise PaymentError("Signer returned a payment with no segCreds") headers = { "Livepeer-Payment": payment.payment, "Livepeer-Segment": payment.seg_creds, @@ -220,7 +221,7 @@ async def _payment_request(self) -> GetPaymentResponse: ) seg_creds = data.get("segCreds") - if seg_creds is not None and not isinstance(seg_creds, str): + if not isinstance(seg_creds, str) or not seg_creds: raise PaymentError( f"GetPayment error: invalid 'segCreds' in response (url={url})" ) diff --git a/tests/test_live_payment_session.py b/tests/test_live_payment_session.py index e96af45..d07bffe 100644 --- a/tests/test_live_payment_session.py +++ b/tests/test_live_payment_session.py @@ -268,6 +268,7 @@ async def _post_json( assert calls[4][1]["orchestrator"] == "new-payment-params" assert session._challenge.payment_url == _PAYMENT_URL + assert session.manifest_id == "manifest-1" async def test_get_signer_info_caches_result(self) -> None: calls: list[tuple[str, dict[str, object]]] = [] diff --git a/tests/test_live_runner.py b/tests/test_live_runner.py index 3d542fd..6be0f73 100644 --- a/tests/test_live_runner.py +++ b/tests/test_live_runner.py @@ -26,7 +26,7 @@ stop_runner_session, create_proxy, ) -from livepeer_gateway.remote_signer import LivePaymentChallenge +from livepeer_gateway.remote_signer import LivePaymentChallenge, LivePaymentSession class TestLiveRunnerHelpers: @@ -631,6 +631,175 @@ def _request_body( assert payment_attempts == 2 +class TestPaymentSessionCallback: + @pytest.fixture + def payment_call(self, request): + url = "https://runner.example.com/call" + stream = getattr(request, "param", False) + reply = ( + ( + SimpleNamespace(close=mock.AsyncMock()), + SimpleNamespace(status=200, headers={}, release=mock.Mock()), + ) + if stream + else _json_data({"ok": True}) + ) + + def challenge(manifest): + return LivepeerHTTPError(402, url, _payment_challenge_body(manifest)) + + transport = mock.AsyncMock(side_effect=[challenge("manifest-1"), reply]) + signer = mock.AsyncMock( + return_value={"payment": "payment", "segCreds": "segment", "state": {}} + ) + with ( + mock.patch.object( + live_runner, "open_stream" if stream else "_request_body", transport + ), + mock.patch.object( + live_runner, + "get_signer_info", + mock.AsyncMock(return_value=SimpleNamespace(address="payer")), + ), + mock.patch("livepeer_gateway.http.post_json", signer), + ): + yield SimpleNamespace( + kwargs=dict( + runner_url=url, + signer_url="https://signer.example.com", + payment_unit="fixed", + stream=stream, + ), + transport=transport, + signer=signer, + reply=reply, + challenge=challenge, + ) + + @pytest.mark.parametrize( + "payment_call,async_callback,payment_unit", + [(False, False, "fixed"), (True, True, "seconds")], + indirect=["payment_call"], + ) + async def test_reports_each_successful_session_before_dispatch( + self, payment_call, async_callback, payment_unit + ) -> None: + call = payment_call + call.kwargs["payment_unit"] = payment_unit + call.signer.side_effect = [SignerRefreshRequired("refresh")] + [ + call.signer.return_value + ] * 2 + call.transport.side_effect = [ + call.challenge("refused"), + call.challenge("manifest-1"), + call.challenge("manifest-2"), + call.reply, + ] + sessions: list[LivePaymentSession] = [] + + def collect(session): + assert ( + call.signer.await_count + == call.transport.await_count + == len(sessions) + 2 + ) + sessions.append(session) + + callback = mock.AsyncMock(side_effect=collect) if async_callback else collect + result = await call_runner(**call.kwargs, on_payment_session=callback) + if call.kwargs["stream"]: + await result.aclose() + assert [s.manifest_id for s in sessions] == ["manifest-1", "manifest-2"] + assert result.payment_session is ( + sessions[-1] if payment_unit == "seconds" else None + ) + assert call.transport.await_count == 4 + + @pytest.mark.parametrize( + "error", + [ + RuntimeError("failed"), + SignerRefreshRequired("failed"), + LivepeerHTTPError(402, "https://callback.example.com"), + ], + ) + async def test_callback_errors_abort_without_retry( + self, payment_call, error + ) -> None: + callback = mock.AsyncMock(side_effect=error) + with pytest.raises(type(error)) as raised: + await call_runner(**payment_call.kwargs, on_payment_session=callback) + assert raised.value is error + assert raised.value.__context__ is None + assert ( + payment_call.transport.await_count == payment_call.signer.await_count == 1 + ) + + @pytest.mark.parametrize( + "outcome", + [ + LivepeerHTTPError(403, "https://signer.example.com"), + {"segCreds": "segment", "state": {}}, + {"payment": "payment", "state": {}}, + {"payment": "payment", "segCreds": None, "state": {}}, + {"payment": "payment", "segCreds": "", "state": {}}, + {"payment": "payment", "segCreds": 123, "state": {}}, + ], + ) + async def test_signer_failure_does_not_report(self, payment_call, outcome) -> None: + if isinstance(outcome, Exception): + payment_call.signer.side_effect = outcome + else: + payment_call.signer.return_value = outcome + callback = mock.Mock() + with pytest.raises(LivepeerGatewayError): + await call_runner(**payment_call.kwargs, on_payment_session=callback) + callback.assert_not_called() + assert payment_call.transport.await_count == 1 + + async def test_paid_request_failure_preserves_notification( + self, payment_call + ) -> None: + call = payment_call + error = LivepeerHTTPError(500, "https://runner.example.com") + call.transport.side_effect = [call.challenge("manifest-1"), error] + sessions: list[LivePaymentSession] = [] + with pytest.raises(LivepeerHTTPError) as raised: + await call_runner(**call.kwargs, on_payment_session=sessions.append) + assert raised.value is error + assert [s.manifest_id for s in sessions] == ["manifest-1"] + + @pytest.mark.parametrize("phase", ["signing", "callback"]) + async def test_cancellation_remains_prompt(self, payment_call, phase) -> None: + call = payment_call + entered, release = asyncio.Event(), asyncio.Event() + + async def block(*_args, **_kwargs): + entered.set() + await release.wait() + return call.signer.return_value + + callback = mock.AsyncMock(side_effect=block if phase == "callback" else None) + if phase == "signing": + call.signer.side_effect = block + task = asyncio.create_task( + call_runner(**call.kwargs, on_payment_session=callback) + ) + try: + await asyncio.wait_for(entered.wait(), timeout=1) + task.cancel() + done, _ = await asyncio.wait([task], timeout=1) + assert task in done + with pytest.raises(asyncio.CancelledError): + await task + assert callback.await_count == (1 if phase == "callback" else 0) + assert call.transport.await_count == 1 + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + def _payment_challenge_body(manifest_id: str) -> str: return json.dumps( { From 49ec8d3acce1f1d82b3773b26a88be6eb6b7b8d2 Mon Sep 17 00:00:00 2001 From: Josh Allmann Date: Wed, 21 Oct 2026 16:09:59 -0700 Subject: [PATCH 2/2] LivePaymentSessionCallback type alias --- src/livepeer_gateway/__init__.py | 2 ++ src/livepeer_gateway/live_runner.py | 7 ++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/livepeer_gateway/__init__.py b/src/livepeer_gateway/__init__.py index 9e4c284..f3fee96 100644 --- a/src/livepeer_gateway/__init__.py +++ b/src/livepeer_gateway/__init__.py @@ -27,6 +27,7 @@ ) from .errors import RunnerRejection from .live_runner import ( + LivePaymentSessionCallback, LiveRunnerCallResult, LiveRunnerCallStream, LiveRunnerGPU, @@ -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 7b06e69..881ba92 100644 --- a/src/livepeer_gateway/live_runner.py +++ b/src/livepeer_gateway/live_runner.py @@ -86,6 +86,7 @@ class LiveRunnerSessionEvent: LiveRunnerSessionCallback = Callable[[LiveRunnerSessionEvent], None | Awaitable[None]] +LivePaymentSessionCallback = Callable[[LivePaymentSession], None | Awaitable[None]] @dataclass(frozen=True) @@ -761,7 +762,7 @@ async def call_runner( payment_unit: str | None = ..., timeout: float = ..., max_payment_challenge_retries: int = ..., - on_payment_session: Callable[[LivePaymentSession], None | Awaitable[None]] | None = ..., + on_payment_session: LivePaymentSessionCallback | None = ..., stream: Literal[False] = False, ) -> LiveRunnerCallResult: pass @@ -779,7 +780,7 @@ async def call_runner( payment_unit: str | None = ..., timeout: float = ..., max_payment_challenge_retries: int = ..., - on_payment_session: Callable[[LivePaymentSession], None | Awaitable[None]] | None = ..., + on_payment_session: LivePaymentSessionCallback | None = ..., stream: Literal[True], ) -> LiveRunnerCallStream: pass @@ -796,7 +797,7 @@ async def call_runner( payment_unit: str | None = None, timeout: float = 5.0, max_payment_challenge_retries: int = 3, - on_payment_session: Callable[[LivePaymentSession], None | Awaitable[None]] | None = None, + 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``).