multi-agent-mux/.agents/skills/lib_py/herdr_config.py

"""Issue #45: the config file MAM's isolated herdr servers start with.

herdr restores a named session's saved layout when its server starts and, by
default, relaunches each saved agent with herdr's own resume command - without
the permission flags MAM adds, and even for agents MAM deleted while the server
was down. `[session] resume_agents_on_restore = false` turns that off. The
shim's headless bootstrap (lib.sh) hands the server `HERDR_CONFIG_PATH` pointing
at the file built here: the user's own herdr config with only that key forced
to false. The user's file is never modified.

  * Source: the file herdr itself would read - HERDR_CONFIG_PATH, else
    $XDG_CONFIG_HOME/herdr/config.toml, else ~/.config/herdr/config.toml. A
    HERDR_CONFIG_PATH naming a file generated here (every MAM pane inherits it
    from the server) is followed back to the source recorded in that file.
  * Output: `mam-server-config.toml` next to the source, so relative paths in
    it (herdr resolves them against the config file's directory) keep their
    meaning. `<.mam>/herdr-server-config.toml` only when that directory is not
    writable. Either way never the source itself and never a file this module
    did not write for the same source: then a name with a hash of the source
    path, and no write at all when that is taken too.
  * The source is parsed, the key set, and the result serialized; the result
    must parse back to exactly the source plus the key (any TOML spelling of
    the key - header, dotted, inline - works). Then herdr's own `config check`
    must accept the written file: herdr drops a file it cannot use whole
    ("using defaults") and would restore again, silently. Anything that fails
    falls back to a file holding only the key, which is checked the same way.

CLI: `python3 -m lib_py.herdr_config server-config [--fallback-dir DIR]
[--herdr PATH]` writes the file and prints its path; warnings go to stderr;
exit 1 when no usable file could be written.
"""
from __future__ import annotations

import argparse
import datetime
import hashlib
import json
import math
import os
import re
import subprocess
import sys
import tempfile
from typing import List, Mapping, Optional, Tuple

MARKER = "# generated by multi-agent-mux (Issue #45) - do not edit"
SOURCE_PREFIX = "# source: "
OUTPUT_NAME = "mam-server-config.toml"
FALLBACK_NAME = "herdr-server-config.toml"
KEY = "resume_agents_on_restore"
MINIMAL = f"[session]\n{KEY} = false\n"
LOSES_SETTINGS = f"the MAM server uses only {KEY} = false, without the rest of your settings"

try:
    import tomllib
except ImportError:  # Python < 3.11: no parser to check a merge with
    tomllib = None

_BARE_KEY = re.compile(r"^[A-Za-z0-9_-]+$")
# Where a pane's server, socket and workspace come from; `config check` must not
# be steered by the caller's pane.
_SCOPE_ENV = ("HERDR_WORKSPACE_ID", "HERDR_PANE_ID", "HERDR_TAB_ID", "HERDR_SOCKET_PATH")


def _generated_source(path: str) -> Optional[str]:
    """The source recorded in a file generated here; '' when it is one without
    a source line; None when it is not a generated file (or unreadable)."""
    try:
        with open(path, "rb") as f:
            first, second = f.readline(), f.readline()
    except OSError:
        return None
    if first.rstrip(b"\n").decode("utf-8", "replace") != MARKER:
        return None
    line = second.rstrip(b"\n").decode("utf-8", "replace")
    return line[len(SOURCE_PREFIX):] if line.startswith(SOURCE_PREFIX) else ""


def user_config_path(env: Mapping[str, str]) -> str:
    """The config file herdr reads under `env`, seen through generated files."""
    explicit = env.get("HERDR_CONFIG_PATH", "")
    if explicit:
        source = _generated_source(explicit)
        if source is None:
            return explicit
        if source:
            return source
    xdg = env.get("XDG_CONFIG_HOME", "")
    if xdg:
        # herdr takes a relative XDG_CONFIG_HOME against its cwd, too.
        return os.path.join(os.path.abspath(xdg), "herdr", "config.toml")
    return os.path.join(env.get("HOME", os.path.expanduser("~")), ".config", "herdr", "config.toml")


# --- TOML writer for what tomllib returns -----------------------------------

def _str(s: str) -> str:
    # json's escapes (\" \\ \b \f \n \r \t \uXXXX) are TOML's; DEL is not
    # escaped by json but is not allowed raw in a TOML string.
    return json.dumps(s, ensure_ascii=False).replace("\x7f", "\\u007f")


def _key(k: str) -> str:
    return k if _BARE_KEY.match(k) else _str(k)


def _value(v) -> str:
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, int):
        return str(v)
    if isinstance(v, float):
        if math.isnan(v):
            return "nan"
        if math.isinf(v):
            return "inf" if v > 0 else "-inf"
        return repr(v)
    if isinstance(v, str):
        return _str(v)
    if isinstance(v, (datetime.datetime, datetime.date, datetime.time)):
        return v.isoformat()
    if isinstance(v, list):
        return "[" + ", ".join(_value(x) for x in v) + "]"
    if isinstance(v, dict):
        return "{" + ", ".join(f"{_key(k)} = {_value(x)}" for k, x in v.items()) + "}"
    raise TypeError(f"unexpected TOML value {type(v).__name__}")


def dumps(doc: dict) -> str:
    """Top-level values, then one [table] per top-level table with its members
    written inline - valid TOML for any document tomllib can return."""
    lines = [f"{_key(k)} = {_value(v)}" for k, v in doc.items() if not isinstance(v, dict)]
    for k, v in doc.items():
        if isinstance(v, dict):
            lines.append(f"[{_key(k)}]")
            lines += [f"{_key(kk)} = {_value(vv)}" for kk, vv in v.items()]
    return "\n".join(lines) + "\n"


def _same_doc(a, b) -> bool:
    """Equality of two parsed TOML documents in which nan equals nan (a valid
    TOML value that a plain == never matches, review R3)."""
    if isinstance(a, float) and isinstance(b, float) and math.isnan(a) and math.isnan(b):
        return True
    if isinstance(a, dict) and isinstance(b, dict):
        return a.keys() == b.keys() and all(_same_doc(a[k], b[k]) for k in a)
    if isinstance(a, list) and isinstance(b, list):
        return len(a) == len(b) and all(_same_doc(x, y) for x, y in zip(a, b))
    return type(a) is type(b) and a == b


def server_config_text(user_text: Optional[str]) -> Tuple[str, List[str]]:
    """(the TOML body for the server, warnings). The body always parses, has
    session.resume_agents_on_restore = false and otherwise equals the source -
    or is MINIMAL, with a warning, when the source cannot be used."""
    if not user_text or not user_text.strip():
        return MINIMAL, []
    if tomllib is None:
        return MINIMAL, [f"this Python has no tomllib (3.11+); {LOSES_SETTINGS}"]
    try:
        user = tomllib.loads(user_text)
    except tomllib.TOMLDecodeError as e:
        return MINIMAL, [f"herdr config does not parse ({e}); {LOSES_SETTINGS}"]
    session = user.get("session", {})
    if not isinstance(session, dict):
        # herdr rejects this file as a whole too.
        return MINIMAL, [f"herdr config has a non-table `session`; {LOSES_SETTINGS}"]
    expected = dict(user, session=dict(session, **{KEY: False}))
    try:
        text = dumps(expected)
        ok = _same_doc(tomllib.loads(text), expected)
    except (TypeError, tomllib.TOMLDecodeError):
        ok = False
    if not ok:
        return MINIMAL, [f"could not rewrite the herdr config faithfully; {LOSES_SETTINGS}"]
    return text, []


# --- herdr's own check ------------------------------------------------------

def herdr_rejects(path: str, herdr: str) -> Optional[str]:
    """None when herdr would use the config at `path`, else herdr's diagnostics.
    herdr 0.9.1 drops a file whole only on a parse or schema error, which it
    reports as "config parse error: ...; using defaults"; other issues (an
    unknown key, an unsupported sound file) cost only that one setting."""
    env = {k: v for k, v in os.environ.items() if k not in _SCOPE_ENV}
    env["HERDR_CONFIG_PATH"] = path
    res = subprocess.run([herdr, "config", "check"], env=env, capture_output=True, text=True, timeout=30)
    out = (res.stdout + res.stderr).strip()
    if res.returncode == 0 or not ("parse error" in out or "using defaults" in out):
        return None
    return out


# --- output -----------------------------------------------------------------

def _same_file(a: str, b: str) -> bool:
    try:
        return os.path.samefile(a, b)
    except OSError:
        return os.path.realpath(a) == os.path.realpath(b)


def output_path(source: str, directory: Optional[str] = None, name: str = OUTPUT_NAME) -> str:
    """`name` in `directory` (default: beside the source), unless that is the
    source itself (aliases included) or a file this module did not write for
    this source; then a per-source name. Every write goes through here."""
    d = directory or os.path.dirname(source)
    tag = hashlib.sha256(os.path.realpath(source).encode()).hexdigest()[:12]
    for n in (name, f"{name[:-len('.toml')]}-{tag}.toml"):
        target = os.path.join(d, n)
        if _same_file(target, source):
            continue
        if os.path.lexists(target) and _generated_source(target) not in ("", source):
            continue
        return target
    raise OSError(f"no free output name in {d}")


def _write_atomic(path: str, text: str, mode: int) -> None:
    d = os.path.dirname(path)
    os.makedirs(d, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=d, prefix=".mam-server-config.", suffix=".tmp")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            f.write(text)
        os.chmod(tmp, mode)
        os.replace(tmp, path)
    except BaseException:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise


def _write_checked(target: str, source: str, body: str, mode: int, herdr: Optional[str],
                   warnings: List[str]) -> bool:
    """Write body (then MINIMAL if herdr rejects it); True when herdr accepts
    what is on disk (or there is no herdr to ask)."""
    for candidate in ((body, MINIMAL) if body != MINIMAL else (MINIMAL,)):
        _write_atomic(target, f"{MARKER}\n{SOURCE_PREFIX}{source}\n{candidate}", mode)
        if not herdr:
            return True
        try:
            why = herdr_rejects(target, herdr)
        except (OSError, subprocess.SubprocessError) as e:
            # Not a pass (review R1): an unchecked merge with a schema error
            # makes herdr drop the whole file - restore back on. The merged
            # body is not used unchecked; MINIMAL is a fixed file herdr 0.9.1
            # accepts (test_real_herdr_accepts_every_result), and returning no
            # config would leave restore on, so it is used with a warning.
            if candidate != MINIMAL:
                warnings.append(f"cannot run `{herdr} config check` ({e}); not using an unchecked merge; "
                                f"{LOSES_SETTINGS}")
                continue
            warnings.append(f"cannot run `{herdr} config check` ({e}); using the fixed minimal config unchecked")
            return True
        if why is None:
            return True
        tail = "" if candidate == MINIMAL else f"; {LOSES_SETTINGS}"
        warnings.append(f"herdr rejects the config built from {source}: {why}{tail}")
    return False


def write_server_config(env: Mapping[str, str], fallback_dir: Optional[str] = None,
                        herdr: Optional[str] = None) -> Tuple[Optional[str], List[str]]:
    """Write the server config; (its path or None, warnings)."""
    source = user_config_path(env)
    warnings: List[str] = []
    user_text: Optional[str] = None
    mode = 0o644
    try:
        with open(source, encoding="utf-8") as f:
            user_text = f.read()
        mode = os.stat(source).st_mode & 0o777
    except FileNotFoundError:
        pass
    except (OSError, UnicodeDecodeError) as e:
        # herdr cannot use it either, so it runs on defaults: same as no file.
        warnings.append(f"cannot read herdr config {source}: {e}; {LOSES_SETTINGS}")
    body, merge_warnings = server_config_text(user_text)
    warnings += merge_warnings
    try:
        target = output_path(source)
        if _write_checked(target, source, body, mode, herdr, warnings):
            return target, warnings
        return None, warnings
    except OSError as e:
        warnings.append(f"cannot write beside {source}: {e}")
    if fallback_dir:
        target = os.path.join(fallback_dir, FALLBACK_NAME)
        try:
            target = output_path(source, fallback_dir, FALLBACK_NAME)
            if _write_checked(target, source, body, mode, herdr, warnings):
                warnings.append(f"wrote {target} instead; relative paths in the herdr config may not resolve")
                return target, warnings
        except OSError as e:
            warnings.append(f"cannot write {target}: {e}")
    return None, warnings


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description="herdr config for MAM's isolated servers (Issue #45)")
    sub = p.add_subparsers(dest="cmd", required=True)
    sc = sub.add_parser("server-config", help="write the server config and print its path")
    sc.add_argument("--fallback-dir", default=None, help="where to write when the config dir is not writable")
    sc.add_argument("--herdr", default=None, help="the real herdr binary, to `config check` the result")
    args = p.parse_args(argv)
    path, warnings = write_server_config(os.environ, args.fallback_dir, args.herdr)
    for w in warnings:
        print(f"WARNING: {w}", file=sys.stderr)
    if not path:
        return 1
    print(path)
    return 0


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