`_submit_fal_request` (image) and `_submit_fal_video_request` (video plugin) translated every managed-gateway 4xx into "This model may not yet be enabled on the Nous Portal's FAL proxy — set FAL_KEY or pick a different model". For HTTP 429 that remediation is wrong: the gateway body is RATE_LIMIT_EXCEEDED with a retryAfter, the model is enabled, and agents reading the message switched models or gave up. On one real install this fired 260 times in a week (17% of image_generate calls). Both surfaces now submit through one shared helper (tools/fal_common.py::submit_managed_fal_with_rate_limit_retry): a 429 whose Retry-After (header, else body error.retryAfter) fits a 30s cap is waited out in interrupt-aware 0.5s slices and resubmitted once under a fresh x-idempotency-key; a second 429, or an unknown/too-long Retry-After, raises a ValueError that names the rate limit and tells the agent to retry later rather than switch models. 429 can no longer reach the "may not yet be enabled" text.
207 lines
10 KiB
Python
207 lines
10 KiB
Python
"""Shared FAL.ai SDK plumbing: lazy import, managed-gateway sync client, small helpers.
|
|
|
|
Stateful pieces (cache globals, ``_managed_fal_client*``, ``_submit_fal_request``)
|
|
intentionally stay on :mod:`tools.image_generation_tool`: it is the patch target for the
|
|
test suites and for ``plugins/image_gen/fal/``'s ``_it`` indirection, so moving the caches
|
|
here would silently defeat ``monkeypatch.setattr(image_tool, "_managed_fal_client", None)``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
import uuid
|
|
from typing import Any, Callable, Dict, Optional, Union
|
|
from urllib.parse import urlencode
|
|
|
|
# A 429 whose Retry-After is longer than this is reported, not waited out inside a tool call.
|
|
MANAGED_FAL_RATE_LIMIT_RETRY_CAP_SECONDS = 30.0
|
|
|
|
|
|
def import_fal_client() -> Any:
|
|
"""Import ``fal_client`` (via ``lazy_deps`` when available); raises ImportError if unavailable.
|
|
|
|
Callers cache the result on their own module global so tests can monkeypatch it.
|
|
"""
|
|
try:
|
|
from tools.lazy_deps import ensure as _lazy_ensure
|
|
_lazy_ensure("image.fal", prompt=False)
|
|
except ImportError:
|
|
pass
|
|
except Exception as exc: # noqa: BLE001 — lazy_deps surfaces install hints
|
|
raise ImportError(str(exc))
|
|
import fal_client # type: ignore # noqa: WPS433 — intentionally lazy
|
|
return fal_client
|
|
|
|
|
|
def _normalize_fal_queue_url_format(queue_run_origin: str) -> str:
|
|
normalized_origin = str(queue_run_origin or "").strip().rstrip("/")
|
|
if not normalized_origin:
|
|
raise ValueError("Managed FAL queue origin is required")
|
|
return f"{normalized_origin}/"
|
|
|
|
|
|
def _extract_http_status(exc: BaseException) -> Optional[int]:
|
|
"""HTTP status from httpx (``.response.status_code``) or fal_client (``.status_code``) exceptions, else None."""
|
|
response = getattr(exc, "response", None)
|
|
if response is not None:
|
|
status = getattr(response, "status_code", None)
|
|
if isinstance(status, int):
|
|
return status
|
|
status = getattr(exc, "status_code", None)
|
|
return status if isinstance(status, int) else None
|
|
|
|
|
|
def _managed_fal_billing_error(exc: BaseException, what: str) -> Optional[str]:
|
|
"""Human-readable tail for a Nous managed-gateway ``BILLING_ERROR`` response, else None.
|
|
|
|
``what`` names the rejected thing ("model", "endpoint"); the wording is shared by the image
|
|
and video callers so the two surfaces never drift.
|
|
"""
|
|
response = getattr(exc, "response", None)
|
|
if response is None:
|
|
return None
|
|
try:
|
|
payload = response.json()
|
|
except Exception: # noqa: BLE001 — diagnostics must not mask the provider error
|
|
return None
|
|
error = payload.get("error") if isinstance(payload, dict) else None
|
|
if not isinstance(error, dict) or error.get("code") != "BILLING_ERROR":
|
|
return None
|
|
details = error.get("details") if isinstance(error.get("details"), dict) else {}
|
|
upstream = details.get("upstreamPayload") if isinstance(details.get("upstreamPayload"), dict) else {}
|
|
code = upstream.get("code") or details.get("chargeIntentErrorCode") or "billing_error"
|
|
detail = upstream.get("error") or "Nous Portal rejected the charge authorization"
|
|
return (
|
|
f"{error.get('message') or 'Charge authorization failed'} (BILLING_ERROR; {code}: {detail}). "
|
|
"This is a Nous Portal billing configuration issue, not a missing local API key. "
|
|
f"The managed route cannot run this {what} until Nous enables its billing meter; "
|
|
"a direct FAL_KEY is an optional bypass."
|
|
)
|
|
|
|
|
|
def _managed_fal_retry_after_seconds(exc: BaseException) -> Optional[float]:
|
|
"""Seconds the managed gateway asked us to wait after a 429: the ``Retry-After`` header,
|
|
else the body's ``error.retryAfter``; None when the status is not 429 or neither is present."""
|
|
response = getattr(exc, "response", None)
|
|
if response is None or _extract_http_status(exc) != 429:
|
|
return None
|
|
headers = getattr(response, "headers", None)
|
|
raw = headers.get("Retry-After") if headers is not None and hasattr(headers, "get") else None
|
|
if raw is None:
|
|
try:
|
|
error = response.json().get("error")
|
|
except Exception: # noqa: BLE001 — a non-JSON 429 body simply has no hint
|
|
return None
|
|
raw = error.get("retryAfter") if isinstance(error, dict) else None
|
|
try:
|
|
return float(raw) if raw is not None else None
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def _managed_fal_rate_limit_message(what: str, name: str, retry_after: Optional[float]) -> str:
|
|
hint = f"retry after {retry_after:g}s" if retry_after is not None else "no Retry-After given"
|
|
return (
|
|
f"Nous Subscription gateway rate-limited {what} '{name}' (HTTP 429; {hint}). "
|
|
"The model is enabled — retry later instead of switching models or setting FAL_KEY."
|
|
)
|
|
|
|
|
|
def submit_managed_fal_with_rate_limit_retry(
|
|
submit: Callable[[Dict[str, str]], Any], *, what: str, name: str,
|
|
):
|
|
"""Call ``submit(headers)`` with a fresh ``x-idempotency-key``; on a 429 whose Retry-After
|
|
fits the cap, wait it out (interrupt-aware) and resubmit ONCE under a new key.
|
|
|
|
A second 429, or one with an unknown/too-long Retry-After, raises ValueError naming the
|
|
rate limit — it must never fall through to the callers' "model may not be enabled" 4xx text,
|
|
which sends agents off to switch models. Every other exception propagates untouched.
|
|
"""
|
|
from tools.interrupt import is_interrupted
|
|
for attempt in (1, 2):
|
|
try:
|
|
return submit({"x-idempotency-key": str(uuid.uuid4())})
|
|
except Exception as exc:
|
|
if _extract_http_status(exc) != 429:
|
|
raise
|
|
retry_after = _managed_fal_retry_after_seconds(exc)
|
|
if attempt == 2 or retry_after is None or retry_after > MANAGED_FAL_RATE_LIMIT_RETRY_CAP_SECONDS:
|
|
raise ValueError(_managed_fal_rate_limit_message(what, name, retry_after)) from exc
|
|
deadline = time.monotonic() + retry_after
|
|
while time.monotonic() < deadline:
|
|
if is_interrupted():
|
|
raise ValueError(_managed_fal_rate_limit_message(what, name, retry_after)) from exc
|
|
time.sleep(min(0.5, max(0.0, deadline - time.monotonic())))
|
|
|
|
|
|
def _require(value: Any, what: str) -> Any:
|
|
if value is None:
|
|
raise RuntimeError(f"{what} is required for managed FAL gateway mode")
|
|
return value
|
|
|
|
|
|
class _ManagedFalSyncClient:
|
|
"""Drives a Nous-managed fal-queue gateway via ``fal_client.SyncClient`` primitives; carries
|
|
its own ``fal_client`` reference so the caller decides which (possibly test-patched) module is used."""
|
|
|
|
def __init__(self, fal_client: Any, *, key: str, queue_run_origin: str):
|
|
sync_client_class = _require(getattr(fal_client, "SyncClient", None), "fal_client.SyncClient")
|
|
client_module = _require(getattr(fal_client, "client", None), "fal_client.client")
|
|
self._queue_url_format = _normalize_fal_queue_url_format(queue_run_origin)
|
|
self._sync_client = sync_client_class(key=key)
|
|
self._http_client = _require(getattr(self._sync_client, "_client", None), "fal_client.SyncClient._client")
|
|
self._maybe_retry_request = getattr(client_module, "_maybe_retry_request", None)
|
|
self._raise_for_status = getattr(client_module, "_raise_for_status", None)
|
|
if self._maybe_retry_request is None or self._raise_for_status is None:
|
|
raise RuntimeError("fal_client.client request helpers are required for managed FAL gateway mode")
|
|
self._request_handle_class = _require(
|
|
getattr(client_module, "SyncRequestHandle", None), "fal_client.client.SyncRequestHandle")
|
|
self._add_hint_header = getattr(client_module, "add_hint_header", None)
|
|
self._add_priority_header = getattr(client_module, "add_priority_header", None)
|
|
self._add_timeout_header = getattr(client_module, "add_timeout_header", None)
|
|
|
|
def submit(
|
|
self, application: str, arguments: Dict[str, Any], *, path: str = "",
|
|
hint: Optional[str] = None, webhook_url: Optional[str] = None, priority: Any = None,
|
|
headers: Optional[Dict[str, str]] = None, start_timeout: Optional[Union[int, float]] = None,
|
|
):
|
|
url = self._queue_url_format + application
|
|
if path:
|
|
url += "/" + path.lstrip("/")
|
|
if webhook_url is not None:
|
|
url += "?" + urlencode({"fal_webhook": webhook_url})
|
|
request_headers = dict(headers or {})
|
|
if hint is not None and self._add_hint_header is not None:
|
|
self._add_hint_header(hint, request_headers)
|
|
if priority is not None:
|
|
if self._add_priority_header is None:
|
|
raise RuntimeError("fal_client.client.add_priority_header is required for priority requests")
|
|
self._add_priority_header(priority, request_headers)
|
|
if start_timeout is not None:
|
|
if self._add_timeout_header is None:
|
|
raise RuntimeError("fal_client.client.add_timeout_header is required for timeout requests")
|
|
self._add_timeout_header(start_timeout, request_headers)
|
|
request_kwargs = {
|
|
"json": arguments,
|
|
"timeout": getattr(self._sync_client, "default_timeout", 120.0),
|
|
"headers": request_headers,
|
|
}
|
|
# The Nous gateway currently records a keyed submission before billing
|
|
# authorization finishes, but cannot replay the resulting error. The
|
|
# SDK's automatic 409 retry therefore replaces the real billing error
|
|
# with an idempotency conflict. Make one attempt when the caller supplied
|
|
# a key; an ambiguous transport failure is safer than a possible duplicate
|
|
# generation or a masked entitlement failure.
|
|
has_idempotency_key = any(
|
|
str(key).lower() == "x-idempotency-key" for key in request_headers
|
|
)
|
|
if has_idempotency_key:
|
|
response = self._http_client.request("POST", url, **request_kwargs)
|
|
else:
|
|
response = self._maybe_retry_request(self._http_client, "POST", url, **request_kwargs)
|
|
self._raise_for_status(response)
|
|
data = response.json()
|
|
return self._request_handle_class(
|
|
request_id=data["request_id"], response_url=data["response_url"],
|
|
status_url=data["status_url"], cancel_url=data["cancel_url"], client=self._http_client)
|