launch-check-studio-v1/agent.py

"""LaunchCheck Studio: SSRF-safe launch readiness audits with user-scoped persistence."""
from __future__ import annotations

import asyncio
import hashlib
import ipaddress
import json
import os
import re
import socket
from dataclasses import dataclass
from datetime import datetime, timezone
from html.parser import HTMLParser
from typing import Annotated, Any
from urllib.parse import urljoin, urlparse, urlunparse

from pydantic import BaseModel, Field

import a2a_pack as a2a
from a2a_pack import (
    A2AAgent,
    AgentDatabase,
    AgentDatabaseEnv,
    AgentDatabaseMigrations,
    AgentPlatformResources,
    EgressPolicy,
    PlatformUserAuth,
    Pricing,
    Resources,
    RunContext,
    State,
)

MAX_REDIRECTS = 4
MAX_RESPONSE_BYTES = 512_000
MAX_ROBOTS_BYTES = 64_000
DNS_TIMEOUT_SECONDS = 2.0
CONNECT_TIMEOUT_SECONDS = 3.0
READ_TIMEOUT_SECONDS = 5.0
TOTAL_TIMEOUT_SECONDS = 12.0
USER_AGENT = "LaunchCheckStudio/0.1 (+https://a2acloud.io)"
DB_CONNECT_OPTIONS = "-c statement_timeout=5000 -c lock_timeout=2000 -c idle_in_transaction_session_timeout=5000"


class LaunchCheckStudioV1Config(BaseModel):
    max_redirects: int = Field(default=MAX_REDIRECTS, ge=0, le=5)
    max_response_bytes: int = Field(default=MAX_RESPONSE_BYTES, ge=32_000, le=1_000_000)


class LaunchCheckStudioV1State(BaseModel):
    storage: str = "managed-postgres"


@dataclass(frozen=True)
class FetchResult:
    final_url: str
    status_code: int | None
    headers: dict[str, str]
    body: bytes
    redirect_chain: list[dict[str, Any]]
    error: dict[str, Any] | None = None


class _PinnedResolver:
    """Resolve one hostname only to addresses validated before the request."""

    def __init__(self, hostname: str, addresses: tuple[str, ...]) -> None:
        self._hostname = hostname.lower().rstrip(".")
        self._addresses = addresses

    async def resolve(
        self,
        host: str,
        port: int = 0,
        family: int = socket.AF_UNSPEC,
    ) -> list[dict[str, Any]]:
        if host.lower().rstrip(".") != self._hostname:
            raise OSError("resolver hostname mismatch")
        return [
            {
                "hostname": host,
                "host": address,
                "port": port,
                "family": (
                    socket.AF_INET6
                    if ipaddress.ip_address(address).version == 6
                    else socket.AF_INET
                ),
                "proto": 0,
                "flags": 0,
            }
            for address in self._addresses
        ]

    async def close(self) -> None:
        return None


class _HeadParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.in_title = False
        self.title_parts: list[str] = []
        self.description: str | None = None
        self.canonical: str | None = None
        self.viewport: str | None = None

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        tag = tag.lower()
        attr = {k.lower(): (v or "") for k, v in attrs}
        if tag == "title":
            self.in_title = True
        elif tag == "meta":
            name = attr.get("name", "").lower()
            prop = attr.get("property", "").lower()
            content = attr.get("content", "").strip()
            if name == "description" and content and self.description is None:
                self.description = content[:500]
            if name == "viewport" and content and self.viewport is None:
                self.viewport = content[:300]
            if prop == "og:description" and content and self.description is None:
                self.description = content[:500]
        elif tag == "link" and attr.get("rel", "").lower() == "canonical":
            href = attr.get("href", "").strip()
            if href and self.canonical is None:
                self.canonical = href[:1000]

    def handle_endtag(self, tag: str) -> None:
        if tag.lower() == "title":
            self.in_title = False

    def handle_data(self, data: str) -> None:
        if self.in_title:
            self.title_parts.append(data)

    @property
    def title(self) -> str | None:
        value = " ".join(" ".join(self.title_parts).split()).strip()
        return value[:300] if value else None


class LaunchCheckStudioV1(A2AAgent[LaunchCheckStudioV1Config, PlatformUserAuth]):
    name = "launch-check-studio-v1"
    description = (
        "LaunchCheck audits public HTTP(S) launch pages with bounded SSRF-safe "
        "network checks and user-scoped Postgres persistence."
    )
    version = "0.1.1"

    config_model = LaunchCheckStudioV1Config
    auth_model = PlatformUserAuth

    state = State.DURABLE
    state_model = LaunchCheckStudioV1State
    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=120)
    egress = EgressPolicy(deny_internet_by_default=False)
    tools_used = ("aiohttp", "psycopg", "postgres")
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic bounded audit; no LLM credential required.",
    )
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name="launch-check-studio-v1-data",
                scope="user",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
            ),
        )
    )

    @a2a.tool(
        description="Run a bounded SSRF-safe launch readiness audit for a public HTTP(S) URL and persist the report.",
        timeout_seconds=60,
        idempotent=False,
        cost_class="network-read",
    )
    async def audit_url(
        self,
        ctx: RunContext[PlatformUserAuth],
        audit_id: Annotated[
            str,
            Field(min_length=1, max_length=128, pattern=r"^[A-Za-z0-9][A-Za-z0-9_.:-]{0,127}$"),
        ],
        url: Annotated[str, Field(min_length=8, max_length=2048)],
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        audit_id = _clean_audit_id(audit_id)
        started_at = _now_iso()
        await ctx.emit_progress("Validating public target and resolving DNS")

        validation = await _validate_public_url(url)
        if not validation["ok"]:
            failure = _failure_result(audit_id, url, validation["code"], validation["message"])
            await _persist_failure_if_possible(tenant, audit_id, url, failure, started_at)
            return failure

        try:
            fetch = await _fetch_bounded(validation["url"], self.config.max_redirects, self.config.max_response_bytes)
            if fetch.error:
                result = _failure_result(audit_id, validation["url"], fetch.error["code"], fetch.error["message"])
                result["target_url"] = validation["url"]
                await _persist_failure_if_possible(tenant, audit_id, validation["url"], result, started_at)
                return result

            await ctx.emit_progress("Analyzing metadata, robots, and security headers")
            robots = await _check_robots(fetch.final_url)
            report = _build_report(fetch, robots)
            receipt = _build_app_receipt(tenant, audit_id, validation["url"], report, started_at)
            payload = {
                "audit_id": audit_id,
                "target_url": validation["url"],
                "ok": True,
                "report": report,
                "receipt": receipt,
                "created_at": started_at,
                "updated_at": _now_iso(),
            }
            await _save_audit_and_receipt(tenant, audit_id, payload, receipt)
            return payload
        except Exception:  # noqa: BLE001 - structured public failure, no raw internals
            result = _failure_result(
                audit_id,
                validation.get("url", url),
                "audit_failed",
                "The audit could not be completed within the bounded execution policy.",
            )
            await _persist_failure_if_possible(tenant, audit_id, validation.get("url", url), result, started_at)
            return result

    @a2a.tool(
        description="Reopen a previously persisted launch-readiness audit for the signed-in user.",
        timeout_seconds=20,
        idempotent=True,
        cost_class="db-read",
    )
    async def get_audit(
        self,
        ctx: RunContext[PlatformUserAuth],
        audit_id: Annotated[
            str,
            Field(min_length=1, max_length=128, pattern=r"^[A-Za-z0-9][A-Za-z0-9_.:-]{0,127}$"),
        ],
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        audit_id = _clean_audit_id(audit_id)
        row = await _load_audit(tenant, audit_id)
        if row is None:
            return {
                "ok": False,
                "code": "audit_not_found",
                "message": "No audit with that id exists for this signed-in user.",
                "audit_id": audit_id,
            }
        return row


def _tenant_key(ctx: RunContext[PlatformUserAuth]) -> str:
    stable_id = getattr(ctx.auth, "user_id", None)
    if stable_id is not None:
        return f"user:{stable_id}"
    sub = str(getattr(ctx.auth, "sub", "") or "").strip()
    if sub:
        return f"user:{sub}"
    raise PermissionError("stable platform identity required")


def _clean_audit_id(value: str) -> str:
    cleaned = str(value or "").strip()
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:-]{0,127}", cleaned):
        raise ValueError("audit_id must be 1-128 safe characters")
    return cleaned


def _now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def _failure_result(audit_id: str, url: str, code: str, message: str) -> dict[str, Any]:
    return {
        "ok": False,
        "code": code,
        "message": message,
        "audit_id": audit_id,
        "target_url": _safe_url_echo(url),
    }


def _safe_url_echo(url: str) -> str:
    try:
        parsed = urlparse(url)
        netloc = parsed.hostname or ""
        if parsed.port:
            netloc = f"{netloc}:{parsed.port}"
        return urlunparse((parsed.scheme, netloc, parsed.path or "", "", parsed.query[:200], ""))[:1000]
    except Exception:  # noqa: BLE001
        return ""


async def _validate_public_url(raw_url: str) -> dict[str, Any]:
    value = str(raw_url or "").strip()
    if len(value) > 2048:
        return {"ok": False, "code": "invalid_url", "message": "URL is too long."}
    parsed = urlparse(value)
    if parsed.scheme.lower() not in {"http", "https"}:
        return {"ok": False, "code": "invalid_scheme", "message": "Only http and https URLs are supported."}
    if parsed.username or parsed.password or "@" in parsed.netloc:
        return {"ok": False, "code": "credentials_not_allowed", "message": "URLs with embedded credentials are not allowed."}
    if not parsed.hostname:
        return {"ok": False, "code": "invalid_url", "message": "URL must include a hostname."}
    host = parsed.hostname.strip().rstrip(".")
    if not host:
        return {"ok": False, "code": "invalid_url", "message": "URL must include a hostname."}
    try:
        port = parsed.port
    except ValueError:
        return {"ok": False, "code": "invalid_url", "message": "URL port is invalid."}
    if port is not None and not (1 <= port <= 65535):
        return {"ok": False, "code": "invalid_url", "message": "URL port is invalid."}
    try:
        host = str(ipaddress.ip_address(host.strip("[]")))
    except ValueError:
        try:
            host = host.encode("idna").decode("ascii").lower()
        except UnicodeError:
            return {"ok": False, "code": "invalid_url", "message": "URL hostname is invalid."}
    addresses = await _resolve_public_addresses(
        host,
        port or (443 if parsed.scheme == "https" else 80),
    )
    if addresses is None:
        return {"ok": False, "code": "disallowed_private_target", "message": "Target resolves to a non-public network address."}
    if not addresses:
        return {"ok": False, "code": "dns_resolution_failed", "message": "Target hostname could not be resolved safely."}
    netloc = f"[{host}]" if ":" in host else host
    if port:
        netloc = f"{netloc}:{port}"
    normalized = urlunparse((parsed.scheme.lower(), netloc, parsed.path or "", "", parsed.query, ""))
    return {
        "ok": True,
        "url": normalized,
        "hostname": host,
        "port": port or (443 if parsed.scheme == "https" else 80),
        "addresses": addresses,
    }


async def _resolve_public_addresses(host: str, port: int) -> tuple[str, ...] | None:
    try:
        ip = ipaddress.ip_address(host.strip("[]"))
        return None if _ip_disallowed(ip) else (str(ip),)
    except ValueError:
        pass
    lowered = host.lower()
    if lowered in {"localhost", "metadata.google.internal"} or lowered.endswith(".localhost"):
        return None
    try:
        infos = await asyncio.wait_for(
            asyncio.get_running_loop().getaddrinfo(host, port, type=socket.SOCK_STREAM),
            timeout=DNS_TIMEOUT_SECONDS,
        )
    except Exception:
        return ()
    addresses: set[str] = set()
    for info in infos:
        sockaddr = info[4]
        if sockaddr:
            addresses.add(str(sockaddr[0]))
    if not addresses:
        return ()
    if any(_ip_disallowed(ipaddress.ip_address(address)) for address in addresses):
        return None
    return tuple(sorted(addresses))


def _ip_disallowed(ip: ipaddress._BaseAddress) -> bool:
    metadata = ipaddress.ip_address("169.254.169.254")
    return bool(
        ip.is_private
        or ip.is_loopback
        or ip.is_link_local
        or ip.is_multicast
        or ip.is_reserved
        or ip.is_unspecified
        or ip == metadata
    )


async def _fetch_bounded(url: str, max_redirects: int, max_bytes: int) -> FetchResult:
    current = url
    chain: list[dict[str, Any]] = []
    for hop in range(max_redirects + 1):
        validation = await _validate_public_url(current)
        if not validation["ok"]:
            return FetchResult(current, None, {}, b"", chain, {"code": validation["code"], "message": validation["message"]})
        try:
            status, headers, body, response_url = await _request_pinned(
                validation,
                max_bytes=max_bytes,
                total_timeout=TOTAL_TIMEOUT_SECONDS,
                connect_timeout=CONNECT_TIMEOUT_SECONDS,
                read_timeout=READ_TIMEOUT_SECONDS,
            )
        except TimeoutError:
            return FetchResult(current, None, {}, b"", chain, {"code": "network_timeout", "message": "The target did not respond within the bounded timeout."})
        except Exception:  # noqa: BLE001 - return a stable public network error
            return FetchResult(current, None, {}, b"", chain, {"code": "network_error", "message": "The target could not be reached safely."})
        if status in {301, 302, 303, 307, 308} and headers.get("location"):
            next_url = urljoin(response_url, headers["location"])
            chain.append({"from": validation["url"], "to": _safe_url_echo(next_url), "status_code": status})
            if hop >= max_redirects:
                return FetchResult(validation["url"], status, headers, body, chain, {"code": "too_many_redirects", "message": "Redirect count exceeded the audit limit."})
            current = next_url
            continue
        return FetchResult(response_url, status, headers, body, chain)
    return FetchResult(current, None, {}, b"", chain, {"code": "too_many_redirects", "message": "Redirect count exceeded the audit limit."})


async def _request_pinned(
    validation: dict[str, Any],
    *,
    max_bytes: int,
    total_timeout: float,
    connect_timeout: float,
    read_timeout: float,
) -> tuple[int, dict[str, str], bytes, str]:
    import aiohttp

    resolver = _PinnedResolver(
        str(validation["hostname"]),
        tuple(str(item) for item in validation["addresses"]),
    )
    connector = aiohttp.TCPConnector(resolver=resolver, use_dns_cache=False)
    timeout = aiohttp.ClientTimeout(
        total=total_timeout,
        connect=connect_timeout,
        sock_connect=connect_timeout,
        sock_read=read_timeout,
    )
    try:
        async with aiohttp.ClientSession(
            connector=connector,
            connector_owner=True,
            timeout=timeout,
            trust_env=False,
            headers={"user-agent": USER_AGENT, "accept-encoding": "identity"},
        ) as client:
            async with client.get(validation["url"], allow_redirects=False) as response:
                body = await _read_limited(response, max_bytes)
                headers = {key.lower(): value[:1000] for key, value in response.headers.items()}
                return response.status, headers, body, str(response.url)
    except asyncio.TimeoutError as exc:
        raise TimeoutError("bounded request timed out") from exc


async def _read_limited(response: Any, max_bytes: int) -> bytes:
    chunks: list[bytes] = []
    total = 0
    async for chunk in response.content.iter_chunked(64 * 1024):
        total += len(chunk)
        if total > max_bytes:
            remaining = max_bytes - (total - len(chunk))
            if remaining > 0:
                chunks.append(chunk[:remaining])
            break
        chunks.append(chunk)
    return b"".join(chunks)


async def _check_robots(final_url: str) -> dict[str, Any]:
    parsed = urlparse(final_url)
    robots_url = urlunparse((parsed.scheme, parsed.netloc, "/robots.txt", "", "", ""))
    validation = await _validate_public_url(robots_url)
    if not validation["ok"]:
        return {"present": False, "status_code": None, "checked_url": _safe_url_echo(robots_url)}
    try:
        status, _headers, body, _response_url = await _request_pinned(
            validation,
            max_bytes=MAX_ROBOTS_BYTES,
            total_timeout=6.0,
            connect_timeout=2.0,
            read_timeout=3.0,
        )
        return {
            "present": status == 200,
            "status_code": status,
            "checked_url": validation["url"],
            "bytes_read": len(body),
        }
    except Exception:  # noqa: BLE001
        return {"present": False, "status_code": None, "checked_url": validation["url"]}


def _build_report(fetch: FetchResult, robots: dict[str, Any]) -> dict[str, Any]:
    text = fetch.body.decode("utf-8", errors="replace")
    parser = _HeadParser()
    parser.feed(text[:200_000])
    headers = fetch.headers
    security = {
        "strict_transport_security": bool(headers.get("strict-transport-security")),
        "content_security_policy": bool(headers.get("content-security-policy")),
        "x_content_type_options": headers.get("x-content-type-options", "").lower() == "nosniff",
        "referrer_policy": bool(headers.get("referrer-policy")),
        "permissions_policy": bool(headers.get("permissions-policy")),
        "x_frame_options": bool(headers.get("x-frame-options")),
    }
    evidence = {
        "status_code": fetch.status_code,
        "final_url": fetch.final_url,
        "redirect_chain": fetch.redirect_chain,
        "https": urlparse(fetch.final_url).scheme == "https",
        "title": parser.title,
        "description": parser.description,
        "canonical": urljoin(fetch.final_url, parser.canonical) if parser.canonical else None,
        "viewport": parser.viewport,
        "robots": robots,
        "security_headers": security,
        "bytes_read": len(fetch.body),
    }
    checks = {
        "reachable": fetch.status_code is not None and 200 <= int(fetch.status_code) < 400,
        "https": evidence["https"],
        "has_title": bool(parser.title),
        "has_description": bool(parser.description),
        "has_canonical": bool(parser.canonical),
        "has_viewport": bool(parser.viewport),
        "robots_present": bool(robots.get("present")),
        "security_headers": security,
    }
    fixes = _recommended_fixes(checks)
    return {
        "reachable": checks["reachable"],
        "final_url": fetch.final_url,
        "checks": checks,
        "evidence": evidence,
        "fixes": fixes,
        "summary": _summary(checks, fixes),
    }


def _recommended_fixes(checks: dict[str, Any]) -> list[dict[str, str]]:
    fixes: list[dict[str, str]] = []
    if not checks["https"]:
        fixes.append({"category": "security", "priority": "high", "fix": "Serve the launch page over HTTPS and redirect HTTP to HTTPS."})
    if not checks["has_title"]:
        fixes.append({"category": "metadata", "priority": "high", "fix": "Add a concise, unique <title> for search and sharing."})
    if not checks["has_description"]:
        fixes.append({"category": "metadata", "priority": "medium", "fix": "Add a meta description that explains the product value proposition."})
    if not checks["has_canonical"]:
        fixes.append({"category": "seo", "priority": "medium", "fix": "Add a canonical link to avoid duplicate indexing ambiguity."})
    if not checks["has_viewport"]:
        fixes.append({"category": "mobile", "priority": "high", "fix": "Add a responsive viewport meta tag for mobile launch traffic."})
    if not checks["robots_present"]:
        fixes.append({"category": "seo", "priority": "low", "fix": "Publish /robots.txt so crawlers receive explicit indexing guidance."})
    security = checks["security_headers"]
    for key, label in {
        "strict_transport_security": "Strict-Transport-Security",
        "content_security_policy": "Content-Security-Policy",
        "x_content_type_options": "X-Content-Type-Options: nosniff",
        "referrer_policy": "Referrer-Policy",
        "permissions_policy": "Permissions-Policy",
    }.items():
        if not security.get(key):
            fixes.append({"category": "security", "priority": "medium", "fix": f"Add {label} to harden the launch page."})
    return fixes


def _summary(checks: dict[str, Any], fixes: list[dict[str, str]]) -> dict[str, Any]:
    total = 7 + len(checks["security_headers"])
    passed = sum(1 for key in ["reachable", "https", "has_title", "has_description", "has_canonical", "has_viewport", "robots_present"] if checks.get(key))
    passed += sum(1 for value in checks["security_headers"].values() if value)
    return {
        "score": round((passed / total) * 100),
        "passed": passed,
        "total": total,
        "fix_count": len(fixes),
        "top_priority": fixes[0]["priority"] if fixes else "none",
    }


def _build_app_receipt(tenant: str, audit_id: str, target_url: str, report: dict[str, Any], started_at: str) -> dict[str, Any]:
    public_payload = {
        "tenant_hash": hashlib.sha256(tenant.encode()).hexdigest()[:16],
        "audit_id": audit_id,
        "target_url": target_url,
        "reachable": report.get("reachable"),
        "started_at": started_at,
    }
    digest = hashlib.sha256(json.dumps(public_payload, sort_keys=True).encode()).hexdigest()
    return {
        "receipt_id": f"launchcheck:{audit_id}:{digest[:16]}",
        "digest": digest,
        "kind": "launchcheck.audit.v1",
        "created_at": _now_iso(),
    }


def _db_url() -> str:
    value = os.environ.get("DATABASE_URL", "").strip()
    if not value:
        raise RuntimeError("database_unavailable")
    return value


def _connect_db() -> Any:
    import psycopg

    return psycopg.connect(_db_url(), options=DB_CONNECT_OPTIONS, connect_timeout=5)


async def _save_audit_and_receipt(tenant: str, audit_id: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
    await asyncio.to_thread(_save_audit_and_receipt_sync, tenant, audit_id, payload, receipt)


def _save_audit_and_receipt_sync(tenant: str, audit_id: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
    from psycopg.types.json import Jsonb

    with _connect_db() as conn:
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT (tenant_key, audit_id) DO UPDATE SET
                        target_url = EXCLUDED.target_url,
                        ok = EXCLUDED.ok,
                        report = EXCLUDED.report,
                        payload = EXCLUDED.payload,
                        updated_at = NOW()
                    """,
                    (tenant, audit_id, payload["target_url"], True, Jsonb(payload["report"]), Jsonb(payload)),
                )
                cur.execute(
                    """
                    INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload)
                    VALUES (%s, %s, %s, %s)
                    ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET
                        payload = EXCLUDED.payload,
                        created_at = NOW()
                    """,
                    (tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)),
                )


async def _persist_failure_if_possible(tenant: str, audit_id: str, target_url: str, payload: dict[str, Any], started_at: str) -> None:
    if not os.environ.get("DATABASE_URL"):
        return
    receipt = {
        "receipt_id": f"launchcheck:{audit_id}:failure:{hashlib.sha256((audit_id + started_at).encode()).hexdigest()[:12]}",
        "kind": "launchcheck.audit_failure.v1",
        "created_at": _now_iso(),
        "code": payload.get("code"),
    }
    saved_payload = {**payload, "receipt": receipt, "created_at": started_at, "updated_at": _now_iso()}
    try:
        await asyncio.to_thread(_persist_failure_sync, tenant, audit_id, target_url, saved_payload, receipt)
    except Exception:  # noqa: BLE001
        return


def _persist_failure_sync(tenant: str, audit_id: str, target_url: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
    from psycopg.types.json import Jsonb

    with _connect_db() as conn:
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT (tenant_key, audit_id) DO UPDATE SET
                        target_url = EXCLUDED.target_url,
                        ok = EXCLUDED.ok,
                        report = EXCLUDED.report,
                        payload = EXCLUDED.payload,
                        updated_at = NOW()
                    """,
                    (tenant, audit_id, _safe_url_echo(target_url), False, Jsonb({}), Jsonb(payload)),
                )
                cur.execute(
                    """
                    INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload)
                    VALUES (%s, %s, %s, %s)
                    ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET payload = EXCLUDED.payload
                    """,
                    (tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)),
                )


async def _load_audit(tenant: str, audit_id: str) -> dict[str, Any] | None:
    try:
        return await asyncio.to_thread(_load_audit_sync, tenant, audit_id)
    except RuntimeError as exc:
        if str(exc) == "database_unavailable":
            return {"ok": False, "code": "database_unavailable", "message": "Persistent storage is not configured for this deployment.", "audit_id": audit_id}
        raise


def _load_audit_sync(tenant: str, audit_id: str) -> dict[str, Any] | None:
    with _connect_db() as conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT payload FROM launch_audits WHERE tenant_key = %s AND audit_id = %s",
                (tenant, audit_id),
            )
            row = cur.fetchone()
            return row[0] if row else None