aios/tools/repair_stale_governed_lease.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import ast
import datetime as dt
import pathlib
import shutil
import sys

UTC = dt.timezone.utc

def die(msg: str) -> None:
    print(f"ERROR: {msg}", file=sys.stderr)
    raise SystemExit(2)

def backup(path: pathlib.Path) -> pathlib.Path:
    stamp = dt.datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    out = path.with_suffix(path.suffix + f".bak.{stamp}")
    shutil.copy2(path, out)
    return out

def replace_once(path: pathlib.Path, old: str, new: str, label: str) -> bool:
    text = path.read_text(encoding="utf-8")
    if new in text:
        print(f"OK already patched: {label}")
        return False
    count = text.count(old)
    if count != 1:
        die(f"{label}: expected exactly 1 match, found {count} in {path}")
    backup_path = backup(path)
    path.write_text(text.replace(old, new, 1), encoding="utf-8")
    print(f"PATCHED {label}; backup={backup_path}")
    return True

def ensure_datetime_import(path: pathlib.Path) -> bool:
    text = path.read_text(encoding="utf-8")
    target = "from datetime import datetime, timezone\n"
    if target in text:
        return False
    lines = text.splitlines(keepends=True)
    insert_at = 0
    if lines and lines[0].startswith("#!"):
        insert_at = 1
    if insert_at < len(lines) and lines[insert_at].startswith("from __future__ import"):
        insert_at += 1
    lines.insert(insert_at, target)
    backup_path = backup(path)
    path.write_text("".join(lines), encoding="utf-8")
    print(f"PATCHED datetime import: {path}; backup={backup_path}")
    return True

def syntax_check(path: pathlib.Path) -> None:
    ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
    print(f"OK syntax: {path}")

def main() -> int:
    ap = argparse.ArgumentParser(description="Repair stale governed writer lease handling.")
    ap.add_argument("--app-dir", default="/opt/aios/current/active-runtime/app")
    args = ap.parse_args()
    app = pathlib.Path(args.app_dir)

    rebind = app / "aios_agent_task_execution_rebind.py"
    worker = app / "aios_online_agent_worker.py"
    bridge = app / "aios_chatgpt_bridge_v2_server.py"
    for p in (rebind, worker, bridge):
        if not p.is_file():
            die(f"missing file: {p}")

    old_rebind = '''        if task["lease_token_hash"] not in (None, "") or task["lease_expires_at"] not in (None, ""):
            raise ExecutionRebindError("EXECUTION_REBIND_ACTIVE_LEASE_DENY", 409)
'''
    new_rebind = '''        # A persisted lease token is not by itself proof of an active lease.
        # Treat an unexpired lease as active; stale lease projection is recoverable.
        lease_expires_at = task["lease_expires_at"]
        if lease_expires_at not in (None, ""):
            try:
                parsed_lease_expires_at = datetime.fromisoformat(
                    str(lease_expires_at).replace("Z", "+00:00")
                ).astimezone(timezone.utc)
            except (TypeError, ValueError) as exc:
                raise ExecutionRebindError("EXECUTION_REBIND_LEASE_EXPIRY_INVALID", 409) from exc
            if parsed_lease_expires_at > datetime.now(timezone.utc):
                raise ExecutionRebindError("EXECUTION_REBIND_ACTIVE_LEASE_DENY", 409)
'''

    old_worker = '''    active = [
        item for item in leases
        if isinstance(item, Mapping)
        and item.get("state") == "ACTIVE"
        and item.get("job_id") == envelope.job_id
        and item.get("branch") == envelope.target_branch
    ]
'''
    new_worker = '''    now = datetime.now(timezone.utc)
    active = []
    for item in leases:
        if (
            not isinstance(item, Mapping)
            or item.get("state") != "ACTIVE"
            or item.get("job_id") != envelope.job_id
            or item.get("branch") != envelope.target_branch
        ):
            continue
        try:
            expires_at = datetime.fromisoformat(
                str(item.get("lease_expires_at") or "").replace("Z", "+00:00")
            ).astimezone(timezone.utc)
        except (TypeError, ValueError):
            continue
        if expires_at > now:
            active.append(item)
'''

    old_bridge = '''            if durable_expires_at <= datetime.now(timezone.utc):
                return self.renew_writer_lease(request)
'''
    new_bridge = '''            # An expired durable lease is recoverable on the same governed
            # task/job/generation. Renew it instead of treating stale ACTIVE state
            # as a durable-binding failure.
            if durable_expires_at <= datetime.now(timezone.utc):
                return self.renew_writer_lease(request)
'''

    changed = []
    if ensure_datetime_import(rebind):
        changed.append(rebind)
    if ensure_datetime_import(worker):
        changed.append(worker)
    if replace_once(rebind, old_rebind, new_rebind, "execution-rebind effective lease expiry"):
        changed.append(rebind)
    if replace_once(worker, old_worker, new_worker, "worker converged-job effective lease expiry"):
        changed.append(worker)
    if replace_once(bridge, old_bridge, new_bridge, "bridge expired-lease recovery contract"):
        changed.append(bridge)

    for p in (rebind, worker, bridge):
        syntax_check(p)

    print("RESULT: PATCH_OK")
    print("CHANGED:", ",".join(dict.fromkeys(str(p) for p in changed)) if changed else "none")
    return 0

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