vrex-flow-engine 0.2.0__py3-none-any.whl

flow_engine/job_store.py

@@ -0,0 +1,28 @@
+"""Per-instance job store.
+
+Extracted from the module-level ``_jobs`` dict in jobs.py so each pool
+instance owns an isolated namespace — jobs from account A never appear in
+account B's job listing.
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from flow_engine.jobs import Job
+
+
+class JobStore:
+    """Simple dict-backed job registry, one per PoolInstance."""
+
+    def __init__(self) -> None:
+        self._jobs: dict[str, "Job"] = {}
+
+    def get(self, job_id: str) -> Optional["Job"]:
+        return self._jobs.get(job_id)
+
+    def put(self, job: "Job") -> None:
+        self._jobs[job.id] = job
+
+    def __len__(self) -> int:
+        return len(self._jobs)

flow_engine/jobs.py

@@ -0,0 +1,245 @@
+"""Async video job runner.
+
+flowboard ran generation through a DB-backed request queue + worker. flow_engine
+replaces that with per-request asyncio tasks and a per-instance JobStore.
+
+The poll loop mirrors flowboard/agent worker `_handle_gen_video`: per-op
+resolution with partial-failure aggregation (some variants can be blocked by
+the content filter while others render). State is process-local; swap for a
+shared store if you scale past one process.
+
+All public functions accept explicit ``job_store``, ``media_store``, and
+``vc_store`` arguments instead of using module globals — this gives each
+PoolInstance full isolation.
+"""
+from __future__ import annotations
+
+import asyncio
+import base64
+import logging
+import os
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from flow_engine.bridge.flow_sdk import FlowSDK
+from flow_engine.job_store import JobStore
+from flow_engine.media_store import MediaStore
+from flow_engine.video_context_store import VideoContextStore
+from flow_engine.video_context import parse_seconds
+
+logger = logging.getLogger(__name__)
+
+# Poll knobs — same budget as flowboard's worker (5 min: 30 cycles × 10s).
+# Poll budget for waiting on Veo results. Default 72 × 10s = 12 min, to tolerate
+# Flow's veo_*_low_priority queue (which can exceed the old 5-min budget and
+# return timeout_waiting_video). Env-tunable without a code change.
+VIDEO_POLL_INTERVAL_S = float(os.getenv("VIDEO_POLL_INTERVAL_S", "10"))
+VIDEO_POLL_MAX_CYCLES = int(os.getenv("VIDEO_POLL_MAX_CYCLES", "72"))
+
+
+@dataclass
+class Job:
+    id: str
+    kind: str  # "video"
+    status: str = "in_progress"  # in_progress | completed | failed
+    created_at: int = field(default_factory=lambda: int(time.time()))
+    # Positional results aligned to dispatch order.
+    media_ids: list[Optional[str]] = field(default_factory=list)
+    slot_errors: list[Optional[str]] = field(default_factory=list)
+    error: Optional[str] = None
+    partial_error: Optional[str] = None
+    raw: dict[str, Any] = field(default_factory=dict)
+
+
+# Generation mode → the FlowSDK dispatch coroutine it maps to.
+_DISPATCH = {
+    "t2v": "gen_video_t2v",
+    "r2v": "gen_video_r2v",
+    "i2v": "gen_video",
+    "extension": "gen_video_extend",
+    "edit": "gen_video_edit",
+}
+
+
+def get(job_id: str, job_store: JobStore) -> Optional[Job]:
+    return job_store.get(job_id)
+
+
+def _new_job(kind: str, job_store: JobStore) -> Job:
+    job = Job(id=f"{kind}_{uuid.uuid4().hex}", kind=kind)
+    job_store.put(job)
+    return job
+
+
+def start_video_job(
+    sdk: FlowSDK,
+    *,
+    mode: str,
+    dispatch_kwargs: dict[str, Any],
+    job_store: JobStore,
+    media_store: MediaStore,
+    vc_store: VideoContextStore,
+) -> Job:
+    """Create a job and kick off dispatch+poll in the background.
+
+    ``mode`` ∈ {t2v, i2v, r2v, extension, edit} selects the FlowSDK dispatch
+    method; ``dispatch_kwargs`` is forwarded to it.
+    """
+    job = _new_job("video", job_store)
+    asyncio.create_task(
+        _run_video_job(job, sdk, mode, dispatch_kwargs, job_store, media_store, vc_store)
+    )
+    return job
+
+
+async def _run_video_job(
+    job: Job,
+    sdk: FlowSDK,
+    mode: str,
+    kwargs: dict[str, Any],
+    job_store: JobStore,
+    media_store: MediaStore,
+    vc_store: VideoContextStore,
+) -> None:
+    try:
+        dispatch = await getattr(sdk, _DISPATCH[mode])(**kwargs)
+    except Exception as exc:  # noqa: BLE001
+        logger.exception("video dispatch crashed")
+        job.status, job.error = "failed", str(exc)[:200]
+        return
+
+    if dispatch.get("error"):
+        job.status = "failed"
+        job.error = str(dispatch["error"])[:200]
+        job.raw = {"dispatch": dispatch}
+        return
+
+    op_names: list[str] = dispatch.get("operation_names") or []
+    if not op_names:
+        job.status, job.error = "failed", "no_operations_returned"
+        job.raw = {"dispatch": dispatch}
+        return
+    workflows = dispatch.get("workflows") or None
+
+    done_by_name = {name: False for name in op_names}
+    entry_by_name: dict[str, dict] = {}
+    op_errors: dict[str, str] = {}
+    last_poll: dict = {}
+    attempts = 0
+
+    while attempts < VIDEO_POLL_MAX_CYCLES and not all(done_by_name.values()):
+        await asyncio.sleep(VIDEO_POLL_INTERVAL_S)
+        attempts += 1
+        last_poll = await sdk.check_async(
+            op_names, workflows=workflows, project_id=kwargs.get("project_id")
+        )
+        if last_poll.get("error"):
+            continue
+        for op in last_poll.get("operations") or []:
+            if not isinstance(op, dict):
+                continue
+            name = op.get("name")
+            if not isinstance(name, str) or done_by_name.get(name, False):
+                continue
+            err = op.get("error")
+            if isinstance(err, str) and err:
+                done_by_name[name] = True
+                op_errors[name] = err
+                continue
+            if op.get("done"):
+                done_by_name[name] = True
+                for e in op.get("media_entries") or []:
+                    if isinstance(e, dict) and e.get("media_id"):
+                        entry_by_name[name] = e
+                        break
+
+    # Unresolved after budget → timeout.
+    for name in op_names:
+        if not done_by_name.get(name) and name not in op_errors:
+            op_errors[name] = "timeout_waiting_video"
+
+    positional_ids: list[Optional[str]] = []
+    slot_errors: list[Optional[str]] = []
+    succeeded: list[dict] = []
+    for name in op_names:
+        e = entry_by_name.get(name)
+        if isinstance(e, dict) and isinstance(e.get("media_id"), str):
+            positional_ids.append(e["media_id"])
+            succeeded.append(e)
+            slot_errors.append(None)
+        else:
+            positional_ids.append(None)
+            slot_errors.append(op_errors.get(name))
+
+    if not any(positional_ids):
+        job.status = "failed"
+        job.error = next(iter(op_errors.values()), "timeout_waiting_video")
+        job.slot_errors = slot_errors
+        job.raw = {"dispatch": dispatch, "last_poll": last_poll}
+        return
+
+    # Ingest what rendered into this instance's MediaStore.
+    url_entries = [e for e in succeeded if isinstance(e, dict) and e.get("url")]
+    if url_entries:
+        try:
+            media_store.ingest_urls(url_entries)
+        except Exception:  # noqa: BLE001
+            logger.exception("ingest_urls failed for video job %s", job.id)
+    for e in succeeded:
+        encoded, mid = e.get("encoded_video"), e.get("media_id")
+        if isinstance(encoded, str) and isinstance(mid, str):
+            try:
+                media_store.ingest_inline_bytes(
+                    mid,
+                    base64.b64decode(encoded, validate=False),
+                    kind="video",
+                    mime="video/mp4",
+                )
+            except Exception:  # noqa: BLE001
+                logger.exception("inline ingest failed for %s", mid)
+
+    # Stash extend/edit context into this instance's VideoContextStore.
+    wf_by_media = {
+        w.get("primary_media_id"): w.get("name")
+        for w in (workflows or [])
+        if isinstance(w, dict)
+    }
+    finished_ids = [
+        e["media_id"]
+        for e in succeeded
+        if isinstance(e, dict) and isinstance(e.get("media_id"), str)
+    ]
+    rich_meta: dict[str, dict] = {}
+    if finished_ids:
+        project_id = kwargs.get("project_id")
+        if isinstance(project_id, str):
+            try:
+                rich_meta = (
+                    await sdk.fetch_video_meta(finished_ids, project_id)
+                ).get("meta") or {}
+            except Exception:  # noqa: BLE001
+                logger.exception("fetch_video_meta failed for job %s", job.id)
+
+    dispatched_scene = dispatch.get("scene_id")
+    for mid in finished_ids:
+        rm = rich_meta.get(mid) or {}
+        vc_store.remember(
+            mid,
+            workflow_id=wf_by_media.get(mid) or rm.get("workflow_id"),
+            operation_id=rm.get("operation_id"),
+            scene_id=dispatched_scene or rm.get("scene_id"),
+            duration_sec=parse_seconds(rm.get("length")),
+        )
+
+    if op_errors:
+        unique = sorted({v for v in op_errors.values()})
+        job.partial_error = (
+            f"{len(op_errors)}/{len(op_names)} variants blocked: {', '.join(unique)}"
+        )
+
+    job.media_ids = positional_ids
+    job.slot_errors = slot_errors
+    job.status = "completed"
+    job.raw = {"dispatch": dispatch, "last_poll": last_poll, "video_meta": rich_meta}

flow_engine/main.py

@@ -0,0 +1,187 @@
+"""flow_engine FastAPI app.
+
+Hosts the extension bridge (WS :9223) and an OpenAI-compatible HTTP surface
+(:8101) for Google Flow image/video generation.
+
+Boot-time checks:
+  - FLOW_ENGINE_API_KEY (or legacy FLOWPROXY_API_KEY) must be non-empty.
+  - WS_HOST must be loopback (the extension WS is unauthenticated by design).
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+
+from fastapi import Depends, FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, Response
+
+from flow_engine.config import API_KEY, WS_HOST
+from flow_engine.openai import images, models, uploads, videos
+from flow_engine.openai._util import require_api_key
+from flow_engine.pool import pool
+from flow_engine.posthog_client import flush_posthog
+
+# The WS server is unauthenticated — refuse to boot on a non-loopback bind.
+if WS_HOST not in ("127.0.0.1", "localhost", "::1"):
+    raise RuntimeError(
+        f"FLOWPROXY_WS_HOST must be loopback (got {WS_HOST!r}); the extension WS "
+        "is unauthenticated by design and must not be network-reachable."
+    )
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Enforce non-empty API key at server startup, not at import time, so that
+    # `pip install -e .` and plain imports never crash (only the running server).
+    if not API_KEY:
+        raise RuntimeError(
+            "FLOW_ENGINE_API_KEY must be set (or legacy FLOWPROXY_API_KEY). "
+            "Set the environment variable before starting flow-engine."
+        )
+
+    from flow_engine.bridge.ws_server import run_ws_server
+
+    ws_task = asyncio.create_task(run_ws_server(), name="ext-ws-server")
+    logger.info("flow_engine started (ws bridge + openai surface)")
+    try:
+        yield
+    finally:
+        ws_task.cancel()
+        await asyncio.gather(ws_task, return_exceptions=True)
+        flush_posthog()  # drain queued events before process exits
+        logger.info("flow_engine stopped")
+
+
+app = FastAPI(title="flow_engine", version="0.0.1", lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.include_router(images.router)
+app.include_router(videos.router)
+app.include_router(uploads.router)
+app.include_router(models.router)
+
+
+@app.get("/api/health")
+def health() -> dict:
+    instances = pool.status()
+    connected = any(i["connected"] for i in instances)
+    return {
+        "ok": True,
+        "extension_connected": connected,
+        "instances": len(instances),
+    }
+
+
+@app.get("/api/pool/status")
+def pool_status(_: str | None = Depends(require_api_key)) -> dict:
+    """Admin view of all registered pool instances (API-key gated)."""
+    return {"instances": pool.status()}
+
+
+@app.post("/api/ext/callback")
+async def ext_callback(
+    request: Request,
+    x_callback_secret: str | None = None,
+) -> dict:
+    """HTTP callback the extension uses to deliver API responses.
+
+    Routes to the correct PoolInstance by searching all instances' pending
+    futures for a matching request id.
+    """
+    # Read the secret from the header manually so the endpoint stays flexible.
+    x_callback_secret = request.headers.get("x-callback-secret")
+    if not x_callback_secret:
+        raise HTTPException(status_code=401, detail="missing X-Callback-Secret")
+
+    # Validate secret against any registered instance — each has its own.
+    matched_instance = None
+    for inst in pool.status():
+        # pool.status() returns dicts; we need the live instances.
+        pass
+    # Access live instances directly from the pool registry.
+    import hmac
+    valid = False
+    for inst_obj in pool._instances.values():
+        if hmac.compare_digest(x_callback_secret, inst_obj.client.callback_secret):
+            valid = True
+            break
+    if not valid:
+        raise HTTPException(status_code=401, detail="invalid callback secret")
+
+    try:
+        payload = await request.json()
+    except Exception:
+        raise HTTPException(status_code=400, detail="invalid json body")
+    if not isinstance(payload, dict) or "id" not in payload:
+        raise HTTPException(status_code=400, detail="missing id")
+
+    resolved = pool.resolve_callback(payload)
+    return {"ok": resolved}
+
+
+@app.get("/media/{media_id}")
+async def serve_media(media_id: str, request: Request) -> Response:
+    """Serve cached media bytes. Routes to the correct instance's MediaStore
+    by trying each registered instance in pool order."""
+    from flow_engine.media import normalize_media_id, is_valid_media_id
+
+    media_id = normalize_media_id(media_id)
+    if not is_valid_media_id(media_id):
+        raise HTTPException(status_code=404, detail="invalid media id")
+
+    # Try per-instance stores first (multi-tenant path).
+    for inst_obj in pool._instances.values():
+        path = await inst_obj.media.ensure_cached(media_id)
+        if path is not None:
+            return FileResponse(path, media_type=inst_obj.media.media_mime(media_id))
+
+    # Fallback: module-level media store (populated by single-tenant code paths
+    # that haven't been migrated yet, or for legacy compatibility).
+    from flow_engine import media as media_service
+    path = await media_service.ensure_cached(media_id)
+    if path is None:
+        raise HTTPException(status_code=404, detail="media not available")
+    return FileResponse(path, media_type=media_service.media_mime(media_id))
+
+
+@app.get("/media/{media_id}/status")
+def media_status(media_id: str) -> dict:
+    from flow_engine.media import normalize_media_id, is_valid_media_id
+
+    media_id = normalize_media_id(media_id)
+    if not is_valid_media_id(media_id):
+        return {"available": False, "has_url": False, "reason": "invalid_id"}
+
+    for inst_obj in pool._instances.values():
+        s = inst_obj.media.status(media_id)
+        if s.get("available") or s.get("has_url"):
+            return s
+
+    from flow_engine import media as media_service
+    return media_service.status(media_id)
+
+
+def main() -> None:
+    import uvicorn
+
+    from flow_engine.config import HTTP_PORT
+
+    uvicorn.run("flow_engine.main:app", host="127.0.0.1", port=HTTP_PORT, reload=False)
+
+
+if __name__ == "__main__":
+    main()

flow_engine/media.py

@@ -0,0 +1,227 @@
+"""Media cache — DB-free port of flowboard's media service.
+
+flowboard persisted media to a SQLModel ``Asset`` table; flow_engine keeps the
+same behaviour with a process-local registry plus an on-disk byte cache under
+``storage/media``. Flow returns short-TTL signed ``flow-content.google`` URLs
+(and, for workflow-mode video, inline base64 MP4), so we fetch/plant bytes
+locally and serve stable ``/media/<id>`` URLs to OpenAI clients.
+
+Registry is in-memory and rebuilt from disk on startup — good enough for a
+single-tenant proxy. Swap for Redis/SQLite if you need persistence across
+restarts or a horizontal pool.
+"""
+from __future__ import annotations
+
+import logging
+import re
+import threading
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+import httpx
+
+from flow_engine.config import STORAGE_DIR
+
+logger = logging.getLogger(__name__)
+
+MEDIA_CACHE_DIR = STORAGE_DIR / "media"
+MEDIA_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+# Media id is a hex-with-dashes UUID in the GCS path.
+_MEDIA_ID_RE = re.compile(r"^[0-9a-fA-F-]{1,64}$")
+
+# Flow serves user-generated media from its `flow-content.google` CDN (signed
+# with short-TTL query params); the response includes the signed URL at
+# `data.media[].image.generatedImage.fifeUrl`.
+_ALLOWED_URL_PREFIXES: tuple[str, ...] = ("https://flow-content.google/",)
+
+_EXT_BY_MIME = {
+    "image/jpeg": ".jpg",
+    "image/png": ".png",
+    "image/webp": ".webp",
+    "image/gif": ".gif",
+    "video/mp4": ".mp4",
+    "video/webm": ".webm",
+}
+
+
+@dataclass
+class Asset:
+    media_id: str
+    url: Optional[str] = None
+    kind: str = "image"
+    local_path: Optional[str] = None
+    mime: Optional[str] = None
+
+
+_registry: dict[str, Asset] = {}
+_lock = threading.Lock()
+
+
+def _url_allowed(url: Any) -> bool:
+    return isinstance(url, str) and any(url.startswith(p) for p in _ALLOWED_URL_PREFIXES)
+
+
+def is_valid_media_id(media_id: str) -> bool:
+    return bool(_MEDIA_ID_RE.fullmatch(media_id or ""))
+
+
+def normalize_media_id(raw: str) -> str:
+    """Accept either ``media/<uuid>`` or the bare uuid."""
+    if raw.startswith("media/"):
+        return raw.split("/", 1)[1]
+    return raw
+
+
+def _cache_glob(media_id: str) -> Optional[Path]:
+    if not is_valid_media_id(media_id):
+        return None
+    for p in MEDIA_CACHE_DIR.glob(f"{media_id}.*"):
+        if p.is_file():
+            return p
+    return None
+
+
+def cached_path(media_id: str) -> Optional[Path]:
+    return _cache_glob(media_id)
+
+
+def _get(media_id: str) -> Optional[Asset]:
+    with _lock:
+        return _registry.get(media_id)
+
+
+def ingest_urls(urls: list[dict[str, Any]]) -> int:
+    """Register Asset rows keyed by uuid media_id. Returns count touched."""
+    touched = 0
+    with _lock:
+        for entry in urls or []:
+            if not isinstance(entry, dict):
+                continue
+            media_id = entry.get("media_id") or entry.get("mediaId")
+            url = entry.get("url")
+            kind = entry.get("mediaType") or entry.get("kind") or "image"
+            if not isinstance(media_id, str) or not is_valid_media_id(media_id):
+                continue
+            if not _url_allowed(url):
+                logger.warning("media: skip non-allowed url for %s: %r", media_id, (url or "")[:80])
+                continue
+            row = _registry.get(media_id)
+            if row is None:
+                _registry[media_id] = Asset(media_id=media_id, url=url, kind=kind)
+            else:
+                row.url = url
+                if not row.kind:
+                    row.kind = kind
+            touched += 1
+    if touched:
+        logger.info("media: ingested %d url(s)", touched)
+    return touched
+
+
+def ingest_inline_bytes(
+    media_id: str, data: bytes, *, kind: str = "video", mime: str = "video/mp4"
+) -> bool:
+    """Cache pre-fetched media bytes (workflow-mode base64 MP4) and mark the
+    Asset as locally available."""
+    if not is_valid_media_id(media_id) or not data:
+        return False
+    ext = _EXT_BY_MIME.get(mime, ".mp4")
+    path = MEDIA_CACHE_DIR / f"{media_id}{ext}"
+    try:
+        path.write_bytes(data)
+    except OSError as exc:
+        logger.error("failed to write inline cache %s: %s", path, exc)
+        return False
+    with _lock:
+        row = _registry.get(media_id) or Asset(media_id=media_id, kind=kind)
+        row.local_path = str(path)
+        row.mime = mime
+        if not row.kind:
+            row.kind = kind
+        _registry[media_id] = row
+    logger.info("media: ingested %d inline bytes for %s", len(data), media_id)
+    return True
+
+
+async def fetch_and_cache(media_id: str) -> Optional[tuple[bytes, str, Path]]:
+    """If the Asset has a URL and no cached file, fetch bytes, cache, return
+    ``(bytes, mime, path)``; else ``None``."""
+    if not is_valid_media_id(media_id):
+        return None
+    row = _get(media_id)
+    if row is None or not row.url:
+        return None
+    url = row.url
+    if not _url_allowed(url):
+        logger.warning("refusing to fetch non-allowed URL: %s", url[:60])
+        return None
+    try:
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            resp = await client.get(url)
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("GCS fetch failed for %s: %s", media_id, exc)
+        return None
+    if resp.status_code != 200:
+        logger.warning("GCS fetch %s returned %d", media_id, resp.status_code)
+        return None
+    mime = resp.headers.get("content-type", "application/octet-stream").split(";")[0].strip()
+    if not (mime.startswith("image/") or mime.startswith("video/")):
+        logger.warning("refusing non-media mime from GCS: %s (media_id=%s)", mime, media_id)
+        return None
+    ext = _EXT_BY_MIME.get(mime, ".bin")
+    path = MEDIA_CACHE_DIR / f"{media_id}{ext}"
+    try:
+        path.write_bytes(resp.content)
+    except OSError as exc:
+        logger.error("failed to write cache %s: %s", path, exc)
+        return None
+    with _lock:
+        r = _registry.get(media_id)
+        if r is not None:
+            r.local_path = str(path)
+            r.mime = mime
+    return resp.content, mime, path
+
+
+async def ensure_cached(media_id: str) -> Optional[Path]:
+    """Return a local file path for the media, fetching from the signed URL
+    on first access. ``None`` if unknown or the fetch fails."""
+    media_id = normalize_media_id(media_id)
+    cached = _cache_glob(media_id)
+    if cached is not None:
+        return cached
+    got = await fetch_and_cache(media_id)
+    return got[2] if got else None
+
+
+def status(media_id: str) -> dict:
+    media_id = normalize_media_id(media_id)
+    if not is_valid_media_id(media_id):
+        return {"available": False, "has_url": False, "reason": "invalid_id"}
+    cached = _cache_glob(media_id)
+    if cached is not None:
+        return {"available": True, "has_url": True, "mime": _mime_from_ext(cached.suffix)}
+    row = _get(media_id)
+    if row is None:
+        return {"available": False, "has_url": False, "reason": "unknown_media"}
+    if row.url:
+        return {"available": False, "has_url": True, "reason": "not_cached_yet"}
+    return {"available": False, "has_url": False, "reason": "no_url_yet"}
+
+
+def media_mime(media_id: str) -> str:
+    media_id = normalize_media_id(media_id)
+    cached = _cache_glob(media_id)
+    if cached is not None:
+        return _mime_from_ext(cached.suffix)
+    row = _get(media_id)
+    return (row.mime if row and row.mime else "application/octet-stream")
+
+
+def _mime_from_ext(ext: str) -> str:
+    for mime, e in _EXT_BY_MIME.items():
+        if e == ext:
+            return mime
+    return "application/octet-stream"