teleo-infrastructure/scripts/collect_leo_restart_survival_proof.py

#!/usr/bin/env python3
"""Collect live Leo restart-survival proof.

This harness intentionally restarts only ``leoclean-gateway.service`` when
``--execute-restart`` is supplied, then proves the gateway came back, canonical
KB counts did not change, and the no-post direct-claim handler still answers.
It never posts to Telegram and never runs production DB apply.
"""

from __future__ import annotations

import argparse
import json
import re
import shlex
import subprocess
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

REPORT_DIR = Path("docs/reports/leo-working-state-20260709")
DEFAULT_OUT = REPORT_DIR / "leo-restart-survival-proof-current.json"
DEFAULT_MARKDOWN_OUT = REPORT_DIR / "leo-restart-survival-proof-current.md"
DEFAULT_SSH_TARGET = "root@77.42.65.182"
DEFAULT_SSH_KEY = Path.home() / ".ssh/livingip_hetzner_20260604_ed25519"
REMOTE_REPO = "/opt/teleo-eval/workspaces/deploy-infra"
REMOTE_HANDLER_REPORT = "/tmp/leo-restart-survival-direct-claim-current.json"
EXPECTED_COUNTS = [
    "public.claims",
    "public.sources",
    "public.claim_evidence",
    "public.claim_edges",
    "kb_stage.kb_proposals",
]

SECRET_PATTERNS = [
    re.compile(r"\b\d{6,}:[A-Za-z0-9_-]{20,}\b"),
    re.compile(r"(Authorization: Bearer )([A-Za-z0-9._-]+)", re.IGNORECASE),
    re.compile(r"((?:api[_-]?key|token|secret|password)[=:]\s*)([A-Za-z0-9._-]+)", re.IGNORECASE),
]


@dataclass(frozen=True)
class CommandResult:
    returncode: int
    stdout: str
    stderr: str


Runner = Callable[[list[str], str, int], CommandResult]


def subprocess_runner(command: list[str], input_text: str, timeout: int) -> CommandResult:
    proc = subprocess.run(command, input=input_text, text=True, capture_output=True, timeout=timeout, check=False)
    return CommandResult(proc.returncode, proc.stdout, proc.stderr)


def redact(text: str) -> str:
    result = text
    for pattern in SECRET_PATTERNS:
        if pattern.groups >= 2:
            result = pattern.sub(r"\1[REDACTED]", result)
        else:
            result = pattern.sub("[REDACTED_TOKEN]", result)
    return result


def parse_last_json_line(text: str) -> dict[str, Any]:
    cleaned = redact(text).strip()
    for line in reversed(cleaned.splitlines()):
        candidate = line.strip()
        if candidate.startswith("{"):
            return json.loads(candidate)
    raise ValueError("command output contained no JSON object line")


def ssh_command(args: argparse.Namespace, remote_command: str) -> list[str]:
    return [
        "ssh",
        "-i",
        str(args.ssh_key),
        "-o",
        "BatchMode=yes",
        "-o",
        "StrictHostKeyChecking=accept-new",
        "-o",
        f"ConnectTimeout={args.connect_timeout}",
        args.ssh_target,
        remote_command,
    ]


def db_counts_sql() -> str:
    return """
select jsonb_build_object(
  'public.claims', (select count(*) from public.claims),
  'public.sources', (select count(*) from public.sources),
  'public.claim_evidence', (select count(*) from public.claim_evidence),
  'public.claim_edges', (select count(*) from public.claim_edges),
  'kb_stage.kb_proposals', (select count(*) from kb_stage.kb_proposals)
)::text;
""".strip()


def remote_readback_command(label: str) -> str:
    script = f"""
import json
import subprocess

def run(command, *, input_text=None):
    return subprocess.run(command, input=input_text, text=True, capture_output=True, check=False)

def kv(text):
    out = {{}}
    for line in text.splitlines():
        if "=" in line:
            key, value = line.split("=", 1)
            out[key] = value.strip()
    return out

service = run([
    "systemctl",
    "show",
    "leoclean-gateway.service",
    "-p", "ActiveState",
    "-p", "SubState",
    "-p", "MainPID",
    "-p", "NRestarts",
    "-p", "ExecMainStartTimestamp",
    "--no-pager",
])
counts = run(
    ["docker", "exec", "-i", "teleo-pg", "psql", "-U", "postgres", "-d", "teleo", "-At", "-q"],
    input_text={db_counts_sql()!r},
)
head = run(["git", "rev-parse", "HEAD"])
stamp = run(["cat", "/opt/teleo-eval/.last-deploy-sha"])
status = run(["git", "status", "--short"])
try:
    parsed_counts = json.loads(counts.stdout.strip()) if counts.returncode == 0 and counts.stdout.strip() else None
except Exception:
    parsed_counts = None
payload = {{
    "label": {label!r},
    "remote_repo": {REMOTE_REPO!r},
    "deploy_head": head.stdout.strip(),
    "last_deploy_sha": stamp.stdout.strip(),
    "git_status_short": status.stdout,
    "service": {{
        "returncode": service.returncode,
        "stderr": service.stderr,
        "state": kv(service.stdout),
    }},
    "db_counts": {{
        "returncode": counts.returncode,
        "stderr": counts.stderr,
        "counts": parsed_counts,
        "raw": counts.stdout.strip(),
    }},
}}
print(json.dumps(payload, sort_keys=True))
""".strip()
    return f"cd {shlex.quote(REMOTE_REPO)} && python3 - <<'PY'\n{script}\nPY"


def restart_command() -> str:
    return "systemctl restart leoclean-gateway.service"


def handler_command(args: argparse.Namespace) -> str:
    return (
        f"cd {shlex.quote(REMOTE_REPO)} && "
        "sudo -u teleo -H /home/teleo/.hermes/hermes-agent/venv/bin/python "
        "scripts/run_leo_direct_claim_handler_remote.py "
        f"--turn-timeout {int(args.handler_turn_timeout)} "
        f"--remote-report-path {shlex.quote(args.remote_handler_report)}"
    )


def run_remote_json(args: argparse.Namespace, runner: Runner, remote_command: str, timeout: int) -> dict[str, Any]:
    proc = runner(ssh_command(args, remote_command), "", timeout)
    if proc.returncode != 0:
        raise RuntimeError(redact(proc.stderr.strip() or proc.stdout.strip() or f"remote command exited {proc.returncode}"))
    return parse_last_json_line(proc.stdout)


def is_active_running(readback: dict[str, Any]) -> bool:
    state = ((readback.get("service") or {}).get("state") or {})
    return state.get("ActiveState") == "active" and state.get("SubState") == "running"


def counts(readback: dict[str, Any]) -> dict[str, Any] | None:
    value = ((readback.get("db_counts") or {}).get("counts"))
    if isinstance(value, dict):
        return value
    return None


def count_keys_present(value: dict[str, Any] | None) -> bool:
    return isinstance(value, dict) and all(key in value for key in EXPECTED_COUNTS)


def wait_for_active(args: argparse.Namespace, runner: Runner) -> dict[str, Any]:
    deadline = time.time() + args.active_timeout
    last: dict[str, Any] | None = None
    while time.time() <= deadline:
        last = run_remote_json(args, runner, remote_readback_command("after_restart_poll"), args.timeout)
        if is_active_running(last):
            return last
        time.sleep(args.poll_interval)
    if last is None:
        raise RuntimeError("service did not produce a post-restart readback")
    return last


def write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def markdown_report(report: dict[str, Any]) -> str:
    before_state = (((report.get("before") or {}).get("service") or {}).get("state") or {})
    after_state = (((report.get("after") or {}).get("service") or {}).get("state") or {})
    handler = report.get("handler") or {}
    lines = [
        "# Leo Restart-Survival Proof",
        "",
        f"Generated UTC: `{report.get('generated_at_utc', '')}`",
        f"Mode: `{report.get('mode', '')}`",
        f"Status: `{report.get('status', '')}`",
        f"Survives restart: `{report.get('survives_restart')}`",
        "",
        "## Safety",
        "",
        f"- Posted to Telegram: `{report.get('posted_to_telegram')}`",
        f"- Production DB apply ran: `{report.get('production_db_apply_ran')}`",
        f"- DB counts changed: `{report.get('db_counts_changed')}`",
        f"- Live profile changed by handler: `{report.get('live_profile_changed_by_handler')}`",
        "",
        "## Service",
        "",
        f"- Before: `ActiveState={before_state.get('ActiveState', '')}`, `SubState={before_state.get('SubState', '')}`, `MainPID={before_state.get('MainPID', '')}`, `NRestarts={before_state.get('NRestarts', '')}`, `ExecMainStartTimestamp={before_state.get('ExecMainStartTimestamp', '')}`",
        f"- After: `ActiveState={after_state.get('ActiveState', '')}`, `SubState={after_state.get('SubState', '')}`, `MainPID={after_state.get('MainPID', '')}`, `NRestarts={after_state.get('NRestarts', '')}`, `ExecMainStartTimestamp={after_state.get('ExecMainStartTimestamp', '')}`",
        "",
        "## Handler Smoke",
        "",
        f"- Handler skipped: `{report.get('handler_skipped')}`",
        f"- Handler pass runtime: `{handler.get('pass_runtime')}`",
        f"- Handler results: `{len(handler.get('results') or [])}`",
        f"- Temp profile removed: `{handler.get('temp_profile_removed')}`",
        "",
        "## Claim Ceiling",
        "",
        "This proves restart-cycle service recovery plus DB-count stability and no-post handler behavior only when `survives_restart=true`. It does not prove Telegram-visible message delivery, production DB apply, or GCP parity.",
        "",
    ]
    return "\n".join(lines)


def collect(args: argparse.Namespace, runner: Runner = subprocess_runner) -> dict[str, Any]:
    report: dict[str, Any] = {
        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
        "mode": "live_vps_leoclean_gateway_restart_survival",
        "ssh_target": args.ssh_target,
        "remote_repo": REMOTE_REPO,
        "remote_handler_report": args.remote_handler_report,
        "posted_to_telegram": False,
        "production_db_apply_ran": False,
        "execute_restart_requested": args.execute_restart,
        "handler_skipped": args.skip_handler,
    }
    try:
        before = run_remote_json(args, runner, remote_readback_command("before_restart"), args.timeout)
        report["before"] = before
        if not args.execute_restart:
            report["status"] = "not_executed_missing_execute_restart"
            report["survives_restart"] = False
            return report

        restart = runner(ssh_command(args, restart_command()), "", args.timeout)
        report["restart_command"] = {
            "returncode": restart.returncode,
            "stdout": redact(restart.stdout),
            "stderr": redact(restart.stderr),
        }
        if restart.returncode != 0:
            report["status"] = "failed_restart_command"
            report["survives_restart"] = False
            return report

        after = wait_for_active(args, runner)
        report["after"] = after
        before_counts = counts(before)
        after_counts = counts(after)
        report["db_counts_changed"] = before_counts != after_counts
        report["db_counts_keys_present"] = count_keys_present(before_counts) and count_keys_present(after_counts)

        handler: dict[str, Any] = {}
        if not args.skip_handler:
            handler = run_remote_json(args, runner, handler_command(args), args.handler_timeout)
        report["handler"] = handler
        report["live_profile_changed_by_handler"] = bool(handler.get("changed_live_profile"))
        report["temp_profile_removed"] = handler.get("temp_profile_removed")

        handler_ok = args.skip_handler or bool(handler.get("pass_runtime"))
        handler_safe = args.skip_handler or (
            handler.get("posted_to_telegram") is False
            and handler.get("mutates_kb_by_harness") is False
            and handler.get("temp_profile_removed") is True
        )
        report["survives_restart"] = bool(
            is_active_running(after)
            and report["db_counts_keys_present"]
            and report["db_counts_changed"] is False
            and handler_ok
            and handler_safe
        )
        report["status"] = "pass" if report["survives_restart"] else "fail"
        return report
    except Exception as exc:
        report["status"] = "error"
        report["survives_restart"] = False
        report["error"] = str(exc)
        return report


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--execute-restart", action="store_true", help="Actually restart leoclean-gateway.service")
    parser.add_argument("--skip-handler", action="store_true", help="Only prove restart and DB-count stability")
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--markdown-out", type=Path, default=DEFAULT_MARKDOWN_OUT)
    parser.add_argument("--ssh-target", default=DEFAULT_SSH_TARGET)
    parser.add_argument("--ssh-key", type=Path, default=DEFAULT_SSH_KEY)
    parser.add_argument("--connect-timeout", type=int, default=10)
    parser.add_argument("--timeout", type=int, default=60)
    parser.add_argument("--active-timeout", type=int, default=90)
    parser.add_argument("--poll-interval", type=float, default=3)
    parser.add_argument("--handler-timeout", type=int, default=2100)
    parser.add_argument("--handler-turn-timeout", type=int, default=300)
    parser.add_argument("--remote-handler-report", default=REMOTE_HANDLER_REPORT)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    report = collect(args)
    write_json(args.out, report)
    args.markdown_out.parent.mkdir(parents=True, exist_ok=True)
    args.markdown_out.write_text(markdown_report(report), encoding="utf-8")
    print(
        json.dumps(
            {
                "out": str(args.out),
                "markdown_out": str(args.markdown_out),
                "status": report.get("status"),
                "survives_restart": report.get("survives_restart"),
                "db_counts_changed": report.get("db_counts_changed"),
                "posted_to_telegram": report.get("posted_to_telegram"),
                "production_db_apply_ran": report.get("production_db_apply_ran"),
            },
            indent=2,
            sort_keys=True,
        )
    )
    if report.get("status") == "not_executed_missing_execute_restart":
        return 2
    return 0 if report.get("survives_restart") else 1


if __name__ == "__main__":
    raise SystemExit(main())