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.
  * The merge edits text (stdlib has no TOML writer) and is then re-parsed: the
    key must read false and everything else must equal the source. A merge that
    fails the check falls back to a file holding only the key - herdr would
    drop a broken file whole and restore again, silently.

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

import argparse
import os
import re
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"

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

_SESSION_HEADER = re.compile(r"^\s*\[\s*session\s*\]\s*(#.*)?$")
_ANY_HEADER = re.compile(r"^\s*\[")
_KEY_LINE = re.compile(rf"^\s*{KEY}\s*=")


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."""
    try:
        with open(path, encoding="utf-8") as f:
            if f.readline().rstrip("\n") != MARKER:
                return None
            line = f.readline().rstrip("\n")
    except OSError:
        return None
    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")


def _without_key(d: dict) -> dict:
    d = dict(d)
    session = dict(d.get("session") or {})
    session.pop(KEY, None)
    if session:
        d["session"] = session
    else:
        d.pop("session", None)
    return d


def _merge_text(user_text: str) -> str:
    lines = user_text.splitlines()
    header = next((i for i, l in enumerate(lines) if _SESSION_HEADER.match(l)), None)
    if header is None:
        sep = "" if user_text.endswith("\n") else "\n"
        return f"{user_text}{sep}\n{MINIMAL}"
    end = next((i for i in range(header + 1, len(lines)) if _ANY_HEADER.match(lines[i])), len(lines))
    key_line = next((i for i in range(header + 1, end) if _KEY_LINE.match(lines[i])), None)
    if key_line is None:
        lines.insert(header + 1, f"{KEY} = false")
    else:
        lines[key_line] = f"{KEY} = false"
    return "\n".join(lines) + "\n"


def server_config_text(user_text: Optional[str]) -> Tuple[str, List[str]]:
    """(the TOML body for the server, warnings). The body always parses and
    always has session.resume_agents_on_restore = false."""
    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+); the MAM server uses only {KEY} = false, "
                         "without the rest of your settings"]
    try:
        user = tomllib.loads(user_text)
    except tomllib.TOMLDecodeError as e:
        return MINIMAL, [f"herdr config does not parse ({e}); the MAM server uses only {KEY} = false"]
    if (user.get("session") or {}).get(KEY) is False:
        return user_text, []
    merged = _merge_text(user_text)
    try:
        result = tomllib.loads(merged)
        ok = result["session"][KEY] is False and _without_key(result) == _without_key(user)
    except (tomllib.TOMLDecodeError, KeyError, TypeError):
        ok = False
    if not ok:
        return MINIMAL, [f"could not set [session] {KEY} in the herdr config (dotted key or inline table?); "
                         f"the MAM server uses only {KEY} = false, without the rest of your settings"]
    return merged, []


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_server_config(env: Mapping[str, str], fallback_dir: 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}")
    body, merge_warnings = server_config_text(user_text)
    warnings += merge_warnings
    text = f"{MARKER}\n{SOURCE_PREFIX}{source}\n{body}"
    target = os.path.join(os.path.dirname(source), OUTPUT_NAME)
    try:
        _write_atomic(target, text, mode)
        return target, warnings
    except OSError as e:
        warnings.append(f"cannot write {target}: {e}")
    if fallback_dir:
        target = os.path.join(fallback_dir, FALLBACK_NAME)
        try:
            _write_atomic(target, text, mode)
            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")
    args = p.parse_args(argv)
    path, warnings = write_server_config(os.environ, args.fallback_dir)
    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())