feat: store async generated images as WebP

app/main.py

@@ -2,29 +2,39 @@ from __future__ import annotations
 
 import asyncio
 import base64
+import binascii
 import hashlib
 import hmac
 import io
 import json
 import logging
 import os
+import re
+import sqlite3
 import time
 import uuid
 from contextlib import asynccontextmanager
+from pathlib import Path
 from typing import Any
+from urllib.parse import quote
 
 import aiosqlite
 import httpx
 from fastapi import FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
-from fastapi.responses import HTMLResponse
+from fastapi.responses import FileResponse, HTMLResponse, Response
 from PIL import Image, ImageOps, UnidentifiedImageError
 
 LOG = logging.getLogger("phproxy")
 DB_PATH = os.getenv("DB_PATH", "./data/proxy.db")
+OUTPUT_DIR = Path(os.getenv("OUTPUT_DIR", "./data/generated")).resolve()
 ADMIN_SECRET = os.getenv("ADMIN_SECRET", "")
 MAX_UPLOAD_BYTES = int(os.getenv("MAX_UPLOAD_BYTES", str(50 * 1024 * 1024)))
 MAX_IMAGES = int(os.getenv("MAX_IMAGES", "8"))
 MAX_IMAGE_PIXELS = int(os.getenv("MAX_IMAGE_PIXELS", "25000000"))
+IMAGE_LINK_TTL_SECONDS = int(os.getenv("IMAGE_LINK_TTL_SECONDS", "3600"))
+OUTPUT_WEBP_QUALITY = int(os.getenv("OUTPUT_WEBP_QUALITY", "90"))
+OUTPUT_RETENTION_SECONDS = int(os.getenv("OUTPUT_RETENTION_SECONDS", "86400"))
+OUTPUT_CLEANUP_INTERVAL_SECONDS = 300
 HTTP_TIMEOUT_SECONDS = float(os.getenv("HTTP_TIMEOUT_SECONDS", "300"))
 CONNECT_TIMEOUT_SECONDS = float(os.getenv("CONNECT_TIMEOUT_SECONDS", "15"))
 MODEL_MAP = {
@@ -47,6 +57,175 @@ def secret_matches(value: str | None, stored_hash: str | None) -> bool:
     return bool(value and stored_hash and hmac.compare_digest(secret_hash(value), stored_hash))
 
 
+def make_image_download_token(job_id: str, client_id: str, image_index: int) -> str:
+    claims = json.dumps(
+        {"job_id": job_id, "client_id": client_id, "image_index": image_index,
+         "expires_at": int(time.time()) + IMAGE_LINK_TTL_SECONDS},
+        separators=(",", ":"),
+    ).encode("utf-8")
+    payload = base64.urlsafe_b64encode(claims).rstrip(b"=").decode("ascii")
+    signature = hmac.new(ADMIN_SECRET.encode("utf-8"), payload.encode("ascii"), hashlib.sha256).hexdigest()
+    return f"{payload}.{signature}"
+
+
+def signed_image_download_url(job_id: str, client_id: str, image_index: int) -> str:
+    token = make_image_download_token(job_id, client_id, image_index)
+    return f"/v1/jobs/{job_id}/images/{image_index}?token={quote(token, safe='')}"
+
+
+def verify_image_download_token(token: str, job_id: str, image_index: int) -> str:
+    try:
+        payload, supplied_signature = token.split(".", 1)
+        expected_signature = hmac.new(
+            ADMIN_SECRET.encode("utf-8"), payload.encode("ascii"), hashlib.sha256
+        ).hexdigest()
+        if not hmac.compare_digest(supplied_signature, expected_signature):
+            raise ValueError("invalid signature")
+        raw_claims = base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))
+        claims = json.loads(raw_claims)
+        if (claims.get("job_id") != job_id or claims.get("image_index") != image_index
+                or not isinstance(claims.get("client_id"), str)):
+            raise ValueError("token scope mismatch")
+        if not isinstance(claims.get("expires_at"), int) or claims["expires_at"] < int(time.time()):
+            raise HTTPException(status_code=410, detail="Image link has expired")
+        return claims["client_id"]
+    except HTTPException:
+        raise
+    except (ValueError, TypeError, UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise HTTPException(status_code=401, detail="Invalid image link") from exc
+
+
+def output_image_path(job_id: str, image_index: int, mime_type: str) -> Path:
+    image_format = IMAGE_FORMAT_TYPES.get(mime_type)
+    if image_format is None or image_index < 0:
+        raise ValueError("Unsupported generated image type or index")
+    job_key = uuid.UUID(job_id).hex
+    return OUTPUT_DIR / f"{job_key}-{image_index}.{image_format.lower()}"
+
+
+def convert_image_to_webp(content: bytes, mime_type: str) -> bytes:
+    if not 1 <= OUTPUT_WEBP_QUALITY <= 100:
+        raise ValueError("OUTPUT_WEBP_QUALITY must be between 1 and 100")
+    Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
+    try:
+        with Image.open(io.BytesIO(content)) as source:
+            actual_mime = IMAGE_FORMATS.get(source.format or "")
+            if actual_mime is None or actual_mime != mime_type:
+                raise ValueError("Generated image content does not match its MIME type")
+            width, height = source.size
+            if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS:
+                raise ValueError("Generated image dimensions exceed configured limits")
+            source.load()
+            image = ImageOps.exif_transpose(source)
+            has_alpha = "A" in image.getbands() or "transparency" in image.info
+            image = image.convert("RGBA" if has_alpha else "RGB")
+            output = io.BytesIO()
+            image.save(output, format="WEBP", quality=OUTPUT_WEBP_QUALITY, method=4)
+            return output.getvalue()
+    except (Image.DecompressionBombError, Image.DecompressionBombWarning) as exc:
+        raise ValueError("Generated image dimensions exceed configured limits") from exc
+    except (UnidentifiedImageError, OSError) as exc:
+        raise ValueError("Provider returned an invalid image") from exc
+
+
+def write_output_file(path: Path, content: bytes) -> None:
+    temp_path = OUTPUT_DIR / f".{path.name}.{uuid.uuid4().hex}.tmp"
+    with temp_path.open("xb") as image_file:
+        os.chmod(temp_path, 0o600)
+        image_file.write(content)
+        image_file.flush()
+        os.fsync(image_file.fileno())
+    os.replace(temp_path, path)
+    os.chmod(path, 0o600)
+
+
+def cleanup_expired_output_files(now: float | None = None) -> int:
+    """Delete generated image files once they have been stored for 24 hours."""
+    cutoff = (time.time() if now is None else now) - OUTPUT_RETENTION_SECONDS
+    generated_name = re.compile(r"[0-9a-f]{32}-[0-9]+\.(?:webp|png|jpe?g)")
+    removed = 0
+    try:
+        paths = list(OUTPUT_DIR.iterdir())
+    except FileNotFoundError:
+        return 0
+    for path in paths:
+        if not generated_name.fullmatch(path.name):
+            continue
+        try:
+            if path.is_file() and path.stat().st_mtime <= cutoff:
+                path.unlink()
+                removed += 1
+        except FileNotFoundError:
+            continue
+        except OSError:
+            LOG.exception("Could not remove expired output image %s", path.name)
+    return removed
+
+
+async def output_cleanup_loop() -> None:
+    while True:
+        await asyncio.sleep(OUTPUT_CLEANUP_INTERVAL_SECONDS)
+        try:
+            removed = await asyncio.to_thread(cleanup_expired_output_files)
+            if removed:
+                LOG.info("Removed %s expired generated image(s)", removed)
+        except Exception:
+            LOG.exception("Generated image cleanup failed")
+
+
+def persist_result_images(job_id: str, result: dict[str, Any]) -> tuple[dict[str, Any], bool, list[Path]]:
+    """Store async image outputs as WebP and return compact result metadata."""
+    compact_result = dict(result)
+    compact_images = []
+    written_paths: list[Path] = []
+    superseded_paths: list[Path] = []
+    changed = False
+    try:
+        for image_index, image in enumerate(result.get("images", [])):
+            image_metadata = dict(image)
+            encoded = image_metadata.pop("b64_json", None)
+            if encoded:
+                mime_type = image_metadata.get("mime_type")
+                source_content = base64.b64decode(encoded, validate=True)
+                webp_content = convert_image_to_webp(source_content, mime_type)
+                path = output_image_path(job_id, image_index, "image/webp")
+                write_output_file(path, webp_content)
+                written_paths.append(path)
+                image_metadata["mime_type"] = "image/webp"
+                changed = True
+            elif image_metadata.get("mime_type") in IMAGE_FORMAT_TYPES:
+                old_path = output_image_path(job_id, image_index, image_metadata["mime_type"])
+                if image_metadata["mime_type"] != "image/webp" and old_path.is_file():
+                    source_content = old_path.read_bytes()
+                    webp_content = convert_image_to_webp(source_content, image_metadata["mime_type"])
+                    path = output_image_path(job_id, image_index, "image/webp")
+                    write_output_file(path, webp_content)
+                    written_paths.append(path)
+                    superseded_paths.append(old_path)
+                    image_metadata["mime_type"] = "image/webp"
+                    changed = True
+            compact_images.append(image_metadata)
+        compact_result["images"] = compact_images
+        return compact_result, changed, superseded_paths
+    except Exception:
+        for path in written_paths:
+            path.unlink(missing_ok=True)
+        raise
+
+
+def compact_sqlite_database() -> None:
+    """Checkpoint and compact SQLite without an outstanding async cursor."""
+    conn = sqlite3.connect(DB_PATH, timeout=30)
+    try:
+        checkpoint = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
+        checkpoint.fetchall()
+        checkpoint.close()
+        vacuum = conn.execute("VACUUM")
+        vacuum.close()
+    finally:
+        conn.close()
+
+
 @asynccontextmanager
 async def db():
     conn = await aiosqlite.connect(DB_PATH, timeout=10)
@@ -61,6 +240,11 @@ async def db():
 
 async def init_db() -> None:
     os.makedirs(os.path.dirname(os.path.abspath(DB_PATH)), exist_ok=True)
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True, mode=0o700)
+    os.chmod(OUTPUT_DIR, 0o700)
+    migrated_result_count = 0
+    cleared_request_count = 0
+    superseded_paths: list[Path] = []
     async with db() as conn:
         await conn.execute("PRAGMA journal_mode=WAL")
         await conn.executescript(
@@ -163,20 +347,80 @@ async def init_db() -> None:
                 LOG.exception("Could not backfill generation history for job %s", row["job_id"])
         await conn.commit()
 
+        cur = await conn.execute(
+            "SELECT job_id,result_json,request_json FROM jobs "
+            "WHERE status IN ('completed','failed')"
+        )
+        for row in await cur.fetchall():
+            if row["result_json"]:
+                try:
+                    result = json.loads(row["result_json"])
+                    compact_result, changed, old_paths = await asyncio.to_thread(
+                        persist_result_images, row["job_id"], result
+                    )
+                    superseded_paths.extend(old_paths)
+                    if changed:
+                        await conn.execute(
+                            "UPDATE jobs SET result_json=? WHERE job_id=?",
+                            (json.dumps(compact_result), row["job_id"]),
+                        )
+                        migrated_result_count += 1
+                except Exception:
+                    # Keep the legacy JSON available if an individual historic image cannot be migrated.
+                    LOG.exception("Could not migrate stored output images for job %s", row["job_id"])
+            if row["request_json"] and row["request_json"] != "{}":
+                await conn.execute("UPDATE jobs SET request_json='{}' WHERE job_id=?", (row["job_id"],))
+                cleared_request_count += 1
+        await conn.commit()
+
+    # Remove old PNG/JPEG outputs only after their WebP replacements are committed.
+    for old_path in superseded_paths:
+        try:
+            old_path.unlink(missing_ok=True)
+        except OSError:
+            LOG.exception("Could not remove superseded output image %s", old_path.name)
+
+    async with db() as conn:
+        cur = await conn.execute("PRAGMA freelist_count")
+        free_page_count = (await cur.fetchone())[0]
+    if migrated_result_count or cleared_request_count or free_page_count > 1000:
+        # Reclaim SQLite pages freed by removing Base64 images and completed request payloads.
+        try:
+            await asyncio.to_thread(compact_sqlite_database)
+        except sqlite3.Error:
+            LOG.exception("Could not compact the SQLite database; service will continue")
+    if migrated_result_count or cleared_request_count:
+        LOG.info("Migrated %s job result(s) to files and cleared %s completed request payload(s)",
+                 migrated_result_count, cleared_request_count)
+
 
 @asynccontextmanager
 async def lifespan(_: FastAPI):
     global http_client
+    cleanup_task: asyncio.Task | None = None
     if not ADMIN_SECRET:
         raise RuntimeError("ADMIN_SECRET must be set before starting the service")
+    if OUTPUT_RETENTION_SECONDS <= 0:
+        raise RuntimeError("OUTPUT_RETENTION_SECONDS must be greater than zero")
     await init_db()
+    removed = await asyncio.to_thread(cleanup_expired_output_files)
+    if removed:
+        LOG.info("Removed %s expired generated image(s) during startup", removed)
     http_client = httpx.AsyncClient(
         timeout=httpx.Timeout(HTTP_TIMEOUT_SECONDS, connect=CONNECT_TIMEOUT_SECONDS),
         limits=httpx.Limits(max_keepalive_connections=50, max_connections=200),
         follow_redirects=False,
     )
-    yield
-    await http_client.aclose()
+    cleanup_task = asyncio.create_task(output_cleanup_loop())
+    try:
+        yield
+    finally:
+        cleanup_task.cancel()
+        try:
+            await cleanup_task
+        except asyncio.CancelledError:
+            pass
+        await http_client.aclose()
 
 
 app = FastAPI(title="Generative AI Proxy", version="0.1.0", lifespan=lifespan)
@@ -511,17 +755,27 @@ async def run_job(job_id: str, client_id: str, request_data: dict[str, Any],
         result.update({"generation_id": job_id,
                        "request_intake_sec": run_metadata["request_intake_sec"],
                        "total_sec": round(completed_at - run_metadata["received_at"], 2)})
+        compact_result, _, superseded_paths = await asyncio.to_thread(
+            persist_result_images, job_id, result
+        )
         async with db() as conn:
-            await conn.execute("UPDATE jobs SET status='completed',result_json=?,updated_at=? WHERE job_id=?",
-                               (json.dumps(result), completed_at, job_id))
+            await conn.execute(
+                "UPDATE jobs SET status='completed',result_json=?,request_json='{}',updated_at=? WHERE job_id=?",
+                (json.dumps(compact_result), completed_at, job_id),
+            )
             await conn.commit()
+        for old_path in superseded_paths:
+            try:
+                old_path.unlink(missing_ok=True)
+            except OSError:
+                LOG.exception("Could not remove superseded output image %s", old_path.name)
         await finish_generation_run(job_id, "completed", completed_at,
                                     result["total_sec"], result=result)
     except HTTPException as exc:
         error = exc.detail
         completed_at = time.time()
         async with db() as conn:
-            await conn.execute("UPDATE jobs SET status='failed',error_json=?,updated_at=? WHERE job_id=?",
+            await conn.execute("UPDATE jobs SET status='failed',request_json='{}',error_json=?,updated_at=? WHERE job_id=?",
                                (json.dumps(error), completed_at, job_id))
             await conn.commit()
         await finish_generation_run(job_id, "failed", completed_at,
@@ -531,7 +785,7 @@ async def run_job(job_id: str, client_id: str, request_data: dict[str, Any],
         completed_at = time.time()
         error = {"message": "Internal generation error"}
         async with db() as conn:
-            await conn.execute("UPDATE jobs SET status='failed',error_json=?,updated_at=? WHERE job_id=?",
+            await conn.execute("UPDATE jobs SET status='failed',request_json='{}',error_json=?,updated_at=? WHERE job_id=?",
                                (json.dumps(error), completed_at, job_id))
             await conn.commit()
         await finish_generation_run(job_id, "failed", completed_at,
@@ -605,6 +859,7 @@ async def generate(request: Request, client_id: str = Form(...), prompt: str = F
 
 @app.get("/v1/jobs/{job_id}")
 async def get_job(job_id: str, client_id: str = Query(...),
+                  include_result: bool = Query(True),
                   x_client_secret: str | None = Header(None)):
     await authorize_client(client_id, x_client_secret)
     async with db() as conn:
@@ -615,12 +870,122 @@ async def get_job(job_id: str, client_id: str = Query(...),
     response: dict[str, Any] = {"job_id": job_id, "status": row["status"],
                                 "created_at": row["created_at"], "updated_at": row["updated_at"]}
     if row["result_json"]:
-        response["result"] = json.loads(row["result_json"])
+        result = json.loads(row["result_json"])
+        if include_result:
+            legacy_result = dict(result)
+            legacy_images = []
+            for index, image in enumerate(result.get("images", [])):
+                image_data = dict(image)
+                if not image_data.get("b64_json") and image_data.get("mime_type") in IMAGE_FORMAT_TYPES:
+                    try:
+                        path = output_image_path(job_id, index, image_data["mime_type"])
+                        if path.is_file():
+                            content = await asyncio.to_thread(path.read_bytes)
+                            image_data["b64_json"] = base64.b64encode(content).decode("ascii")
+                    except (ValueError, OSError):
+                        LOG.exception("Could not read stored output image for job %s, image %s",
+                                      job_id, index)
+                legacy_images.append(image_data)
+            legacy_result["images"] = legacy_images
+            response["result"] = legacy_result
+        else:
+            # Keep status polling small; image bytes are fetched through the binary route.
+            result_metadata = {key: value for key, value in result.items() if key != "images"}
+            result_metadata["images"] = []
+            for index, image in enumerate(result.get("images", [])):
+                image_metadata = {key: value for key, value in image.items()
+                                  if key not in {"b64_json", "url"}}
+                image_is_local = bool(image.get("b64_json"))
+                if not image_is_local and image.get("mime_type") in IMAGE_FORMAT_TYPES:
+                    try:
+                        path = output_image_path(job_id, index, image.get("mime_type"))
+                        image_is_local = path.is_file()
+                    except ValueError:
+                        LOG.warning("Stored output image has unsupported MIME type for job %s, image %s",
+                                    job_id, index)
+                if image_is_local:
+                    image_metadata["download_url"] = signed_image_download_url(
+                        job_id, client_id, index
+                    )
+                elif image.get("url"):
+                    image_metadata["url"] = image["url"]
+                result_metadata["images"].append(image_metadata)
+            response["result"] = result_metadata
     if row["error_json"]:
         response["error"] = json.loads(row["error_json"])
     return response
 
 
+@app.get("/v1/jobs/{job_id}/images/{image_index}")
+async def get_job_image(job_id: str, image_index: int, client_id: str | None = Query(None),
+                        token: str | None = Query(None),
+                        x_client_secret: str | None = Header(None)):
+    """Return raw image bytes via a short-lived signed link or client-secret auth."""
+    if token:
+        token_client_id = verify_image_download_token(token, job_id, image_index)
+        if client_id and client_id != token_client_id:
+            raise HTTPException(status_code=401, detail="Invalid image link")
+        client_id = token_client_id
+    elif client_id:
+        await authorize_client(client_id, x_client_secret)
+    else:
+        raise HTTPException(status_code=401, detail="An image link or client authentication is required")
+
+    async with db() as conn:
+        cur = await conn.execute(
+            "SELECT j.status,j.result_json,c.is_active FROM jobs j "
+            "JOIN clients c ON c.client_id=j.client_id WHERE j.job_id=? AND j.client_id=?",
+            (job_id, client_id),
+        )
+        row = await cur.fetchone()
+    if not row or not row["is_active"]:
+        raise HTTPException(status_code=404, detail="Job not found")
+    if row["status"] != "completed" or not row["result_json"]:
+        raise HTTPException(status_code=409, detail="Job has no completed result")
+    images = json.loads(row["result_json"]).get("images", [])
+    if image_index < 0 or image_index >= len(images):
+        raise HTTPException(status_code=404, detail="Image not found")
+    image = images[image_index]
+    mime_type = image.get("mime_type")
+    if mime_type in IMAGE_FORMAT_TYPES:
+        path = output_image_path(job_id, image_index, mime_type)
+        if path.is_file():
+            extension = IMAGE_FORMAT_TYPES[mime_type].lower()
+            return FileResponse(
+                path,
+                media_type=mime_type,
+                filename=f"generated-{image_index}.{extension}",
+                content_disposition_type="inline",
+                headers={
+                    "Cache-Control": "private, no-store",
+                    "X-Content-Type-Options": "nosniff",
+                    "Referrer-Policy": "no-referrer",
+                },
+            )
+
+    # Keep legacy completed jobs accessible if a file migration was not possible.
+    encoded = image.get("b64_json")
+    if encoded:
+        try:
+            content = base64.b64decode(encoded, validate=True)
+        except (ValueError, TypeError, binascii.Error) as exc:
+            LOG.error("Stored output image is invalid for job %s, image %s", job_id, image_index)
+            raise HTTPException(status_code=500, detail="Stored image is invalid") from exc
+        return Response(
+            content=content,
+            media_type=mime_type if mime_type in IMAGE_FORMAT_TYPES else "application/octet-stream",
+            headers={
+                "Content-Disposition": f'inline; filename="generated-{image_index}"',
+                "Cache-Control": "private, no-store",
+                "X-Content-Type-Options": "nosniff",
+                "Referrer-Policy": "no-referrer",
+            },
+        )
+    if mime_type in IMAGE_FORMAT_TYPES:
+        raise HTTPException(status_code=410, detail="Image is no longer available")
+    raise HTTPException(status_code=409, detail="Image is available only from its upstream URL")
+
+
 @app.get("/api/admin/state")
 async def admin_state(x_admin_secret: str | None = Header(None)):
     auth_admin(x_admin_secret)