Lab 7: Cancellation and Compensation

"""Start AgentRun with one tool call, optionally cancel it, wait, then print the annotated history.

    python starter.py                                 # 7.1 / 7.2: goal "deploy ..." -> the GPU saga; register_endpoint fails
    python starter.py --tool crunch --cancel-after 4  # 7.3: a 20 s tool with no heartbeat; cancel while it runs
    python starter.py --tool crunch --cancel-after 4 --heartbeat
    python starter.py --tool crunch --cancel-after 4 --heartbeat --shield
"""
import argparse
import asyncio
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import TASK_QUEUE, connect, start_workflow  # noqa: E402
from temporalio.api.enums.v1 import EventType  # noqa: E402
from temporalio.client import Client, WorkflowFailureError  # noqa: E402
from workflows import AgentRun, ToolOptions  # noqa: E402


async def history(client: Client, workflow_id: str) -> None:
    handle = client.get_workflow_handle(workflow_id)
    async for ev in handle.fetch_history_events():
        name = EventType.Name(ev.event_type).removeprefix("EVENT_TYPE_")
        detail = ""
        if ev.HasField("activity_task_scheduled_event_attributes"):
            detail = ev.activity_task_scheduled_event_attributes.activity_type.name
        elif ev.HasField("activity_task_started_event_attributes"):
            detail = f"attempt={ev.activity_task_started_event_attributes.attempt}"
        elif ev.HasField("activity_task_failed_event_attributes"):
            detail = ev.activity_task_failed_event_attributes.failure.message
        elif ev.HasField("activity_task_timed_out_event_attributes"):
            detail = ev.activity_task_timed_out_event_attributes.failure.message
        elif ev.HasField("workflow_task_started_event_attributes"):
            detail = f"identity={ev.workflow_task_started_event_attributes.identity}"
        elif ev.HasField("workflow_execution_failed_event_attributes"):
            detail = ev.workflow_execution_failed_event_attributes.failure.message
        which = ev.WhichOneof("attributes")
        if which in ("activity_task_completed_event_attributes", "activity_task_canceled_event_attributes",
                     "activity_task_failed_event_attributes", "activity_task_timed_out_event_attributes",
                     "activity_task_cancel_requested_event_attributes"):
            detail = f"{detail} (scheduled by event {getattr(ev, which).scheduled_event_id})".strip()
        print(f"{ev.event_id:4d}  {name:38s} {detail}")


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--id", default=os.environ.get("WORKFLOW_ID", "agent-42"))
    ap.add_argument("--tool", choices=["deploy_model", "crunch"], default="deploy_model")
    ap.add_argument("--heartbeat", action="store_true", help="execute_tool heartbeats every chunk")
    ap.add_argument("--shield", action="store_true", help="AgentRun shields the in-flight execute_tool")
    ap.add_argument("--chunks", type=int, default=20, help="length of the crunch tool in seconds")
    ap.add_argument("--cancel-after", type=float, default=0, help="seconds after start to call handle.cancel()")
    ap.add_argument("--no-wait", action="store_true")
    args = ap.parse_args()

    goal = "deploy the fine-tuned model" if args.tool == "deploy_model" else "crunch the eval set"
    options = ToolOptions(heartbeat=args.heartbeat, shield=args.shield, chunks=args.chunks)

    client = await connect()
    handle = await start_workflow(client, AgentRun.run, args=[goal, options], id=args.id, task_queue=TASK_QUEUE)
    print(f"started {handle.id} run={handle.result_run_id} tool={args.tool} options={options}")
    print(f"open http://localhost:8233/namespaces/default/workflows/{handle.id}")
    if args.no_wait:
        return

    if args.cancel_after > 0:
        await asyncio.sleep(args.cancel_after)
        print(f"status before cancel: {await handle.query(AgentRun.status)}")
        await handle.cancel()
        print("handle.cancel() sent — a WorkflowExecutionCancelRequested Event; the Workflow sees it at its next await")

    try:
        print("result:", await handle.result())
    except WorkflowFailureError as e:
        print(f"execution closed with {type(e.cause).__name__}: {e.cause}")
    desc = await handle.describe()
    print(f"close status: {desc.status.name}")
    print(f"final status(): {await handle.query(AgentRun.status)}")
    print("\nhistory:")
    await history(client, handle.id)


if __name__ == "__main__":
    asyncio.run(main())

"""Lab 7 self-check against the time-skipping test server:  python -m pytest tests/ -q

Activities run in real time even under time skipping, so the fake work is scaled down
(`workflows.SPEED`) and the crunch tool is a few chunks long.
"""
from __future__ import annotations

import asyncio
import os
import sys
import uuid

import pytest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, ".."))
sys.path.insert(0, os.path.join(HERE, "..", ".."))

import workflows  # noqa: E402
from temporalio.api.enums.v1 import EventType  # noqa: E402
from temporalio.client import WorkflowFailureError  # noqa: E402
from temporalio.exceptions import ActivityError, CancelledError  # noqa: E402
from temporalio.testing import WorkflowEnvironment  # noqa: E402
from temporalio.worker import Worker  # noqa: E402
from workflows import (  # noqa: E402
    AgentRun,
    ToolOptions,
    allocate_sandbox,
    call_llm,
    execute_tool,
    register_endpoint,
    release_gpu,
    release_sandbox,
    reserve_gpu,
    start_model,
    stop_model,
    unregister_endpoint,
)

ACTIVITIES = [call_llm, execute_tool, reserve_gpu, allocate_sandbox, start_model, register_endpoint,
              unregister_endpoint, stop_model, release_sandbox, release_gpu]


@pytest.fixture(autouse=True)
def fast_activities(monkeypatch):
    monkeypatch.setattr(workflows, "SPEED", 0.05)      # 1 s of fake work -> 50 ms
    workflows.FINISHED_TOOLS.clear()


async def _annotated(handle) -> list[str]:
    """['ACTIVITY_TASK_SCHEDULED reserve_gpu', 'ACTIVITY_TASK_COMPLETED', ...] in history order."""
    out = []
    async for ev in handle.fetch_history_events():
        name = EventType.Name(ev.event_type).removeprefix("EVENT_TYPE_")
        if ev.HasField("activity_task_scheduled_event_attributes"):
            name += " " + ev.activity_task_scheduled_event_attributes.activity_type.name
        out.append(name)
    return out


def _scheduled(names: list[str]) -> list[str]:
    return [n.split(" ", 1)[1] for n in names if n.startswith("ACTIVITY_TASK_SCHEDULED ")]


def test_saga_runs_compensations_in_reverse_when_register_endpoint_fails() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            tq = f"lab-07-test-{uuid.uuid4()}"
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=ACTIVITIES):
                handle = await env.client.start_workflow(
                    AgentRun.run, args=["deploy the fine-tuned model", ToolOptions()],
                    id=f"lab07-test-{uuid.uuid4()}", task_queue=tq)
                with pytest.raises(WorkflowFailureError) as exc:
                    await handle.result()
                assert isinstance(exc.value.cause, ActivityError)          # register_endpoint, retries spent
                assert (await handle.describe()).status.name == "FAILED"
                names = await _annotated(handle)
                assert _scheduled(names) == [
                    "call_llm", "reserve_gpu", "allocate_sandbox", "start_model", "register_endpoint",
                    "stop_model", "release_sandbox", "release_gpu",          # reverse order; no unregister_endpoint
                ]
                assert names.count("ACTIVITY_TASK_FAILED") == 1              # one failure Event after 3 attempts
                assert names.index("ACTIVITY_TASK_FAILED") < names.index("ACTIVITY_TASK_SCHEDULED stop_model")
                status = await handle.query(AgentRun.status)
                assert status["status"] == "failed" and status["rolled_back"] == ["stop_model", "release_sandbox", "release_gpu"]

    asyncio.run(go())


def test_cancel_without_heartbeats_the_tool_runs_to_the_end_anyway() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            tq = f"lab-07-test-{uuid.uuid4()}"
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=ACTIVITIES):
                handle = await env.client.start_workflow(
                    AgentRun.run, args=["crunch the eval set", ToolOptions(heartbeat=False, chunks=30)],
                    id=f"lab07-test-{uuid.uuid4()}", task_queue=tq)
                while (await handle.query(AgentRun.status))["in_flight"] != "crunch":
                    await asyncio.sleep(0.05)
                await handle.cancel()
                with pytest.raises(WorkflowFailureError) as exc:
                    await handle.result()
                assert isinstance(exc.value.cause, CancelledError)
                assert (await handle.describe()).status.name == "CANCELED"
                names = await _annotated(handle)
                assert "ACTIVITY_TASK_CANCEL_REQUESTED" in names
                assert "ACTIVITY_TASK_CANCELED" not in names                 # nobody could tell the Activity
                assert _scheduled(names)[-1] == "release_sandbox"            # cleanup ran, unshielded, in `finally`
                assert (await handle.query(AgentRun.status))["rolled_back"] == ["release_sandbox"]
                # ... and the tool is still running on the Worker. Wait for it: it finishes every chunk.
                for _ in range(200):
                    if workflows.FINISHED_TOOLS:
                        break
                    await asyncio.sleep(0.05)
                assert workflows.FINISHED_TOOLS == [f"{handle.id}:1:crunch"]

    asyncio.run(go())


def test_cancel_with_heartbeats_stops_the_tool_and_cleanup_still_runs(monkeypatch) -> None:
    # The cancel rides back on the next heartbeat, up to 0.8 x HEARTBEAT_TIMEOUT later; the run must
    # still be open then for the acknowledgement to land in *its* history, so the cleanup is slow here.
    monkeypatch.setattr(workflows, "UNDO_SECONDS", 80)      # x SPEED 0.05 = 4 s

    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            tq = f"lab-07-test-{uuid.uuid4()}"
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=ACTIVITIES):
                handle = await env.client.start_workflow(
                    AgentRun.run, args=["crunch the eval set", ToolOptions(heartbeat=True, chunks=200)],
                    id=f"lab07-test-{uuid.uuid4()}", task_queue=tq)
                while (await handle.query(AgentRun.status))["in_flight"] != "crunch":
                    await asyncio.sleep(0.05)
                await handle.cancel()
                with pytest.raises(WorkflowFailureError):
                    await handle.result()
                assert (await handle.describe()).status.name == "CANCELED"
                names = await _annotated(handle)
                assert "ACTIVITY_TASK_CANCEL_REQUESTED" in names
                assert "ACTIVITY_TASK_CANCELED" in names                     # the Activity acknowledged the cancel
                assert _scheduled(names)[-1] == "release_sandbox"
                assert workflows.FINISHED_TOOLS == []                        # it did not run to the end

    asyncio.run(go())


def test_shielded_tool_finishes_first_then_the_run_closes_cancelled() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            tq = f"lab-07-test-{uuid.uuid4()}"
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=ACTIVITIES):
                handle = await env.client.start_workflow(
                    AgentRun.run, args=["crunch the eval set", ToolOptions(heartbeat=True, shield=True, chunks=30)],
                    id=f"lab07-test-{uuid.uuid4()}", task_queue=tq)
                while (await handle.query(AgentRun.status))["in_flight"] != "crunch":
                    await asyncio.sleep(0.05)
                await handle.cancel()
                with pytest.raises(WorkflowFailureError):
                    await handle.result()
                assert (await handle.describe()).status.name == "CANCELED"
                names = await _annotated(handle)
                assert "ACTIVITY_TASK_CANCEL_REQUESTED" not in names         # the shield kept the cancel off the Activity
                assert names.index("ACTIVITY_TASK_SCHEDULED execute_tool") < names.index("WORKFLOW_EXECUTION_CANCEL_REQUESTED")
                assert _scheduled(names)[-2:] == ["execute_tool", "release_sandbox"]
                assert names.count("ACTIVITY_TASK_COMPLETED") == 4           # llm, sandbox, the tool, the release
                assert workflows.FINISHED_TOOLS == [f"{handle.id}:1:crunch"]

    asyncio.run(go())

import asyncio
import logging
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import TASK_QUEUE, connect  # noqa: E402
from temporalio import activity  # noqa: E402
from temporalio.worker import Worker  # noqa: E402
from workflows import (  # noqa: E402
    AgentRun,
    allocate_sandbox,
    call_llm,
    execute_tool,
    register_endpoint,
    release_gpu,
    release_sandbox,
    reserve_gpu,
    start_model,
    stop_model,
    unregister_endpoint,
)


def _one_line_failures(record: logging.LogRecord) -> bool:
    """register_endpoint fails on purpose, three times; keep the SDK's line, drop its traceback."""
    if record.getMessage().startswith("Completing activity as failed"):
        record.exc_info = None
    return True


async def main() -> None:
    # The Activities narrate what they do; this lab is read from the Worker's terminal as much as from history.
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S")
    logging.getLogger("temporalio").setLevel(logging.WARNING)
    logging.getLogger("temporalio.activity").setLevel(logging.INFO)     # our activity.logger lines
    logging.getLogger("temporalio.activity").addFilter(_one_line_failures)
    activity.logger.activity_info_on_message = False                    # no context dict after every line
    client = await connect()
    worker = Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=[AgentRun],
        activities=[call_llm, execute_tool, reserve_gpu, allocate_sandbox, start_model, register_endpoint,
                    unregister_endpoint, stop_model, release_sandbox, release_gpu],
    )
    print(f"worker pid={os.getpid()} polling {TASK_QUEUE!r}; kill -9 {os.getpid()} to crash it")
    await worker.run()


if __name__ == "__main__":
    asyncio.run(main())

"""Lab 7: AgentRun, stage 7 — the GPU saga, and cancellation that actually stops things.

Two additions. (1) `execute_tool`'s GPU lane is four Activities orchestrated by
Workflow code — `reserve_gpu -> allocate_sandbox -> start_model -> register_endpoint` —
and after each success AgentRun appends the matching undo to `self.compensations`.
That list is Workflow state: when a later step fails, the undos run in reverse, and
a Worker that dies halfway through the rollback is replaced by one that replays the
list and finishes it. (2) `handle.cancel()` is Temporal's cancellation, delivered
once as `asyncio.CancelledError` at the Workflow's next await. The rollback runs in
`finally`, after that delivery, with no shield. What a cancel does to the Activity
that was in flight depends on two things: whether the Workflow propagates the cancellation to it at
all (7.3's shield does not), and, when it does, whether the Activity heartbeats — the request rides
back on a heartbeat response.
"""
from __future__ import annotations

import asyncio
import hashlib
import os
import sys
from dataclasses import dataclass, field
from datetime import timedelta

from temporalio import activity, workflow
from temporalio.common import RetryPolicy
from temporalio.exceptions import ActivityError, ApplicationError, CancelledError

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import TASK_QUEUE  # noqa: E402,F401

SPEED = 1.0                                   # scale for every fake Activity's sleep (tests shrink it)
STEP_TIMEOUT = timedelta(seconds=10)          # forward steps and undos; also how long a dead Worker's attempt stays "running" (7.2)
FORWARD_RETRY = RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3)
UNDO_RETRY = RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=10)
UNDO_SECONDS = 3                              # each undo takes this long: the window to kill the Worker in 7.2
TOOL_CHUNKS = 20                              # execute_tool works in 1-second chunks
HEARTBEAT_TIMEOUT = timedelta(seconds=3)      # heartbeats are throttled to 0.8x this, so a cancel lands within ~2.4 s


@dataclass
class ToolCall:
    key: str
    tool: str                                 # "deploy_model" (GPU lane) | "crunch" (CPU lane)
    args: dict = field(default_factory=dict)


@dataclass
class ToolResult:
    key: str
    output: str


@dataclass
class ToolOptions:
    """Lab knobs for 7.3, passed to run() so the Workflow stays deterministic."""
    heartbeat: bool = False                   # execute_tool calls activity.heartbeat() each chunk
    shield: bool = False                      # AgentRun wraps the in-flight execute_tool in asyncio.shield
    chunks: int = TOOL_CHUNKS


@dataclass
class Step:
    kind: str                                 # "llm" | "tool" | "finish"
    prompt: str = ""
    tool: str = ""


async def _work(seconds: float) -> None:
    await asyncio.sleep(seconds * SPEED)


@activity.defn
async def call_llm(prompt: str, key: str) -> str:
    return f"[model answer to: {prompt[:48]}]"


# ---- the GPU lane, forward: each step returns the identifier its undo will need ----

@activity.defn
async def reserve_gpu(key: str) -> str:
    await _work(1)
    gpu = f"gpu-for-{key.split(':')[0]}"          # the reservation id; release_gpu is keyed on it
    activity.logger.info("reserve_gpu -> %s", gpu)
    return gpu


@activity.defn
async def allocate_sandbox(gpu: str) -> str:
    await _work(1)
    box = f"sandbox-on-{gpu}"
    activity.logger.info("allocate_sandbox -> %s", box)
    return box


@activity.defn
async def start_model(box: str) -> str:
    await _work(1)
    model = f"model-in-{box}"
    activity.logger.info("start_model -> %s", model)
    return model


@activity.defn
async def register_endpoint(model: str) -> str:
    await _work(0.2)
    activity.logger.info("register_endpoint attempt %d for %s: rejected", activity.info().attempt, model)
    raise ApplicationError(f"endpoint registry rejected the name for {model}")   # retryable; the policy ends it


# ---- the GPU lane, reverse: each undo may run more than once ----
# The stable resource id is what *lets* the external system make the undo idempotent; it does not
# make it idempotent. Releasing gpu-123 twice is safe only if the allocator treats the second
# release as a no-op rather than an error — or, worse, as releasing a gpu-123 that has since been
# handed to someone else. These fake Activities have no external state to get wrong.

async def _undo(what: str, resource: str) -> None:
    activity.logger.info("%s %s: attempt %d, starting", what, resource, activity.info().attempt)
    await _work(UNDO_SECONDS)
    activity.logger.info("%s %s: done", what, resource)


@activity.defn
async def unregister_endpoint(endpoint: str) -> None:
    await _undo("unregister_endpoint", endpoint)


@activity.defn
async def stop_model(model: str) -> None:
    await _undo("stop_model", model)


@activity.defn
async def release_sandbox(box: str) -> None:
    await _undo("release_sandbox", box)


@activity.defn
async def release_gpu(gpu: str) -> None:
    await _undo("release_gpu", gpu)


# ---- the CPU lane: one long tool ----

FINISHED_TOOLS: list[str] = []                # in-process marker for the tests: which tool calls ran to the end


@activity.defn
async def execute_tool(call: ToolCall) -> ToolResult:
    """A long tool in 1-second chunks. With `heartbeat` off it never talks to the Service while it
    runs, so nothing can tell it the run was cancelled. With it on, the cancellation rides back on a
    heartbeat response and the SDK raises CancelledError at the next await."""
    chunks = int(call.args.get("chunks", TOOL_CHUNKS))
    heartbeat = bool(call.args.get("heartbeat", False))
    i = 0
    try:
        for i in range(1, chunks + 1):
            if heartbeat:
                activity.heartbeat(i)             # progress out, cancellation (if any) back
            await _work(1)
            activity.logger.info("execute_tool %s: chunk %d/%d", call.tool, i, chunks)
    except asyncio.CancelledError:
        activity.logger.info("execute_tool %s: cancelled at chunk %d/%d", call.tool, i, chunks)
        raise                                     # re-raise, or the Activity does not appear cancelled
    FINISHED_TOOLS.append(call.key)
    activity.logger.info("execute_tool %s: finished all %d chunks", call.tool, chunks)
    return ToolResult(key=call.key, output=f"{call.tool}: {chunks} chunks done")


@workflow.defn
class AgentRun:
    def __init__(self) -> None:
        self.step = 0
        self.status = "running"                    # running | done | failed | cancelled
        self.goal = ""
        self.context_summary = ""
        self.attempts = 0
        self.compensations: list = []              # [(undo_activity, arg)] — Workflow state, so replay rebuilds it
        self.options = ToolOptions()
        self.in_flight = ""
        self.rolled_back: list[str] = []

    def plan(self) -> Step:
        if self.step == 0:
            return Step("llm", prompt=f"Plan one tool call for: {self.goal}")
        if self.step == 1:
            return Step("tool", tool="deploy_model" if self.goal.startswith("deploy") else "crunch")
        return Step("finish")

    @workflow.run
    async def run(self, goal: str, options: ToolOptions | None = None) -> str:
        self.goal = goal
        self.options = options or ToolOptions()
        try:
            while self.status == "running":
                step = self.plan()
                if step.kind == "finish":
                    self.status = "done"
                    break
                if step.kind == "llm":
                    key = f"{workflow.info().workflow_id}:{self.step}:{hashlib.sha256(step.prompt.encode()).hexdigest()[:12]}"
                    self.context_summary = await workflow.execute_activity(
                        call_llm, args=[step.prompt, key], start_to_close_timeout=timedelta(seconds=10)
                    )
                else:
                    call = ToolCall(key=f"{workflow.info().workflow_id}:{self.step}:{step.tool}", tool=step.tool,
                                    args={"heartbeat": self.options.heartbeat, "chunks": self.options.chunks})
                    result = await (self.gpu_lane(call) if step.tool == "deploy_model" else self.cpu_lane(call))
                    self.context_summary += f"\n{result.output}"
                self.step += 1
        except asyncio.CancelledError:
            self.status = "cancelled"              # Temporal's cancel, delivered once at the await above
            raise                                  # ... after `finally` has run the rollback
        except ActivityError as e:
            if isinstance(e.cause, CancelledError):
                self.status = "cancelled"          # the same cancel, but it landed while an Activity was in
            else:                                  # flight: it arrives wrapped as that Activity's cancellation
                self.status = "failed"             # a step failed for good (its retries are spent)
            raise
        finally:
            if self.status in ("cancelled", "failed"):
                await self.rollback()              # runs normally: the one cancellation was already delivered
        return self.context_summary

    async def gpu_lane(self, call: ToolCall) -> ToolResult:
        """Four steps, each recorded; the undo list grows with every success."""
        self.compensations = []
        gpu = await workflow.execute_activity(reserve_gpu, call.key, start_to_close_timeout=STEP_TIMEOUT)
        self.compensations.append((release_gpu, gpu))
        box = await workflow.execute_activity(allocate_sandbox, gpu, start_to_close_timeout=STEP_TIMEOUT)
        self.compensations.append((release_sandbox, box))
        model = await workflow.execute_activity(start_model, box, start_to_close_timeout=STEP_TIMEOUT)
        self.compensations.append((stop_model, model))
        endpoint = await workflow.execute_activity(                       # fails after FORWARD_RETRY is spent
            register_endpoint, model, start_to_close_timeout=STEP_TIMEOUT, retry_policy=FORWARD_RETRY
        )
        self.compensations.append((unregister_endpoint, endpoint))
        return ToolResult(key=call.key, output=f"deployed at {endpoint}")

    async def cpu_lane(self, call: ToolCall) -> ToolResult:
        """A sandbox, then one long tool inside it. The sandbox is the resource a cancel must give back."""
        self.compensations = []
        box = await workflow.execute_activity(allocate_sandbox, call.key, start_to_close_timeout=STEP_TIMEOUT)
        self.compensations.append((release_sandbox, box))
        self.in_flight = call.tool
        tool = workflow.start_activity(              # a handle: the Activity is scheduled now, awaited below
            execute_tool, call,
            start_to_close_timeout=timedelta(seconds=self.options.chunks + 30),
            heartbeat_timeout=HEARTBEAT_TIMEOUT if self.options.heartbeat else None,
            retry_policy=RetryPolicy(maximum_attempts=1),
        )
        if self.options.shield:
            try:
                result = await asyncio.shield(tool)  # the cancel reaches this await, not the Activity
            except asyncio.CancelledError:
                result = await tool                  # let the in-flight work finish, then close as cancelled
                self.context_summary += f"\n{result.output}"
                raise
        else:
            result = await tool                      # cancelling this await cancels the Activity too
        self.in_flight = ""
        return result

    async def rollback(self) -> None:
        """The saga's reverse path. Each undo is an Activity with its own retry policy."""
        for undo, arg in reversed(self.compensations):
            await workflow.execute_activity(undo, arg, start_to_close_timeout=STEP_TIMEOUT, retry_policy=UNDO_RETRY)
            self.rolled_back.append(undo.__name__)
        self.compensations = []

    @workflow.query
    def status(self) -> dict:
        return {
            "step": self.step,
            "status": self.status,
            "in_flight": self.in_flight,
            "compensations": [undo.__name__ for undo, _ in self.compensations],
            "rolled_back": self.rolled_back,
        }