flow-extension/local_server_example.py

#!/usr/bin/env python3
"""Flow Auto-Pilot local queue.

Source of work:  queue.jsonl   — one task per line: {"prompt": "...", "kind": "video"|"image", "model": "Veo 3.1 - Fast"}
                                  (id optional; auto-assigned by line number)
Output:          results.jsonl — one record per finished task: {id, status, assets:[urls], error, ts}

Behaviour:
  * /next-task  -> hands out the next PENDING task, video-first (rinse the pricey credits first).
  * /task-completed -> appends to results.jsonl, marks done.
  * On startup, tasks already in results.jsonl are skipped -> safe to restart during a month-long run.

Harvested `assets` are googleusercontent URLs; they need your logged-in session to download,
so this just records them. ponytail: fetch+save the blobs from the extension later if you want local files.
"""
import base64, json, os, time
from http.server import HTTPServer, BaseHTTPRequestHandler

HERE = os.path.dirname(os.path.abspath(__file__))
QUEUE_FILE = os.path.join(HERE, "queue.jsonl")
RESULTS_FILE = os.path.join(HERE, "results.jsonl")
DOWNLOADS = os.path.join(HERE, "downloads")

EXT = {"video/mp4": "mp4", "image/png": "png", "image/jpeg": "jpg", "image/webp": "webp", "image/gif": "gif"}


def save_asset(rec):
    """Write downloaded bytes to downloads/<category>/<id>_<idx>.<ext> (+ dest copy)."""
    data = base64.b64decode(rec.get("dataBase64", ""))
    ct = (rec.get("contentType") or "").split(";")[0].strip()
    ext = EXT.get(ct, "mp4" if "video" in ct else "png")
    name = f'{rec.get("id", "asset")}_{rec.get("idx", 0)}.{ext}'
    outdir = os.path.join(DOWNLOADS, rec.get("category") or "misc")
    os.makedirs(outdir, exist_ok=True)
    path = os.path.join(outdir, name)
    with open(path, "wb") as f:
        f.write(data)
    dest = rec.get("dest")            # game asked for a copy in its own folder
    if dest:
        os.makedirs(dest, exist_ok=True)
        with open(os.path.join(dest, name), "wb") as f:
            f.write(data)
    return path


def enqueue_task(rec):
    """Append a game-requested task to the queue (persisted, picked up on next poll)."""
    t = {
        "id": rec.get("id") or f'req_{int(time.time() * 1000)}',
        "kind": rec.get("kind", "image"),
        "model": rec.get("model") or ("Veo 3.1 - Fast" if rec.get("kind") == "video" else "Nano Banana 2"),
        "prompt": rec["prompt"],
        "category": rec.get("category", "request"),
    }
    if rec.get("dest"):
        t["dest"] = rec["dest"]
    with open(QUEUE_FILE, "a") as f:
        f.write(json.dumps(t) + "\n")
    return t

# Fallback demo queue if queue.jsonl is absent (so it runs out of the box).
DEMO = [
    {"prompt": "Cinematic synthwave keyboard warrior fighting digital demons, 8k", "kind": "video", "model": "Veo 3.1 - Fast"},
    {"prompt": "Op-shop record crate, warm 90s film grain, product hero shot", "kind": "image", "model": "Nano Banana 2"},
]

def load_queue():
    rows = []
    if os.path.exists(QUEUE_FILE):
        with open(QUEUE_FILE) as f:
            for line in f:
                line = line.strip()
                if line:
                    rows.append(json.loads(line))
    else:
        rows = list(DEMO)
    # assign stable ids by position; video first so the expensive credits burn first
    for i, r in enumerate(rows):
        r.setdefault("id", f"task_{i:05d}")
        r.setdefault("kind", "image")
    rows.sort(key=lambda r: 0 if r.get("kind") == "video" else 1)
    return rows

def done_ids():
    """Ids that must NOT be handed out again.

    'failed' IS counted: a timeout may have burned credits invisibly, so nothing is retried
    automatically. To re-run a failed task, delete its line from results.jsonl (requeue.py).

    'ratelimited'/'requeue' are skipped belt-and-braces: /task-completed already refuses to
    persist those rows (they stay PENDING by design), so they can't reach this file. This is
    defence in depth against someone appending a row by hand — not a bug fix.
    """
    ids = set()
    if os.path.exists(RESULTS_FILE):
        with open(RESULTS_FILE) as f:
            for line in f:
                line = line.strip()
                if line:
                    row = json.loads(line)
                    if row.get("status") != "ratelimited":
                        ids.add(row.get("id"))
    return ids

class Handler(BaseHTTPRequestHandler):
    def _send(self, code=200, body=None):
        self.send_response(code)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        if body is not None:
            self.wfile.write(json.dumps(body).encode())

    def do_OPTIONS(self):
        self._send(200, {})

    def do_GET(self):
        if self.path != "/next-task":
            return self._send(404, {"error": "not found"})
        done = done_ids()
        for t in load_queue():
            # in-flight entries expire after 15 min: a hard-refreshed tab that never
            # reported back must get its task re-served, not stranded
            if t["id"] not in done and (t["id"] not in INFLIGHT or time.time() - INFLIGHT[t["id"]] > 900):
                INFLIGHT[t["id"]] = time.time()
                remaining = sum(1 for x in load_queue() if x["id"] not in done)
                print(f"[SERVER] -> {t['id']} ({t.get('kind')}) | {remaining} left")
                return self._send(200, t)
        return self._send(204)  # queue drained

    def do_POST(self):
        n = int(self.headers.get("Content-Length", 0))
        rec = json.loads(self.rfile.read(n).decode() or "{}")

        if self.path == "/task-completed":
            INFLIGHT.pop(rec.get("id"), None)
            if rec.get("status") not in ("ratelimited", "requeue"):   # both stay PENDING, retried after cooldown/reload
                rec["ts"] = time.time()
                with open(RESULTS_FILE, "a") as f:
                    f.write(json.dumps(rec) + "\n")
            print(f"[SERVER] <- {rec.get('id')} {str(rec.get('status')).upper()} | {len(rec.get('assets') or [])} asset(s)")
            return self._send(200, {"ok": True})

        if self.path == "/save":
            path = save_asset(rec)
            print(f"[SERVER] saved {rec.get('id')} -> {path}")
            return self._send(200, {"ok": True, "path": path})

        if self.path == "/enqueue":
            t = enqueue_task(rec)
            print(f"[SERVER] enqueued {t['id']} ({t['kind']})")
            return self._send(200, {"ok": True, "id": t["id"]})

        return self._send(404, {"error": "not found"})

    def log_message(self, *a):  # quiet the default noise
        pass

INFLIGHT = {}

def selftest():
    # video-first ordering + resume-skip
    global load_queue
    _orig = load_queue
    load_queue = lambda: [  # noqa: E731
        {"id": "a", "kind": "image"}, {"id": "b", "kind": "video"}]
    assert [t["id"] for t in load_queue()] == ["a", "b"]  # loader keeps as-is here
    q = sorted(load_queue(), key=lambda r: 0 if r["kind"] == "video" else 1)
    assert q[0]["id"] == "b", "video must be handed out first"
    load_queue = _orig
    print("selftest ok")

def run(port=8017):
    # ponytail: loopback ONLY. /save base64-writes to an attacker-supplied absolute `dest`
    # with no auth — bound to 0.0.0.0 that is arbitrary file write as this user (drop a
    # ~/Library/LaunchAgents plist → code execution) for anyone on the same wifi.
    # The extension only ever calls localhost. FLOW_BIND=<tailscale-ip> if a remote box
    # must enqueue — a tailnet IP is routed over utun, so the LAN still can't reach it.
    # Never set it back to "" / 0.0.0.0.
    host = os.environ.get("FLOW_BIND", "127.0.0.1")
    print(f"--- Flow queue on http://{host}:{port} ---")
    print(f"queue:   {QUEUE_FILE if os.path.exists(QUEUE_FILE) else '(demo — no queue.jsonl)'}")
    print(f"pending: {sum(1 for t in load_queue() if t['id'] not in done_ids())}")
    HTTPServer((host, port), Handler).serve_forever()

if __name__ == "__main__":
    import sys
    if "--selftest" in sys.argv:
        selftest()
    else:
        try:
            run()
        except KeyboardInterrupt:
            print("\nstopped")