"""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, }