Blame for src/lib/engine.ts

/**
 * Runs a workflow: walks the execution order, resolves each node's inputs
 * from upstream results or its own widgets, calls the Space, and stores the
 * normalised outputs on the node.
 */
import { executionOrder } from "./groups";
import { callEndpoint, fileFromUrl, GradioError, resolveFileUrl, uploadFile, type GradioFileData } from "./gradio";
import { getLocalFile, isLocalRef } from "./files";
import { canConnect, getDef, REROUTE_TYPE } from "./registry";
import type { CanvasEdge, CanvasNode, CanvasNodeData, FileRef, NodeDef, OutputValue, Port, SpaceRef } from "./types";

export interface RunOptions {
  nodes: CanvasNode[];
  edges: CanvasEdge[];
  order: string[];
  /** When set, only these nodes and what feeds them run. */
  targets?: string[];
  token: string | null;
  signal: AbortSignal;
  patch: (id: string, data: Partial<CanvasNodeData>) => void;
  flow: (fromId: string | null) => void;
  step: (completed: number, total: number, currentId: string | null) => void;
  /** Called before each Space call: resolves once the Space runs, throws when it cannot. */
  ensureAwake?: (spaceId: string, report: (text: string) => void) => Promise<void>;
}

export interface RunResult {
  ok: boolean;
  ran: number;
  total: number;
  error?: string;
  failedNode?: string;
  thumb?: string;
  cancelled?: boolean;
}

export const hashString = (s: string) => {
  let h = 5381;
  for (let i = 0; i < s.length; i++) h = ((h << 5) + h + s.charCodeAt(i)) | 0;
  return (h >>> 0).toString(36);
};

const isFileRef = (v: unknown): v is FileRef => !!v && typeof v === "object" && typeof (v as FileRef).url === "string" && !("path" in (v as object));
const isFileData = (v: unknown): v is GradioFileData =>
  !!v && typeof v === "object" && (typeof (v as GradioFileData).path === "string" || typeof (v as GradioFileData).url === "string") && ("meta" in (v as object) || "orig_name" in (v as object) || "path" in (v as object));

/* Outputs → FileRef / primitives ---------------------------------------- */

const toRef = (fd: GradioFileData, ref: SpaceRef): FileRef => ({
  url: resolveFileUrl(ref.host, ref.apiPrefix, fd),
  name: fd.orig_name ?? fd.path?.split("/").pop(),
  mime: fd.mime_type,
});

/** Gradio wraps outputs of `gr.update(...)` as `{ __type__: "update", value, visible, ... }`. */
const isUpdate = (v: unknown): v is { __type__: string; value?: unknown; visible?: boolean } =>
  !!v && typeof v === "object" && !Array.isArray(v) && (v as { __type__?: unknown }).__type__ === "update";

export function normaliseOutput(raw: unknown, port: Port, ref: SpaceRef): OutputValue {
  if (raw == null) return null;
  if (isUpdate(raw)) {
    if (raw.visible === false || !("value" in raw)) return null;
    return normaliseOutput(raw.value, port, ref);
  }
  if (isFileData(raw)) return toRef(raw, ref);
  if (typeof raw === "object" && !Array.isArray(raw)) {
    const o = raw as Record<string, unknown>;
    if (isFileData(o.video)) return toRef(o.video, ref);
    if (isFileData(o.composite)) return toRef(o.composite, ref);
    if (isFileData(o.background)) return toRef(o.background, ref);
    if (isFileData(o.image)) return toRef(o.image, ref);
    return o;
  }
  if (Array.isArray(raw)) {
    const refs = raw.map((x) => {
      if (isFileData(x)) return toRef(x, ref);
      if (x && typeof x === "object" && isFileData((x as { image?: unknown }).image)) return toRef((x as { image: GradioFileData }).image, ref);
      if (Array.isArray(x) && isFileData(x[0])) return toRef(x[0], ref);
      return null;
    });
    if (refs.length && refs.every(Boolean)) return refs as FileRef[];
    return raw;
  }
  if (port.type === "text" && typeof raw !== "string") return typeof raw === "number" || typeof raw === "boolean" ? raw : JSON.stringify(raw);
  return raw as OutputValue;
}

/* Values → Gradio payload ----------------------------------------------- */

type Uploader = (file: File) => Promise<GradioFileData>;

async function toFileData(v: unknown, upload: Uploader): Promise<GradioFileData | null> {
  if (v == null || v === "") return null;
  if (isFileRef(v)) return toFileData(v.url, upload);
  if (Array.isArray(v)) return toFileData(v[v.length - 1], upload);
  if (typeof v !== "string") return null;
  if (isLocalRef(v)) {
    const f = getLocalFile(v);
    if (!f) throw new Error("The picked file is no longer available. Choose it again.");
    return upload(f);
  }
  if (v.startsWith("data:")) {
    const blob = await (await fetch(v)).blob();
    return upload(new File([blob], `upload.${blob.type.split("/")[1] ?? "bin"}`, { type: blob.type }));
  }
  if (/^https?:\/\//.test(v)) return fileFromUrl(v);
  throw new Error(`"${v.slice(0, 40)}" is not a URL or a file.`);
}

async function toFileList(v: unknown, upload: Uploader): Promise<GradioFileData[]> {
  if (v == null || v === "") return [];
  if (Array.isArray(v)) return (await Promise.all(v.map((x) => toFileData(x, upload)))).filter((x): x is GradioFileData => !!x);
  const one = await toFileData(v, upload);
  return one ? [one] : [];
}

const parseJson = (v: unknown) => {
  if (typeof v !== "string") return v;
  try {
    return JSON.parse(v);
  } catch {
    return v;
  }
};

export async function toPayload(v: unknown, port: Port, upload: Uploader): Promise<unknown> {
  const c = port.component ?? "";
  switch (c) {
    case "video": {
      const fd = await toFileData(v, upload);
      return fd ? { video: fd, subtitles: null } : null;
    }
    case "imageeditor":
    case "sketchpad":
    case "paint":
    case "imagemask": {
      const fd = await toFileData(v, upload);
      return fd ? { background: fd, layers: [], composite: null, id: null } : null;
    }
    case "gallery":
      return (await toFileList(v, upload)).map((image) => ({ image, caption: null }));
    case "imageslider": {
      const list = await toFileList(v, upload);
      return list.length ? [list[0], list[1] ?? null] : null;
    }
    case "checkboxgroup":
    case "dataframe":
    case "json":
      return parseJson(v);
    case "dropdown":
    case "radio": {
      const w = port.widget;
      if (w?.kind === "textarea") return parseJson(v);
      // Map the displayed option back to the Space's typed choice (e.g. 16 not "16").
      if (w?.kind === "select" && w.raw) {
        const i = w.options.indexOf(String(v));
        if (i >= 0) return w.raw[i];
      }
      break;
    }
  }
  switch (port.type) {
    case "image":
    case "audio":
    case "file":
    case "video":
      return toFileData(v, upload);
    case "gallery":
      return toFileList(v, upload);
    case "number": {
      const n = typeof v === "number" ? v : Number(v);
      return Number.isFinite(n) ? n : null;
    }
    case "boolean":
      return typeof v === "boolean" ? v : v === "true";
    case "text":
      if (isFileRef(v)) return v.url;
      return v == null ? "" : typeof v === "string" ? v : typeof v === "object" ? JSON.stringify(v) : String(v);
    case "json":
      return parseJson(v);
    default:
      return v ?? null;
  }
}

/* The run --------------------------------------------------------------- */

const upstreamClosure = (edges: CanvasEdge[], targets: string[]) => {
  const keep = new Set<string>();
  const stack = [...targets];
  while (stack.length) {
    const id = stack.pop()!;
    if (keep.has(id)) continue;
    keep.add(id);
    for (const e of edges) if (e.target === id) stack.push(e.source);
  }
  return keep;
};

const describeProgress = (p: import("./gradio").CallProgress): { statusText?: string; progress?: number; progressEstimated?: boolean } => {
  switch (p.kind) {
    case "queued":
      return { statusText: p.rank != null && p.queueSize ? `Queue ${p.rank + 1} of ${p.queueSize}${p.eta ? ` · ~${Math.round(p.eta)}s` : ""}` : "Queued…", progress: 0 };
    case "starting":
      return { statusText: p.eta ? `Running · ~${Math.round(p.eta)}s` : "Running…" };
    case "progress":
      return { statusText: p.desc || "Running…", progress: p.progress ?? undefined };
    case "log":
      return { statusText: p.log };
    case "generating":
      return { statusText: "Streaming…" };
  }
};

const inputValue = (node: CanvasNode, def: NodeDef, port: Port, edge: CanvasEdge | undefined, results: Map<string, Record<string, OutputValue>>): { value: unknown; missing: boolean } => {
  if (edge) {
    const out = results.get(edge.source);
    if (!out) return { value: null, missing: true };
    return { value: out[edge.sourceHandle ?? ""] ?? null, missing: false };
  }
  void def;
  const v = node.data.values[port.id];
  if (port.widget?.kind === "file" && (v === "" || v == null)) return { value: null, missing: !port.optional };
  return { value: v ?? null, missing: false };
};

export interface RunIssue {
  nodeId: string;
  title: string;
  message: string;
}

/** Problems that would make a run fail before it starts: required inputs with no value and no link. */
export function findRunIssues(nodes: CanvasNode[], edges: CanvasEdge[], order: string[], targets?: string[]): RunIssue[] {
  let ids = executionOrder(nodes, edges, order);
  if (targets?.length) {
    const keep = upstreamClosure(edges, targets);
    ids = ids.filter((id) => keep.has(id));
  }
  const byId = new Map(nodes.map((n) => [n.id, n]));
  const issues: RunIssue[] = [];
  for (const id of ids) {
    const node = byId.get(id);
    if (!node || node.data.mode !== "always") continue;
    const def = getDef(node.data.defType);
    const title = node.data.title ?? def.title;
    if (def.emits) {
      const v = node.data.values[def.emits];
      if ((v == null || v === "") && def.outputs[0]?.type !== "text") issues.push({ nodeId: id, title, message: "Pick a file or paste a URL" });
      continue;
    }
    for (const port of def.inputs) {
      if (port.optional) continue;
      const edge = edges.find((e) => e.target === id && e.targetHandle === port.id);
      if (edge) continue;
      const v = node.data.values[port.id];
      const empty = port.widget ? inputValue(node, def, port, undefined, new Map()).missing : v == null || v === "";
      if (empty) issues.push({ nodeId: id, title, message: `“${port.label || port.id}” needs a value or a link from another node` });
    }
  }
  return issues;
}

export async function runGraph(o: RunOptions): Promise<RunResult> {
  const { nodes, edges, token, signal, patch, flow, step } = o;
  let order = executionOrder(nodes, edges, o.order);
  if (o.targets?.length) {
    const keep = upstreamClosure(edges, o.targets);
    order = order.filter((id) => keep.has(id));
  }
  const byId = new Map(nodes.map((n) => [n.id, n]));
  const results = new Map<string, Record<string, OutputValue>>();
  // Nodes outside this run still provide their last outputs.
  for (const n of nodes) if (!order.includes(n.id) && n.data.outputs && n.data.status !== "error") results.set(n.id, n.data.outputs);
  const failed = new Set<string>();
  const total = order.length;
  for (const id of order) patch(id, { status: "queued", progress: 0, statusText: undefined, error: undefined });
  step(0, total, null);

  let thumb: string | undefined;
  let firstError: { id: string; message: string } | undefined;
  let completed = 0;

  for (const id of order) {
    if (signal.aborted) return { ok: false, ran: completed, total, cancelled: true, thumb };
    const node = byId.get(id);
    if (!node) continue;
    const def = getDef(node.data.defType);
    step(completed, total, id);
    const incoming = edges.filter((e) => e.target === id);
    const inputs: Record<string, unknown> = {};
    let skipped: string | null = null;
    for (const port of def.inputs) {
      const edge = incoming.find((e) => e.targetHandle === port.id);
      if (edge && failed.has(edge.source)) {
        skipped = `Skipped: ${byId.get(edge.source)?.data.title ?? getDef(byId.get(edge.source)?.data.defType ?? "").title} failed`;
        break;
      }
      const { value, missing } = inputValue(node, def, port, edge, results);
      if (missing && !edge && !port.optional) {
        skipped = null;
        patch(id, { status: "error", error: `“${port.label || port.id}” needs a value or a connection.`, statusText: undefined });
        failed.add(id);
        break;
      }
      if (missing && edge) {
        skipped = "Skipped: input not available";
        break;
      }
      inputs[port.id] = value;
    }
    if (failed.has(id)) {
      firstError ??= { id, message: `${node.data.title ?? def.title}: missing input` };
      completed++;
      continue;
    }
    if (skipped) {
      patch(id, { status: "idle", statusText: skipped, progress: 0 });
      failed.add(id);
      completed++;
      continue;
    }

    // Input nodes just emit their widget value.
    if (def.emits) {
      const v = node.data.values[def.emits];
      const out: Record<string, OutputValue> = { value: typeof v === "string" && v === "" && def.outputs[0]?.type !== "text" ? null : (v as OutputValue) };
      if (out.value == null && def.outputs[0]?.type !== "text") {
        patch(id, { status: "error", error: "Pick a file or paste a URL." });
        failed.add(id);
        firstError ??= { id, message: `${node.data.title ?? def.title}: no file` };
        completed++;
        continue;
      }
      results.set(id, out);
      patch(id, { status: "done", progress: 1, outputs: out, statusText: undefined });
      completed++;
      step(completed, total, id);
      continue;
    }

    if (def.type === REROUTE_TYPE) {
      const out: Record<string, OutputValue> = { out: (inputs.in as OutputValue) ?? null };
      results.set(id, out);
      patch(id, { status: "done", progress: 1, outputs: out, statusText: undefined });
      completed++;
      continue;
    }

    if (node.data.mode === "bypass") {
      const out: Record<string, OutputValue> = {};
      for (const op of def.outputs) {
        const ip = def.inputs.find((p) => canConnect(p.type, op.type) && inputs[p.id] != null);
        out[op.id] = ip ? (inputs[ip.id] as OutputValue) : null;
      }
      results.set(id, out);
      patch(id, { status: "done", progress: 1, outputs: out, statusText: "Bypassed" });
      completed++;
      continue;
    }

    const ref = def.space;
    if (!ref) {
      patch(id, { status: "error", error: "This node's Space is not loaded." });
      failed.add(id);
      firstError ??= { id, message: `${def.title}: Space not loaded` };
      completed++;
      continue;
    }

    const cacheKey = hashString(JSON.stringify([def.type, inputs]));
    if (node.data.cacheKey === cacheKey && node.data.outputs && node.data.status !== "error") {
      results.set(id, node.data.outputs);
      patch(id, { status: "done", progress: 1, statusText: "Cached", error: undefined });
      completed++;
      step(completed, total, id);
      continue;
    }

    patch(id, { status: "running", progress: 0, progressEstimated: undefined, statusText: "Uploading…" });
    flow(null);
    // Estimated progress, like Gradio's own UI: the queue's `process_starts` message carries the
    // endpoint's average process time, and we run a clock against it until real progress arrives.
    let eta: number | null = null;
    let startedAt = 0;
    let realProgress = false;
    let generic = true; // statusText is ours (safe to overwrite with "~Ns left")
    let ticker: ReturnType<typeof setInterval> | undefined;
    const tick = () => {
      if (realProgress || !startedAt || !eta) return;
      const elapsed = Date.now() - startedAt;
      const left = Math.max(0, eta - elapsed);
      patch(id, {
        // Linear to the estimate, then creep toward 97% so a slow run never looks finished.
        progress: elapsed <= eta ? 0.85 * (elapsed / eta) : 0.85 + 0.12 * (1 - Math.exp(1 - elapsed / eta)),
        progressEstimated: true,
        ...(generic ? { statusText: left > 0 ? `Running · ~${Math.ceil(left / 1000)}s left` : "Running · almost done…" } : {}),
      });
    };
    try {
      if (o.ensureAwake) await o.ensureAwake(ref.id, (text) => patch(id, { statusText: text }));
      if (signal.aborted) throw new DOMException("Aborted", "AbortError");
      const upload = (file: File) => uploadFile(ref.host, ref.apiPrefix, file, token, signal);
      const data = await Promise.all(def.inputs.map((p) => toPayload(inputs[p.id], p, upload)));
      patch(id, { statusText: "Sending…" });
      ticker = setInterval(tick, 250);
      const raw = await callEndpoint(ref, data, {
        token,
        signal,
        onProgress: (p) => {
          const d = describeProgress(p);
          if (p.kind === "queued") {
            startedAt = 0; // don't count queue wait toward the run
            generic = true;
          } else if (p.kind === "starting") {
            startedAt = Date.now();
            eta = p.eta && p.eta > 0 ? p.eta * 1000 : null;
            generic = true;
            if (eta) return tick();
          } else if (p.kind === "progress" && p.progress != null) {
            realProgress = true;
            d.progressEstimated = false;
            generic = false;
          } else if (p.kind === "log" || (p.kind === "progress" && p.desc)) {
            generic = false;
          }
          patch(id, d);
          if (p.kind === "generating" && p.data) {
            const partial: Record<string, OutputValue> = {};
            def.outputs.forEach((op, i) => (partial[op.id] = normaliseOutput(p.data![i], op, ref)));
            patch(id, { outputs: partial });
          }
        },
      });
      const out: Record<string, OutputValue> = {};
      def.outputs.forEach((op, i) => (out[op.id] = normaliseOutput(raw[i], op, ref)));
      results.set(id, out);
      for (const op of def.outputs) {
        const v = out[op.id];
        if (op.type === "image" && v) thumb = Array.isArray(v) ? (v[v.length - 1] as FileRef)?.url : isFileRef(v) ? v.url : thumb;
        else if (op.type === "gallery" && Array.isArray(v) && v.length) thumb = (v[0] as FileRef)?.url ?? thumb;
      }
      patch(id, { status: "done", progress: 1, progressEstimated: undefined, outputs: out, cacheKey, statusText: undefined, error: undefined });
      flow(id);
    } catch (e) {
      if (signal.aborted) {
        patch(id, { status: "idle", progress: 0, statusText: undefined });
        return { ok: false, ran: completed, total, cancelled: true, thumb };
      }
      const message = e instanceof GradioError ? e.message : e instanceof TypeError ? "Could not reach the Space (network or CORS)." : e instanceof Error ? e.message : String(e);
      patch(id, { status: "error", error: message, statusText: undefined, progress: 0, progressEstimated: undefined });
      failed.add(id);
      firstError ??= { id, message: `${node.data.title ?? def.title}: ${message}` };
    } finally {
      clearInterval(ticker);
    }
    completed++;
    step(completed, total, id);
  }
  flow(null);
  return { ok: !firstError, ran: completed, total, error: firstError?.message, failedNode: firstError?.id, thumb };
}