deer-flow/scripts/dev_multi_instance.py

#!/usr/bin/env python3
"""Helpers for ``scripts/dev_multi_instance.sh``, the two-Gateway local harness.

The shell entry point owns containers and processes. This module owns the parts
that decide whether the pair can start and whether it behaves as one cluster:

* ``render-config`` overlays the multi-instance prerequisites on a base
  ``config.yaml`` so the result passes the Gateway's multi-instance startup gate
  (``app.gateway.deps._enforce_postgres_for_multi_worker``) and leaves every
  other section of the developer's config alone.
* ``secrets`` generates the secrets every Gateway process must share.
* ``nginx-conf`` renders a minimal round-robin nginx front for both Gateways.
* ``check`` runs the scriptable cross-instance acceptance checks.

It needs PyYAML and httpx, so run it with the backend environment::

    uv run --no-sync --project backend python scripts/dev_multi_instance.py --help
"""

from __future__ import annotations

import argparse
import base64
import copy
import hashlib
import json
import os
import secrets
import sys
import time
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# Each Gateway exports its own value; the shared config.yaml references it so
# one file serves both processes (the Helm chart does the same with a per-Pod
# emptyDir mounted at one path).
RETRIEVAL_INDEX_ENV_VAR = "DEERFLOW_MI_RETRIEVAL_INDEX_PATH"

# Every Gateway process must see identical values. DEER_FLOW_CREDENTIALS_KEY is
# the at-rest credentials key: without one shared value a declared
# multi-instance pair with channel_connections enabled refuses to start.
SHARED_SECRET_NAMES = (
    "AUTH_JWT_SECRET",
    "DEER_FLOW_INTERNAL_AUTH_TOKEN",
    "DEER_FLOW_CREDENTIALS_KEY",
    "DEERFLOW_MI_POSTGRES_PASSWORD",
)

GENERATED_HEADER = "# Generated by scripts/dev_multi_instance.sh on every `up`; edit the base config instead.\n# Base config: {base}\n"


# ---------------------------------------------------------------------------
# Config overlay
# ---------------------------------------------------------------------------


def _section(config: dict[str, Any], name: str) -> dict[str, Any]:
    value = config.get(name)
    if not isinstance(value, dict):
        value = {}
        config[name] = value
    return value


def _is_browser_tool(tool: Any) -> bool:
    if not isinstance(tool, Mapping):
        return False
    return str(tool.get("name") or "").startswith("browser_") or "browser_automation" in str(tool.get("use") or "")


def _is_local_aio_sandbox(sandbox: Mapping[str, Any]) -> bool:
    """AIO sandboxes run as named Docker containers on this host unless a provisioner owns them."""
    return "aio_sandbox" in str(sandbox.get("use") or "") and not sandbox.get("provisioner_url")


def sandbox_container_prefix(state_dir: Path | str) -> str:
    """Docker name prefix for the local AIO sandboxes of one harness state directory.

    Stable for a state directory (warm sandboxes survive restarts) and distinct
    from the provider default, so the harness never adopts or reaps the
    sandbox containers of `make dev` or of another harness.
    """
    digest = hashlib.sha256(os.path.abspath(state_dir).encode("utf-8")).hexdigest()[:10]
    return f"deer-flow-mi-{digest}-sandbox"


def build_multi_instance_config(
    base: Mapping[str, Any] | None,
    *,
    database_url: str,
    redis_url: str,
    sandbox_container_prefix: str,
    keep_channels: bool = False,
) -> tuple[dict[str, Any], list[str]]:
    """Return ``(config, notes)``: *base* plus every multi-instance prerequisite.

    Everything the harness writes must land in its own Postgres, Redis, home
    and containers, never in the developer's own deployment, so every
    persistence endpoint or data root the base config names is redirected.
    ``notes`` describes each change beyond the always-applied keys, so the
    harness can tell the developer what differs from their own config.
    """
    config: dict[str, Any] = copy.deepcopy(dict(base or {}))
    notes: list[str] = []

    _section(config, "deployment")["multi_instance"] = True

    database = _section(config, "database")
    database["backend"] = "postgres"
    database["postgres_url"] = database_url
    cache = database.get("checkpoint_cache")
    if isinstance(cache, dict) and (cache.get("type") == "redis" or cache.get("redis_url")):
        cache["redis_url"] = redis_url
        notes.append("pointed database.checkpoint_cache at the harness Redis")

    _section(config, "run_events")["backend"] = "db"
    _section(config, "run_ownership")["heartbeat_enabled"] = True

    bridge = _section(config, "stream_bridge")
    bridge["type"] = "redis"
    bridge["redis_url"] = redis_url

    checkpointer = config.pop("checkpointer", None)
    if checkpointer is not None:
        # An explicit checkpointer section wins over `database` for LangGraph
        # checkpoints and the Store: a sqlite/memory one would split threads per
        # process, a postgres one would write into the developer's database.
        kind = checkpointer.get("type") if isinstance(checkpointer, Mapping) else None
        notes.append(f"removed checkpointer (type {kind!r}); the harness Postgres holds checkpoints and the Store")

    scheduler = config.get("scheduler")
    if isinstance(scheduler, dict) and scheduler.get("enabled") and not scheduler.get("multi_instance"):
        scheduler["multi_instance"] = True
        notes.append("set scheduler.multi_instance: true (the scheduler is enabled; the gate refuses a per-process scheduler)")

    tools = config.get("tools")
    if isinstance(tools, list):
        removed = [str(tool.get("name")) for tool in tools if _is_browser_tool(tool)]
        if removed:
            config["tools"] = [tool for tool in tools if not _is_browser_tool(tool)]
            notes.append(f"removed process-local browser tools: {', '.join(removed)}")

    sandbox = config.get("sandbox")
    if isinstance(sandbox, dict):
        ownership = sandbox.get("ownership")
        if ownership is not None:
            # An explicit section stops ownership from inheriting the bridge's
            # Redis, so its endpoint must be rewritten; timing settings stay.
            ownership = dict(ownership) if isinstance(ownership, Mapping) else {}
            previous = f"type {ownership.get('type', 'memory')!r}" + (f", {ownership['redis_url']}" if ownership.get("redis_url") else "")
            ownership.update(type="redis", redis_url=redis_url)
            sandbox["ownership"] = ownership
            notes.append(f"pointed sandbox.ownership at the harness Redis (was {previous})")
        if _is_local_aio_sandbox(sandbox):
            previous_prefix = sandbox.get("container_prefix") or "deer-flow-sandbox"
            sandbox["container_prefix"] = sandbox_container_prefix
            notes.append(f"set sandbox.container_prefix: {sandbox_container_prefix} (was {previous_prefix!r}) so the harness never adopts or reaps other sandbox containers")

    if not keep_channels:
        # UI-connectable providers can start transports after credentials arrive,
        # so disable them as well as the static YAML providers. ``leader_lease``
        # and the scalar connection settings are not providers.
        for section_name in ("channels", "channel_connections"):
            channels = config.get(section_name)
            if not isinstance(channels, dict):
                continue
            disabled = sorted(name for name, value in channels.items() if name != "leader_lease" and isinstance(value, dict) and value.get("enabled"))
            for name in disabled:
                channels[name]["enabled"] = False
            if disabled:
                notes.append(f"disabled IM channels under {section_name}: {', '.join(disabled)}: the harness keeps your bots offline unless asked (set DEERFLOW_MI_KEEP_CHANNELS=1 to keep them; the leader lease then gives each bot to one Gateway)")

    memory = _section(config, "memory")
    manager_class = memory.get("manager_class")
    if manager_class in (None, "", "deermem"):
        backend_config = memory.get("backend_config")
        if not isinstance(backend_config, dict):
            backend_config = {}
            memory["backend_config"] = backend_config
        # Empty storage_path means the harness DEER_FLOW_HOME. Both spellings go:
        # a legacy top-level memory.storage_path is migrated into backend_config on load.
        dropped = [f"{label}={value!r}" for label, value in (("memory.storage_path", memory.pop("storage_path", None)), ("memory.backend_config.storage_path", backend_config.pop("storage_path", None))) if value]
        if dropped:
            notes.append(f"dropped {', '.join(dropped)}: DeerMem data stays in the harness home")
        backend_config["retrieval_index_path"] = f"${RETRIEVAL_INDEX_ENV_VAR}"
    else:
        notes.append(f"memory.manager_class={manager_class!r} is not DeerMem: the harness uses that backend as configured, shared with your own deployment")

    blob_storage = config.get("blob_storage")
    blob_config = blob_storage.get("backend_config") if isinstance(blob_storage, dict) else None
    if isinstance(blob_config, dict) and blob_config.get("root"):
        notes.append(f"dropped blob_storage.backend_config.root={blob_config.pop('root')!r}: blobs stay in the harness home")

    return config, notes


def _write_private_text(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(f".{path.name}.tmp")
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w", encoding="utf-8") as fh:
        fh.write(text)
    os.chmod(tmp, 0o600)
    os.replace(tmp, path)


def render_config(base_path: Path, out_path: Path, *, database_url: str, redis_url: str, state_dir: Path, keep_channels: bool = False) -> list[str]:
    """Write the generated multi-instance config (mode 0600) and return its notes."""
    import yaml

    base = yaml.safe_load(Path(base_path).read_text(encoding="utf-8-sig")) or {}
    if not isinstance(base, dict):
        raise SystemExit(f"{base_path} is not a YAML mapping")
    config, notes = build_multi_instance_config(
        base,
        database_url=database_url,
        redis_url=redis_url,
        sandbox_container_prefix=sandbox_container_prefix(state_dir),
        keep_channels=keep_channels,
    )
    body = yaml.safe_dump(config, sort_keys=False, allow_unicode=True, default_flow_style=False)
    _write_private_text(Path(out_path), GENERATED_HEADER.format(base=base_path) + body)
    return notes


# ---------------------------------------------------------------------------
# Shared secrets
# ---------------------------------------------------------------------------


def generate_shared_secrets() -> dict[str, str]:
    return {
        "AUTH_JWT_SECRET": secrets.token_urlsafe(32),
        "DEER_FLOW_INTERNAL_AUTH_TOKEN": secrets.token_urlsafe(32),
        # A raw urlsafe-base64 32-byte Fernet key, not a passphrase.
        "DEER_FLOW_CREDENTIALS_KEY": base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("ascii"),
        "DEERFLOW_MI_POSTGRES_PASSWORD": secrets.token_hex(16),
    }


def read_env_file(path: Path) -> dict[str, str]:
    values: dict[str, str] = {}
    for line in Path(path).read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        values[key.strip()] = value.strip()
    return values


def ensure_secrets_file(path: Path) -> dict[str, str]:
    """Create the shell-sourceable secrets file once (``O_EXCL``, 0600) and return it."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    values = generate_shared_secrets()
    try:
        fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    except FileExistsError:
        return read_env_file(path)
    with os.fdopen(fd, "w", encoding="utf-8") as fh:
        fh.write("# Shared by every Gateway of the dev_multi_instance harness. Do not commit.\n")
        for name in SHARED_SECRET_NAMES:
            fh.write(f"{name}={values[name]}\n")
    return values


# ---------------------------------------------------------------------------
# nginx front
# ---------------------------------------------------------------------------

_PROXY_HEADERS = """            proxy_set_header Host $http_host;
            proxy_set_header X-Real-IP $remote_addr;
            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
            proxy_set_header X-Forwarded-Proto $scheme;"""


def _nginx_path(path: Path) -> str:
    """Render *path* as an nginx double-quoted string literal.

    Quoting keeps whitespace and ``;``/``{`` inside one argument. ``$`` cannot be
    escaped (nginx interpolates it in directives such as ``access_log``) and
    control characters cannot be spelled, so both are refused.
    """
    text = str(path)
    if "$" in text or any(ord(char) < 32 for char in text):
        raise ValueError(f"nginx cannot reference {text!r}: choose a state directory without '$' or control characters")
    return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"'


def render_nginx_conf(*, listen_port: int, gateway_ports: Sequence[int], frontend_port: int, state_dir: Path) -> str:
    """Return a minimal nginx config that round-robins ``/api`` over every Gateway."""
    state_dir = Path(state_dir)
    run_dir, log_dir, temp_dir = state_dir / "run", state_dir / "logs", state_dir / "nginx" / "temp"
    servers = "\n".join(f"        server 127.0.0.1:{port};" for port in gateway_ports)
    return f"""# Generated by scripts/dev_multi_instance.sh; regenerated on every `up`.
worker_processes 1;
pid {_nginx_path(run_dir / "nginx.pid")};
error_log {_nginx_path(log_dir / "nginx-error.log")} warn;

events {{
    worker_connections 1024;
}}

http {{
    access_log {_nginx_path(log_dir / "nginx-access.log")};
    client_body_temp_path {_nginx_path(temp_dir / "client_body")};
    proxy_temp_path {_nginx_path(temp_dir / "proxy")};
    fastcgi_temp_path {_nginx_path(temp_dir / "fastcgi")};
    uwsgi_temp_path {_nginx_path(temp_dir / "uwsgi")};
    scgi_temp_path {_nginx_path(temp_dir / "scgi")};

    upstream deerflow_gateways {{
{servers}
    }}

    upstream deerflow_frontend {{
        server 127.0.0.1:{frontend_port};
    }}

    map $http_upgrade $connection_upgrade {{
        default upgrade;
        ''      '';
    }}

    server {{
        listen 127.0.0.1:{listen_port};
        server_name _;

        # Which Gateway answered: lets `check` and developers see the balancing.
        add_header X-DeerFlow-Upstream $upstream_addr always;

        proxy_http_version 1.1;
        proxy_buffering off;
        proxy_cache off;
        proxy_request_buffering off;
        client_max_body_size 101M;
        proxy_connect_timeout 10s;
        proxy_send_timeout 600s;
        proxy_read_timeout 600s;

        location /api/langgraph/ {{
            rewrite ^/api/langgraph/(.*) /api/$1 break;
{_PROXY_HEADERS}
            proxy_set_header Connection '';
            proxy_pass http://deerflow_gateways;
        }}

        location /api/ {{
{_PROXY_HEADERS}
            proxy_set_header Connection '';
            proxy_pass http://deerflow_gateways;
        }}

        location /health {{
{_PROXY_HEADERS}
            proxy_pass http://deerflow_gateways;
        }}

        location / {{
{_PROXY_HEADERS}
            proxy_set_header Upgrade $http_upgrade;
            proxy_set_header Connection $connection_upgrade;
            proxy_pass http://deerflow_frontend;
        }}
    }}
}}
"""


# ---------------------------------------------------------------------------
# Cross-instance checks
# ---------------------------------------------------------------------------

PASS, FAIL, SKIP = "PASS", "FAIL", "SKIP"


@dataclass
class CheckResult:
    name: str
    status: str
    detail: str


@dataclass
class SSEEvent:
    id: str | None
    event: str | None
    data: str


def parse_sse_events(lines: Iterable[str]) -> Iterator[SSEEvent]:
    """Yield server-sent events from decoded lines (blank line = dispatch)."""
    event_id: str | None = None
    event_name: str | None = None
    data: list[str] = []
    for raw in lines:
        line = raw.rstrip("\r\n")
        if not line:
            if event_name is not None or data or event_id is not None:
                yield SSEEvent(event_id, event_name, "\n".join(data))
            event_id, event_name, data = None, None, []
            continue
        if line.startswith(":"):
            continue
        name, _, value = line.partition(":")
        value = value[1:] if value.startswith(" ") else value
        if name == "id":
            event_id = value
        elif name == "event":
            event_name = value
        elif name == "data":
            data.append(value)
    if event_name is not None or data or event_id is not None:
        yield SSEEvent(event_id, event_name, "\n".join(data))


def exit_code(results: Iterable[CheckResult]) -> int:
    return 1 if any(result.status == FAIL for result in results) else 0


class CheckFailed(Exception):
    """A check's expectation did not hold; the message is the reported detail."""


class CheckSkipped(Exception):
    """A check cannot run in this setup; the message says why."""


@dataclass
class Harness:
    """State shared by the checks: two Gateways, an optional nginx, one session."""

    gateways: list[tuple[str, str]]
    nginx_url: str | None
    state_dir: Path
    with_llm: bool
    timeout: float
    internal_token: str = ""
    cookies: dict[str, str] = field(default_factory=dict)
    user: dict[str, Any] = field(default_factory=dict)
    is_admin: bool = False
    thread_id: str | None = None
    models: list[Any] | None = None

    def __post_init__(self) -> None:
        import httpx

        # Every target is loopback: never route it through an HTTP(S)_PROXY from the shell.
        self.http = httpx.Client(timeout=httpx.Timeout(self.timeout, read=max(self.timeout, 60.0)), follow_redirects=False, trust_env=False)

    # -- request helpers -----------------------------------------------------

    def headers(self, *, mutating: bool = False) -> dict[str, str]:
        headers: dict[str, str] = {}
        if self.cookies:
            headers["Cookie"] = "; ".join(f"{name}={value}" for name, value in self.cookies.items())
        if mutating and "csrf_token" in self.cookies:
            headers["X-CSRF-Token"] = self.cookies["csrf_token"]
        return headers

    def request(self, method: str, base: str, path: str, *, expect: Iterable[int] = (200,), session: bool = True, **kwargs: Any):
        mutating = method.upper() in {"POST", "PUT", "PATCH", "DELETE"}
        headers = {**(self.headers(mutating=mutating) if session else {}), **kwargs.pop("headers", {})}
        response = self.http.request(method, f"{base}{path}", headers=headers, **kwargs)
        self.http.cookies.clear()  # cookies travel explicitly, never through the jar
        if response.status_code not in tuple(expect):
            raise CheckFailed(f"{method} {base}{path} -> HTTP {response.status_code}: {response.text[:300]}")
        return response

    @property
    def a(self) -> str:
        return self.gateways[0][1]

    @property
    def b(self) -> str:
        return self.gateways[1][1]

    # -- session ---------------------------------------------------------------

    def _adopt_session(self, response) -> None:
        for name in ("access_token", "csrf_token"):
            value = response.cookies.get(name)
            if value:
                self.cookies[name] = value
        self.http.cookies.clear()
        if "access_token" not in self.cookies:
            raise CheckFailed(f"{response.request.url} set no access_token cookie")

    def ensure_session(self) -> str:
        """Log in as the harness's throwaway user, creating it on first use."""
        creds_path = self.state_dir / "check-user.json"
        if creds_path.is_file():
            creds = json.loads(creds_path.read_text(encoding="utf-8"))
            response = self.request("POST", self.a, "/api/v1/auth/login/local", data={"username": creds["email"], "password": creds["password"]}, expect=(200, 401))
            if response.status_code == 200:
                self._adopt_session(response)
                self.is_admin = bool(creds.get("admin"))
                return f"logged in as {creds['email']} on A"
        creds = {"email": f"mi-check-{secrets.token_hex(4)}@example.com", "password": secrets.token_urlsafe(18)}
        setup = self.request("GET", self.a, "/api/v1/auth/setup-status").json()
        if setup.get("needs_setup"):
            response = self.request("POST", self.a, "/api/v1/auth/initialize", json=creds, expect=(201,))
            creds["admin"] = True
        else:
            response = self.request("POST", self.a, "/api/v1/auth/register", json=creds, expect=(201,))
            creds["admin"] = False
        self._adopt_session(response)
        self.is_admin = creds["admin"]
        _write_private_text(creds_path, json.dumps(creds))
        role = "admin" if self.is_admin else "user"
        return f"created throwaway {role} {creds['email']} on A (credentials: {creds_path})"


def _check_health(h: Harness) -> str:
    seen = []
    for name, base in h.gateways:
        body = h.request("GET", base, "/health/ready").json()
        if body.get("status") != "ready" or body.get("stream_bridge") != "ok":
            raise CheckFailed(f"{name} /health/ready: {body}")
        seen.append(f"{name}: database={body.get('database')} stream_bridge={body.get('stream_bridge')}")
    return "; ".join(seen)


def _check_session(h: Harness) -> str:
    created = h.ensure_session()
    me_a = h.request("GET", h.a, "/api/v1/auth/me").json()
    me_b = h.request("GET", h.b, "/api/v1/auth/me").json()
    if me_a.get("id") != me_b.get("id"):
        raise CheckFailed(f"A sees user {me_a.get('id')}, B sees {me_b.get('id')}")
    h.user = me_a
    return f"{created}; B accepts A's session (shared AUTH_JWT_SECRET)"


def _check_internal_token(h: Harness) -> str:
    if not h.internal_token:
        raise CheckSkipped("no DEER_FLOW_INTERNAL_AUTH_TOKEN in the secrets file")
    bodies = []
    for _name, base in h.gateways:
        # Negative control first: without a session, a wrong token must be refused.
        h.request("GET", base, "/api/models", session=False, headers={"X-DeerFlow-Internal-Token": "not-the-token"}, expect=(401,))
        bodies.append(h.request("GET", base, "/api/models", session=False, headers={"X-DeerFlow-Internal-Token": h.internal_token}).json())
    models = bodies[0].get("models", []) if isinstance(bodies[0], dict) else []
    h.models = models
    return f"both Gateways accept the shared DEER_FLOW_INTERNAL_AUTH_TOKEN ({len(models)} model(s) configured)"


def _require_thread(h: Harness) -> str:
    if h.thread_id is None:
        raise CheckSkipped("no thread (thread check did not pass)")
    return h.thread_id


def _check_thread(h: Harness) -> str:
    created = h.request("POST", h.a, "/api/threads", json={"metadata": {"source": "dev_multi_instance check"}}).json()
    thread_id = created["thread_id"]
    fetched = h.request("GET", h.b, f"/api/threads/{thread_id}").json()
    if fetched.get("thread_id") != thread_id:
        raise CheckFailed(f"B returned {fetched}")
    h.thread_id = thread_id
    return f"thread {thread_id} created on A is readable on B"


def _check_upload(h: Harness) -> str:
    thread_id = _require_thread(h)
    name = f"mi-check-{secrets.token_hex(4)}.txt"
    payload = f"multi-instance upload {secrets.token_hex(16)}\n".encode()
    uploaded = h.request("POST", h.a, f"/api/threads/{thread_id}/uploads", files={"files": (name, payload, "text/plain")}).json()
    info = next((item for item in uploaded.get("files", []) if item.get("filename") == name), None)
    if info is None:
        raise CheckFailed(f"A did not report {name}: {uploaded}")
    listed = h.request("GET", h.b, f"/api/threads/{thread_id}/uploads/list").json()
    if name not in {item.get("filename") for item in listed.get("files", [])}:
        raise CheckFailed(f"B does not list {name}: {listed}")
    body = h.request("GET", h.b, info["artifact_url"]).content
    if body != payload:
        raise CheckFailed(f"B served {len(body)} bytes that differ from the {len(payload)} uploaded on A")
    return f"uploaded on A; B lists it and serves identical bytes via {info['artifact_url']}"


def _check_skill_toggle(h: Harness) -> str:
    if not h.is_admin:
        raise CheckSkipped("the check user is not an admin (an admin already existed); PUT /api/skills is admin-only")
    skills = h.request("GET", h.a, "/api/skills").json().get("skills", [])
    skill = next((s for s in skills if s.get("category") == "public"), None)
    if skill is None:
        raise CheckSkipped("no public skill to toggle")
    name, original = skill["name"], bool(skill.get("enabled", True))
    h.request("PUT", h.a, f"/api/skills/{name}", json={"enabled": not original})
    try:
        started = time.monotonic()
        deadline = started + 30
        while True:
            seen = {s["name"]: s.get("enabled") for s in h.request("GET", h.b, "/api/skills").json().get("skills", [])}
            if seen.get(name) is (not original):
                return f"skill {name!r} toggled on A is {'enabled' if not original else 'disabled'} on B after {time.monotonic() - started:.1f}s (restored)"
            if time.monotonic() > deadline:
                raise CheckFailed(f"B still reports {name!r} enabled={seen.get(name)} 30s after A toggled it")
            time.sleep(1)
    finally:
        h.request("PUT", h.a, f"/api/skills/{name}", json={"enabled": original})


def _check_nginx(h: Harness) -> str:
    if not h.nginx_url:
        raise CheckSkipped("nginx is not running (up --no-nginx, or nginx not installed)")
    expected = {base.split("://", 1)[1] for _name, base in h.gateways}
    upstreams = set()
    for _ in range(2 * len(h.gateways) + 2):
        response = h.request("GET", h.nginx_url, "/health")
        upstreams.add(response.headers.get("x-deerflow-upstream", ""))
    if not expected <= upstreams:
        raise CheckFailed(f"nginx reached {sorted(upstreams)}, expected every Gateway {sorted(expected)}")
    detail = f"round-robin reached {', '.join(sorted(expected))}"
    if h.thread_id:
        for _ in range(len(h.gateways) * 2):
            h.request("GET", h.nginx_url, f"/api/threads/{h.thread_id}")
        detail += f"; thread {h.thread_id} readable through every upstream"
    return detail


def bounded_lines(lines: Iterable[str], *, deadline: float, what: str, clock: Callable[[], float] = time.monotonic) -> Iterator[str]:
    """Yield raw SSE lines until *deadline*, counting comment heartbeats too.

    The production stream keeps an idle live run open with ``: heartbeat``
    comments, which the event parser drops and which also reset httpx's read
    timeout, so the budget must be enforced per raw line, not per event.
    """
    for line in lines:
        if clock() > deadline:
            raise CheckFailed(f"{what} exceeded its time budget without reaching `end`")
        yield line


def _read_stream(h: Harness, base: str, path: str, *, last_event_id: str | None = None, budget: float = 90.0) -> list[SSEEvent]:
    headers = {**h.headers(), "Accept": "text/event-stream"}
    if last_event_id:
        headers["Last-Event-ID"] = last_event_id
    events: list[SSEEvent] = []
    deadline = time.monotonic() + budget
    try:
        with h.http.stream("GET", f"{base}{path}", headers=headers) as response:
            if response.status_code != 200:
                response.read()
                raise CheckFailed(f"GET {base}{path} -> HTTP {response.status_code}: {response.text[:300]}")
            for event in parse_sse_events(bounded_lines(response.iter_lines(), deadline=deadline, what=f"GET {base}{path} ({budget:g}s budget)")):
                events.append(event)
                if event.event == "end":
                    break
    finally:
        h.http.cookies.clear()
    return events


def _check_sse_resume(h: Harness) -> str:
    thread_id = _require_thread(h)
    if not h.with_llm:
        # Only a confirmed empty model list makes the run free (it fails before any LLM call).
        if h.models is None:
            raise CheckSkipped("model discovery did not succeed, so a run might call an LLM; rerun with `check --with-llm`")
        if h.models:
            raise CheckSkipped(f"{len(h.models)} model(s) configured: a run would call the LLM; rerun with `check --with-llm`")
    body = {"assistant_id": "lead_agent", "input": {"messages": [{"role": "user", "content": "Reply with the single word: pong"}]}}
    run = h.request("POST", h.a, f"/api/threads/{thread_id}/runs", json=body).json()
    run_id = run["run_id"]
    path = f"/api/threads/{thread_id}/runs/{run_id}/join"
    full = _read_stream(h, h.b, path)
    ids = [event.id for event in full if event.id]
    if not full or full[-1].event != "end":
        raise CheckFailed(f"joining on B did not reach `end`: {[e.event for e in full]}")
    if len(ids) < 2:
        raise CheckSkipped(f"run {run_id} produced only {len(ids)} event id(s); nothing to resume from")
    resumed = _read_stream(h, h.a, path, last_event_id=ids[0])
    resumed_ids = [event.id for event in resumed if event.id]
    if resumed_ids != ids[1:]:
        raise CheckFailed(f"resuming on A after {ids[0]} returned ids {resumed_ids}, expected {ids[1:]}")
    status = h.request("GET", h.b, f"/api/threads/{thread_id}/runs/{run_id}").json().get("status")
    kind = "no-model (fails fast, no LLM call)" if h.models == [] else "LLM"
    return f"{kind} run {run_id} created on A (status {status}); B replayed {len(full)} events; A resumed from Last-Event-ID {ids[0]} with the remaining {len(resumed_ids)}"


CHECKS: list[tuple[str, Callable[[Harness], str]]] = [
    ("health/ready on both (stream_bridge ok)", _check_health),
    ("session from A accepted by B", _check_session),
    ("shared internal auth token", _check_internal_token),
    ("thread created on A visible on B", _check_thread),
    ("upload on A, list + download on B", _check_upload),
    ("skill toggle on A seen by B", _check_skill_toggle),
    ("nginx balances both Gateways", _check_nginx),
    ("SSE join on B, Last-Event-ID resume on A", _check_sse_resume),
]

MANUAL_CHECKS = (
    "`stop b --kill` while A has an active run: A's run must not turn `error`",
    "cancel a run on B from A: takes effect within about one heartbeat (~10s)",
    "change an MCP filesystem server's args on A: B's local-bash allowlist follows",
    "edit a skill on A: B's system-prompt skills section changes within 30s",
    "IM channel gates (one leader per platform, failover, runtime-config): once channel leader election lands",
)


def run_checks(h: Harness, *, out=sys.stdout) -> list[CheckResult]:
    results: list[CheckResult] = []
    for name, check in CHECKS:
        try:
            result = CheckResult(name, PASS, check(h))
        except CheckSkipped as exc:
            result = CheckResult(name, SKIP, str(exc))
        except CheckFailed as exc:
            result = CheckResult(name, FAIL, str(exc))
        except Exception as exc:  # connection errors, unexpected payloads
            result = CheckResult(name, FAIL, f"{type(exc).__name__}: {exc}")
        results.append(result)
        print(f"[{result.status}] {result.name}: {result.detail}", file=out, flush=True)
    counts = {status: sum(1 for r in results if r.status == status) for status in (PASS, FAIL, SKIP)}
    print(f"\n{counts[PASS]} passed, {counts[FAIL]} failed, {counts[SKIP]} skipped", file=out)
    print('Still manual (backend/docs/CONFIGURATION.md, "Local two-Gateway harness"):', file=out)
    for item in MANUAL_CHECKS:
        print(f"  - {item}", file=out)
    return results


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def _parse_gateway(value: str) -> tuple[str, str]:
    name, sep, url = value.partition("=")
    if not sep or not name or not url:
        raise argparse.ArgumentTypeError("expected NAME=URL, e.g. A=http://127.0.0.1:8001")
    return name, url.rstrip("/")


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="command", required=True)

    render = sub.add_parser("render-config", help="write the multi-instance config.yaml")
    render.add_argument("--base", required=True, type=Path)
    render.add_argument("--out", required=True, type=Path)
    render.add_argument("--database-url", required=True)
    render.add_argument("--redis-url", required=True)
    render.add_argument("--state-dir", required=True, type=Path)
    render.add_argument("--keep-channels", action="store_true")

    secrets_cmd = sub.add_parser("secrets", help="create the shared secrets file if missing")
    secrets_cmd.add_argument("--out", required=True, type=Path)

    nginx = sub.add_parser("nginx-conf", help="write the round-robin nginx config")
    nginx.add_argument("--out", required=True, type=Path)
    nginx.add_argument("--state-dir", required=True, type=Path)
    nginx.add_argument("--listen-port", required=True, type=int)
    nginx.add_argument("--gateway-port", required=True, type=int, action="append")
    nginx.add_argument("--frontend-port", type=int, default=3000)

    check = sub.add_parser("check", help="run the cross-instance checks")
    check.add_argument("--gateway", required=True, type=_parse_gateway, action="append", help="NAME=URL; pass exactly two")
    check.add_argument("--nginx", default=None)
    check.add_argument("--state-dir", required=True, type=Path)
    check.add_argument("--with-llm", action="store_true", help="allow checks that start a run on a configured model")
    check.add_argument("--timeout", type=float, default=15.0)

    args = parser.parse_args(argv)

    if args.command == "render-config":
        for note in render_config(args.base, args.out, database_url=args.database_url, redis_url=args.redis_url, state_dir=args.state_dir, keep_channels=args.keep_channels):
            print(f"  note: {note}")
        return 0
    if args.command == "secrets":
        ensure_secrets_file(args.out)
        return 0
    if args.command == "nginx-conf":
        _write_private_text(args.out, render_nginx_conf(listen_port=args.listen_port, gateway_ports=args.gateway_port, frontend_port=args.frontend_port, state_dir=args.state_dir))
        return 0
    if args.command == "check":
        if len(args.gateway) != 2:
            parser.error("check needs exactly two --gateway NAME=URL values")
        secrets_path = args.state_dir / "secrets.env"
        token = read_env_file(secrets_path).get("DEER_FLOW_INTERNAL_AUTH_TOKEN", "") if secrets_path.is_file() else ""
        harness = Harness(gateways=args.gateway, nginx_url=(args.nginx or "").rstrip("/") or None, state_dir=args.state_dir, with_llm=args.with_llm, timeout=args.timeout, internal_token=token)
        return exit_code(run_checks(harness))
    return 2


if __name__ == "__main__":
    sys.exit(main())