Image Requests That Hang After the Data Has Arrived

import json
import time
import base64
import queue
import threading
import requests

MIN_BYTES = 1024          # check 4: anything smaller cannot be an image


def _pump(raw, q):
    """Background thread: its only job is to push chunks onto the queue."""
    try:
        for chunk in raw.stream(65536, decode_content=True):
            q.put(chunk)
        q.put(None)                       # the server finished normally
    except Exception as exc:              # transport error, re-raised by main
        q.put(exc)


def _try_parse(buf):
    if len(buf) < MIN_BYTES:                    # check 4
        return None
    if not buf.rstrip().endswith(b"}"):         # check 6: cheap pre-filter
        return None
    try:
        return json.loads(buf.decode("utf-8"))  # check 7: the final verdict
    except ValueError:
        return None


def generate_image(url, headers, payload,
                   ttfb_timeout=180, term_grace=5, body_timeout=60):
    """Non-streaming image generation with client-side completion.

    ttfb_timeout: waiting for the first byte. The upstream is generating during
                  this phase — slow is not the same as broken, so leave room.
    term_grace:   how long to wait for the terminating signal once the body is
                  complete. If it never comes, finish the response ourselves.
    body_timeout: how long the whole body may take after the first byte.
    """
    resp = requests.post(url, headers=headers, json=payload,
                         stream=True, timeout=(10, ttfb_timeout))
    resp.raise_for_status()                     # check 2

    if resp.headers.get("Content-Length"):      # check 3
        return resp.json()                      # length declared: let the library finish

    q = queue.Queue()
    threading.Thread(target=_pump, args=(resp.raw, q), daemon=True).start()

    buf = bytearray()
    deadline = time.monotonic() + body_timeout

    while True:
        try:
            item = q.get(timeout=term_grace)
        except queue.Empty:                     # nothing new within the grace period
            obj = _try_parse(buf)
            if obj is not None:
                resp.close()                    # body is complete: finish it ourselves
                return obj
            if time.monotonic() > deadline:
                raise TimeoutError("incomplete body and no new data for a long time")
            continue                            # still incomplete, keep waiting

        if item is None:                        # healthy path: the server finished it
            break
        if isinstance(item, Exception):
            raise item
        buf += item

    obj = _try_parse(buf)
    if obj is None:
        raise ValueError("Incomplete response")
    return obj


def extract_image(obj):
    for cand in obj.get("candidates", []):
        for part in cand.get("content", {}).get("parts", []):
            if "inlineData" in part:
                return base64.b64decode(part["inlineData"]["data"])
    return None

const TERM_GRACE_MS = 5000;       // how long to wait for the terminating signal
const TOTAL_TIMEOUT_MS = 180000;  // total timeout: must cover upstream generation
const MIN_BYTES = 1024;

async function generateImage(url, headers, payload) {
  const resp = await fetch(url, {
    method: "POST",
    headers: { ...headers, "Content-Type": "application/json" },
    body: JSON.stringify(payload),
    signal: AbortSignal.timeout(TOTAL_TIMEOUT_MS),
  });
  if (!resp.ok) throw new Error(`HTTP ${resp.status}`);        // check 2
  if (resp.headers.get("content-length")) return resp.json();  // check 3

  const reader = resp.body.getReader();
  const chunks = [];
  let total = 0;

  while (true) {
    // Promise.race caps how long we wait for the next chunk
    const next = reader.read();
    const timer = new Promise((r) => setTimeout(() => r("GRACE"), TERM_GRACE_MS));
    const winner = await Promise.race([next, timer]);

    if (winner === "GRACE") {
      const parsed = tryParse(chunks, total);
      if (parsed) {                 // body is complete: server never finished
        reader.cancel().catch(() => {});
        return parsed;
      }
      continue;                     // body still incomplete, keep waiting
    }

    const { done, value } = winner;
    if (done) break;                // healthy path: the server finished it
    chunks.push(value);
    total += value.length;
  }

  const parsed = tryParse(chunks, total);
  if (!parsed) throw new Error("Incomplete response");
  return parsed;
}

function tryParse(chunks, total) {
  if (total < MIN_BYTES) return null;                 // check 4
  const buf = new Uint8Array(total);
  let off = 0;
  for (const c of chunks) { buf.set(c, off); off += c.length; }
  const text = new TextDecoder().decode(buf).trimEnd();
  if (!text.endsWith("}")) return null;               // check 6
  try { return JSON.parse(text); } catch { return null; }   // check 7
}