From 5b3a1ef23bf3320d36ef4fbfae2516b2f34cff78 Mon Sep 17 00:00:00 2001 From: seanhanca Date: Fri, 26 Jun 2026 20:21:33 -0700 Subject: [PATCH 1/2] fix(byoc): send real capability + model_id in live payment request `_create_byoc_payment` sent the BYOC `capability` only as a bare JSON field that the remote signer dropped (its RemotePaymentRequest had no such field) and sent no model id at all. Combined with the hardcoded `type:"lv2v"`, the signer's create_signed_ticket metering event recorded pipeline=live-video-to-video / model_id=unknown for every BYOC job. Thread the REAL capability (already present) and the model id into the `/generate-live-payment` body so the signer can attribute usage to the true pipeline + model. `type` stays "lv2v" so fee/pixel routing is unchanged. A new `_extract_model_id` helper locates the model in the request payload (`model_id`/`model`), falling back to job parameters; when no model id is found the field is omitted so behavior is byte-identical to today (the signer falls back). Adds tests covering the wire contract (capability + model_id present for a nano-banana request), model extraction precedence, and the backward-compatible omit-when-absent path. --- src/livepeer_gateway/byoc.py | 48 +++++- tests/test_byoc_usage_attribution.py | 209 +++++++++++++++++++++++++++ 2 files changed, 255 insertions(+), 2 deletions(-) create mode 100644 tests/test_byoc_usage_attribution.py diff --git a/src/livepeer_gateway/byoc.py b/src/livepeer_gateway/byoc.py index 4d27781..ec70488 100644 --- a/src/livepeer_gateway/byoc.py +++ b/src/livepeer_gateway/byoc.py @@ -144,6 +144,36 @@ def audio_url(self) -> Optional[str]: # Header building # --------------------------------------------------------------------------- +# Keys, in priority order, under which a BYOC request conventionally carries +# the specific provider model id (distinct from the coarse capability name). +_MODEL_ID_KEYS = ("model_id", "model") + + +def _extract_model_id( + payload: Optional[dict[str, Any]], + parameters: Optional[dict[str, Any]] = None, +) -> str: + """Best-effort extraction of the provider model id from a BYOC request. + + The *model id* is the specific provider model the worker runs (e.g. the + value billed/attributed per generation), as opposed to the *capability* + name (e.g. ``nano-banana``) which selects the orchestrator pipeline. It is + conventionally carried in the request ``payload`` under ``model_id`` or + ``model``; some capabilities place it in the job ``parameters`` instead. + + Returns an empty string when no model id can be determined so callers can + fall back to today's behavior (the signer then derives/defaults it). + """ + for source in (payload, parameters): + if not isinstance(source, dict): + continue + for key in _MODEL_ID_KEYS: + value = source.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + def _create_byoc_payment( *, orch_origin: str, @@ -151,6 +181,7 @@ def _create_byoc_payment( livepeer_hdr: str, signer_url: str, signer_headers: Optional[dict[str, str]] = None, + model_id: str = "", timeout: float = 30.0, ) -> dict[str, str]: """ @@ -194,11 +225,22 @@ def _create_byoc_payment( signer_origin = _http_origin(signer_url) payment_url = f"{signer_origin}/generate-live-payment" - payment_body = json.dumps({ + # `type` stays "lv2v" so the remote signer keeps reusing the live + # fee/pixel routing for BYOC jobs (BYOC has no per-job pixel count). The + # additive `capability` + `model_id` fields below carry the REAL usage + # attribution so the signer's `create_signed_ticket` metering event + # records the true pipeline + model instead of `live-video-to-video` / + # `unknown`. Both are backward-compatible: an older signer ignores the + # extra fields, and an empty `model_id` is simply omitted so behavior is + # byte-identical to today. + payment_payload: dict[str, Any] = { "orchestrator": orch_info_b64, "type": "lv2v", "capability": capability, - }).encode("utf-8") + } + if model_id: + payment_payload["model_id"] = model_id + payment_body = json.dumps(payment_payload).encode("utf-8") payment_headers = { "Content-Type": "application/json", "Livepeer-Capability": capability, @@ -390,6 +432,7 @@ def submit_byoc_job( livepeer_hdr=livepeer_hdr, signer_url=signer_url, signer_headers=signer_headers, + model_id=_extract_model_id(req.payload, req.parameters), timeout=http_timeout, ) headers.update(payment_headers) @@ -673,6 +716,7 @@ def submit_training_job( livepeer_hdr=livepeer_hdr, signer_url=signer_url, signer_headers=signer_headers, + model_id=_extract_model_id(byoc_req.payload, byoc_req.parameters), timeout=http_timeout, ) headers.update(payment_headers) diff --git a/tests/test_byoc_usage_attribution.py b/tests/test_byoc_usage_attribution.py new file mode 100644 index 0000000..afabe06 --- /dev/null +++ b/tests/test_byoc_usage_attribution.py @@ -0,0 +1,209 @@ +""" +Unit tests for BYOC usage attribution (real capability + model_id). + +Root cause being fixed: ``_create_byoc_payment`` used to send only a bare +``capability`` field (which the remote signer dropped) and no model id, so the +signer's ``create_signed_ticket`` metering event recorded +``pipeline=live-video-to-video`` / ``model_id=unknown`` for every BYOC job. + +These tests prove the additive wire contract on the gateway side: +- the ``/generate-live-payment`` body now carries the REAL ``capability`` and + ``model_id`` for a representative BYOC capability (e.g. ``nano-banana``); +- ``type`` stays ``"lv2v"`` (fee/pixel routing is unchanged); +- when no model id can be determined the field is omitted (backward-compatible, + byte-identical to today so the signer falls back). +""" +from __future__ import annotations + +import json +from contextlib import contextmanager +from unittest.mock import MagicMock, patch +from urllib.request import Request + +from livepeer_gateway.byoc import ( + ByocJobRequest, + _create_byoc_payment, + _extract_model_id, +) + + +def _stub_orch_info(): + info = MagicMock() + tp = MagicMock() + tp.face_value = b"\x01\x00" # non-zero → payment required + info.ticket_params = tp + info.HasField = lambda field: field == "ticket_params" + info.SerializeToString = lambda: b"stub-orch-info-protobuf" + return info + + +@contextmanager +def _mock_signer(*, payment_body=b'{"payment":"TICKETS","segCreds":"SEG"}'): + """Mock the signer HTTP call and capture the request sent to it.""" + captured: list[Request] = [] + + class _MockResponse: + def __init__(self, body: bytes): + self._body = body + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def _fake_urlopen(req, *args, **kwargs): + captured.append(req) + return _MockResponse(payment_body) + + with patch("livepeer_gateway.byoc.urlopen", side_effect=_fake_urlopen), patch( + "livepeer_gateway.orch_info.get_orch_info", + side_effect=lambda *a, **k: _stub_orch_info(), + ): + yield captured + + +def _signer_body(captured: list[Request]) -> dict: + signer_reqs = [ + r for r in captured if "generate-live-payment" in r.full_url + ] + assert len(signer_reqs) == 1, f"expected 1 signer call, got {len(signer_reqs)}" + return json.loads(signer_reqs[0].data.decode("utf-8")) + + +# --------------------------------------------------------------------------- +# _extract_model_id +# --------------------------------------------------------------------------- + + +def test_extract_model_id_from_payload_model_id_key(): + assert _extract_model_id({"model_id": "google/nano-banana"}) == "google/nano-banana" + + +def test_extract_model_id_from_payload_model_key(): + assert _extract_model_id({"prompt": "a cat", "model": "flux-dev"}) == "flux-dev" + + +def test_extract_model_id_prefers_payload_over_parameters(): + assert ( + _extract_model_id({"model_id": "from-payload"}, {"model_id": "from-params"}) + == "from-payload" + ) + + +def test_extract_model_id_falls_back_to_parameters(): + assert _extract_model_id({"prompt": "x"}, {"model": "from-params"}) == "from-params" + + +def test_extract_model_id_empty_when_absent(): + assert _extract_model_id({"prompt": "no model here"}) == "" + assert _extract_model_id(None, None) == "" + assert _extract_model_id({"model_id": " "}) == "" # whitespace-only ignored + + +# --------------------------------------------------------------------------- +# _create_byoc_payment wire contract +# --------------------------------------------------------------------------- + + +def test_payment_body_carries_real_capability_and_model_id(): + """A representative BYOC capability sends the REAL capability + model_id.""" + with _mock_signer() as captured: + result = _create_byoc_payment( + orch_origin="https://orch.test:8936", + capability="nano-banana", + livepeer_hdr="", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + model_id="google/nano-banana", + ) + + body = _signer_body(captured) + assert body["capability"] == "nano-banana" + assert body["model_id"] == "google/nano-banana" + # type stays lv2v so fee/pixel routing is unchanged + assert body["type"] == "lv2v" + assert body["orchestrator"] # base64 orch info present + # payment headers still returned as before + assert result["Livepeer-Payment"] == "TICKETS" + assert result["Livepeer-Segment"] == "SEG" + + +def test_payment_body_omits_model_id_when_absent(): + """Backward-compatible: no model id → field omitted (signer falls back).""" + with _mock_signer() as captured: + _create_byoc_payment( + orch_origin="https://orch.test:8936", + capability="nano-banana", + livepeer_hdr="", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + # model_id omitted → defaults to "" + ) + + body = _signer_body(captured) + assert body["capability"] == "nano-banana" + assert "model_id" not in body, "empty model_id must be omitted for byte-identical body" + assert body["type"] == "lv2v" + + +# --------------------------------------------------------------------------- +# End-to-end through submit_byoc_job +# --------------------------------------------------------------------------- + + +def test_submit_byoc_job_threads_model_id_from_payload(): + """submit_byoc_job extracts the model id from the request payload and + forwards it (and the capability) to the signer payment request.""" + from livepeer_gateway import byoc as byoc_mod + + captured: list[Request] = [] + + 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, *a): + return False + + def _fake_urlopen(req, *args, **kwargs): + captured.append(req) + url = req.full_url if hasattr(req, "full_url") else req.get_full_url() + if "generate-live-payment" in url: + return _MockResponse(b'{"payment":"TICKETS","segCreds":"SEG"}') + if "sign-byoc-job" in url: + return _MockResponse(b'{"sender":"0xabc","signature":"0xsig"}') + # orchestrator /process/request/ + return _MockResponse(b'{"images":[{"url":"https://img"}]}', 200) + + req = ByocJobRequest( + capability="nano-banana", + payload={"prompt": "a dragon", "model_id": "google/nano-banana"}, + ) + + with patch("livepeer_gateway.byoc.urlopen", side_effect=_fake_urlopen), patch( + "livepeer_gateway.orch_info.get_orch_info", + side_effect=lambda *a, **k: _stub_orch_info(), + ): + byoc_mod.submit_byoc_job( + req, + orch_url="https://orch.test:8936", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + ) + + body = _signer_body(captured) + assert body["capability"] == "nano-banana" + assert body["model_id"] == "google/nano-banana" + assert body["type"] == "lv2v" From 57e9e1683534d3d4cf36436a0c97937a5d6a7f6a Mon Sep 17 00:00:00 2001 From: seanhanca Date: Tue, 30 Jun 2026 15:53:50 -0700 Subject: [PATCH 2/2] fix(byoc): attribute training payment to top-level model_id + harden wire body Addresses review feedback on PR #33: - submit_training_job built the signer attribution model_id from the merged byoc_req.payload ({"model_id": req.model_id, **req.params}), so a stray `model_id` inside req.params could override the real training model and diverge from the orchestrator body (which uses req.model_id). Resolve it from the explicit top-level req.model_id, falling back to req.params, so payment attribution matches the orchestrator body. (CodeRabbit, byoc.py:719) - _create_byoc_payment now defensively strips (and ignores non-string) model_id so whitespace-only values are omitted from the wire body, keeping it byte-identical to today when no real model id is present. (Copilot, byoc.py:242) - Add a submit_training_job regression test asserting the signer body uses the top-level training model_id, not the nested params value. (CodeRabbit nitpick) --- src/livepeer_gateway/byoc.py | 11 ++++- tests/test_byoc_usage_attribution.py | 65 ++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 1 deletion(-) diff --git a/src/livepeer_gateway/byoc.py b/src/livepeer_gateway/byoc.py index ec70488..09961c3 100644 --- a/src/livepeer_gateway/byoc.py +++ b/src/livepeer_gateway/byoc.py @@ -238,6 +238,11 @@ def _create_byoc_payment( "type": "lv2v", "capability": capability, } + # Defensively strip (and ignore non-strings): callers normally pass a value + # already cleaned by `_extract_model_id`, but a direct caller could pass a + # whitespace-only or non-string value, which must still be omitted so the + # wire body stays byte-identical to today and the signer falls back. + model_id = model_id.strip() if isinstance(model_id, str) else "" if model_id: payment_payload["model_id"] = model_id payment_body = json.dumps(payment_payload).encode("utf-8") @@ -716,7 +721,11 @@ def submit_training_job( livepeer_hdr=livepeer_hdr, signer_url=signer_url, signer_headers=signer_headers, - model_id=_extract_model_id(byoc_req.payload, byoc_req.parameters), + # Attribute payment to the explicit top-level training + # model_id (matching the orchestrator body above), not the + # merged payload where a `model_id` inside req.params could + # otherwise override it and diverge from the real model. + model_id=_extract_model_id({"model_id": req.model_id}, req.params), timeout=http_timeout, ) headers.update(payment_headers) diff --git a/tests/test_byoc_usage_attribution.py b/tests/test_byoc_usage_attribution.py index afabe06..08f7ad3 100644 --- a/tests/test_byoc_usage_attribution.py +++ b/tests/test_byoc_usage_attribution.py @@ -207,3 +207,68 @@ def _fake_urlopen(req, *args, **kwargs): assert body["capability"] == "nano-banana" assert body["model_id"] == "google/nano-banana" assert body["type"] == "lv2v" + + +# --------------------------------------------------------------------------- +# End-to-end through submit_training_job +# --------------------------------------------------------------------------- + + +def test_submit_training_job_uses_top_level_model_id_over_params(): + """submit_training_job must attribute payment to the explicit top-level + training ``model_id`` (which also builds the orchestrator body), not a + ``model_id`` nested in ``params`` that would otherwise win the payload + merge and send the wrong model to the signer.""" + from livepeer_gateway import byoc as byoc_mod + from livepeer_gateway.byoc import ByocTrainingRequest + + captured: list[Request] = [] + + 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, *a): + return False + + def _fake_urlopen(req, *args, **kwargs): + captured.append(req) + url = req.full_url if hasattr(req, "full_url") else req.get_full_url() + if "generate-live-payment" in url: + return _MockResponse(b'{"payment":"TICKETS","segCreds":"SEG"}') + if "sign-byoc-job" in url: + return _MockResponse(b'{"sender":"0xabc","signature":"0xsig"}') + # orchestrator /process/train/ + return _MockResponse(b'{"job_id":"train-123","status":"submitted"}', 200) + + # Top-level training model differs from a stray params["model_id"]. + req = ByocTrainingRequest( + capability="flux-lora-trainer", + model_id="fal-ai/flux-lora-fast-training", + params={"model_id": "wrong/params-model", "trigger_word": "TOK"}, + ) + + with patch("livepeer_gateway.byoc.urlopen", side_effect=_fake_urlopen), patch( + "livepeer_gateway.orch_info.get_orch_info", + side_effect=lambda *a, **k: _stub_orch_info(), + ): + byoc_mod.submit_training_job( + req, + orch_url="https://orch.test:8936", + signer_url="https://signer.test", + signer_headers={"Authorization": "Bearer sk_test"}, + ) + + body = _signer_body(captured) + assert body["capability"] == "flux-lora-trainer" + # The explicit training model_id wins — NOT the nested params value. + assert body["model_id"] == "fal-ai/flux-lora-fast-training" + assert body["type"] == "lv2v"