fix(byoc): send real capability + model_id for usage attribution

src/livepeer_gateway/byoc.py

@@ -144,13 +144,44 @@ 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,
     capability: str,
     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,27 @@ 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")
+    }
+    # 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")
     payment_headers = {
         "Content-Type": "application/json",
         "Livepeer-Capability": capability,
@@ -390,6 +437,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 +721,11 @@ def submit_training_job(
                     livepeer_hdr=livepeer_hdr,
                     signer_url=signer_url,
                     signer_headers=signer_headers,
+                    # 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)

tests/test_byoc_usage_attribution.py

@@ -0,0 +1,274 @@
+"""
+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/<cap>
+        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"
+
+
+# ---------------------------------------------------------------------------
+# 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/<cap>
+        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"