出圖請求傳完卻不返回的相容處理

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

MIN_BYTES = 1024          # 判定 4：比這還小不可能是一張圖


def _pump(raw, q):
    """後臺執行緒：只負責把收到的塊塞進佇列。"""
    try:
        for chunk in raw.stream(65536, decode_content=True):
            q.put(chunk)
        q.put(None)                       # 服務端正常收尾了
    except Exception as exc:              # 傳輸異常，交給主執行緒拋
        q.put(exc)


def _try_parse(buf):
    if len(buf) < MIN_BYTES:                    # 判定 4
        return None
    if not buf.rstrip().endswith(b"}"):         # 判定 6：極便宜的預篩
        return None
    try:
        return json.loads(buf.decode("utf-8"))  # 判定 7：最終裁決
    except ValueError:
        return None


def generate_image(url, headers, payload,
                   ttfb_timeout=180, term_grace=5, body_timeout=60):
    """非流式出圖，帶主動收尾。

    ttfb_timeout: 等首位元組。這一段是上游在生成，慢不等於故障，要留夠。
    term_grace:   資料到齊後，再等多久結束訊號。等不到就主動收尾。
    body_timeout: 首位元組之後，整個響應體最多再花多久。
    """
    resp = requests.post(url, headers=headers, json=payload,
                         stream=True, timeout=(10, ttfb_timeout))
    resp.raise_for_status()                     # 判定 2

    if resp.headers.get("Content-Length"):      # 判定 3
        return resp.json()                      # 長度已宣告，交給庫自己結束

    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:                     # 寬限期內沒有任何新資料
            obj = _try_parse(buf)
            if obj is not None:
                resp.close()                    # 資料已到齊，主動收尾
                return obj
            if time.monotonic() > deadline:
                raise TimeoutError("響應不完整，且長時間沒有新資料")
            continue                            # 還不完整，繼續等

        if item is None:                        # 正常路徑：服務端自己收了尾
            break
        if isinstance(item, Exception):
            raise item
        buf += item

    obj = _try_parse(buf)
    if obj is None:
        raise ValueError("響應不完整")
    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;       // 資料到齊後再等多久結束訊號
const TOTAL_TIMEOUT_MS = 180000;  // 總超時：必須留夠上游生成的時間
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}`);        // 判定 2
  if (resp.headers.get("content-length")) return resp.json();  // 判定 3

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

  while (true) {
    // 用 Promise.race 給「等下一塊資料」加一個寬限期上限
    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) {                 // 資料已完整，是服務端沒收尾
        reader.cancel().catch(() => {});
        return parsed;
      }
      continue;                     // 資料還不完整，繼續等
    }

    const { done, value } = winner;
    if (done) break;                // 正常路徑：服務端收了尾
    chunks.push(value);
    total += value.length;
  }

  const parsed = tryParse(chunks, total);
  if (!parsed) throw new Error("響應不完整");
  return parsed;
}

function tryParse(chunks, total) {
  if (total < MIN_BYTES) return null;                 // 判定 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;               // 判定 6
  try { return JSON.parse(text); } catch { return null; }   // 判定 7
}