Webhook receiver

import sqlite3

db = sqlite3.connect("jobs.db", isolation_level=None, check_same_thread=False)
db.execute(
    """CREATE TABLE IF NOT EXISTS jobs (
        request_id TEXT PRIMARY KEY,
        state      TEXT NOT NULL,      -- pending | processing | completed | failed | nsfw | canceled | error
        updated_at TEXT NOT NULL DEFAULT (datetime('now'))
    )"""
)


def record_pending(request_id: str) -> None:
    db.execute("INSERT OR IGNORE INTO jobs (request_id, state) VALUES (?, 'pending')", (request_id,))


def claim(request_id: str) -> bool:
    """Atomically take a job for processing. False if it's done or someone else has it."""
    record_pending(request_id)
    cursor = db.execute(
        "UPDATE jobs SET state = 'processing', updated_at = datetime('now') "
        "WHERE request_id = ? AND state IN ('pending', 'error')",
        (request_id,),
    )
    return cursor.rowcount == 1


def finish(request_id: str, state: str) -> None:
    db.execute(
        "UPDATE jobs SET state = ?, updated_at = datetime('now') WHERE request_id = ?", (state, request_id)
    )


def pending_older_than(seconds: int) -> list[str]:
    rows = db.execute(
        "SELECT request_id FROM jobs WHERE state IN ('pending', 'error') "
        "AND updated_at < datetime('now', ?)",
        (f"-{seconds} seconds",),
    )
    return [row[0] for row in rows]

import asyncio
import mimetypes
import os
import re
from contextlib import asynccontextmanager
from pathlib import Path

import httpx
from fastapi import BackgroundTasks, FastAPI, Request, Response

import store

API = os.environ.get("WINDPAINT_API_URL", "https://api.windpaint.ai").rstrip("/") + "/v1"
KEY = os.environ["WINDPAINT_API_KEY"]
OUT = Path("outputs")
OUT.mkdir(exist_ok=True)
TERMINAL = {"completed", "failed", "nsfw", "canceled"}
UUID = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$")


async def process(request_id: str) -> None:
    if not store.claim(request_id):
        return  # already handled, or being handled
    try:
        async with httpx.AsyncClient(
            base_url=API, headers={"Authorization": f"Bearer {KEY}"}, timeout=60
        ) as api:
            response = await api.get(f"/generation/requests/{request_id}/status")
            if response.status_code == 404:
                store.finish(request_id, "ignored")  # not a job in your organization
                return
            response.raise_for_status()
            status = response.json()

            if status["status"] not in TERMINAL:
                store.finish(request_id, "pending")  # early or spoofed call; reconcile will retry
                return

            if status["status"] == "completed":
                for i, output in enumerate(status["outputs"]):
                    signed = await api.get(f"/generation/assets/{output['id']}/url")
                    signed.raise_for_status()
                    url = signed.json()["data"]["url"]
                    async with httpx.AsyncClient(timeout=300) as plain:  # no Authorization on signed URLs
                        media = await plain.get(url)
                        media.raise_for_status()
                    ext = mimetypes.guess_extension(output["content_type"]) or ""
                    (OUT / f"{request_id}-{i}{ext}").write_bytes(media.content)
                print(f"{request_id}: saved {len(status['outputs'])} output(s), "
                      f"{status['credits']['actual']} credits")
            else:
                print(f"{request_id}: {status['status']} ({status['error']})")
            store.finish(request_id, status["status"])
    except Exception:
        store.finish(request_id, "error")  # picked up again by the reconcile loop
        raise


async def reconcile_forever() -> None:
    """Catch jobs whose webhook never arrived: webhooks are sent once and not retried."""
    while True:
        await asyncio.sleep(60)
        for request_id in store.pending_older_than(120):
            try:
                await process(request_id)
            except Exception as exc:
                print(f"{request_id}: reconcile failed: {exc!r}")


@asynccontextmanager
async def lifespan(app: FastAPI):
    task = asyncio.create_task(reconcile_forever())
    yield
    task.cancel()


app = FastAPI(lifespan=lifespan)


@app.post("/windpaint/webhook")
async def webhook(request: Request, background: BackgroundTasks) -> Response:
    try:
        payload = await request.json()
    except ValueError:
        return Response(status_code=400)
    request_id = payload.get("request_id") if isinstance(payload, dict) else None
    if not isinstance(request_id, str) or not UUID.match(request_id):
        return Response(status_code=400)
    # Answer within the 10 s timeout; do the work after responding.
    background.add_task(process, request_id)
    return Response(status_code=204)

import os
import sys

import httpx

import store

API = os.environ.get("WINDPAINT_API_URL", "https://api.windpaint.ai").rstrip("/") + "/v1"

response = httpx.post(
    f"{API}/generation/capabilities/image.generate",
    headers={"Authorization": f"Bearer {os.environ['WINDPAINT_API_KEY']}"},
    json={
        "prompt": sys.argv[1],
        "aspect_ratio": "16:9",
        "webhook_url": os.environ["WEBHOOK_URL"],
    },
    timeout=60,
)
response.raise_for_status()
job = response.json()
store.record_pending(job["request_id"])  # the reconcile loop now owns this job too
print(f"submitted {job['request_id']} ({job['credits_estimate']} credits held)")