Add payment session callback for runner calls

CHANGELOG.md

@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ### Added
 
+- `call_runner(on_payment_session=...)` callback for each new payment session
+- `LivePaymentSession.manifest_id` for payment attribution
 - Flexible request body types on call_runner and selectors.
 
 ### Changed

src/livepeer_gateway/__init__.py

@@ -27,6 +27,7 @@
 )
 from .errors import RunnerRejection
 from .live_runner import (
+    LivePaymentSessionCallback,
     LiveRunnerCallResult,
     LiveRunnerCallStream,
     LiveRunnerGPU,
@@ -79,6 +80,7 @@
     "LiveRunnerSessionEvent",
     "LivePaymentChallenge",
     "LivePaymentSession",
+    "LivePaymentSessionCallback",
     "LiveRunnerProxy",
     "RunnerPayload",
     "StreamPayload",

src/livepeer_gateway/live_runner.py

@@ -31,7 +31,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,
@@ -89,6 +88,7 @@ class LiveRunnerSessionEvent:
 
 
 LiveRunnerSessionCallback = Callable[[LiveRunnerSessionEvent], None | Awaitable[None]]
+LivePaymentSessionCallback = Callable[[LivePaymentSession], None | Awaitable[None]]
 
 
 @dataclass(frozen=True)
@@ -768,6 +768,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
@@ -786,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
@@ -803,6 +805,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``).
@@ -811,6 +814,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``.
 
@@ -852,15 +859,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
@@ -871,8 +879,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.
@@ -1012,31 +1024,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

@@ -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})"
             )

tests/test_call_runner_uploads.py

@@ -33,8 +33,8 @@
 @pytest.fixture
 def signer(monkeypatch):
     monkeypatch.setattr(live_runner, "get_signer_info", mock.AsyncMock(return_value=SimpleNamespace(address="payer")))
-    payment = mock.AsyncMock(return_value=(None, SimpleNamespace(payment="paid", seg_creds="segment")))
-    monkeypatch.setattr(live_runner, "_get_runner_payment", payment)
+    payment = mock.AsyncMock(return_value={"payment": "paid", "segCreds": "segment", "state": {}})
+    monkeypatch.setattr("livepeer_gateway.http.post_json", payment)
     return payment
 
 
@@ -207,6 +207,7 @@ async def test_early_payment_challenge_does_not_touch_source(
 ):
     source = Source(b"one", b"two")
     attempts = 0
+    sessions = []
 
     async def handler(reader, writer):
         nonlocal attempts
@@ -219,6 +220,7 @@ async def handler(reader, writer):
             assert await reader.read() == b""
             return
         assert request_headers["Livepeer-Payment"] == "paid"
+        assert [session.manifest_id for session in sessions] == [CHALLENGE["manifest_id"]]
         assert source.iterations == source.reads == 0
         await allow_upload(writer)
         assert await body(reader, request_headers) == b"onetwo"
@@ -234,6 +236,7 @@ async def handler(reader, writer):
                 payload=StreamPayload(source=source, retry_buffer=False),
                 signer_url="http://signer.example",
                 payment_unit="fixed",
+                on_payment_session=sessions.append,
                 stream=stream,
             ),
             2,
@@ -244,10 +247,9 @@ async def handler(reader, writer):
     assert source.reads == 3
     assert not source.closed
     signer.assert_awaited_once()
-    challenge = signer.call_args.args[0]
-    assert challenge.payment_params == CHALLENGE["payment_params"]
-    assert challenge.manifest_id == CHALLENGE["manifest_id"]
-    assert challenge.payment_url == CHALLENGE["payment_url"]
+    payment_request = signer.call_args.args[1]
+    assert payment_request["orchestrator"] == CHALLENGE["payment_params"]
+    assert payment_request["ManifestID"] == CHALLENGE["manifest_id"]
 
 
 @pytest.mark.parametrize("status,kind", [(402, "nonseekable"), (503, "iterator"), (307, "file")])

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]]] = []