fix(music): poll the job the gateway hands back, on both rails

blockrun_llm/image.py

@@ -42,8 +42,8 @@
     payment_mode,
     raise_for_api_key_402,
     resolve_api_key,
-    resolve_poll_url,
 )
+from .jobs import poll_until_completed
 from .tx_log import paid_request_error_prefix
 from .types import APIError, ImageResponse, PaymentError, retry_after_of
 from .validation import (
@@ -409,106 +409,25 @@ def _handle_payment_and_retry(
             retry_after=retry_after_of(retry_response),
         )
 
-    def _absolute_url(self, url: str) -> str:
-        """Resolve a relative ``poll_url`` against the configured API host.
-
-        Server-returned poll URLs look like ``/api/v1/images/generations/<id>``;
-        our ``self.api_url`` already ends with ``/api`` so we strip it once
-        to avoid double-prefixing.
-        """
-        if url.startswith(("http://", "https://")):
-            return url
-        return resolve_poll_url(url, self.api_url, self.api_key)
-
     def _poll_until_completed(
         self,
         submit_resp: httpx.Response,
         payment_payload: str | None,
     ) -> ImageResponse:
-        """Poll the gateway's ``poll_url`` with the same PAYMENT-SIGNATURE
-        until the upstream returns the finished image.
-
-        Settlement happens on the first ``status=completed`` poll, so
-        timeout = no spend. Returns the parsed :class:`ImageResponse`.
-        """
-        import time as _time
-
-        try:
-            submit_data = submit_resp.json()
-        except Exception:
-            submit_data = {}
-
-        poll_url_rel = submit_data.get("poll_url")
-        job_id = submit_data.get("id")
-        if not poll_url_rel:
-            raise APIError(
-                "Slow-path 202 missing poll_url",
-                202,
-                {"response": submit_data},
-            )
-
-        poll_url = self._absolute_url(poll_url_rel)
-        # A signature exists only on the wallet rail; the key rides on the
-        # client's default headers.
-        poll_headers = {"PAYMENT-SIGNATURE": payment_payload} if payment_payload else {}
-        deadline = _time.monotonic() + self.IMAGE_POLL_BUDGET_SECONDS
-        last_status = submit_data.get("status", "queued")
-
-        while _time.monotonic() < deadline:
-            _time.sleep(self.IMAGE_POLL_INTERVAL_SECONDS)
-
-            poll_resp = self._client.get(poll_url, headers=poll_headers)
-            try:
-                poll_data = poll_resp.json()
-            except Exception:
-                poll_data = {}
-            last_status = poll_data.get("status", last_status)
-
-            if poll_resp.status_code == 402:
-                # Account rail: a 402 is the account being out of credit, not a
-                # challenge to sign. Nothing here can sign, so say so plainly.
-                raise_for_api_key_402(poll_resp, self.api_key)
-                # Settlement failed on this poll — surface the gateway reason.
-                raise build_payment_rejected_error(poll_resp)
-
-            if last_status == "failed":
-                raise APIError(
-                    f"Image generation failed upstream: {poll_data.get('error', 'unknown')}",
-                    poll_resp.status_code,
-                    sanitize_error_response(poll_data if isinstance(poll_data, dict) else {}),
-                    retry_after=retry_after_of(poll_resp),
-                )
-
-            if poll_resp.status_code == 200 and last_status == "completed":
-                return ImageResponse(**poll_data)
-
-            if poll_resp.status_code in (202, 504):
-                # 202 = still queued/in_progress; 504 = transient upstream
-                # hiccup. Both retriable inside the budget.
-                continue
-
-            if poll_resp.status_code != 200:
-                try:
-                    error_body = poll_resp.json()
-                except Exception:
-                    error_body = {"error": "Request failed"}
-                raise APIError(
-                    f"Image poll failed: HTTP {poll_resp.status_code}",
-                    poll_resp.status_code,
-                    sanitize_error_response(error_body),
-                    retry_after=retry_after_of(poll_resp),
-                )
-
-        raise APIError(
-            (
-                f"Image generation did not complete within "
-                f"{self.IMAGE_POLL_BUDGET_SECONDS:.0f}s "
-                f"(last status: {last_status}). Settlement only happens on "
-                "completion, so no payment was taken."
-            ),
-            504,
-            {"id": job_id, "last_status": last_status},
+        """Poll the gateway's ``poll_url`` until the upstream returns the
+        finished image. Settlement happens on the first ``status=completed``
+        poll, so timeout = no spend. Shared with music — see jobs.py."""
+        data = poll_until_completed(
+            self._client,
+            submit_resp,
+            payment_payload,
+            api_url=self.api_url,
+            api_key=self.api_key,
+            interval_seconds=self.IMAGE_POLL_INTERVAL_SECONDS,
+            budget_seconds=self.IMAGE_POLL_BUDGET_SECONDS,
+            label="Image",
         )
+        return ImageResponse(**data)
 
     @property
     def payment_mode(self) -> str:

blockrun_llm/jobs.py

@@ -0,0 +1,123 @@
+"""Polling for the gateway's async media jobs.
+
+A slow generation answers ``202`` with a ``poll_url`` and settles on the first
+poll that observes ``completed``, so a poll that times out has cost nothing.
+Images and music share this loop: the same statuses, the same settlement rule,
+the same two rails. One copy, so the two clients cannot drift apart on how a
+job ends.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import httpx
+
+from .apikey import raise_for_api_key_402, resolve_poll_url
+from .types import APIError, retry_after_of
+from .validation import build_payment_rejected_error, sanitize_error_response
+
+
+def absolute_poll_url(url: str, api_url: str, api_key: str | None) -> str:
+    """Resolve a relative ``poll_url`` against the configured API host.
+
+    Server-returned poll URLs look like ``/api/v1/images/generations/<id>``;
+    ``api_url`` already ends with ``/api`` on the wallet rail, and the account
+    rail serves the same route without that prefix.
+    """
+    if url.startswith(("http://", "https://")):
+        return url
+    return resolve_poll_url(url, api_url, api_key)
+
+
+def poll_until_completed(
+    client: httpx.Client,
+    submit_resp: httpx.Response,
+    payment_payload: str | None,
+    *,
+    api_url: str,
+    api_key: str | None,
+    interval_seconds: float,
+    budget_seconds: float,
+    label: str,
+) -> dict[str, Any]:
+    """Poll ``poll_url`` until the job completes; return the completed body.
+
+    ``payment_payload`` is the create's PAYMENT-SIGNATURE on the wallet rail
+    (the job is bound to that wallet and settles against it) and ``None`` on
+    the account rail, where the key rides on the client's default headers.
+    ``label`` names the product in errors ("Image", "Music").
+    """
+    try:
+        submit_data = submit_resp.json()
+    except Exception:
+        submit_data = {}
+
+    poll_url_rel = submit_data.get("poll_url")
+    job_id = submit_data.get("id")
+    if not poll_url_rel:
+        raise APIError("Slow-path 202 missing poll_url", 202, {"response": submit_data})
+
+    poll_url = absolute_poll_url(poll_url_rel, api_url, api_key)
+    poll_headers = {"PAYMENT-SIGNATURE": payment_payload} if payment_payload else {}
+    deadline = time.monotonic() + budget_seconds
+    last_status = submit_data.get("status", "queued")
+
+    while time.monotonic() < deadline:
+        time.sleep(interval_seconds)
+
+        poll_resp = client.get(poll_url, headers=poll_headers)
+        try:
+            poll_data = poll_resp.json()
+        except Exception:
+            poll_data = {}
+        last_status = poll_data.get("status", last_status)
+
+        if poll_resp.status_code == 402:
+            # Account rail: a 402 is the account being out of credit, not a
+            # challenge to sign. Nothing here can sign, so say so plainly.
+            raise_for_api_key_402(poll_resp, api_key)
+            # Settlement failed on this poll — surface the gateway reason.
+            raise build_payment_rejected_error(poll_resp)
+
+        if last_status == "failed":
+            raise APIError(
+                f"{label} generation failed upstream: {poll_data.get('error', 'unknown')}",
+                poll_resp.status_code,
+                sanitize_error_response(poll_data if isinstance(poll_data, dict) else {}),
+                retry_after=retry_after_of(poll_resp),
+            )
+
+        if poll_resp.status_code == 200 and last_status == "completed":
+            tx_hash = poll_resp.headers.get("x-payment-receipt")
+            if tx_hash and "txHash" not in poll_data:
+                poll_data["txHash"] = tx_hash
+            return poll_data
+
+        if poll_resp.status_code in (202, 504):
+            # 202 = still queued/in_progress; 504 = transient upstream
+            # hiccup. Both retriable inside the budget.
+            continue
+
+        if poll_resp.status_code != 200:
+            try:
+                error_body = poll_resp.json()
+            except Exception:
+                error_body = {"error": "Request failed"}
+            raise APIError(
+                f"{label} poll failed: HTTP {poll_resp.status_code}",
+                poll_resp.status_code,
+                sanitize_error_response(error_body),
+                retry_after=retry_after_of(poll_resp),
+            )
+
+    raise APIError(
+        (
+            f"{label} generation did not complete within {budget_seconds:.0f}s "
+            f"(last status: {last_status}). Settlement only happens on "
+            "completion, so no payment was taken."
+        ),
+        504,
+        {"id": job_id, "last_status": last_status},
+    )

blockrun_llm/music.py

@@ -46,6 +46,7 @@
     raise_for_api_key_402,
     resolve_api_key,
 )
+from .jobs import poll_until_completed
 from .tx_log import paid_request_error_prefix
 from .types import APIError, MusicResponse, PaymentError, retry_after_of
 from .validation import (
@@ -71,6 +72,11 @@ class MusicClient:
     DEFAULT_API_URL = "https://blockrun.ai/api"
     DEFAULT_MODEL = "minimax/music-2.5+"
     DEFAULT_TIMEOUT = 210.0  # music gen takes 1-3 min
+    # A track takes one to three minutes and the gateway answers 202 +
+    # poll_url at once, so the wait happens here, poll by poll. Settlement is
+    # on the completed poll: a budget that runs out has cost nothing.
+    MUSIC_POLL_INTERVAL_SECONDS = 5.0
+    MUSIC_POLL_BUDGET_SECONDS = 300.0
 
     def __init__(
         self,
@@ -194,6 +200,12 @@ def _request_with_payment(self, endpoint: str, body: dict[str, Any]) -> MusicRes
             raise_for_api_key_402(response, self.api_key)
             return self._handle_payment_and_retry(url, body, response)
 
+        # Account rail: the key already paid, so the job's 202 comes on the
+        # FIRST post. Music is never fast enough to finish inline, so without
+        # this branch every API-key music request raised "API error: 202".
+        if self.api_key and response.status_code == 202:
+            return self._poll_until_completed(response, None)
+
         if response.status_code != 200:
             try:
                 error_body = response.json()
@@ -263,6 +275,11 @@ def _handle_payment_and_retry(
             raise_for_api_key_402(retry_response, self.api_key)
             raise PaymentError("Payment was rejected. Check your wallet balance.")
 
+        if retry_response.status_code == 202:
+            # The signed create is queued: replay the same signature on each
+            # poll — the job is bound to this wallet — and settle on completion.
+            return self._poll_until_completed(retry_response, payment_payload)
+
         if retry_response.status_code != 200:
             try:
                 error_body = retry_response.json()
@@ -285,6 +302,21 @@ def _handle_payment_and_retry(
 
         return MusicResponse(**data)
 
+    def _poll_until_completed(
+        self, submit_resp: httpx.Response, payment_payload: str | None
+    ) -> MusicResponse:
+        data = poll_until_completed(
+            self._client,
+            submit_resp,
+            payment_payload,
+            api_url=self.api_url,
+            api_key=self.api_key,
+            interval_seconds=self.MUSIC_POLL_INTERVAL_SECONDS,
+            budget_seconds=self.MUSIC_POLL_BUDGET_SECONDS,
+            label="Music",
+        )
+        return MusicResponse(**data)
+
     @property
     def payment_mode(self) -> str:
         """Which rail this client pays on: ``"apikey"`` or ``"wallet"``.