dify/api/dev/cached_console_openapi.py

"""Local, idle-expiring OpenAPI export cache. Contract bytes never touch the filesystem.

The POSIX server owns a private Unix socket and an empty lifecycle lock. It caches
the document in RAM. A single-threaded snapshot preloads only dependency modules,
never controllers. Each refresh imports controllers in a fresh fork; edits to a
preloaded dependency re-exec the snapshot before exporting. Unsupported platforms
or unsafe snapshots use the ordinary fresh-process exporter.
"""

from __future__ import annotations

import contextlib
import hashlib
import importlib.machinery
import io
import json
import os
import select
import signal
import socket
import subprocess
import sys
import tempfile
import time
import traceback
from pathlib import Path
from typing import Any, Never

API_ROOT = Path(__file__).resolve().parents[1]
SOURCE_ROOTS = (API_ROOT, API_ROOT.parent / "dify-agent/src")
PACKAGE_INPUTS = (API_ROOT / "pyproject.toml", API_ROOT / "uv.lock", API_ROOT.parent / "dify-agent/pyproject.toml")
TIMEOUT = 90
IDLE_TIMEOUT = 120
MAX_BYTES = 64 * 1024 * 1024


def source_imports() -> None:
    """Read editable backend modules from source, preserving installed-library bytecode caches."""
    original = importlib.machinery.SourceFileLoader.get_code

    def get_code(self: importlib.machinery.SourceFileLoader, fullname: str):
        path = Path(self.path)
        if any(path.is_relative_to(root) for root in SOURCE_ROOTS) and ".venv" not in path.parts:
            return self.source_to_code(self.get_data(self.path), self.path)
        return original(self, fullname)

    importlib.machinery.SourceFileLoader.get_code = get_code
    sys.dont_write_bytecode = True


def source_state() -> dict[str, tuple[int, int, int]]:
    """Match the Console descriptor's inputs, excluding tests and local environments."""
    files = [*PACKAGE_INPUTS, Path.cwd() / ".env"]
    for root in SOURCE_ROOTS:
        for directory, children, names in os.walk(root):
            children[:] = [name for name in children if name not in {".venv", ".git", "__pycache__"}]
            if Path(directory) == API_ROOT:
                children[:] = [name for name in children if name != "tests"]
            files.extend(Path(directory) / name for name in names if name.endswith(".py"))
    return {str(file): file_stamp(file) for file in files}


def file_stamp(file: Path) -> tuple[int, int, int]:
    try:
        stat = file.stat()
        return stat.st_mtime_ns, stat.st_ctime_ns, stat.st_size
    except FileNotFoundError:
        return 0, 0, 0


def fresh_export() -> bytes:
    result = subprocess.run(
        [
            sys.executable,
            "-B",
            "-c",
            "import sys;sys.path.insert(0,sys.argv[1]);"
            "from dev.cached_console_openapi import source_imports;source_imports();"
            "from dev.export_console_openapi import main;main()",
            str(API_ROOT),
        ],
        capture_output=True,
        timeout=TIMEOUT,
        check=False,
    )
    if result.returncode:
        raise RuntimeError(result.stderr.decode(errors="replace")[-16000:])
    return result.stdout


def single_threaded() -> bool:
    """Check native threads too: threading.active_count misses C library pools."""
    try:
        if sys.platform.startswith("linux"):
            return len(os.listdir("/proc/self/task")) == 1
        if sys.platform == "darwin":
            # Darwin -M always prints its thread header, even with -o pid=.
            output = subprocess.check_output(["ps", "-M", "-p", str(os.getpid())], timeout=5)
            return len(output.strip().splitlines()) == 2
    except (OSError, subprocess.SubprocessError):
        return False
    return False


def preload() -> tuple[bool, dict[str, tuple[int, int, int]]]:
    """Load schema dependencies without retaining route registries or controller classes."""
    import urllib.request

    getproxies = urllib.request.getproxies
    try:
        if sys.platform == "darwin":
            # macOS SystemConfiguration starts native threads. Resolve its proxy
            # settings in a helper, preserving the values without forking those threads.
            proxies = json.loads(
                subprocess.check_output(
                    [sys.executable, "-c", "import json,urllib.request;print(json.dumps(urllib.request.getproxies()))"],
                    timeout=10,
                )
            )
            urllib.request.getproxies = lambda: proxies
        sys.path.insert(0, str(API_ROOT))
        from dev.generate_swagger_specs import apply_runtime_defaults

        apply_runtime_defaults()
        # Keep runtime composition (and libraries with native thread pools) out
        # of the snapshot. These pure model dependencies account for much of startup.
        import dify_agent.protocol.schemas

        import core.mcp.types
        import graphon.nodes.parameter_extractor.entities
        import models
        import services.knowledge_fs.product_dto

        _ = core.mcp.types, dify_agent.protocol.schemas, graphon.nodes.parameter_extractor.entities
        _ = models, services.knowledge_fs.product_dto
        safe = not any(name.startswith("controllers.") for name in sys.modules) and single_threaded()
    except Exception:
        # A cache optimization must not prevent a normal export from working.
        safe = False
    finally:
        urllib.request.getproxies = getproxies
    dependencies = {}
    for module in tuple(sys.modules.values()):
        filename = getattr(module, "__file__", None)  # guard-ignore: no-new-getattr -- optional module source metadata
        if isinstance(filename, str):
            dependencies[str(Path(filename).resolve())] = file_stamp(Path(filename))
    for filename in (*PACKAGE_INPUTS, Path.cwd() / ".env"):
        dependencies[str(filename)] = file_stamp(filename)
    return safe, dependencies


def export_worker(write_fd: int) -> Never:
    """Run the isolated exporter and keep its pipe open through process exit."""
    try:
        from dev.export_console_openapi import main

        output = io.StringIO()
        with contextlib.redirect_stdout(output):
            main()
        payload = b"OK\n" + output.getvalue().encode()
    except BaseException:
        payload = b"ERROR\n" + traceback.format_exc().encode()
    try:
        with os.fdopen(write_fd, "wb") as pipe:
            pipe.write(payload)
            pipe.flush()
            # Keep the pipe open through process exit. Closing it first lets
            # the parent kill a successful worker while exit instrumentation
            # (for example, subprocess coverage) is still flushing its data.
            os._exit(0)
    finally:
        os._exit(1)


def fork_export() -> bytes:
    """Export from a clean controller registry; the parent snapshot is never mutated."""
    read_fd, write_fd = os.pipe()
    pid = os.fork()
    if pid == 0:
        os.close(read_fd)
        export_worker(write_fd)
    os.close(write_fd)
    chunks: list[bytes] = []
    size = 0
    deadline = time.monotonic() + TIMEOUT
    try:
        while True:
            if not select.select([read_fd], [], [], max(0, deadline - time.monotonic()))[0]:
                raise TimeoutError("OpenAPI refresh timed out")
            chunk = os.read(read_fd, 65536)
            if not chunk:
                break
            size += len(chunk)
            if size > MAX_BYTES:
                raise ValueError("OpenAPI export exceeds the response limit")
            chunks.append(chunk)
        result = b"".join(chunks)
        status, _, payload = result.partition(b"\n")
        if status != b"OK":
            raise RuntimeError(payload.decode(errors="replace") or "OpenAPI worker exited without a response")
        return payload
    finally:
        os.close(read_fd)
        # A worker can fail after closing its pipe. Never leave it running or unreaped.
        with contextlib.suppress(ProcessLookupError):
            os.kill(pid, signal.SIGKILL)
        os.waitpid(pid, 0)


def socket_path() -> Path:
    directory = Path(tempfile.gettempdir()) / f"dify-openapi-{os.getuid()}"
    # macOS Unix socket paths have a small length limit.
    if len(str(directory)) > 55:
        directory = Path("/tmp") / f"dify-openapi-{os.getuid()}"
    directory.mkdir(mode=0o700, exist_ok=True)
    stat = directory.lstat()
    if directory.is_symlink() or stat.st_uid != os.getuid() or stat.st_mode & 0o077:
        raise PermissionError("OpenAPI IPC directory must be private to the current user")
    environment = {
        key: value
        for key, value in os.environ.items()
        if key not in {"_", "PWD", "OLDPWD", "SHLVL", "UV_RUN_RECURSION_DEPTH"}
    }
    key = hashlib.sha256(json.dumps([str(API_ROOT), os.getcwd(), sys.executable, environment], sort_keys=True).encode())
    key.update(Path(__file__).read_bytes())
    return directory / f"{key.hexdigest()[:24]}.sock"


def serve(path: Path, inherited: list[str]) -> None:
    import fcntl

    if inherited:
        listener = socket.socket(fileno=int(inherited[0]))
        lock_fd = int(inherited[1])
    else:
        lock_fd = os.open(path.with_suffix(".lock"), os.O_CREAT | os.O_RDWR, 0o600)
        try:
            fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            os.close(lock_fd)
            return
        path.unlink(missing_ok=True)
        listener = socket.socket(socket.AF_UNIX)
        listener.bind(str(path))
        listener.listen(16)
    listener.settimeout(IDLE_TIMEOUT)
    baseline = source_state()
    safe, dependencies = preload()
    snapshot_dirty = source_state() != baseline
    cached_state: dict[str, tuple[int, int, int]] | None = None
    cached = b""
    try:
        while True:
            try:
                connection, _ = listener.accept()
            except TimeoutError:
                return
            with connection:
                connection.settimeout(TIMEOUT)
                try:
                    request = b""
                    while not request.endswith(b"\n") and len(request) < 16:
                        part = connection.recv(16 - len(request))
                        if not part:
                            break
                        request += part
                    if request == b"stop\n":
                        connection.sendall(b"OK\n")
                        return
                    if request != b"export\n":
                        continue
                    current = source_state()
                    stale = (
                        snapshot_dirty
                        or set(current) != set(baseline)
                        or any(file_stamp(Path(file)) != stamp for file, stamp in dependencies.items())
                    )
                    if stale:
                        # No reload(sys.modules): restart all preloaded dependencies.
                        connection.sendall(b"RETRY\n")
                        connection.close()
                        os.set_inheritable(listener.fileno(), True)
                        os.set_inheritable(lock_fd, True)
                        os.execv(
                            sys.executable,
                            [sys.executable, __file__, "--serve", str(path), str(listener.fileno()), str(lock_fd)],
                        )
                    if current != cached_state:
                        output = fork_export() if safe and single_threaded() else fresh_export()
                        # Do not publish a contract assembled while source files changed.
                        if source_state() != current:
                            connection.sendall(b"RETRY\n")
                            continue
                        payload: Any = json.loads(output)
                        if not isinstance(payload, dict) or set(payload) != {"console", "fast"}:
                            raise ValueError("Invalid Console export response")
                        cached, cached_state = output, current
                    connection.sendall(b"OK\n" + cached)
                except (BrokenPipeError, ConnectionResetError):
                    pass
                except Exception:
                    with contextlib.suppress(OSError):
                        connection.sendall(b"ERROR\n" + traceback.format_exc().encode()[-16000:])
    finally:
        listener.close()
        path.unlink(missing_ok=True)
        os.close(lock_fd)


def request_export(*, stop: bool = False) -> bytes:
    if not hasattr(os, "fork"):
        return b"" if stop else fresh_export()
    path = socket_path()
    deadline = time.monotonic() + TIMEOUT
    launched = False
    while time.monotonic() < deadline:
        try:
            with socket.socket(socket.AF_UNIX) as connection:
                connection.settimeout(max(0.1, deadline - time.monotonic()))
                connection.connect(str(path))
                connection.sendall(b"stop\n" if stop else b"export\n")
                chunks: list[bytes] = []
                size = 0
                while chunk := connection.recv(65536):
                    size += len(chunk)
                    if size > MAX_BYTES:
                        raise ValueError("OpenAPI response exceeds the limit")
                    chunks.append(chunk)
            status, _, payload = b"".join(chunks).partition(b"\n")
            if status == b"OK":
                return payload
            if status == b"ERROR":
                raise RuntimeError(payload.decode(errors="replace"))
            if status != b"RETRY":
                raise RuntimeError("Incomplete OpenAPI cache response")
        except (FileNotFoundError, ConnectionRefusedError):
            if stop:
                return b""
            if not launched:
                subprocess.Popen(
                    [sys.executable, __file__, "--serve", str(path)],
                    start_new_session=True,
                    stdin=subprocess.DEVNULL,
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
                launched = True
        time.sleep(0.03)
    raise TimeoutError("OpenAPI cache did not respond in time")


if __name__ == "__main__":
    # Timestamp-based .pyc files can serve stale code after same-size edits within one second.
    source_imports()
    if sys.argv[1:2] == ["--serve"]:
        serve(Path(sys.argv[2]), sys.argv[3:])
    elif sys.argv[1:] == ["--stop"]:
        request_export(stop=True)
    else:
        sys.stdout.buffer.write(request_export())