from __future__ import annotations import math import re from typing import Any from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from app.generation.domain.enums import ( GenerationModality, WorkerCancellationStatus, WorkerErrorCategory, WorkerHealthStatus, WorkerJobStatus, WorkerReadinessStatus, ) _EXTERNAL_ID = re.compile(r"^[A-Za-z0-9._:-]{1,255}$") _SHA256 = re.compile(r"^[0-9a-f]{64}$") _SAFE_FILENAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,254}$") _BEARER = re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/=-]+") _HTTP_URL = re.compile(r"(?i)\bhttps?://[^\s\"'<>]+") _ASSIGNED_SECRET = re.compile( r"(?i)(access[_-]?token|refresh[_-]?token|id[_-]?token|client[_-]?secret|" r"authorization|api[_-]?key|password|secret|credential)" r"([\"']?\s*[:=]\s*[\"']?)([^\"'\s,&}]+)" ) _SENSITIVE_METADATA_PARTS = frozenset( { "access_token", "refresh_token", "id_token", "token", "secret", "authorization", "cookie", "password", "api_key", "credential", "url", "uri", } ) _MAX_METADATA_DEPTH = 8 _MAX_METADATA_ITEMS = 256 _MAX_METADATA_STRING_LENGTH = 8_192 _DROP = object() def _safe_metadata(value: Any, *, depth: int = 0) -> Any: """Return bounded JSON-safe metadata with credential-like content removed.""" if depth > _MAX_METADATA_DEPTH: return _DROP if value is None or isinstance(value, bool) or isinstance(value, int): return value if isinstance(value, float): return value if math.isfinite(value) else _DROP if isinstance(value, str): without_bearer = _BEARER.sub("Bearer [REDACTED]", value) without_urls = _HTTP_URL.sub("[REDACTED_URL]", without_bearer) redacted = _ASSIGNED_SECRET.sub( lambda match: f"{match.group(1)}{match.group(2)}[REDACTED]", without_urls, ) return redacted[:_MAX_METADATA_STRING_LENGTH] if isinstance(value, dict): result: dict[str, object] = {} for key, item in list(value.items())[:_MAX_METADATA_ITEMS]: if not isinstance(key, str) or len(key) > 255: continue normalized_key = key.lower().replace("-", "_") if any(part in normalized_key for part in _SENSITIVE_METADATA_PARTS): continue cleaned = _safe_metadata(item, depth=depth + 1) if cleaned is not _DROP: result[key] = cleaned return result if isinstance(value, (list, tuple)): result: list[object] = [] for item in list(value)[:_MAX_METADATA_ITEMS]: cleaned = _safe_metadata(item, depth=depth + 1) if cleaned is not _DROP: result.append(cleaned) return result return _DROP def safe_worker_metadata(value: Any) -> Any: """Remove secret-like fields and non-JSON values before persistence. Worker metadata is useful for diagnostics, but it is never a credential store. This helper is deliberately conservative and is applied both at worker-transport parsing and at persistence boundaries. """ cleaned = _safe_metadata(value) return {} if cleaned is _DROP and isinstance(value, dict) else cleaned def _safe_metadata_dict(value: Any) -> dict[str, object]: cleaned = safe_worker_metadata(value) if not isinstance(cleaned, dict): raise ValueError("worker metadata must be a JSON object") return cleaned class WorkerModelInfo(BaseModel): """One worker-discovered model; never an operator-configured endpoint.""" model_config = ConfigDict(extra="forbid") id: str = Field(min_length=1, max_length=255) name: str = Field(min_length=1, max_length=255) media_types: list[GenerationModality] = Field(default_factory=list) metadata: dict[str, object] = Field(default_factory=dict) @field_validator("metadata", mode="before") @classmethod def clean_metadata(cls, value: Any) -> dict[str, object]: return _safe_metadata_dict(value) class WorkerInfo(BaseModel): """Verified non-secret identity returned by a configured worker.""" model_config = ConfigDict(extra="forbid") id: str = Field(min_length=1, max_length=255) name: str = Field(min_length=1, max_length=255) media_types: list[GenerationModality] = Field(default_factory=list) models: list[WorkerModelInfo] = Field(default_factory=list) status: WorkerHealthStatus = WorkerHealthStatus.UNKNOWN metadata: dict[str, object] = Field(default_factory=dict) @field_validator("models") @classmethod def unique_models(cls, values: list[WorkerModelInfo]) -> list[WorkerModelInfo]: if len({value.id for value in values}) != len(values): raise ValueError("worker metadata contains duplicate model IDs") return values @field_validator("metadata", mode="before") @classmethod def clean_metadata(cls, value: Any) -> dict[str, object]: return _safe_metadata_dict(value) class WorkerHealth(BaseModel): """Normalised liveness result. It never implies model availability.""" model_config = ConfigDict(extra="forbid") status: WorkerHealthStatus metadata: dict[str, object] = Field(default_factory=dict) @field_validator("metadata", mode="before") @classmethod def clean_metadata(cls, value: Any) -> dict[str, object]: return _safe_metadata_dict(value) class WorkerReadiness(BaseModel): """Normalised inference readiness result.""" model_config = ConfigDict(extra="forbid") status: WorkerReadinessStatus model_loaded: bool = False model_ids: list[str] = Field(default_factory=list) metadata: dict[str, object] = Field(default_factory=dict) @field_validator("model_ids") @classmethod def valid_model_ids(cls, values: list[str]) -> list[str]: if len(values) != len(set(values)): raise ValueError("worker readiness model_ids must be unique") for value in values: if not value or len(value) > 255 or any(character.isspace() for character in value): raise ValueError("worker readiness model_ids contain an invalid value") return values @field_validator("metadata", mode="before") @classmethod def clean_metadata(cls, value: Any) -> dict[str, object]: return _safe_metadata_dict(value) class WorkerOutput(BaseModel): """A worker-issued descriptor, never a filesystem path or arbitrary URL.""" model_config = ConfigDict(extra="forbid") output_type: GenerationModality mime_type: str = Field( min_length=3, max_length=255, pattern=r"^[a-z0-9!#$&^_.+-]+/[a-z0-9!#$&^_.+-]+$", ) provider_output_id: str = Field(min_length=1, max_length=255) # A worker may expose a download endpoint only under its configured origin. # RemoteWorkerClient rejects absolute URLs, traversal, queries, and fragments. download_path: str = Field(min_length=2, max_length=2048) filename: str | None = Field(default=None, max_length=255) sha256: str | None = Field(default=None, max_length=64) byte_size: int | None = Field(default=None, ge=0) metadata: dict[str, object] = Field(default_factory=dict) @field_validator("provider_output_id") @classmethod def valid_provider_output_id(cls, value: str) -> str: if _EXTERNAL_ID.fullmatch(value) is None: raise ValueError("provider_output_id contains unsupported characters") return value @field_validator("download_path") @classmethod def valid_download_path(cls, value: str) -> str: if ( not value.startswith("/") or "//" in value or "\\" in value or "%" in value or "?" in value or "#" in value or any(part in {"", ".", ".."} for part in value.split("/")[1:]) ): raise ValueError("download_path must be a safe absolute worker-relative path") return value @field_validator("filename") @classmethod def valid_filename(cls, value: str | None) -> str | None: if value is not None and _SAFE_FILENAME.fullmatch(value) is None: raise ValueError("filename must not contain a path") return value @field_validator("sha256") @classmethod def valid_sha256(cls, value: str | None) -> str | None: if value is not None and _SHA256.fullmatch(value) is None: raise ValueError("sha256 must be a lowercase SHA-256 hex digest") return value @field_validator("metadata", mode="before") @classmethod def clean_metadata(cls, value: Any) -> dict[str, object]: return _safe_metadata_dict(value) @model_validator(mode="after") def media_type_matches_output_type(self) -> "WorkerOutput": if not self.mime_type.startswith(f"{self.output_type.value}/"): raise ValueError("output MIME type does not match its generation modality") return self class WorkerJob(BaseModel): """Provider-neutral remote job representation.""" model_config = ConfigDict(extra="forbid") external_job_id: str = Field(min_length=1, max_length=255) status: WorkerJobStatus output: WorkerOutput | None = None error_category: WorkerErrorCategory | None = None error_code: str | None = Field(default=None, max_length=100) error_message: str | None = Field(default=None, max_length=500) metadata: dict[str, object] = Field(default_factory=dict) @field_validator("external_job_id") @classmethod def valid_external_job_id(cls, value: str) -> str: if _EXTERNAL_ID.fullmatch(value) is None: raise ValueError("external_job_id contains unsupported characters") return value @field_validator("error_code") @classmethod def valid_error_code(cls, value: str | None) -> str | None: if value is not None and _EXTERNAL_ID.fullmatch(value) is None: raise ValueError("worker error_code contains unsupported characters") return value @field_validator("error_message", mode="before") @classmethod def clean_error_message(cls, value: Any) -> str | None: if value is None: return None if not isinstance(value, str): raise ValueError("worker error_message must be a string") cleaned = safe_worker_metadata(value) if not isinstance(cleaned, str): # Defensive: strings are JSON-safe. raise ValueError("worker error_message is invalid") return cleaned @model_validator(mode="after") def completed_job_requires_output(self) -> "WorkerJob": if self.status is WorkerJobStatus.COMPLETED and self.output is None: raise ValueError("completed worker jobs must include an output descriptor") return self @field_validator("metadata", mode="before") @classmethod def clean_metadata(cls, value: Any) -> dict[str, object]: return _safe_metadata_dict(value) class WorkerCancellationResult(BaseModel): model_config = ConfigDict(extra="forbid") status: WorkerCancellationStatus metadata: dict[str, object] = Field(default_factory=dict) @field_validator("metadata", mode="before") @classmethod def clean_metadata(cls, value: Any) -> dict[str, object]: return _safe_metadata_dict(value)