feat(guard): read the Zerto task, and ask before guarding unknown tools #6
+6
-2
@@ -34,8 +34,12 @@ A VPG whose status is MeetingSLA or a NotMeetingSLA variant, and whose substatus
|
|||||||
_Avoid_: healthy, in sync, Protecting (as status 0)
|
_Avoid_: healthy, in sync, Protecting (as status 0)
|
||||||
|
|
||||||
**Mutating catalog**:
|
**Mutating catalog**:
|
||||||
The opt-in list of MCP tools that must call `zerto_guard_before_mutate` first. Unlisted tools pass through. Users add entries; the starter list is not the whole world.
|
The opt-in list of MCP tools that must call `zerto_guard_before_mutate` first. Paired with `read_only_tools`, the list known not to change a guest. A tool in neither is **unknown**, which is the normal case: the agent asks the human whether to checkpoint rather than assuming either way.
|
||||||
_Avoid_: denylist, hold-everything
|
_Avoid_: denylist, hold-everything, treating unknown as safe
|
||||||
|
|
||||||
|
**Zerto task**:
|
||||||
|
Write operations return a task id, not a result. A 200 means queued. `GET /v1/tasks/{id}` carries the real outcome in `Status.State`: 1 InProgress, 4 Failed, 5 Stopped, 6 Completed (terminal is 4/5/6). A second tagged-checkpoint insert fired at a VPG while the first runs comes back Failed, which is only visible if the task is read.
|
||||||
|
_Avoid_: treating HTTP 200 as success
|
||||||
|
|
||||||
**Official ZVM MCP**:
|
**Official ZVM MCP**:
|
||||||
HPE Zerto 10.9 MCP (`ZVM.MCP`): inventory, VPG settings, failover test. Not in the demo path. This PoC is one server.
|
HPE Zerto 10.9 MCP (`ZVM.MCP`): inventory, VPG settings, failover test. Not in the demo path. This PoC is one server.
|
||||||
|
|||||||
@@ -6,10 +6,11 @@ If the loop works, these tools are the delta to add to official ZVM MCP (`ZVM.MC
|
|||||||
|
|
||||||
## What it does
|
## What it does
|
||||||
|
|
||||||
|
0. `zerto_check_tool` — before running anything against a VM: `read_only` (go), `mutating` (guard first), or `unknown` (**ask the human** whether to checkpoint).
|
||||||
1. `zerto_find_protection` — VM name, hostname, or vmIdentifier to exactly one VM and every VPG. Zero or two-plus VMs: stop.
|
1. `zerto_find_protection` — VM name, hostname, or vmIdentifier to exactly one VM and every VPG. Zero or two-plus VMs: stop.
|
||||||
2. `zerto_create_tagged_checkpoint` / `zerto_guard_before_mutate` — same tag on every protecting VPG, wait until listed. The name records which agent and what it is doing: `ai:<agent> | <action> | vm=<vm> | change=<id> | <utc>`.
|
2. `zerto_create_tagged_checkpoint` / `zerto_guard_before_mutate` — same tag on every protecting VPG, wait until listed. The name records which agent and what it is doing: `ai:<agent> | <action> | vm=<vm> | change=<id> | <utc>`.
|
||||||
3. `zerto_recover_file` — FLR after a human sets `confirmed=true`. Linux and Windows guest paths. Locally replicated VPGs only: FLR runs at the VPG's recovery site. Reports its own unmount; `zerto_list_flr_sessions` / `zerto_end_flr_session` find and reap a mount orphaned by a crashed recovery.
|
3. `zerto_recover_file` — FLR after a human sets `confirmed=true`. Linux and Windows guest paths. Locally replicated VPGs only: FLR runs at the VPG's recovery site. Reports its own unmount; `zerto_list_flr_sessions` / `zerto_end_flr_session` find and reap a mount orphaned by a crashed recovery.
|
||||||
4. Mutating catalog — opt-in list of MCP tools that must be guarded. Unlisted tools pass through. Users add entries. The starter list covers Linux (`ssh`, `ansible`) and Windows (`winrm`, `powershell`, `smb`), and is illustrative, not exhaustive.
|
4. Two catalogs — `mutating_tools` (guard first) and `read_only_tools` (safe). Both cover Linux (`ssh`, `ansible`) and Windows (`winrm`, `powershell`, `smb`), and both are illustrative, not exhaustive. A tool in neither is **unknown, not safe**: ask the human.
|
||||||
|
|
||||||
Official ZVM MCP already has inventory and failover test. It does not insert tagged checkpoints or run FLR.
|
Official ZVM MCP already has inventory and failover test. It does not insert tagged checkpoints or run FLR.
|
||||||
|
|
||||||
|
|||||||
@@ -42,5 +42,27 @@
|
|||||||
"vm_arg": "host",
|
"vm_arg": "host",
|
||||||
"notes": "Writes a file into a Windows share. Changes the guest without a shell."
|
"notes": "Writes a file into a Windows share. Changes the guest without a shell."
|
||||||
}
|
}
|
||||||
|
],
|
||||||
|
"read_only_tools": [
|
||||||
|
{
|
||||||
|
"server": "ssh",
|
||||||
|
"tool": "read_file",
|
||||||
|
"notes": "Reads a file from the guest. Does not change it."
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"server": "ssh",
|
||||||
|
"tool": "stat",
|
||||||
|
"notes": "File metadata only."
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"server": "winrm",
|
||||||
|
"tool": "get_item",
|
||||||
|
"notes": "Windows read of a file or registry value."
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"server": "powershell",
|
||||||
|
"tool": "get_content",
|
||||||
|
"notes": "Reads file contents on a Windows guest."
|
||||||
|
}
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,16 +11,33 @@ You talk to **one** MCP: `zerto_rewind_mcp`. Do not also require official ZVM MC
|
|||||||
|
|
||||||
## Loop (mandatory)
|
## Loop (mandatory)
|
||||||
|
|
||||||
Before **every** guest-mutating tool call:
|
Before **every** tool call that might touch a guest:
|
||||||
|
|
||||||
1. Take the hostname / VM name / Zerto `vmIdentifier` from the tool args.
|
1. Take the hostname / VM name / Zerto `vmIdentifier` from the tool args.
|
||||||
2. Call `zerto_guard_before_mutate` with `change_id` and `action` (or `zerto_find_protection` then `zerto_create_tagged_checkpoint`).
|
2. Call `zerto_check_tool(server, tool, vm)`. Act on the verdict:
|
||||||
3. If `ok` is not true: **stop**. Do not mutate.
|
|
||||||
4. Then run the mutating call.
|
|
||||||
|
|
||||||
Reads skip the guard.
|
| verdict | what you do |
|
||||||
|
|---|---|
|
||||||
|
| `read_only` | Run the tool. No checkpoint. |
|
||||||
|
| `mutating` | Guard first. Do not ask — it is already known to change the guest. |
|
||||||
|
| `unknown` | **Ask the human.** Do not assume it is safe, and do not silently guard. |
|
||||||
|
|
||||||
Unlisted MCP tools pass through. If you are about to change a protected VM with a tool that is not in the catalog, call `zerto_add_mutating_tool` (server, tool, `vm_arg`) and then guard.
|
3. For `mutating`, or for `unknown` where the human said yes: call
|
||||||
|
`zerto_guard_before_mutate` with `change_id` and `action`.
|
||||||
|
4. If `ok` is not true: **stop**. Do not mutate.
|
||||||
|
5. Then run the call.
|
||||||
|
|
||||||
|
`unknown` is the normal case, not an edge case. The catalogs are short and the
|
||||||
|
world of tools is not, so most tools are unclassified. Unknown means *nobody has
|
||||||
|
said this is read-only* — it does not mean safe. Put the decision to the human:
|
||||||
|
|
||||||
|
> `winrm/run_ps` is not a known read-only command. It may change `web01`.
|
||||||
|
> Insert a Zerto tagged checkpoint first so this is reversible?
|
||||||
|
|
||||||
|
If they say yes, guard, then run. If they say no, run it and tell them plainly
|
||||||
|
that it is not reversible through Zerto. If they want it remembered, call
|
||||||
|
`zerto_add_mutating_tool` (server, tool, `vm_arg`) so it is guarded
|
||||||
|
automatically next time instead of asking again.
|
||||||
|
|
||||||
## find_protection outcomes
|
## find_protection outcomes
|
||||||
|
|
||||||
|
|||||||
@@ -36,12 +36,42 @@ def entry_from_dict(raw: dict[str, Any]) -> CatalogEntry:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
MUTATING = "mutating"
|
||||||
|
READ_ONLY = "read_only"
|
||||||
|
UNKNOWN = "unknown"
|
||||||
|
|
||||||
|
|
||||||
class MutatingCatalog:
|
class MutatingCatalog:
|
||||||
def __init__(self, entries: list[CatalogEntry] | None = None, path: Path | None = None):
|
def __init__(
|
||||||
|
self,
|
||||||
|
entries: list[CatalogEntry] | None = None,
|
||||||
|
path: Path | None = None,
|
||||||
|
read_only: list[CatalogEntry] | None = None,
|
||||||
|
):
|
||||||
self._entries: dict[tuple[str, str], CatalogEntry] = {}
|
self._entries: dict[tuple[str, str], CatalogEntry] = {}
|
||||||
|
self._read_only: dict[tuple[str, str], CatalogEntry] = {}
|
||||||
self.path = path
|
self.path = path
|
||||||
for entry in entries or []:
|
for entry in entries or []:
|
||||||
self._entries[entry.key()] = entry
|
self._entries[entry.key()] = entry
|
||||||
|
for entry in read_only or []:
|
||||||
|
self._read_only[entry.key()] = entry
|
||||||
|
|
||||||
|
def read_only(self) -> list[CatalogEntry]:
|
||||||
|
return sorted(self._read_only.values(), key=lambda e: e.key())
|
||||||
|
|
||||||
|
def classify(self, server: str, tool: str) -> str:
|
||||||
|
"""mutating / read_only / unknown.
|
||||||
|
|
||||||
|
unknown is the important one. It does not mean safe: it means nobody has
|
||||||
|
said this tool is read-only, so the agent must ask a human before
|
||||||
|
changing a protected VM with it.
|
||||||
|
"""
|
||||||
|
key = (server.lower(), tool.lower())
|
||||||
|
if key in self._entries:
|
||||||
|
return MUTATING
|
||||||
|
if key in self._read_only:
|
||||||
|
return READ_ONLY
|
||||||
|
return UNKNOWN
|
||||||
|
|
||||||
def list(self) -> list[CatalogEntry]:
|
def list(self) -> list[CatalogEntry]:
|
||||||
return sorted(self._entries.values(), key=lambda e: e.key())
|
return sorted(self._entries.values(), key=lambda e: e.key())
|
||||||
@@ -68,6 +98,9 @@ class MutatingCatalog:
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, data: dict[str, Any], path: Path | None = None) -> MutatingCatalog:
|
def from_config(cls, data: dict[str, Any], path: Path | None = None) -> MutatingCatalog:
|
||||||
raw = data.get("mutating_tools") or []
|
entries = [entry_from_dict(item) for item in data.get("mutating_tools") or []]
|
||||||
entries = [entry_from_dict(item) for item in raw]
|
read_only = [
|
||||||
return cls(entries, path=path)
|
entry_from_dict({**item, "vm_arg": item.get("vm_arg") or "-"})
|
||||||
|
for item in data.get("read_only_tools") or []
|
||||||
|
]
|
||||||
|
return cls(entries, path=path, read_only=read_only)
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ from typing import Any
|
|||||||
|
|
||||||
from zerto_rewind_mcp.client import ZertoClient, ZertoError
|
from zerto_rewind_mcp.client import ZertoClient, ZertoError
|
||||||
from zerto_rewind_mcp.protection import FindResult
|
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
|
from zerto_rewind_mcp.util import pick
|
||||||
|
|
||||||
_CTRL = re.compile(r"[\x00-\x1f\x7f]+")
|
_CTRL = re.compile(r"[\x00-\x1f\x7f]+")
|
||||||
@@ -108,16 +109,20 @@ async def tag_vpgs(
|
|||||||
"message": result.message,
|
"message": result.message,
|
||||||
"find": result.as_dict(),
|
"find": result.as_dict(),
|
||||||
}
|
}
|
||||||
# Insert then wait, one VPG at a time. Measured on 10.x: tagged checkpoint
|
# Insert then wait for the TASK, one VPG at a time. The POST returns a task
|
||||||
# inserts fired back to back at the same VPG are silently dropped -- the POST
|
# id, not a result: 200 only means queued. Measured on 10.x, a second insert
|
||||||
# returns 200 and queues a task, but only the first checkpoint ever appears.
|
# fired at the same VPG while the first is running comes back Failed (state
|
||||||
# Do not turn this loop into an asyncio.gather().
|
# 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]] = []
|
tagged: list[dict[str, Any]] = []
|
||||||
skipped = [v.as_dict() for v in result.vm.vpgs if not v.can_tag]
|
skipped = [v.as_dict() for v in result.vm.vpgs if not v.can_tag]
|
||||||
errors: list[str] = []
|
errors: list[str] = []
|
||||||
for vpg in result.taggable_vpgs:
|
for vpg in result.taggable_vpgs:
|
||||||
try:
|
try:
|
||||||
insert = await client.insert_checkpoint(vpg.vpg_identifier, tag)
|
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)
|
row = await wait_for_tag(client, vpg.vpg_identifier, tag)
|
||||||
tagged.append(
|
tagged.append(
|
||||||
{
|
{
|
||||||
@@ -125,7 +130,8 @@ async def tag_vpgs(
|
|||||||
"vpg_name": vpg.vpg_name,
|
"vpg_name": vpg.vpg_name,
|
||||||
"checkpoint_id": checkpoint_id(row),
|
"checkpoint_id": checkpoint_id(row),
|
||||||
"tag": checkpoint_tag(row) or tag,
|
"tag": checkpoint_tag(row) or tag,
|
||||||
"insert_result": insert if not isinstance(insert, dict) else "ok",
|
"task_id": task_id,
|
||||||
|
"task_state": task["label"],
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
except ZertoError as exc:
|
except ZertoError as exc:
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ from typing import Any
|
|||||||
|
|
||||||
from mcp.server.fastmcp import FastMCP
|
from mcp.server.fastmcp import FastMCP
|
||||||
|
|
||||||
from zerto_rewind_mcp.catalog import CatalogEntry, MutatingCatalog
|
from zerto_rewind_mcp.catalog import MUTATING, READ_ONLY, CatalogEntry, MutatingCatalog
|
||||||
from zerto_rewind_mcp.checkpoints import checkpoint_id, checkpoint_tag, make_tag, tag_vpgs
|
from zerto_rewind_mcp.checkpoints import checkpoint_id, checkpoint_tag, make_tag, tag_vpgs
|
||||||
from zerto_rewind_mcp.client import ZertoClient, ZertoError
|
from zerto_rewind_mcp.client import ZertoClient, ZertoError
|
||||||
from zerto_rewind_mcp.config import load_catalog, load_config
|
from zerto_rewind_mcp.config import load_catalog, load_config
|
||||||
@@ -219,6 +219,69 @@ async def zerto_guard_before_mutate(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mcp.tool(
|
||||||
|
name="zerto_check_tool",
|
||||||
|
annotations={
|
||||||
|
"title": "Is this tool safe to run on a protected VM?",
|
||||||
|
"readOnlyHint": True,
|
||||||
|
"destructiveHint": False,
|
||||||
|
"idempotentHint": True,
|
||||||
|
"openWorldHint": False,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
async def zerto_check_tool(server: str, tool: str, vm: str = "") -> str:
|
||||||
|
"""Classify a tool before you run it against a VM. Call this first.
|
||||||
|
|
||||||
|
Returns one of:
|
||||||
|
|
||||||
|
read_only - known not to change the guest. Run it. No checkpoint.
|
||||||
|
mutating - known to change the guest. Guard first, do not ask.
|
||||||
|
unknown - NOT known to be read-only. Do not assume it is safe and do not
|
||||||
|
silently guard either. ASK THE HUMAN whether to insert a
|
||||||
|
tagged checkpoint before running it, then do what they say.
|
||||||
|
|
||||||
|
unknown is the normal case for any tool nobody has classified yet, which is
|
||||||
|
most of them. The catalogs are short; the world of tools is not.
|
||||||
|
"""
|
||||||
|
verdict = get_catalog().classify(server, tool)
|
||||||
|
payload: dict[str, Any] = {
|
||||||
|
"ok": True,
|
||||||
|
"server": server,
|
||||||
|
"tool": tool,
|
||||||
|
"vm": vm or None,
|
||||||
|
"verdict": verdict,
|
||||||
|
}
|
||||||
|
target = vm or "<the VM this tool targets>"
|
||||||
|
if verdict == READ_ONLY:
|
||||||
|
payload["checkpoint_required"] = False
|
||||||
|
payload["next"] = "Run the tool. Reads do not need a checkpoint."
|
||||||
|
elif verdict == MUTATING:
|
||||||
|
payload["checkpoint_required"] = True
|
||||||
|
payload["next"] = (
|
||||||
|
f"Call zerto_guard_before_mutate(query={target!r}, change_id=..., "
|
||||||
|
"action=<what you are about to do>) and only run the tool if it "
|
||||||
|
"returns ok true."
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
payload["checkpoint_required"] = None
|
||||||
|
payload["ask_user"] = True
|
||||||
|
payload["question"] = (
|
||||||
|
f"{server}/{tool} is not a known read-only command. It may change "
|
||||||
|
f"{target}. Insert a Zerto tagged checkpoint first so this is "
|
||||||
|
"reversible?"
|
||||||
|
)
|
||||||
|
payload["next"] = (
|
||||||
|
"Ask the human the question above and wait for an answer. Do not "
|
||||||
|
"guess. If they say yes, call zerto_guard_before_mutate("
|
||||||
|
f"query={target!r}, change_id=..., action=<what you are about to "
|
||||||
|
"do>) and run the tool only if it returns ok true. If they say no, "
|
||||||
|
"run the tool without a checkpoint and tell them it is not "
|
||||||
|
"reversible through Zerto. If they want this remembered, call "
|
||||||
|
"zerto_add_mutating_tool so it is guarded automatically next time."
|
||||||
|
)
|
||||||
|
return _dump(payload)
|
||||||
|
|
||||||
|
|
||||||
@mcp.tool(
|
@mcp.tool(
|
||||||
name="zerto_list_mutating_tools",
|
name="zerto_list_mutating_tools",
|
||||||
annotations={
|
annotations={
|
||||||
@@ -230,9 +293,22 @@ async def zerto_guard_before_mutate(
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
async def zerto_list_mutating_tools() -> str:
|
async def zerto_list_mutating_tools() -> str:
|
||||||
"""Tools that must be guarded. Unlisted tools pass through."""
|
"""Both catalogs. A tool in neither is unknown: ask the human before mutating."""
|
||||||
entries = [e.as_dict() for e in get_catalog().list()]
|
cat = get_catalog()
|
||||||
return _dump({"ok": True, "count": len(entries), "mutating_tools": entries})
|
entries = [e.as_dict() for e in cat.list()]
|
||||||
|
read_only = [e.as_dict() for e in cat.read_only()]
|
||||||
|
return _dump(
|
||||||
|
{
|
||||||
|
"ok": True,
|
||||||
|
"count": len(entries),
|
||||||
|
"mutating_tools": entries,
|
||||||
|
"read_only_tools": read_only,
|
||||||
|
"unlisted": (
|
||||||
|
"A tool in neither list is unknown, not safe. Call "
|
||||||
|
"zerto_check_tool and ask the human before changing a protected VM."
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@mcp.tool(
|
@mcp.tool(
|
||||||
|
|||||||
@@ -0,0 +1,90 @@
|
|||||||
|
"""Zerto async task polling.
|
||||||
|
|
||||||
|
Most write operations return a task id, not a result. The POST returning 200
|
||||||
|
only means the task was queued. Whether it actually happened is in
|
||||||
|
GET /v1/tasks/{id}. Skipping that read is how a failed checkpoint insert looks
|
||||||
|
like a success.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from zerto_rewind_mcp.client import ZertoClient, ZertoError
|
||||||
|
|
||||||
|
# Zerto task states (9.0 Tasks API; mirrored by Zerto's published quickstarts).
|
||||||
|
FIRST_UNUSED_VALUE = 0
|
||||||
|
IN_PROGRESS = 1
|
||||||
|
WAITING_FOR_USER_INPUT = 2
|
||||||
|
PAUSED = 3
|
||||||
|
FAILED = 4
|
||||||
|
STOPPED = 5
|
||||||
|
COMPLETED = 6
|
||||||
|
CANCELLING = 7
|
||||||
|
|
||||||
|
TASK_STATE_LABELS: dict[int, str] = {
|
||||||
|
FIRST_UNUSED_VALUE: "FirstUnusedValue",
|
||||||
|
IN_PROGRESS: "InProgress",
|
||||||
|
WAITING_FOR_USER_INPUT: "WaitingForUserInput",
|
||||||
|
PAUSED: "Paused",
|
||||||
|
FAILED: "Failed",
|
||||||
|
STOPPED: "Stopped",
|
||||||
|
COMPLETED: "Completed",
|
||||||
|
CANCELLING: "Cancelling",
|
||||||
|
}
|
||||||
|
TERMINAL_STATES = {FAILED, STOPPED, COMPLETED}
|
||||||
|
|
||||||
|
|
||||||
|
def task_id_from(payload: Any) -> str:
|
||||||
|
"""A write endpoint returns the task id as a bare (often quoted) string."""
|
||||||
|
if isinstance(payload, str):
|
||||||
|
value = payload.strip().strip('"')
|
||||||
|
if value:
|
||||||
|
return value
|
||||||
|
if isinstance(payload, dict):
|
||||||
|
for key in ("TaskIdentifier", "taskIdentifier", "Identifier", "identifier", "id"):
|
||||||
|
if payload.get(key):
|
||||||
|
return str(payload[key]).strip().strip('"')
|
||||||
|
raise ZertoError(f"Could not read a task id from {payload!r}")
|
||||||
|
|
||||||
|
|
||||||
|
def task_state(payload: Any) -> int | None:
|
||||||
|
"""Status is {"State": n, "Progress": n}, but can be a bare int."""
|
||||||
|
if not isinstance(payload, dict):
|
||||||
|
return None
|
||||||
|
status = payload.get("Status")
|
||||||
|
if isinstance(status, dict):
|
||||||
|
status = status.get("State")
|
||||||
|
return status if isinstance(status, bool) is False and isinstance(status, int) else None
|
||||||
|
|
||||||
|
|
||||||
|
def task_label(state: int | None) -> str:
|
||||||
|
return TASK_STATE_LABELS.get(state, f"Unknown({state})") if state is not None else "Unknown"
|
||||||
|
|
||||||
|
|
||||||
|
async def wait_for_task(
|
||||||
|
client: ZertoClient,
|
||||||
|
task_id: str,
|
||||||
|
*,
|
||||||
|
timeout_s: float = 120.0,
|
||||||
|
interval_s: float = 3.0,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Poll until the task reaches a terminal state. Raise unless it Completed."""
|
||||||
|
deadline = asyncio.get_event_loop().time() + timeout_s
|
||||||
|
last: dict[str, Any] = {}
|
||||||
|
while asyncio.get_event_loop().time() < deadline:
|
||||||
|
last = await client.get_task(task_id)
|
||||||
|
state = task_state(last)
|
||||||
|
if state in TERMINAL_STATES:
|
||||||
|
if state == COMPLETED:
|
||||||
|
return {"task_id": task_id, "state": state, "label": task_label(state)}
|
||||||
|
raise ZertoError(
|
||||||
|
f"Zerto task {task_id} ({last.get('Type')}) finished as "
|
||||||
|
f"{task_label(state)}. The operation did not happen."
|
||||||
|
)
|
||||||
|
await asyncio.sleep(interval_s)
|
||||||
|
raise ZertoError(
|
||||||
|
f"Zerto task {task_id} ({last.get('Type')}) still "
|
||||||
|
f"{task_label(task_state(last))} after {timeout_s:.0f}s."
|
||||||
|
)
|
||||||
@@ -50,3 +50,22 @@ def test_example_config_covers_windows_and_linux():
|
|||||||
# every entry must name the arg holding the VM, or the guard cannot resolve one
|
# every entry must name the arg holding the VM, or the guard cannot resolve one
|
||||||
for entry in cat.list():
|
for entry in cat.list():
|
||||||
assert entry.vm_arg, f"{entry.server}/{entry.tool} has no vm_arg"
|
assert entry.vm_arg, f"{entry.server}/{entry.tool} has no vm_arg"
|
||||||
|
|
||||||
|
|
||||||
|
def test_classify_three_way():
|
||||||
|
data = {
|
||||||
|
"mutating_tools": [{"server": "winrm", "tool": "run_ps", "vm_arg": "host"}],
|
||||||
|
"read_only_tools": [{"server": "ssh", "tool": "read_file"}],
|
||||||
|
}
|
||||||
|
cat = MutatingCatalog.from_config(data)
|
||||||
|
assert cat.classify("winrm", "run_ps") == "mutating"
|
||||||
|
assert cat.classify("SSH", "Read_File") == "read_only"
|
||||||
|
# unknown is not safe; it means nobody classified it
|
||||||
|
assert cat.classify("anything", "else") == "unknown"
|
||||||
|
|
||||||
|
|
||||||
|
def test_read_only_entries_need_no_vm_arg():
|
||||||
|
data = {"read_only_tools": [{"server": "ssh", "tool": "stat"}]}
|
||||||
|
cat = MutatingCatalog.from_config(data)
|
||||||
|
assert cat.classify("ssh", "stat") == "read_only"
|
||||||
|
assert [e.tool for e in cat.read_only()] == ["stat"]
|
||||||
|
|||||||
@@ -0,0 +1,69 @@
|
|||||||
|
import asyncio
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from zerto_rewind_mcp.client import ZertoError
|
||||||
|
from zerto_rewind_mcp.tasks import (
|
||||||
|
COMPLETED,
|
||||||
|
FAILED,
|
||||||
|
task_id_from,
|
||||||
|
task_label,
|
||||||
|
task_state,
|
||||||
|
wait_for_task,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_id_from_shapes():
|
||||||
|
# a write endpoint returns the id as a bare, often quoted, string
|
||||||
|
assert task_id_from('"abc.def"') == "abc.def"
|
||||||
|
assert task_id_from("abc.def") == "abc.def"
|
||||||
|
assert task_id_from({"TaskIdentifier": "t1"}) == "t1"
|
||||||
|
with pytest.raises(ZertoError):
|
||||||
|
task_id_from(None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_state_handles_nested_and_bare():
|
||||||
|
assert task_state({"Status": {"State": 6, "Progress": 100}}) == COMPLETED
|
||||||
|
assert task_state({"Status": 4}) == FAILED
|
||||||
|
assert task_state({"Status": None}) is None
|
||||||
|
assert task_state("nope") is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_labels():
|
||||||
|
assert task_label(COMPLETED) == "Completed"
|
||||||
|
assert task_label(FAILED) == "Failed"
|
||||||
|
|
||||||
|
|
||||||
|
class _TaskClient:
|
||||||
|
def __init__(self, states):
|
||||||
|
self.states = list(states)
|
||||||
|
self.calls = 0
|
||||||
|
|
||||||
|
async def get_task(self, task_id):
|
||||||
|
self.calls += 1
|
||||||
|
state = self.states.pop(0) if self.states else self.states
|
||||||
|
return {"Type": "InsertTaggedCP", "Status": {"State": state, "Progress": 100}}
|
||||||
|
|
||||||
|
|
||||||
|
def test_wait_for_task_returns_on_completed():
|
||||||
|
client = _TaskClient([1, 1, COMPLETED])
|
||||||
|
out = asyncio.run(wait_for_task(client, "t1", interval_s=0))
|
||||||
|
assert out["state"] == COMPLETED
|
||||||
|
assert out["label"] == "Completed"
|
||||||
|
assert client.calls == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_wait_for_task_raises_on_failed():
|
||||||
|
# this is the case that used to look like success: 200 queued, task Failed
|
||||||
|
client = _TaskClient([1, FAILED])
|
||||||
|
with pytest.raises(ZertoError) as err:
|
||||||
|
asyncio.run(wait_for_task(client, "t1", interval_s=0))
|
||||||
|
assert "Failed" in str(err.value)
|
||||||
|
assert "did not happen" in str(err.value)
|
||||||
|
|
||||||
|
|
||||||
|
def test_wait_for_task_times_out_while_in_progress():
|
||||||
|
client = _TaskClient([1] * 50)
|
||||||
|
with pytest.raises(ZertoError) as err:
|
||||||
|
asyncio.run(wait_for_task(client, "t1", timeout_s=0.05, interval_s=0.01))
|
||||||
|
assert "InProgress" in str(err.value)
|
||||||
Reference in New Issue
Block a user