coglet 0.23.1: scalar Path outputs intermittently returned as a one-element array (["https://…"]) #3205

import tempfile

from cog import BasePredictor, Path


class Predictor(BasePredictor):
    def predict(self, size_kb: int = 512) -> Path:
        print("Generating output...")
        out = Path(tempfile.mkdtemp()) / "output.wav"
        out.write_bytes(b"\0" * (size_kb * 1024))
        print("Generated output")
        return out

"""Minimal repro: coglet webhooks for a scalar `-> Path` output.

Runs a real coglet server against a local stub that
* accepts output uploads with `202 Accepted` + `Location` (as in coglet's own tests)
* records every webhook in arrival order.

Usage: public_repro.py PYTHON_FROM_VENV_WITH_COG [--runs N]
"""

import argparse
import json
import os
import subprocess
import sys
import threading
import time
import urllib.request
import uuid
from collections import Counter, defaultdict
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

HERE = os.path.dirname(os.path.abspath(__file__))
STUB_PORT, COG_PORT = 4998, 5056
TERMINAL = {"succeeded", "failed", "canceled"}
lock = threading.Lock()
webhooks = defaultdict(list)
prediction_id_headers = Counter()


def shape(v):
    return "null" if v is None else f"list[{len(v)}]" if isinstance(v, list) else type(v).__name__


class Stub(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, *a):
        pass

    def _read(self):
        return self.rfile.read(int(self.headers.get("Content-Length", 0)))

    def _reply(self, code, headers=()):
        self.send_response(code)
        for k, v in headers:
            self.send_header(k, v)
        self.send_header("Content-Length", "0")
        self.end_headers()

    def do_PUT(self):  # output upload
        self._read()
        with lock:
            prediction_id_headers[bool(self.headers.get("X-Prediction-Id"))] += 1
        name = self.path.rsplit("/", 1)[-1]
        self._reply(202, [("Location", f"https://example.com/outputs/{uuid.uuid4().hex}/{name}")])

    def do_POST(self):  # webhook
        p = json.loads(self._read())
        self._reply(200)
        with lock:
            webhooks[p["id"]].append((p["status"], shape(p.get("output"))))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("python")
    ap.add_argument("--runs", type=int, default=50)
    args = ap.parse_args()

    srv = ThreadingHTTPServer(("127.0.0.1", STUB_PORT), Stub)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    venv_bin = os.path.dirname(os.path.abspath(args.python))
    env = dict(os.environ, PORT=str(COG_PORT), COG_PREDICT_TYPE_STUB="predict.py:Predictor",
               PATH=f"{venv_bin}{os.pathsep}{os.environ['PATH']}")
    proc = subprocess.Popen([args.python, "-m", "cog.server.http", "--upload-url",
                             f"http://127.0.0.1:{STUB_PORT}/upload/"],
                            cwd=HERE, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        for _ in range(300):
            try:
                with urllib.request.urlopen(f"http://127.0.0.1:{COG_PORT}/health-check") as r:
                    if json.load(r)["status"] == "READY":
                        break
            except Exception:
                pass
            time.sleep(0.2)
        else:
            sys.exit("coglet did not become READY")

        ver = subprocess.run([args.python, "-c", "import coglet;print(coglet.__version__)"],
                             capture_output=True, text=True).stdout.strip()
        print(f"coglet {ver}, Output schema: {{type: string, format: uri}}, {args.runs} runs")
        stats = Counter()
        for i in range(args.runs):
            pid = f"p{i:03d}{uuid.uuid4().hex[:8]}"
            req = urllib.request.Request(
                f"http://127.0.0.1:{COG_PORT}/predictions/{pid}", method="PUT",
                data=json.dumps({"id": pid, "input": {},
                                 "webhook": f"http://127.0.0.1:{STUB_PORT}/webhook"}).encode(),
                headers={"Content-Type": "application/json", "Prefer": "respond-async"})
            urllib.request.urlopen(req).read()
            deadline = time.time() + 30
            while time.time() < deadline:
                with lock:
                    if any(s in TERMINAL for s, _ in webhooks[pid]):
                        break
                time.sleep(0.02)
            time.sleep(0.3)  # let any straggler webhooks land
            with lock:
                seq = list(webhooks[pid])
            t = next(n for n, (s, _) in enumerate(seq) if s in TERMINAL)
            array_nonterminal = any(s not in TERMINAL and o.startswith("list") for s, o in seq)
            after_terminal = len(seq) > t + 1
            stats["non-terminal webhook had output=list[1]"] += array_nonterminal
            stats["a non-terminal webhook arrived AFTER the terminal one"] += after_terminal
            stats[f"last webhook received: {seq[-1]}"] += 1
            if i < 2 or (after_terminal and stats["shown"] < 2):
                stats["shown"] += after_terminal
                print(f"  {pid} webhooks in arrival order: {seq}")
        stats.pop("shown", None)
        print(f"  uploads with X-Prediction-Id header: {prediction_id_headers[True]}/{sum(prediction_id_headers.values())}")
        for k, v in stats.items():
            print(f"  {k}: {v}/{args.runs}")
    finally:
        proc.terminate()
        proc.wait(10)
        srv.shutdown()


if __name__ == "__main__":
    main()