zerto-ai-rewind/src/zerto_rewind_mcp/checkpoints.py

"""Tagged checkpoint insert and wait-until-listed."""

from __future__ import annotations

import asyncio
import re
from datetime import UTC, datetime
from typing import Any

from zerto_rewind_mcp.client import ZertoClient, ZertoError
from zerto_rewind_mcp.protection import FindResult
from zerto_rewind_mcp.tasks import task_id_from, wait_for_task
from zerto_rewind_mcp.util import pick

_CTRL = re.compile(r"[\x00-\x1f\x7f]+")
_WS = re.compile(r"\s+")

# CheckpointName is the only field POST /v1/vpgs/{id}/checkpoints accepts
# (VpgInsertTagCheckpointDataApi has exactly one property). Who the agent is
# and what it is about to do therefore have to live inside the name itself.
# Measured on 10.x: names of at least 400 chars are accepted, and spaces,
# slashes, parentheses, '=' and '|' all survive. We still cap at TAG_MAX_LEN
# so the name stays readable in the Zerto UI checkpoint list.
TAG_MAX_LEN = 250


def _clean(value: Any, limit: int) -> str:
    """One field of a checkpoint name: single-line, no separator collisions."""
    text = _CTRL.sub(" ", str(value or ""))
    # Zerto appends "; Used for File Level Restore" to its own tags, and we use
    # " | " as our field separator. Keep both out of user-supplied text.
    text = text.replace(";", ",").replace("|", "/")
    return _WS.sub(" ", text).strip()[:limit]


def make_tag(
    agent: str,
    change_id: str,
    when: datetime | None = None,
    *,
    action: str | None = None,
    vm_name: str | None = None,
) -> str:
    """Build the checkpoint name.

    Says which agent, what it is about to do, to which VM, under what change id:

        ai:claude | edit /etc/app.conf | vm=jp-ubuntu | change=chg-99 | 20260921T150405Z

    action is free text from the caller describing the pending mutation.
    """
    stamp = (when or datetime.now(UTC)).strftime("%Y%m%dT%H%M%SZ")
    parts = [f"ai:{_clean(agent, 40) or 'agent'}"]
    cleaned_action = _clean(action, 120)
    if cleaned_action:
        parts.append(cleaned_action)
    cleaned_vm = _clean(vm_name, 60)
    if cleaned_vm:
        parts.append(f"vm={cleaned_vm}")
    parts.append(f"change={_clean(change_id, 80) or 'change'}")
    parts.append(stamp)
    return " | ".join(parts)[:TAG_MAX_LEN]


def checkpoint_tag(row: dict[str, Any]) -> str:
    value = pick(row, "Tag", "tag", "CheckpointName", "checkpointName") or ""
    return str(value)


def checkpoint_id(row: dict[str, Any]) -> str:
    value = pick(
        row, "CheckpointId", "checkpointId", "checkpointIdentifier", "CheckpointIdentifier"
    )
    return str(value) if value is not None else ""


async def wait_for_tag(
    client: ZertoClient,
    vpg_identifier: str,
    tag: str,
    *,
    timeout_s: float = 45.0,
    interval_s: float = 1.5,
) -> dict[str, Any]:
    deadline = asyncio.get_event_loop().time() + timeout_s
    last: list[dict[str, Any]] = []
    while asyncio.get_event_loop().time() < deadline:
        last = await client.list_checkpoints(vpg_identifier)
        for row in last:
            if checkpoint_tag(row) == tag:
                return row
        await asyncio.sleep(interval_s)
    raise ZertoError(
        f"Tagged checkpoint {tag!r} did not appear on VPG {vpg_identifier} "
        f"within {timeout_s:.0f}s. Do not mutate. "
        "If the protected site is Azure or AWS, tagged checkpoints are not supported."
    )


async def tag_vpgs(
    client: ZertoClient,
    result: FindResult,
    tag: str,
) -> dict[str, Any]:
    if result.outcome != "ok" or result.vm is None:
        return {
            "ok": False,
            "tag": tag,
            "message": result.message,
            "find": result.as_dict(),
        }
    # Insert then wait for the TASK, one VPG at a time. The POST returns a task
    # id, not a result: 200 only means queued. Measured on 10.x, a second insert
    # fired at the same VPG while the first is running comes back Failed (state
    # 4) from GET /v1/tasks/{id} while the first is Completed (6). That failure
    # is reported, not silent -- but only if you read the task. Do not turn this
    # loop into an asyncio.gather().
    tagged: list[dict[str, Any]] = []
    skipped = [v.as_dict() for v in result.vm.vpgs if not v.can_tag]
    errors: list[str] = []
    for vpg in result.taggable_vpgs:
        try:
            insert = await client.insert_checkpoint(vpg.vpg_identifier, tag)
            task_id = task_id_from(insert)
            task = await wait_for_task(client, task_id)
            row = await wait_for_tag(client, vpg.vpg_identifier, tag)
            tagged.append(
                {
                    "vpg_identifier": vpg.vpg_identifier,
                    "vpg_name": vpg.vpg_name,
                    "checkpoint_id": checkpoint_id(row),
                    "tag": checkpoint_tag(row) or tag,
                    "task_id": task_id,
                    "task_state": task["label"],
                }
            )
        except ZertoError as exc:
            errors.append(f"{vpg.vpg_name} ({vpg.vpg_identifier}): {exc}")
    ok = bool(tagged) and not errors
    message = (
        f"Tagged {len(tagged)} VPG(s) with {tag!r}."
        if ok
        else f"Checkpoint failed on {len(errors)} VPG(s). Refuse the change. " + "; ".join(errors)
    )
    if skipped and ok:
        message += " Some VPGs were skipped (syncing or not Protecting)."
    return {
        "ok": ok,
        "tag": tag,
        "vm": result.vm.as_dict(),
        "tagged": tagged,
        "skipped": skipped,
        "errors": errors,
        "message": message,
    }