feat(runner): report each signed payment session from call_runner

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

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

src/livepeer_gateway/live_runner.py

@@ -33,6 +33,7 @@
     GetPaymentResponse,
     LivePaymentChallenge,
     LivePaymentSession,
+    _SignedPaymentError,
     _freeze_headers,
     get_signer_info,
 )
@@ -87,6 +88,12 @@ class LiveRunnerSessionEvent:
 
 
 LiveRunnerSessionCallback = Callable[[LiveRunnerSessionEvent], None | Awaitable[None]]
+LivePaymentSessionCallback = Callable[[LivePaymentSession], None | Awaitable[None]]
+
+
+@dataclass(frozen=True)
+class _PaymentSessionCallbackFailure:
+    error: BaseException
 
 
 @dataclass(frozen=True)
@@ -762,6 +769,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 +787,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 +804,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 +813,11 @@ 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 payment string 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. Cancellation waits for
+    an in-progress signer request and callback; callback errors still propagate.
+
     ``application/json`` and ``+json`` types parse into ``result.data``; anything else
     (an image, ndjson) comes back unparsed in ``result.content`` + ``result.content_type``.
 
@@ -833,14 +848,18 @@ 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_outcome = await _await_runner_payment(
+                    payment_session,
+                    on_payment_session=on_payment_session,
                 )
             except SignerRefreshRequired as e:
                 if attempt + 1 >= attempts:
@@ -852,8 +871,15 @@ async def call_runner(
                 )
                 challenge = None
                 continue
+            if isinstance(payment_outcome, _PaymentSessionCallbackFailure):
+                raise payment_outcome.error
+            payment = payment_outcome
+            if not payment.seg_creds:
+                raise LivepeerGatewayError(
+                    "Live runner payment response missing segCreds"
+                )
             request_headers["Livepeer-Payment"] = payment.payment
-            request_headers["Livepeer-Segment"] = payment.seg_creds or ""
+            request_headers["Livepeer-Segment"] = cast(str, payment.seg_creds)
             session_id = challenge.manifest_id
 
             # Metered pricing bills for as long as the work runs.
@@ -940,6 +966,77 @@ async def call_runner(
     raise LivepeerGatewayError("Live runner call exhausted payment challenge retries")
 
 
+async def _await_runner_payment(
+    payment_session: LivePaymentSession,
+    *,
+    on_payment_session: LivePaymentSessionCallback | None,
+) -> GetPaymentResponse | _PaymentSessionCallbackFailure:
+    if on_payment_session is None:
+        return await payment_session.get_payment()
+
+    async def _pay_and_report() -> GetPaymentResponse | _PaymentSessionCallbackFailure:
+        try:
+            payment_outcome = await payment_session.get_payment()
+        except _SignedPaymentError as error:
+            payment_outcome = error
+        callback_error = await _notify_payment_session(
+            on_payment_session,
+            payment_session,
+        )
+        if callback_error is not None:
+            return _PaymentSessionCallbackFailure(callback_error)
+        if isinstance(payment_outcome, _SignedPaymentError):
+            raise payment_outcome
+        return payment_outcome
+
+    payment_task = asyncio.create_task(_pay_and_report())
+    waiter = asyncio.gather(payment_task, return_exceptions=True)
+    cancelled = False
+    while not waiter.done():
+        try:
+            await asyncio.shield(waiter)
+        # Reporting must finish even if the caller cancels more than once.
+        except asyncio.CancelledError:  # NOSONAR
+            cancelled = True
+
+    return _runner_payment_result(payment_task, cancelled=cancelled)
+
+
+def _runner_payment_result(
+    payment_task: asyncio.Task[
+        GetPaymentResponse | _PaymentSessionCallbackFailure
+    ],
+    *,
+    cancelled: bool,
+) -> GetPaymentResponse | _PaymentSessionCallbackFailure:
+    if payment_task.cancelled():
+        raise asyncio.CancelledError
+    error = payment_task.exception()
+    if error is not None:
+        if cancelled:
+            raise asyncio.CancelledError
+        raise error
+    result = payment_task.result()
+    if isinstance(result, _PaymentSessionCallbackFailure):
+        return result
+    if cancelled:
+        raise asyncio.CancelledError
+    return result
+
+
+async def _notify_payment_session(
+    callback: LivePaymentSessionCallback,
+    payment_session: LivePaymentSession,
+) -> BaseException | None:
+    try:
+        result = callback(payment_session)
+        if inspect.isawaitable(result):
+            await result
+    except (Exception, asyncio.CancelledError) as error:
+        return error
+    return None
+
+
 def _parse_runner_payment_challenge(error: LivepeerHTTPError) -> LivePaymentChallenge:
     try:
         data = json.loads(error.body)
@@ -988,31 +1085,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,

src/livepeer_gateway/remote_signer.py

@@ -18,6 +18,11 @@
 # Must stay under the signer's opening payment of 10 seconds.
 PAYMENT_INTERVAL_S = 3.0
 
+
+class _SignedPaymentError(PaymentError):
+    """The signer returned a payment string but an unusable response."""
+
+
 @dataclass(frozen=True)
 class GetPaymentResponse:
     payment: str
@@ -131,6 +136,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)
@@ -220,14 +230,17 @@ async def _payment_request(self) -> GetPaymentResponse:
             )
 
         seg_creds = data.get("segCreds")
+        # A non-string segment is unusable and carries no cursor. A blank one
+        # (missing or "") can still advance the signer's state, so store that
+        # before the caller refuses to send it.
         if seg_creds is not None and not isinstance(seg_creds, str):
-            raise PaymentError(
+            raise _SignedPaymentError(
                 f"GetPayment error: invalid 'segCreds' in response (url={url})"
             )
 
         state = data.get("state")
         if not isinstance(state, dict):
-            raise PaymentError(
+            raise _SignedPaymentError(
                 f"Remote signer response missing 'state' object (url={url})"
             )
 

tests/test_live_payment_session.py

@@ -7,6 +7,7 @@
 
 from livepeer_gateway.errors import (
     LivepeerHTTPError,
+    PaymentError,
     SignerRefreshRequired,
 )
 from livepeer_gateway.remote_signer import (
@@ -268,6 +269,61 @@ 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"
+
+    @pytest.mark.parametrize("seg_creds", [None, ""])
+    async def test_blank_segment_renewal_sends_the_new_cursor(
+        self, seg_creds: str | None
+    ) -> None:
+        calls: list[dict[str, object]] = []
+        responses = [
+            {
+                "payment": "payment-1",
+                "segCreds": "segment-1",
+                "state": {"cursor": "previous"},
+            },
+            {
+                "payment": "payment-2",
+                "segCreds": seg_creds,
+                "state": {"cursor": "advanced"},
+            },
+            {
+                "payment": "payment-3",
+                "segCreds": "segment-3",
+                "state": {"cursor": "later"},
+            },
+        ]
+
+        async def _post_json(
+            url: str,
+            payload: dict[str, object],
+            *,
+            headers: dict[str, str] | None = None,
+            timeout: float = 5.0,
+        ) -> dict[str, object]:
+            del url, headers, timeout
+            calls.append(payload)
+            return responses[len(calls) - 1]
+
+        post_empty = mock.AsyncMock()
+        with (
+            mock.patch("livepeer_gateway.http.post_json", side_effect=_post_json),
+            mock.patch("livepeer_gateway.http._post_empty", post_empty),
+        ):
+            session = LivePaymentSession(
+                "https://signer.example.com",
+                type="live",
+                challenge=_challenge(),
+            )
+            await session.send_payment()
+            with pytest.raises(PaymentError, match="no segCreds"):
+                await session.send_payment()
+            await session.send_payment()
+
+        assert post_empty.await_count == 2
+        assert "state" not in calls[0]
+        assert calls[1]["state"] == {"cursor": "previous"}
+        assert calls[2]["state"] == {"cursor": "advanced"}
 
     async def test_get_signer_info_caches_result(self) -> None:
         calls: list[tuple[str, dict[str, object]]] = []