shared/gradio/gradio_queue_focus_patch.py

"""Runtime monkey patch that injects a Gradio background-scheduler fix into templates."""

from __future__ import annotations

import sys
from pathlib import Path
from typing import Callable

ENABLE_GRADIO_FOCUS_QUEUE_MONKEYPATCH = True
GRADIO_FOCUS_QUEUE_MONKEYPATCH_VERBOSE = False

_PATCH_SENTINEL = "window.__gradioFocusQueuePatch"
_TARGET_TEMPLATES = {"frontend/index.html", "frontend/share.html"}
_PROJECT_ROOT = Path(__file__).resolve().parents[2]


def get_javascript() -> str:
    if not ENABLE_GRADIO_FOCUS_QUEUE_MONKEYPATCH:
        return ""
    verbose = "true" if GRADIO_FOCUS_QUEUE_MONKEYPATCH_VERBOSE else "false"
    return f"""
(function () {{
  if (typeof window === "undefined" || window.__gradioFocusQueuePatch) {{
    return;
  }}

  const nativeRequestAnimationFrame = window.requestAnimationFrame.bind(window);
  const nativeCancelAnimationFrame = window.cancelAnimationFrame.bind(window);
  const channel = typeof MessageChannel === "function" ? new MessageChannel() : null;

  function isElementVisible(element) {{
    if (!element) {{
      return false;
    }}
    const rect = element.getBoundingClientRect();
    const style = window.getComputedStyle(element);
    return style.display !== "none" && style.visibility !== "hidden" && rect.width > 0 && rect.height > 0;
  }}

  function getVideoGenTab() {{
    return Array.from(document.querySelectorAll('[role="tab"]')).find((element) => element.textContent.trim() === "Video Generator") || null;
  }}

  function getVideoGenPanel() {{
    const tab = getVideoGenTab();
    const panelId = tab?.getAttribute("aria-controls");
    const panel = panelId ? document.getElementById(panelId) : null;
    return isElementVisible(panel) ? panel : null;
  }}

  function isVideoGenActive() {{
    return getVideoGenPanel() !== null;
  }}

  const patch = {{
    enabled: true,
    forceBackground: false,
    verbose: {verbose},
    nextId: 1,
    pending: new Map(),
    queue: [],
    nativeRequestAnimationFrame,
    nativeCancelAnimationFrame,
    isBackground() {{
      if (this.forceBackground) {{
        return true;
      }}
      try {{
        return document.visibilityState !== "visible" || !document.hasFocus();
      }} catch (_error) {{
        return false;
      }}
    }},
    shouldPatch() {{
      return this.enabled && this.isBackground() && isVideoGenActive();
    }},
    shouldPatchAnimationFrame(stack, callback) {{
      if (!this.shouldPatch()) {{
        return false;
      }}
      const stackText = String(stack || "");
      const source = String(callback || "");
      const isBlocksDispatch = stackText.includes("/assets/Blocks-") && source.includes("Jt(");
      const isCoreFlush = stackText.includes("/assets/index-") && source.includes("l.update(") && source.includes("ge.length") && source.includes("f.props[v.prop]=j");
      return isBlocksDispatch || isCoreFlush;
    }}
  }};

  function cancelSynthetic(id) {{
    const job = patch.pending.get(id);
    if (!job) {{
      return false;
    }}
    job.canceled = true;
    patch.pending.delete(id);
    return true;
  }}

  function dispatchSynthetic(job) {{
    if (!job || job.canceled) {{
      return;
    }}
    patch.pending.delete(job.id);
    try {{
      if (job.kind === "raf") {{
        job.callback(window.performance.now());
      }}
    }} catch (error) {{
      window.setTimeout(() => {{
        throw error;
      }}, 0);
    }}
  }}

  function flushOne() {{
    const job = patch.queue.shift();
    dispatchSynthetic(job);
  }}

  function scheduleSyntheticAnimationFrame(callback) {{
    if (!channel) {{
      return nativeRequestAnimationFrame(callback);
    }}
    const id = -patch.nextId++;
    const job = {{ id, kind: "raf", callback, args: null, canceled: false }};
    patch.pending.set(id, job);
    patch.queue.push(job);
    channel.port2.postMessage(id);
    if (patch.verbose) {{
      console.debug("[Gradio] focus queue synthetic animation frame", id);
    }}
    return id;
  }}

  if (channel) {{
    channel.port1.onmessage = flushOne;
  }}

  window.__gradioFocusQueuePatch = patch;

  window.requestAnimationFrame = function (callback) {{
    if (typeof callback !== "function") {{
      return nativeRequestAnimationFrame(callback);
    }}
    if (!patch.shouldPatch()) {{
      return nativeRequestAnimationFrame(callback);
    }}
    const stack = new Error().stack || "";
    if (!patch.shouldPatchAnimationFrame(stack, callback)) {{
      return nativeRequestAnimationFrame(callback);
    }}
    return scheduleSyntheticAnimationFrame(callback);
  }};

  window.cancelAnimationFrame = function (id) {{
    if (cancelSynthetic(id)) {{
      return;
    }}
    return nativeCancelAnimationFrame(id);
  }};

  console.info("[Gradio] focus queue patch installed");
}})();
"""


def _inject_script(template_source: str) -> str:
    if _PATCH_SENTINEL in template_source:
        return template_source
    script_tag = f"\n\t\t<script>\n{get_javascript()}\n\t\t</script>\n"
    module_tag = '<script type="module"'
    insert_at = template_source.find(module_tag)
    if insert_at != -1:
        return template_source[:insert_at] + script_tag + template_source[insert_at:]
    head_close = template_source.find("</head>")
    if head_close != -1:
        return template_source[:head_close] + script_tag + template_source[head_close:]
    return template_source + script_tag


def install() -> bool:
    if not ENABLE_GRADIO_FOCUS_QUEUE_MONKEYPATCH:
        return False
    argv0 = Path(sys.argv[0]).name.lower() if sys.argv and sys.argv[0] else ""
    cwd = Path.cwd().resolve()
    if cwd != _PROJECT_ROOT and _PROJECT_ROOT not in cwd.parents and argv0 != "wgp.py":
        return False
    import gradio.routes as gradio_routes

    templates = getattr(gradio_routes, "templates", None)
    loader = getattr(getattr(templates, "env", None), "loader", None)
    if loader is None:
        return False
    if getattr(loader, "_focus_queue_patch_installed", False):
        return True

    original_get_source: Callable = loader.get_source

    def patched_get_source(environment, template):
        source, filename, uptodate = original_get_source(environment, template)
        if template in _TARGET_TEMPLATES:
            source = _inject_script(source)
        return source, filename, uptodate

    loader.get_source = patched_get_source
    loader._focus_queue_patch_installed = True
    loader._focus_queue_patch_original_get_source = original_get_source
    templates.env.cache.clear()
    return True


__all__ = [
    "ENABLE_GRADIO_FOCUS_QUEUE_MONKEYPATCH",
    "GRADIO_FOCUS_QUEUE_MONKEYPATCH_VERBOSE",
    "get_javascript",
    "install",
]