Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import asyncio | |
| import ipaddress | |
| from collections.abc import AsyncIterator, Awaitable, Callable | |
| from contextlib import asynccontextmanager | |
| from pathlib import Path | |
| from typing import Any | |
| from urllib.parse import urlparse | |
| import httpx | |
| from pydantic import SecretStr | |
| from app.generation.domain.enums import ( | |
| GenerationModality, | |
| WorkerCancellationStatus, | |
| WorkerErrorCategory, | |
| WorkerHealthStatus, | |
| WorkerJobStatus, | |
| WorkerReadinessStatus, | |
| ) | |
| from app.generation.domain.errors import GenerationOutputError, GenerationWorkerError | |
| from app.generation.domain.retry import GenerationRetryPolicy | |
| from app.generation.domain.runtime import ( | |
| WorkerCancellationResult, | |
| WorkerHealth, | |
| WorkerInfo, | |
| WorkerJob, | |
| WorkerModelInfo, | |
| WorkerOutput, | |
| WorkerReadiness, | |
| safe_worker_metadata, | |
| ) | |
| class RemoteWorkerClient: | |
| """Strict HTTP client for a trusted, configured MediaRouter worker. | |
| The constructor is deliberately internal-facing: no REST, MCP, SDK, n8n, | |
| or browser payload may supply its base URL or token. Redirects and proxy | |
| environment variables are disabled, endpoint paths are fixed/validated, | |
| and no request/response body is logged. | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| base_url: str, | |
| bearer_token: SecretStr | str | None, | |
| connect_timeout_seconds: float, | |
| request_timeout_seconds: float, | |
| read_timeout_seconds: float, | |
| retry_policy: GenerationRetryPolicy, | |
| http_client: httpx.AsyncClient | None = None, | |
| sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, | |
| ) -> None: | |
| self.base_url = self._validate_base_url(base_url) | |
| self._bearer_token = ( | |
| bearer_token.get_secret_value() | |
| if isinstance(bearer_token, SecretStr) | |
| else bearer_token | |
| ) | |
| self._request_timeout_seconds = request_timeout_seconds | |
| self.retry_policy = retry_policy | |
| self._sleep = sleep | |
| self._client = http_client or httpx.AsyncClient( | |
| timeout=httpx.Timeout( | |
| connect=connect_timeout_seconds, | |
| read=read_timeout_seconds, | |
| write=request_timeout_seconds, | |
| pool=connect_timeout_seconds, | |
| ), | |
| follow_redirects=False, | |
| trust_env=False, | |
| headers={"User-Agent": "mediarouter-generation-runtime/1"}, | |
| ) | |
| self._owns_client = http_client is None | |
| async def aclose(self) -> None: | |
| if self._owns_client: | |
| await self._client.aclose() | |
| async def health(self) -> WorkerHealth: | |
| payload = await self._request_json("GET", "/health", idempotent=True) | |
| return WorkerHealth( | |
| status=self._health_status(payload.get("status")), | |
| metadata=self._metadata(payload, known={"status"}), | |
| ) | |
| async def ready(self) -> WorkerReadiness: | |
| payload = await self._request_json( | |
| "GET", "/ready", idempotent=True, readiness_endpoint=True | |
| ) | |
| model_ids = payload.get("model_ids") | |
| if model_ids is None: | |
| model = payload.get("model") | |
| if model is None: | |
| model_ids = [] | |
| elif isinstance(model, str): | |
| model_ids = [model] | |
| else: | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned invalid readiness model IDs.", | |
| ) | |
| elif not isinstance(model_ids, list) or not all( | |
| isinstance(value, str) for value in model_ids | |
| ): | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned invalid readiness model IDs.", | |
| ) | |
| model_loaded = payload.get("model_loaded", False) | |
| if not isinstance(model_loaded, bool): | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned invalid readiness metadata.", | |
| ) | |
| return WorkerReadiness( | |
| status=self._readiness_status(payload.get("status")), | |
| model_loaded=model_loaded, | |
| model_ids=model_ids, | |
| metadata=self._metadata( | |
| payload, known={"status", "model_loaded", "model_ids", "model"} | |
| ), | |
| ) | |
| async def info(self) -> WorkerInfo: | |
| payload = await self._request_json("GET", "/v1/info", idempotent=True) | |
| identifier = payload.get("id") | |
| name = payload.get("name") | |
| if ( | |
| not isinstance(identifier, str) | |
| or not identifier | |
| or not isinstance(name, str) | |
| or not name | |
| ): | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned invalid model metadata.", | |
| ) | |
| try: | |
| media_types = self._media_types(payload) | |
| models = self._worker_models(payload, fallback_id=identifier, fallback_name=name) | |
| except ValueError as exc: | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned an unsupported media type.", | |
| ) from exc | |
| return WorkerInfo( | |
| id=identifier, | |
| name=name, | |
| media_types=list(dict.fromkeys(media_types)), | |
| models=models, | |
| status=self._health_status(payload.get("status")), | |
| metadata=self._metadata( | |
| payload, | |
| known={"id", "name", "type", "media_types", "models", "status"}, | |
| ), | |
| ) | |
| async def submit( | |
| self, *, payload: dict[str, object], idempotency_key: str | |
| ) -> WorkerJob: | |
| if not idempotency_key.strip(): | |
| raise self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation submission requires an idempotency key.", | |
| ) | |
| response = await self._request_json( | |
| "POST", | |
| "/v1/generate", | |
| json_payload=payload, | |
| headers={"Idempotency-Key": idempotency_key}, | |
| idempotent=True, | |
| expected_statuses={200, 202}, | |
| ) | |
| return self._worker_job(response) | |
| async def submit_form( | |
| self, | |
| *, | |
| fields: dict[str, str], | |
| idempotency_key: str, | |
| idempotent: bool, | |
| ) -> dict[str, object]: | |
| """Submit a strict scalar form to a worker that does not need a file. | |
| Some workers use ``multipart/form-data`` only when an optional input | |
| asset is supplied, but still require form fields for text-only work. | |
| This provider-neutral primitive keeps that transport detail out of | |
| adapters without falling back to an incompatible JSON request body. | |
| ``idempotent`` remains explicit because workers can accept a job | |
| without exposing any request-idempotency protocol. | |
| """ | |
| if not idempotency_key.strip(): | |
| raise self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation submission requires an idempotency key.", | |
| ) | |
| if not fields or any( | |
| not isinstance(key, str) or not isinstance(value, str) | |
| for key, value in fields.items() | |
| ): | |
| raise self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation submission fields are invalid.", | |
| ) | |
| return await self._request_json( | |
| "POST", | |
| "/v1/generate", | |
| data=fields, | |
| headers={"Idempotency-Key": idempotency_key}, | |
| idempotent=idempotent, | |
| expected_statuses={200, 202}, | |
| ) | |
| async def submit_multipart( | |
| self, | |
| *, | |
| fields: dict[str, str], | |
| file_field: str, | |
| file_path: Path, | |
| filename: str, | |
| mime_type: str, | |
| idempotency_key: str, | |
| idempotent: bool, | |
| ) -> dict[str, object]: | |
| """Submit one canonical local input file as multipart data. | |
| This is a transport primitive rather than a model-specific API. The | |
| source file is opened by the server from a verified canonical asset; | |
| it is never a client filesystem path. ``httpx`` streams the file | |
| object while encoding multipart data, so large media is not loaded | |
| into memory. A caller may opt out of automatic retries when its | |
| worker does not offer submission idempotency. | |
| """ | |
| if not idempotency_key.strip(): | |
| raise self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation submission requires an idempotency key.", | |
| ) | |
| if not file_field or not filename or not mime_type: | |
| raise self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation submission file metadata is invalid.", | |
| ) | |
| source_input = file_path.expanduser() | |
| if source_input.is_symlink(): | |
| raise self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation input asset is unavailable.", | |
| ) | |
| source = source_input.resolve() | |
| if not source.is_file(): | |
| raise self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation input asset is unavailable.", | |
| ) | |
| try: | |
| # The WAN worker intentionally has no idempotency key support. | |
| # Its adapter sets ``idempotent=False``, preventing an uncertain | |
| # network failure from causing a second expensive GPU submission. | |
| with source.open("rb") as stream: | |
| return await self._request_json( | |
| "POST", | |
| "/v1/generate", | |
| data=fields, | |
| files={file_field: (filename, stream, mime_type)}, | |
| headers={"Idempotency-Key": idempotency_key}, | |
| idempotent=idempotent, | |
| expected_statuses={200, 202}, | |
| ) | |
| except GenerationWorkerError: | |
| raise | |
| except OSError as exc: | |
| raise self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation input asset could not be read.", | |
| ) from exc | |
| async def get_job_payload(self, external_job_id: str) -> dict[str, object]: | |
| """Return a fixed worker job response for adapter-specific parsing.""" | |
| path = f"/v1/jobs/{self._safe_external_id(external_job_id)}" | |
| return await self._request_json("GET", path, idempotent=True, job_endpoint=True) | |
| async def cancel_job_payload( | |
| self, | |
| external_job_id: str, | |
| *, | |
| expected_statuses: set[int] | None = None, | |
| ) -> dict[str, object]: | |
| """Call the fixed cancellation endpoint without provider parsing.""" | |
| path = f"/v1/jobs/{self._safe_external_id(external_job_id)}/cancel" | |
| return await self._request_json( | |
| "POST", | |
| path, | |
| idempotent=True, | |
| expected_statuses=expected_statuses or {200, 202, 204}, | |
| job_endpoint=True, | |
| ) | |
| async def get_job(self, external_job_id: str) -> WorkerJob: | |
| return self._worker_job(await self.get_job_payload(external_job_id)) | |
| async def cancel(self, external_job_id: str) -> WorkerCancellationResult: | |
| payload = await self.cancel_job_payload(external_job_id) | |
| raw_status = str(payload.get("status", "")).strip().lower() | |
| if raw_status in {"cancelled", "canceled"}: | |
| status = WorkerCancellationStatus.CANCELLED | |
| elif raw_status in {"requested", "cancel_requested", "cancellation_requested"}: | |
| status = WorkerCancellationStatus.REQUESTED | |
| elif raw_status in {"unsupported", "not_supported"}: | |
| status = WorkerCancellationStatus.UNSUPPORTED | |
| elif not raw_status: | |
| # A successful 204 has no representation of whether a running | |
| # GPU operation actually stopped. It can only mean the worker | |
| # accepted the cancellation request, never that it completed it. | |
| status = WorkerCancellationStatus.REQUESTED | |
| else: | |
| status = WorkerCancellationStatus.FAILED | |
| return WorkerCancellationResult( | |
| status=status, | |
| metadata=self._metadata(payload, known={"status"}), | |
| ) | |
| async def retrieve_output(self, external_job_id: str) -> WorkerOutput: | |
| job = await self.get_job(external_job_id) | |
| if job.status is not WorkerJobStatus.COMPLETED or job.output is None: | |
| raise GenerationOutputError("Generation output is not ready.") | |
| return job.output | |
| async def stream_output(self, output: WorkerOutput) -> AsyncIterator[AsyncIterator[bytes]]: | |
| """Stream a worker-owned relative output path without buffering it. | |
| The caller must write into a controlled MediaRouter staging location, | |
| verify the optional checksum, then register it through | |
| ``CanonicalAssetService``. No worker filesystem path is ever trusted. | |
| """ | |
| path = self._safe_worker_path(output.download_path) | |
| context, response = await self._open_stream(path) | |
| try: | |
| yield response.aiter_bytes() | |
| finally: | |
| await context.__aexit__(None, None, None) | |
| async def _request_json( | |
| self, | |
| method: str, | |
| path: str, | |
| *, | |
| json_payload: dict[str, object] | None = None, | |
| data: dict[str, str] | None = None, | |
| files: Any | None = None, | |
| headers: dict[str, str] | None = None, | |
| idempotent: bool, | |
| expected_statuses: set[int] | None = None, | |
| readiness_endpoint: bool = False, | |
| job_endpoint: bool = False, | |
| ) -> dict[str, object]: | |
| expected = expected_statuses or {200} | |
| safe_path = self._safe_worker_path(path) | |
| retry_number = 0 | |
| while True: | |
| try: | |
| response = await asyncio.wait_for( | |
| self._client.request( | |
| method, | |
| self._url_for(safe_path), | |
| headers=self._headers(headers), | |
| json=json_payload, | |
| data=data, | |
| files=files, | |
| ), | |
| timeout=self._request_timeout_seconds, | |
| ) | |
| if response.status_code not in expected: | |
| raise self._response_error( | |
| response.status_code, | |
| readiness_endpoint=readiness_endpoint, | |
| job_endpoint=job_endpoint, | |
| ) | |
| try: | |
| payload = response.json() if response.content else {} | |
| except (ValueError, UnicodeDecodeError) as exc: | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned an invalid JSON response.", | |
| ) from exc | |
| if not isinstance(payload, dict): | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned an invalid response shape.", | |
| ) | |
| return payload | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception as exc: | |
| error = self._normalise_exception( | |
| exc, | |
| readiness_endpoint=readiness_endpoint, | |
| job_endpoint=job_endpoint, | |
| ) | |
| decision = self.retry_policy.decide( | |
| category=error.category, | |
| http_status=error.http_status, | |
| retry_number=retry_number, | |
| idempotent=idempotent, | |
| ) | |
| if not decision.retryable: | |
| raise error from None | |
| retry_number += 1 | |
| await self._sleep(decision.delay_seconds) | |
| async def _open_stream( | |
| self, path: str | |
| ) -> tuple[Any, httpx.Response]: | |
| # httpx exposes stream() as an async context manager. It is kept | |
| # private to this method so all callers close it in a finally block. | |
| retry_number = 0 | |
| while True: | |
| context = self._client.stream( | |
| "GET", self._url_for(path), headers=self._headers(None) | |
| ) | |
| try: | |
| response = await asyncio.wait_for( | |
| context.__aenter__(), timeout=self._request_timeout_seconds | |
| ) | |
| if response.status_code != 200: | |
| raise self._response_error(response.status_code) | |
| return context, response | |
| except asyncio.CancelledError: | |
| await context.__aexit__(None, None, None) | |
| raise | |
| except Exception as exc: | |
| await context.__aexit__(None, None, None) | |
| error = self._normalise_exception(exc) | |
| decision = self.retry_policy.decide( | |
| category=error.category, | |
| http_status=error.http_status, | |
| retry_number=retry_number, | |
| idempotent=True, | |
| ) | |
| if not decision.retryable: | |
| raise error from None | |
| retry_number += 1 | |
| await self._sleep(decision.delay_seconds) | |
| def _worker_job(self, payload: dict[str, object]) -> WorkerJob: | |
| raw_status = str(payload.get("status", "")).strip().lower() | |
| statuses = { | |
| "queued": WorkerJobStatus.QUEUED, | |
| "running": WorkerJobStatus.RUNNING, | |
| "processing": WorkerJobStatus.RUNNING, | |
| "completed": WorkerJobStatus.COMPLETED, | |
| "succeeded": WorkerJobStatus.COMPLETED, | |
| "failed": WorkerJobStatus.FAILED, | |
| "cancelled": WorkerJobStatus.CANCELLED, | |
| "canceled": WorkerJobStatus.CANCELLED, | |
| } | |
| job_id = payload.get("job_id", payload.get("external_job_id")) | |
| if not isinstance(job_id, str) or raw_status not in statuses: | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned an invalid job response.", | |
| ) | |
| raw_output = payload.get("output") | |
| try: | |
| output = self._worker_output(raw_output) if isinstance(raw_output, dict) else None | |
| return WorkerJob( | |
| external_job_id=job_id, | |
| status=statuses[raw_status], | |
| output=output, | |
| error_category=self._worker_error_category(payload, statuses[raw_status]), | |
| error_code=self._worker_error_code(payload), | |
| error_message=None, | |
| metadata=self._metadata( | |
| payload, | |
| known={"job_id", "external_job_id", "status", "output", "error"}, | |
| ), | |
| ) | |
| except (TypeError, ValueError) as exc: | |
| raise self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned invalid job metadata.", | |
| ) from exc | |
| def _worker_output(self, payload: dict[str, object]) -> WorkerOutput: | |
| type_value = payload.get("output_type", payload.get("type")) | |
| mime_type = payload.get("mime_type") | |
| provider_output_id = payload.get("provider_output_id", payload.get("id")) | |
| download_path = payload.get("download_path") | |
| if not all( | |
| isinstance(value, str) | |
| for value in (type_value, mime_type, provider_output_id, download_path) | |
| ): | |
| raise GenerationOutputError() | |
| filename = payload.get("filename") | |
| sha256 = payload.get("sha256") | |
| byte_size = payload.get("byte_size") | |
| if filename is not None and not isinstance(filename, str): | |
| raise GenerationOutputError() | |
| if sha256 is not None and not isinstance(sha256, str): | |
| raise GenerationOutputError() | |
| if byte_size is not None and ( | |
| not isinstance(byte_size, int) or isinstance(byte_size, bool) | |
| ): | |
| raise GenerationOutputError() | |
| metadata = self._metadata( | |
| payload, | |
| known={ | |
| "output_type", | |
| "type", | |
| "mime_type", | |
| "provider_output_id", | |
| "id", | |
| "download_path", | |
| "filename", | |
| "sha256", | |
| "byte_size", | |
| }, | |
| ) | |
| return WorkerOutput( | |
| output_type=type_value, | |
| mime_type=mime_type, | |
| provider_output_id=provider_output_id, | |
| download_path=self._safe_worker_path(download_path), | |
| filename=filename, | |
| sha256=sha256, | |
| byte_size=byte_size, | |
| metadata=metadata, | |
| ) | |
| def _media_types(payload: dict[str, object]) -> list[GenerationModality]: | |
| values = payload.get("media_types") | |
| if values is None: | |
| legacy_type = payload.get("type") | |
| if legacy_type is None: | |
| values = [] | |
| elif isinstance(legacy_type, str): | |
| values = [legacy_type] | |
| else: | |
| raise ValueError("worker media type is invalid") | |
| elif not isinstance(values, list) or not all(isinstance(value, str) for value in values): | |
| raise ValueError("worker media types are invalid") | |
| return [GenerationModality(value) for value in values] | |
| def _worker_models( | |
| self, payload: dict[str, object], *, fallback_id: str, fallback_name: str | |
| ) -> list[WorkerModelInfo]: | |
| raw_models = payload.get("models") | |
| if raw_models is None: | |
| return [ | |
| WorkerModelInfo( | |
| id=fallback_id, | |
| name=fallback_name, | |
| media_types=list(dict.fromkeys(self._media_types(payload))), | |
| ) | |
| ] | |
| if isinstance(raw_models, dict): | |
| # A compact single-model worker may expose named model variants as | |
| # a JSON object instead of a list of independently selectable | |
| # models. It is still one discovered top-level model; retain the | |
| # safe variant map as metadata for the concrete adapter to verify. | |
| variants = safe_worker_metadata(raw_models) | |
| if not isinstance(variants, dict): | |
| raise ValueError("worker model variants must be an object") | |
| return [ | |
| WorkerModelInfo( | |
| id=fallback_id, | |
| name=fallback_name, | |
| media_types=list(dict.fromkeys(self._media_types(payload))), | |
| metadata={"variants": variants}, | |
| ) | |
| ] | |
| if not isinstance(raw_models, list) or not raw_models: | |
| raise ValueError("worker models must be a non-empty list") | |
| models: list[WorkerModelInfo] = [] | |
| for raw_model in raw_models: | |
| if not isinstance(raw_model, dict): | |
| raise ValueError("worker model must be an object") | |
| model_id = raw_model.get("id") | |
| model_name = raw_model.get("name") | |
| if not isinstance(model_id, str) or not model_id: | |
| raise ValueError("worker model ID is invalid") | |
| if not isinstance(model_name, str) or not model_name: | |
| raise ValueError("worker model name is invalid") | |
| models.append( | |
| WorkerModelInfo( | |
| id=model_id, | |
| name=model_name, | |
| media_types=list(dict.fromkeys(self._media_types(raw_model))), | |
| metadata=self._metadata( | |
| raw_model, known={"id", "name", "type", "media_types"} | |
| ), | |
| ) | |
| ) | |
| return models | |
| def _worker_error_code(payload: dict[str, object]) -> str | None: | |
| raw_error = payload.get("error") | |
| if not isinstance(raw_error, dict): | |
| return None | |
| code = raw_error.get("code") | |
| return code if isinstance(code, str) and len(code) <= 100 else None | |
| def _worker_error_category( | |
| payload: dict[str, object], status: WorkerJobStatus | |
| ) -> WorkerErrorCategory | None: | |
| raw_error = payload.get("error") | |
| if isinstance(raw_error, dict): | |
| category = raw_error.get("category") | |
| if isinstance(category, str): | |
| try: | |
| return WorkerErrorCategory(category) | |
| except ValueError: | |
| pass | |
| # A terminal worker failure is an inference failure unless the worker | |
| # explicitly supplied a supported, non-secret category. | |
| return WorkerErrorCategory.INFERENCE_ERROR if status is WorkerJobStatus.FAILED else None | |
| def _normalise_exception( | |
| self, | |
| exc: Exception, | |
| *, | |
| readiness_endpoint: bool = False, | |
| job_endpoint: bool = False, | |
| ) -> GenerationWorkerError: | |
| if isinstance(exc, GenerationWorkerError): | |
| return exc | |
| if isinstance(exc, asyncio.TimeoutError) or isinstance( | |
| exc, (httpx.ReadTimeout, httpx.WriteTimeout, httpx.PoolTimeout) | |
| ): | |
| return self._error(WorkerErrorCategory.TIMEOUT, "Generation worker request timed out.") | |
| if isinstance(exc, (httpx.ConnectTimeout, httpx.NetworkError)): | |
| return self._error( | |
| WorkerErrorCategory.WORKER_UNAVAILABLE, | |
| "Generation worker is unavailable.", | |
| ) | |
| if isinstance(exc, httpx.RequestError): | |
| return self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker request failed.", | |
| ) | |
| del readiness_endpoint, job_endpoint | |
| # Never propagate implementation exception text: httpx exceptions can | |
| # include request URLs and caller implementations can include secrets. | |
| return self._error( | |
| WorkerErrorCategory.UNKNOWN_ERROR, | |
| "Generation worker operation failed unexpectedly.", | |
| ) | |
| def _response_error( | |
| self, | |
| status_code: int, | |
| *, | |
| readiness_endpoint: bool = False, | |
| job_endpoint: bool = False, | |
| ) -> GenerationWorkerError: | |
| if status_code in {400, 422}: | |
| return self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation worker rejected the request.", | |
| http_status=status_code, | |
| ) | |
| if status_code == 401: | |
| return self._error( | |
| WorkerErrorCategory.AUTHENTICATION_ERROR, | |
| "Generation worker authentication failed.", | |
| http_status=status_code, | |
| ) | |
| if status_code == 403: | |
| return self._error( | |
| WorkerErrorCategory.AUTHORIZATION_ERROR, | |
| "Generation worker authorization failed.", | |
| http_status=status_code, | |
| ) | |
| if status_code == 404 and job_endpoint: | |
| return self._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation worker job was not found.", | |
| http_status=status_code, | |
| ) | |
| if status_code == 429: | |
| return self._error( | |
| WorkerErrorCategory.RATE_LIMITED, | |
| "Generation worker is rate limited.", | |
| http_status=status_code, | |
| ) | |
| if status_code == 503 and readiness_endpoint: | |
| return self._error( | |
| WorkerErrorCategory.WORKER_NOT_READY, | |
| "Generation worker is not ready.", | |
| http_status=status_code, | |
| ) | |
| if status_code in {502, 503, 504}: | |
| return self._error( | |
| WorkerErrorCategory.WORKER_UNAVAILABLE, | |
| "Generation worker is temporarily unavailable.", | |
| http_status=status_code, | |
| ) | |
| if status_code >= 500: | |
| return self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker failed to process the request.", | |
| http_status=status_code, | |
| ) | |
| return self._error( | |
| WorkerErrorCategory.PROVIDER_ERROR, | |
| "Generation worker returned an unsupported response.", | |
| http_status=status_code, | |
| ) | |
| def _error( | |
| category: WorkerErrorCategory, | |
| message: str, | |
| *, | |
| http_status: int | None = None, | |
| ) -> GenerationWorkerError: | |
| return GenerationWorkerError( | |
| category=category, | |
| message=message, | |
| retryable=False, | |
| http_status=http_status, | |
| ) | |
| def _headers(self, extra: dict[str, str] | None) -> dict[str, str]: | |
| headers = dict(extra or {}) | |
| if self._bearer_token: | |
| headers["Authorization"] = f"Bearer {self._bearer_token}" | |
| return headers | |
| def _url_for(self, path: str) -> str: | |
| return f"{self.base_url}{path}" | |
| def _safe_external_id(value: str) -> str: | |
| if not value or len(value) > 255 or any( | |
| character not in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._:-" | |
| for character in value | |
| ): | |
| raise RemoteWorkerClient._error( | |
| WorkerErrorCategory.INVALID_REQUEST, | |
| "Generation worker job identifier is invalid.", | |
| ) | |
| return value | |
| def _safe_worker_path(value: str) -> str: | |
| parsed = urlparse(value) | |
| if ( | |
| not value.startswith("/") | |
| or parsed.scheme | |
| or parsed.netloc | |
| or parsed.query | |
| or parsed.fragment | |
| or "\\" in value | |
| or "%" in value | |
| or "//" in value | |
| or any(part in {"", ".", ".."} for part in value.split("/")[1:]) | |
| ): | |
| raise GenerationOutputError("Generation worker returned an unsafe endpoint path.") | |
| return value | |
| def _validate_base_url(value: str) -> str: | |
| parsed = urlparse(value.strip()) | |
| if ( | |
| parsed.scheme not in {"https", "http"} | |
| or not parsed.hostname | |
| or parsed.username | |
| or parsed.password | |
| or parsed.query | |
| or parsed.fragment | |
| ): | |
| raise ValueError("Generation worker base URL must be an absolute HTTP(S) origin.") | |
| host = parsed.hostname.lower() | |
| try: | |
| address = ipaddress.ip_address(host) | |
| except ValueError: | |
| address = None | |
| is_loopback = host == "localhost" or (address is not None and address.is_loopback) | |
| if address is not None and not address.is_global and not is_loopback: | |
| raise ValueError("Generation worker base URL uses a prohibited address.") | |
| if parsed.scheme == "http" and not is_loopback: | |
| raise ValueError("Generation workers require HTTPS outside local development.") | |
| path = parsed.path.rstrip("/") | |
| if path and ( | |
| "\\" in path | |
| or "%" in path | |
| or "//" in path | |
| or any(part in {"", ".", ".."} for part in path.split("/")[1:]) | |
| ): | |
| raise ValueError("Generation worker base URL contains an unsafe path.") | |
| return f"{parsed.scheme}://{parsed.netloc}{path}" | |
| def _health_status(value: object) -> WorkerHealthStatus: | |
| normalized = str(value or "").strip().lower() | |
| if normalized in {"ok", "healthy", "ready"}: | |
| return WorkerHealthStatus.HEALTHY | |
| if normalized in {"starting", "loading", "initializing"}: | |
| return WorkerHealthStatus.STARTING | |
| if normalized in {"unavailable", "offline"}: | |
| return WorkerHealthStatus.UNAVAILABLE | |
| if normalized in {"unhealthy", "failed", "error"}: | |
| return WorkerHealthStatus.UNHEALTHY | |
| return WorkerHealthStatus.UNKNOWN | |
| def _readiness_status(value: object) -> WorkerReadinessStatus: | |
| normalized = str(value or "").strip().lower() | |
| if normalized == "ready": | |
| return WorkerReadinessStatus.READY | |
| if normalized in {"starting", "loading", "initializing"}: | |
| return WorkerReadinessStatus.STARTING | |
| if normalized in { | |
| "not_ready", | |
| "unavailable", | |
| "offline", | |
| "unhealthy", | |
| "failed", | |
| "error", | |
| }: | |
| return WorkerReadinessStatus.UNAVAILABLE | |
| return WorkerReadinessStatus.UNKNOWN | |
| def _metadata(payload: dict[str, object], *, known: set[str]) -> dict[str, object]: | |
| data = safe_worker_metadata( | |
| {key: value for key, value in payload.items() if key not in known} | |
| ) | |
| return data if isinstance(data, dict) else {} | |