diff --git a/src/livepeer_gateway/byoc.py b/src/livepeer_gateway/byoc.py index 4d27781..970d2b1 100644 --- a/src/livepeer_gateway/byoc.py +++ b/src/livepeer_gateway/byoc.py @@ -36,11 +36,12 @@ import ssl import uuid from dataclasses import dataclass, field -from typing import Any, Optional, Sequence +from typing import Any, Literal, Optional, Sequence from urllib.error import HTTPError, URLError from urllib.parse import urlparse from urllib.request import Request, urlopen +from .capabilities import byoc_capabilities_from_app from .orchestrator import _http_origin, _parse_http_url, discover_orchestrators from .errors import LivepeerGatewayError, NoOrchestratorAvailableError, OrchestratorRejection @@ -144,6 +145,22 @@ def audio_url(self) -> Optional[str]: # Header building # --------------------------------------------------------------------------- +_LEGACY_DAYDREAM_SIGNER_HOST = "signer.daydream.live" + + +def _payment_type_for_signer(signer_url: str) -> Literal["byoc", "lv2v"]: + """ + Select payment payload shape for the target signer. + + Legacy Daydream signer expects type:lv2v + string capability. + pymthouse DMZ (and other modern signers) expect type:byoc + capabilities proto. + """ + hostname = (urlparse(signer_url).hostname or "").lower() + if hostname == _LEGACY_DAYDREAM_SIGNER_HOST: + return "lv2v" + return "byoc" + + def _create_byoc_payment( *, orch_origin: str, @@ -173,10 +190,20 @@ def _create_byoc_payment( parsed = urlparse(orch_origin) grpc_url = f"https://{parsed.hostname}:8935" + payment_type = _payment_type_for_signer(signer_url) + byoc_caps = ( + byoc_capabilities_from_app(capability) + if payment_type == "byoc" + else None + ) + + # Capabilities on orch discovery are required for type:byoc so TicketParams + # use PriceInfoForCaps (per-cap wei/sec) when ByocPerCapPricing is enabled. info = get_orch_info( grpc_url, signer_url=signer_url, signer_headers=signer_headers, + capabilities=byoc_caps, ) # Check if orch has a price set — if price is 0, skip payment @@ -189,16 +216,22 @@ def _create_byoc_payment( _LOG.info("BYOC orch has no ticket_params, skipping payment") return {} - # Step 2: Generate payment via signer + # Step 2: Generate payment via signer — shape depends on signer host. orch_info_b64 = base64.b64encode(info.SerializeToString()).decode("ascii") + payment_payload: dict[str, Any] = { + "orchestrator": orch_info_b64, + "type": payment_type, + } + if payment_type == "byoc" and byoc_caps is not None: + payment_payload["capabilities"] = base64.b64encode( + byoc_caps.SerializeToString() + ).decode("ascii") + elif payment_type == "lv2v": + payment_payload["capability"] = capability signer_origin = _http_origin(signer_url) payment_url = f"{signer_origin}/generate-live-payment" - payment_body = json.dumps({ - "orchestrator": orch_info_b64, - "type": "lv2v", - "capability": capability, - }).encode("utf-8") + payment_body = json.dumps(payment_payload).encode("utf-8") payment_headers = { "Content-Type": "application/json", "Livepeer-Capability": capability, @@ -759,8 +792,8 @@ def refresh_training_payment( Args: job_id: training job_id assigned by orch on submit. orch_url: orchestrator URL accepting the job. - capability: capability name (needed by signer to pick correct - ticket_params). + capability: capability/app name encoded into the BYOC capabilities + proto sent to the signer (not a separate string field). signer_url: remote signer with /generate-live-payment. signer_headers: pass-through headers (notably Authorization Bearer of the user). diff --git a/src/livepeer_gateway/capabilities.py b/src/livepeer_gateway/capabilities.py index 27e5688..a400dc0 100644 --- a/src/livepeer_gateway/capabilities.py +++ b/src/livepeer_gateway/capabilities.py @@ -45,6 +45,7 @@ class CapabilityId(IntEnum): IMAGE_TO_TEXT = 34 LIVE_VIDEO_TO_VIDEO = 35 TEXT_TO_SPEECH = 36 + BYOC = 37 CAPABILITY_ID_TO_NAME: dict[int, str] = { -2: "Invalid", @@ -85,9 +86,18 @@ class CapabilityId(IntEnum): 34: "Image to text", 35: "Live video to video", 36: "Text to speech", + 37: "byoc", } +def byoc_capabilities_from_app(app: str) -> Optional[lp_rpc_pb2.Capabilities]: + """Build BYOC capability constraints from a capability/app name.""" + app = app.strip() + if not app: + return None + return build_capabilities(CapabilityId.BYOC, app) + + def capability_name(cap_id: int) -> str: return CAPABILITY_ID_TO_NAME.get(cap_id, "Unknown capability") diff --git a/tests/test_byoc_payment.py b/tests/test_byoc_payment.py new file mode 100644 index 0000000..356cd9f --- /dev/null +++ b/tests/test_byoc_payment.py @@ -0,0 +1,186 @@ +""" +Unit tests for _create_byoc_payment orch-discovery capabilities threading. + +When ByocPerCapPricing is enabled on the signer, TicketParams must be issued +via PriceInfoForCaps — that requires passing capabilities into get_orch_info(). +""" + +from __future__ import annotations + +import base64 +import json +from contextlib import contextmanager +from unittest.mock import MagicMock, patch +from urllib.request import Request + +from livepeer_gateway.byoc import _create_byoc_payment, _payment_type_for_signer +from livepeer_gateway.capabilities import CapabilityId, byoc_capabilities_from_app + + +def _stub_orch_info(): + info = MagicMock() + tp = MagicMock() + tp.face_value = b"\x01\x00" + info.ticket_params = tp + info.HasField = lambda field: field == "ticket_params" + info.SerializeToString = lambda: b"stub-orch-info-protobuf" + return info + + +@contextmanager +def _mock_payment_http(*, get_orch_info_side_effect): + captured: list[Request] = [] + get_orch_info_calls: list[dict] = [] + + class _MockResponse: + def __init__(self, body: bytes, status: int = 200): + self._body = body + self.status = status + self.headers = {} + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def _fake_get_orch_info(*args, **kwargs): + get_orch_info_calls.append(kwargs) + return get_orch_info_side_effect(*args, **kwargs) + + def _fake_urlopen(req, *args, **kwargs): + captured.append(req) + return _MockResponse( + json.dumps({"payment": "TICKETS_B64", "segCreds": "SEG_B64"}).encode() + ) + + with patch("livepeer_gateway.byoc.urlopen", side_effect=_fake_urlopen), \ + patch("livepeer_gateway.orch_info.get_orch_info", side_effect=_fake_get_orch_info): + yield captured, get_orch_info_calls + + +def test_create_byoc_payment_passes_capabilities_to_get_orch_info(): + capability = "flux-schnell" + expected_caps = byoc_capabilities_from_app(capability) + + with _mock_payment_http(get_orch_info_side_effect=lambda *a, **k: _stub_orch_info()) as ( + _reqs, + orch_calls, + ): + result = _create_byoc_payment( + orch_origin="https://byoc-staging-1.daydream.monster:8936", + capability=capability, + livepeer_hdr="ignored", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + ) + + assert result["Livepeer-Payment"] == "TICKETS_B64" + assert len(orch_calls) == 1 + passed_caps = orch_calls[0]["capabilities"] + assert passed_caps is not None + assert passed_caps.capacities[int(CapabilityId.BYOC)] == 1 + assert ( + "flux-schnell" + in passed_caps.constraints.PerCapability[int(CapabilityId.BYOC)].models + ) + assert passed_caps.SerializeToString() == expected_caps.SerializeToString() + + +def test_create_byoc_payment_includes_capabilities_on_signer_payload(): + capability = "flux-dev" + + with _mock_payment_http(get_orch_info_side_effect=lambda *a, **k: _stub_orch_info()) as ( + reqs, + _orch_calls, + ): + _create_byoc_payment( + orch_origin="https://byoc-staging-1.daydream.monster:8936", + capability=capability, + livepeer_hdr="ignored", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + ) + + signer_req = [r for r in reqs if "generate-live-payment" in r.full_url][0] + payload = json.loads(signer_req.data.decode("utf-8")) + assert payload["type"] == "byoc" + assert "capabilities" in payload + + caps = byoc_capabilities_from_app(capability) + assert payload["capabilities"] == base64.b64encode( + caps.SerializeToString() + ).decode("ascii") + + +def test_payment_type_for_signer_legacy_daydream(): + assert _payment_type_for_signer("https://signer.daydream.live") == "lv2v" + assert _payment_type_for_signer("https://signer.daydream.live/generate-live-payment") == "lv2v" + + +def test_payment_type_for_signer_pymthouse_dmz(): + assert ( + _payment_type_for_signer( + "https://pymthouse-production.up.railway.app" + ) + == "byoc" + ) + assert _payment_type_for_signer("https://signer.test") == "byoc" + + +def test_create_byoc_payment_legacy_daydream_signer_uses_lv2v(): + capability = "flux-schnell" + + with _mock_payment_http(get_orch_info_side_effect=lambda *a, **k: _stub_orch_info()) as ( + reqs, + orch_calls, + ): + _create_byoc_payment( + orch_origin="https://byoc-staging-1.daydream.monster:8936", + capability=capability, + livepeer_hdr="ignored", + signer_url="https://signer.daydream.live", + signer_headers={"Authorization": "Bearer sk_test"}, + ) + + assert len(orch_calls) == 1 + assert orch_calls[0]["capabilities"] is None + + signer_req = [r for r in reqs if "generate-live-payment" in r.full_url][0] + payload = json.loads(signer_req.data.decode("utf-8")) + assert payload["type"] == "lv2v" + assert payload["capability"] == capability + assert "capabilities" not in payload + + +def test_create_byoc_payment_pymthouse_signer_uses_byoc(): + capability = "flux-dev" + + with _mock_payment_http(get_orch_info_side_effect=lambda *a, **k: _stub_orch_info()) as ( + reqs, + orch_calls, + ): + _create_byoc_payment( + orch_origin="https://byoc-staging-1.daydream.monster:8936", + capability=capability, + livepeer_hdr="ignored", + signer_url="https://pymthouse-production.up.railway.app", + signer_headers={"Authorization": "Bearer app.pmth_test"}, + ) + + assert len(orch_calls) == 1 + passed_caps = orch_calls[0]["capabilities"] + assert passed_caps is not None + assert ( + capability + in passed_caps.constraints.PerCapability[int(CapabilityId.BYOC)].models + ) + + signer_req = [r for r in reqs if "generate-live-payment" in r.full_url][0] + payload = json.loads(signer_req.data.decode("utf-8")) + assert payload["type"] == "byoc" + assert "capabilities" in payload + assert "capability" not in payload diff --git a/tests/test_capabilities.py b/tests/test_capabilities.py new file mode 100644 index 0000000..e821e28 --- /dev/null +++ b/tests/test_capabilities.py @@ -0,0 +1,20 @@ +"""Unit tests for BYOC capability protobuf helpers.""" + +from __future__ import annotations + +from livepeer_gateway.capabilities import ( + CapabilityId, + byoc_capabilities_from_app, +) + + +def test_byoc_capabilities_from_app_builds_constraint(): + caps = byoc_capabilities_from_app("flux-schnell") + assert caps is not None + assert caps.capacities[int(CapabilityId.BYOC)] == 1 + assert "flux-schnell" in caps.constraints.PerCapability[int(CapabilityId.BYOC)].models + + +def test_byoc_capabilities_from_app_empty_returns_none(): + assert byoc_capabilities_from_app("") is None + assert byoc_capabilities_from_app(" ") is None