Lab 8: Testing and Replay Tests

"""Export a Workflow Execution's Event History as JSON, the format `Replayer` reads.

    python export_history.py agent-42                      # -> tests/histories/agent-42.json
    python export_history.py agent-42 --run-id <run-id>    # one specific run of the chain
    python export_history.py agent-42 --out some/path.json

This is the Python form of

    temporal workflow show --workflow-id agent-42 --output json > tests/histories/agent-42.json

for a box without the Temporal CLI. Both produce the same `temporal.api.history.v1.History`
JSON (the Web UI's "Download" button does too), and `WorkflowHistory.from_json` reads any of them.
"""
from __future__ import annotations

import argparse
import asyncio
import os
import sys
from pathlib import Path

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

HISTORIES = Path(__file__).parent / "tests" / "histories"


async def export(workflow_id: str, run_id: str | None, out: Path) -> int:
    client = await connect()
    handle = client.get_workflow_handle(workflow_id, run_id=run_id)
    history = await handle.fetch_history()          # WorkflowHistory: the events, decoded
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(history.to_json())                # proto JSON, same shape as the CLI's
    return len(history.events)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("workflow_id")
    ap.add_argument("--run-id", default=None)
    ap.add_argument("--out", default=None, help=f"default: {HISTORIES}/<workflow_id>.json")
    ns = ap.parse_args()
    out = Path(ns.out) if ns.out else HISTORIES / f"{ns.workflow_id}.json"
    n = asyncio.run(export(ns.workflow_id, ns.run_id, out))
    print(f"wrote {n} events to {out}")


if __name__ == "__main__":
    main()

# 8.5 — Answers

Do not open this before you have written your five answers down.

History: `tests/histories/forensic_lab08-agent-42-broken.json` (Workflow ID `lab08-agent-42-broken`,
16 events, recorded 2026-09-11 on Temporal Server 1.31.2 / temporalio 1.32.0). Event numbers
below are the `eventId` fields in that file.

## The events, annotated

```
 1  WorkflowExecutionStarted       AgentRun, task queue lab-08, input "summarize the Temporal docs"
 2  WorkflowTaskScheduled
 3  WorkflowTaskStarted            identity worker-a
 4  WorkflowTaskCompleted
 5  ActivityTaskScheduled          call_llm   retry: 3 attempts, 1 s → ×2.0 ; no heartbeat timeout
 6  ActivityTaskStarted            attempt 1, identity worker-c, lastFailure: none
 7  ActivityTaskCompleted          "TOOL: fetch temporal-docs"
 8  WorkflowTaskScheduled
 9  WorkflowTaskStarted            identity worker-a
10  WorkflowTaskCompleted
11  ActivityTaskScheduled          execute_tool   input key "lab08-agent-42-broken:1:d3f35deb8827", tool fetch
                                   retry: maximumAttempts 3, initialInterval 1 s, backoff 1.0 ; heartbeatTimeout 5 s
12  ActivityTaskStarted            attempt 3, identity worker-c,
                                   lastFailure: "activity Heartbeat timeout" (TIMEOUT_TYPE_HEARTBEAT, source Server)
13  ActivityTaskFailed             ApplicationError type "BadToolArguments", non_retryable,
                                   message "tool rejected the arguments: target must be a URL",
                                   retryState RETRY_STATE_NON_RETRYABLE_FAILURE
14  WorkflowTaskScheduled          task queue "worker-a-<uuid>" kind STICKY (normalName lab-08)
15  WorkflowTaskTimedOut           TIMEOUT_TYPE_SCHEDULE_TO_START   ← nobody polled worker-a's sticky queue
16  WorkflowTaskScheduled          task queue lab-08 kind NORMAL     ← last event; nobody has picked it up
```

## 1. Which Activity attempt failed?

Attempt **3** of `execute_tool` (event 12: `attempt: 3`) failed with the non-retryable
`BadToolArguments` (event 13). Attempts 1 and 2 also failed — you know that only because the
one `ActivityTaskStarted` the history contains says `attempt: 3`. Intermediate attempts leave
no events of their own. The Service generates `ActivityTaskStarted` when it dispatches the
Task, but withholds it from history until the Activity Execution reaches a terminal outcome —
so history carries one such Event, describing the attempt that closed it, instead of one per
retry. While the Activity was still pending, the
attempt counter lived only in `temporal workflow describe` (Pending Activities), not in history.

## 2. Was it retried?

Yes, twice. Event 11 carries the RetryPolicy (`maximumAttempts: 3`); event 12 says attempt 3
ran. Why attempt 2 ended is recorded — `lastFailure` on event 12 is a **heartbeat timeout**:
attempt 2 stopped heartbeating (heartbeatTimeout 5 s) and the Service gave up on it. Why
attempt 1 ended is **not recorded anywhere**: `lastFailure` holds only the immediately
preceding attempt's failure. Will it be retried again? No: event 13's
`retryState: RETRY_STATE_NON_RETRYABLE_FAILURE` — the Activity is closed. (Had attempt 3 been
a retryable error, the RetryPolicy was exhausted anyway: `RETRY_STATE_MAXIMUM_ATTEMPTS_REACHED`.)

## 3. Did the side effect possibly occur?

**Possibly, yes — and the history cannot say.** Three attempts of `fetch temporal-docs` ran.
History names the Worker for the *terminal* attempt only: event 12 says attempt 3 ran on
worker-c. Attempts 1 and 2 left no `ActivityTaskStarted` of their own, so which Worker ran them
is not recorded anywhere — do not assume it was the same one. Attempt 3 was rejected before doing anything (`BadToolArguments` is the tool
refusing its input). Attempt 2 *started* — an attempt can only time out on heartbeat after
it has started — and then went silent; a heartbeat timeout means the Service stopped hearing
from it, not that the code stopped running. The Worker may have completed the fetch after the
Service had already written it off. (This history's `lastFailure` carries no
`lastHeartbeatDetails`, so how far attempt 2 got is not recorded either.) Attempt 1: unknown. The only handle on the question is the idempotency key in event 11's input,
`lab08-agent-42-broken:1:d3f35deb8827`: if the tool's system of record has an entry for that
key, the effect happened once; if it does not, it did not. Temporal records that the attempts
*ran*; whether they *did* something is in your database, not here.

## 4. What state is the Workflow in now?

**Open, with a Workflow Task scheduled that no Worker has picked up** (event 16). Note the
pair 14–15: the task was first offered to worker-a's sticky queue, timed out on
schedule-to-start because worker-a is gone, and was re-offered on the normal queue `lab-08`.

Split that answer in two, because the exercise gave you no source and the two halves have
different epistemic status.

**Temporal execution state — determinable from history alone.** The execution is open. No
Activity is pending (`describe` shows Pending Activities: 0 — the Activity closed at event 13).
A Workflow Task is scheduled on the normal queue and no Worker has started it. Event 13 has not
been applied to any Workflow code yet; that is what the pending task is for.

**Application state — not determinable from history alone.** Event History records inputs,
Activity results, timers, Signals and Commands. It does not serialize `self.step`,
`self.context_summary` or `self.attempts`. Those exist only as the product of *history + the
Workflow Definition*, reconstructed by replay. Hand this history to compatible code and it
rebuilds `goal` = "summarize the Temporal docs", `step` = 1, `context_summary` =
"\nTOOL: fetch temporal-docs", `attempts` = 0, `status` = "running" — but that is the code
talking, not the history. With the Worker image gone, the honest answer is the first paragraph
plus "the rest needs the code".

## 5. What happens mechanically if a compatible Worker starts now?

**Predictable from history alone:** it polls `lab-08`, receives the pending Workflow Task,
replays events 1–13, and at event 13 the `execute_activity` call raises `ActivityError` (cause
`ApplicationError("BadToolArguments")`). **Not predictable without the source:** what the
Workflow catches, what it mutates, and therefore which Command comes next. From the history alone you can
say: the code will issue *some* Command; the Service will accept it if it is consistent with
events 1–13. Predict: either the Workflow fails/compensates, or it records the error and asks
the model again.

**What actually happened** (observed 2026-09-11 by starting `worker.py` on `lab-08`, events
16–49 of the continued execution):

```
17  WorkflowTaskStarted        a fresh identity (pid 3811222)
19  ActivityTaskScheduled      call_llm   prompt now contains "RESULT: error BadToolArguments"
21  ActivityTaskCompleted      "TOOL: fetch temporal-docs"        ← the model asks for the same tool again
25  ActivityTaskScheduled      execute_tool  key "lab08-agent-42-broken:3:d3f35deb8827"   ← step 3, a NEW key
27  ActivityTaskCompleted      ok "fetch: temporal-docs done"
...                            summarize, FINAL
49  WorkflowExecutionCompleted "the agent fetched and summarized the docs"
```

`status` query afterwards: `attempts: 1, step: 7, status: "done"`. The code recorded the
failure (`RESULT: error BadToolArguments`), asked the model what to do, and the model said
"fetch" again. The second fetch ran under a **different key** (`:3:` instead of `:1:`) because
the key is `workflow_id:step:hash` and the step had moved on. So if attempt 2 of the first
call did complete the fetch (question 3), the returning Worker fetched it **twice** — the
idempotency key protects a *retry of the same Activity*, not a *re-plan by the Workflow*. That
is the design bug the forensic exercise is there to make you notice: the key that dedupes a
tool call should be derived from what the call *is* (tool + args), not from where in the
loop it happened.

## The two things the history does not record

1. Why attempt 1 failed (only the last failure survives).
2. Whether attempts 1 and 2 had a side effect (only your system of record knows).

Both are things your own Activities must record if you want to answer them next time.

"""8.5 — reproduce the broken AgentRun execution on your dev server and export its history.

    TASK_QUEUE=lab-08 python forensic/make_broken_history.py     # ~10 s; leaves the execution OPEN

Do not read past this docstring until you have written your five answers. The whole point of
the exercise is that the history alone must answer them.
"""
from __future__ import annotations

import asyncio
import logging
import os
import sys
from datetime import datetime, timezone
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
from common import TASK_QUEUE, connect  # noqa: E402
from temporalio import activity  # noqa: E402
from temporalio.api.enums.v1 import EventType, TaskQueueType  # noqa: E402
from temporalio.api.taskqueue.v1 import TaskQueue  # noqa: E402
from temporalio.api.workflowservice.v1 import DescribeTaskQueueRequest  # noqa: E402
from temporalio.client import WorkflowExecutionStatus  # noqa: E402
from temporalio.runtime import LoggingConfig, Runtime, TelemetryConfig, TelemetryFilter  # noqa: E402
from temporalio.exceptions import ApplicationError  # noqa: E402
from temporalio.worker import Worker  # noqa: E402
from workflows import AgentRun, ToolCall, ToolResult, call_llm  # noqa: E402

WORKFLOW_ID = os.environ.get("WORKFLOW_ID", "agent-42-broken")
OUT = Path(__file__).resolve().parent.parent / "tests" / "histories" / f"forensic_{WORKFLOW_ID}.json"


@activity.defn(name="execute_tool")
async def sabotaged_execute_tool(call: ToolCall) -> ToolResult:
    attempt = activity.info().attempt
    if attempt == 1:
        raise ApplicationError("sandbox not ready", type="ToolTransientError")   # retryable
    if attempt == 2:
        activity.heartbeat(f"{call.tool}:1/3")      # got this far ...
        await asyncio.sleep(3600)                    # ... then stopped heartbeating
    raise ApplicationError("tool rejected the arguments: target must be a URL",
                           type="BadToolArguments", non_retryable=True)


# The sabotage is supposed to fail; the SDK would log each failure with a full traceback, which
# reads like the script breaking. Errors of its own still print.
logging.getLogger("temporalio").setLevel(logging.ERROR)
Runtime.set_default(Runtime(telemetry=TelemetryConfig(logging=LoggingConfig(
    filter=TelemetryFilter(core_level="ERROR", other_level="ERROR")))))


async def live_workflow_pollers(client) -> list[str]:
    """Workers that polled this queue for Workflow Tasks in the last 90 s."""
    r = await client.workflow_service.describe_task_queue(DescribeTaskQueueRequest(
        namespace=client.namespace, task_queue=TaskQueue(name=TASK_QUEUE),
        task_queue_type=TaskQueueType.TASK_QUEUE_TYPE_WORKFLOW))
    now = datetime.now(timezone.utc)
    return [p.identity for p in r.pollers
            if (now - p.last_access_time.ToDatetime().replace(tzinfo=timezone.utc)).total_seconds() <= 90]


async def main() -> None:
    client = await connect()
    # The exercise needs the run left with no Workflow Worker. One still running from 8.2 on the same
    # queue would take the run over and finish it, and the history would be a different story.
    others = await live_workflow_pollers(client)
    if others:
        raise SystemExit(f"A Worker is still polling {TASK_QUEUE!r} for Workflow Tasks: {', '.join(others)}.\n"
                         f"Stop the Worker from 8.2 (Ctrl-C, or Crash Worker in the browser), wait about"
                         f" 90 s for its last poll to age out, and run this again.")
    # Two Workers, so the Workflow Worker can disappear while the Activity Worker keeps going.
    workflow_worker = Worker(client, task_queue=TASK_QUEUE, workflows=[AgentRun], identity="worker-a")
    activity_worker = Worker(client, task_queue=TASK_QUEUE, activities=[call_llm, sabotaged_execute_tool],
                             identity="worker-c")
    wf_task = asyncio.create_task(workflow_worker.run())
    act_task = asyncio.create_task(activity_worker.run())
    handle = await client.start_workflow(AgentRun.run, "summarize the Temporal docs",
                                         id=WORKFLOW_ID, task_queue=TASK_QUEUE)
    print(f"started {handle.id}")

    async def tool_scheduled() -> bool:
        async for e in handle.fetch_history_events():
            if (e.HasField("activity_task_scheduled_event_attributes")
                    and e.activity_task_scheduled_event_attributes.activity_type.name == "execute_tool"):
                return True
        return False

    while not await tool_scheduled():
        await asyncio.sleep(0.2)
    await workflow_worker.shutdown()                 # no Workflow Worker from here on
    await wf_task
    print("workflow worker gone; letting the tool fail its attempts")

    async def tool_closed() -> bool:
        async for e in handle.fetch_history_events():
            if e.HasField("activity_task_failed_event_attributes"):
                return True
        return False

    while not await tool_closed():
        await asyncio.sleep(0.5)
    await asyncio.sleep(1.0)                          # let the Service write the tail
    await activity_worker.shutdown()
    act_task.cancel()

    # The Workflow Task that delivers the failure went to worker-a's sticky queue. Wait for it to time
    # out there and be rescheduled on the shared queue (~10 s), so the file ends where the story does.
    print("waiting for the sticky Workflow Task to time out (~10 s)")
    for _ in range(40):
        types = [e.event_type async for e in handle.fetch_history_events()]
        if types[-2:] == [EventType.EVENT_TYPE_WORKFLOW_TASK_TIMED_OUT, EventType.EVENT_TYPE_WORKFLOW_TASK_SCHEDULED]:
            break
        await asyncio.sleep(0.5)
    history = await handle.fetch_history()
    desc = await handle.describe()
    if desc.status != WorkflowExecutionStatus.RUNNING:
        raise SystemExit(f"{handle.id} is {desc.status.name}, not open: another Worker on {TASK_QUEUE!r} took it over.\n"
                         f"Stop every Worker on that queue and rerun with a fresh id, e.g. WORKFLOW_ID={handle.id}-2.")
    OUT.write_text(history.to_json())
    print(f"wrote {len(history.events)} events to {OUT}; {handle.id} is still open — do not start a Worker yet")


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

[pytest]
# Tests are plain `async def`; conftest.py runs them with asyncio.run so no plugin is required
# (pytest-asyncio, if installed, is left in its default strict mode and does not interfere).
testpaths = tests

import asyncio
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import TASK_QUEUE, connect, show_history, start_workflow  # noqa: E402
from workflows import AgentRun  # noqa: E402

# Override when several people share one dev server: WORKFLOW_ID=lab08-agent-42 python starter.py
WORKFLOW_ID = os.environ.get("WORKFLOW_ID", "agent-42")


async def main() -> None:
    client = await connect()
    handle = await start_workflow(client,
        AgentRun.run, "summarize the Temporal docs", id=WORKFLOW_ID, task_queue=TASK_QUEUE
    )
    print(f"started {handle.id} run={handle.result_run_id}; open http://localhost:8233")
    print("result:", await handle.result())
    await show_history(client, handle.id)


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

"""Test plumbing: import the lab's modules, run `async def` tests without a plugin."""
from __future__ import annotations

import asyncio
import inspect
import os
import sys

import pytest

LAB_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, LAB_DIR)                                 # workflows.py
sys.path.insert(0, os.path.dirname(LAB_DIR))                # common/


@pytest.hookimpl(tryfirst=True)
def pytest_pyfunc_call(pyfuncitem):
    """Run coroutine tests on a fresh event loop. stdlib only — no pytest-asyncio needed."""
    if inspect.iscoroutinefunction(pyfuncitem.obj):
        kwargs = {name: pyfuncitem.funcargs[name] for name in pyfuncitem._fixtureinfo.argnames}
        asyncio.run(pyfuncitem.obj(**kwargs))
        return True
    return None


@pytest.fixture(autouse=True)
def _fresh_tool_effects():
    """Each test starts with an empty tool system-of-record."""
    import workflows
    workflows.reset_effects()
    yield
    workflows.reset_effects()

"""A Worker to kill. 8.4 starts it as a subprocess against the test server and SIGKILLs it
mid-Activity; the Service must then hand the Activity attempt to the next Worker."""
from __future__ import annotations

import asyncio
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from temporalio.client import Client  # noqa: E402
from temporalio.worker import Worker  # noqa: E402

import fakes  # noqa: E402
from workflows import AgentRun  # noqa: E402


async def main() -> None:
    client = await Client.connect(os.environ["TEMPORAL_ADDRESS"], namespace=os.environ["TEMPORAL_NAMESPACE"])
    worker = Worker(client, task_queue=os.environ["TASK_QUEUE"], workflows=[AgentRun],
                    activities=[fakes.fake_call_llm, fakes.slow_execute_tool], identity=f"crashable-{os.getpid()}",
                    max_cached_workflows=int(os.environ.get("MAX_CACHED_WORKFLOWS", "1000")))
    print("ready", flush=True)
    await worker.run()


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

"""Fake Activities for Workflow tests. Same *names* and signatures as the real ones — the
Workflow schedules `call_llm` by name, so a Worker that registers these runs them instead.
No model, no tool, no network: each fake is a rule plus a counter the test can read."""
from __future__ import annotations

import asyncio

from temporalio import activity
from temporalio.exceptions import ApplicationError

from workflows import ToolCall, ToolResult

calls = {"call_llm": 0, "execute_tool": 0}
tool_keys: list[str] = []


def reset() -> None:
    calls["call_llm"] = calls["execute_tool"] = 0
    cancelled["execute_tool"] = 0
    tool_keys.clear()


@activity.defn(name="call_llm")
async def fake_call_llm(prompt: str, key: str) -> str:
    """The stand-in model: ask for one tool call, then answer."""
    calls["call_llm"] += 1
    if "RESULT: ok" in prompt:
        return "FINAL: 42"
    return "TOOL: fetch the-answer"


@activity.defn(name="execute_tool")
async def fake_execute_tool(call: ToolCall) -> ToolResult:
    calls["execute_tool"] += 1
    tool_keys.append(call.key)
    return ToolResult(key=call.key, ok=True, output="42")


@activity.defn(name="execute_tool")
async def failing_execute_tool(call: ToolCall) -> ToolResult:
    """A tool that never works: `ok=False` every time. Tests the give-up branch."""
    calls["execute_tool"] += 1
    return ToolResult(key=call.key, ok=False, output="sandbox refused")


@activity.defn(name="execute_tool")
async def slow_execute_tool(call: ToolCall) -> ToolResult:
    """A tool that runs long and heartbeats: cancellation reaches it at the next heartbeat."""
    calls["execute_tool"] += 1
    tool_keys.append(call.key)
    if call.tool.startswith("undo_"):
        return ToolResult(key=call.key, ok=True, output="undone")
    try:
        for i in range(200):
            activity.heartbeat(i)
            await asyncio.sleep(0.05)
    except asyncio.CancelledError:
        cancelled["execute_tool"] += 1            # delivered on a heartbeat; re-raise it
        raise
    return ToolResult(key=call.key, ok=True, output="finally")


@activity.defn(name="execute_tool")
async def hanging_execute_tool(call: ToolCall) -> ToolResult:
    """A tool that hangs and never heartbeats: only a timeout ends it."""
    calls["execute_tool"] += 1
    await asyncio.sleep(3600)
    return ToolResult(key=call.key, ok=True, output="unreachable")


@activity.defn(name="execute_tool")
async def bad_arguments_execute_tool(call: ToolCall) -> ToolResult:
    """Fails non-retryably: the Workflow sees one ActivityError, no retries."""
    calls["execute_tool"] += 1
    raise ApplicationError("tool rejected the arguments", type="BadToolArguments", non_retryable=True)

cancelled = {"execute_tool": 0}


@activity.defn(name="call_llm")
async def fake_call_llm_two_tools(prompt: str, key: str) -> str:
    """Like the stand-in in workflows.py: fetch, then summarize, then answer."""
    calls["call_llm"] += 1
    n = prompt.count("RESULT: ok")
    return ["TOOL: fetch the-answer", "TOOL: summarize the-answer", "FINAL: 42"][min(n, 2)]

"""8.1 — Workflow tests in the time-skipping environment, plus plan() with no environment."""
from __future__ import annotations

import uuid
from datetime import timedelta

import pytest
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

import fakes
from workflows import CHECKPOINT_TIMEOUT, MAX_TOOL_FAILURES, AgentRun, Step


@pytest.fixture(autouse=True)
def _reset_fakes():
    fakes.reset()


# ---------------------------------------------------------------- plan(): no environment at all
def test_plan_first_step_asks_the_model():
    agent = AgentRun()
    agent.goal = "summarize the Temporal docs"
    step = agent.plan()
    assert step == Step("llm", prompt="Goal: summarize the Temporal docs\nMake a plan.")


def test_plan_turns_a_tool_answer_into_a_tool_step():
    agent = AgentRun()
    agent.step, agent.context_summary = 1, "\nTOOL: fetch temporal-docs"
    assert agent.plan() == Step("tool", tool="fetch", args={"target": "temporal-docs"})


def test_plan_finishes_on_final_and_gives_up_after_three_failures():
    agent = AgentRun()
    agent.step, agent.context_summary = 3, "\nFINAL: done"
    assert agent.plan().kind == "finish"
    agent.context_summary, agent.attempts = "\nRESULT: error x", MAX_TOOL_FAILURES
    assert agent.plan().kind == "finish"          # the give-up branch, no model asked


# ---------------------------------------------------------------- the loop, Activities mocked
async def test_agent_loop_calls_one_tool_then_answers():
    tq = f"lab08-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                          activities=[fakes.fake_call_llm, fakes.fake_execute_tool]):
            result = await env.client.execute_workflow(
                AgentRun.run, "find the answer", id=f"lab08-test-{uuid.uuid4()}", task_queue=tq)
    assert result == "42"
    assert fakes.calls == {"call_llm": 2, "execute_tool": 1}
    # The key the Workflow handed the tool: workflow_id:step:hash — step 1 is the tool step.
    assert fakes.tool_keys[0].split(":")[1] == "1"


async def test_agent_gives_up_after_three_failed_tool_steps():
    tq = f"lab08-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                          activities=[fakes.fake_call_llm, fakes.failing_execute_tool]):
            result = await env.client.execute_workflow(
                AgentRun.run, "find the answer", id=f"lab08-test-{uuid.uuid4()}", task_queue=tq)
    assert result == f"gave up after {MAX_TOOL_FAILURES} failed tool steps"
    assert fakes.calls["execute_tool"] == MAX_TOOL_FAILURES


async def test_paused_agent_expires_after_72_hours_in_under_a_second():
    """The 72-hour checkpoint: paused before its first step, never resumed."""
    tq = f"lab08-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                          activities=[fakes.fake_call_llm, fakes.fake_execute_tool]):
            handle = await env.client.start_workflow(
                AgentRun.run, "needs approval", id=f"lab08-test-{uuid.uuid4()}", task_queue=tq,
                start_signal="pause")               # delivered before the first Workflow Task
            result = await handle.result()
            assert result == "expired: no resume within 72h"
            # The history really contains a 72-hour timer; the test did not wait for it.
            timers = [e for e in (await handle.fetch_history()).events
                      if e.HasField("timer_started_event_attributes")]
            assert len(timers) == 1
            assert timers[0].timer_started_event_attributes.start_to_fire_timeout.ToTimedelta() == CHECKPOINT_TIMEOUT
            assert fakes.calls["call_llm"] == 0     # it never got past the checkpoint
    assert CHECKPOINT_TIMEOUT == timedelta(hours=72)


async def test_resume_signal_releases_the_checkpoint():
    tq = f"lab08-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                          activities=[fakes.fake_call_llm, fakes.fake_execute_tool]):
            handle = await env.client.start_workflow(
                AgentRun.run, "needs approval", id=f"lab08-test-{uuid.uuid4()}", task_queue=tq,
                start_signal="pause")
            # Time skipping only advances while the Workflow is blocked *and* nobody holds the
            # clock: a query then a signal is the human coming back a day later.
            assert (await handle.query(AgentRun.status))["status"] == "paused"
            await handle.signal(AgentRun.resume)
            assert await handle.result() == "42"

"""8.2 — Replay tests: every history in tests/histories/ must replay under the current AgentRun.

A replay test is a Worker coming back — against a file, in CI, before any real Worker sees the
new code. `Replayer` feeds the recorded events to the Workflow code and checks that the code
issues the same Commands the history recorded; a mismatch is the non-determinism error a real
Worker would raise on its first Workflow Task."""
from __future__ import annotations

import asyncio
from pathlib import Path

import pytest
from temporalio import workflow
from temporalio.client import WorkflowHistory
from temporalio.worker import Replayer

from workflows import AgentRun, AgentState

HISTORIES = Path(__file__).parent / "histories"


def load(path: Path) -> WorkflowHistory:
    # The filename is the Workflow ID (export_history.py names files that way), so from_json
    # gets it for nothing. `temporal workflow show --output json` and the UI download match.
    return WorkflowHistory.from_json(path.stem, path.read_text())


@pytest.mark.parametrize("path", sorted(HISTORIES.glob("*.json")), ids=lambda p: p.stem)
async def test_history_replays_under_current_code(path: Path):
    await Replayer(workflows=[AgentRun]).replay_workflow(load(path))   # raises on non-determinism


# ---------------------------------------------------------------- the failure, on purpose
@workflow.defn(name="AgentRun")            # same Workflow Type: the history says "AgentRun"
class AgentRunWithAnExtraTimer(AgentRun):
    """The change 9.1 warns about, in its smallest form: one new Command before the first
    Activity. Every recorded history says "schedule call_llm" where this code says "start a
    timer"; the replay must fail on every one of them."""

    @workflow.run
    async def run(self, goal: str, state: AgentState | None = None) -> str:
        await asyncio.sleep(1)                 # NEW: a Command the histories never recorded
        return await super().run(goal, state)


@pytest.mark.parametrize("path", sorted(HISTORIES.glob("*.json")), ids=lambda p: p.stem)
async def test_changed_workflow_fails_replay(path: Path):
    result = await Replayer(workflows=[AgentRunWithAnExtraTimer]).replay_workflow(
        load(path), raise_on_replay_failure=False)
    assert result.replay_failure is not None, "the changed code replayed clean — it should not"
    assert isinstance(result.replay_failure, workflow.NondeterminismError)
    # The message names the event the code could not match: the first ActivityTaskScheduled.
    assert "ActivityTaskScheduled" in str(result.replay_failure)

"""8.3 — Activity tests with ActivityEnvironment: heartbeats, cancellation, idempotency by key."""
from __future__ import annotations

import asyncio

import pytest
from temporalio.exceptions import ApplicationError
from temporalio.testing import ActivityEnvironment

import workflows
from workflows import ToolCall, execute_tool


def call(key: str = "agent-test:3:9f3c", tool: str = "fetch") -> ToolCall:
    return ToolCall(key=key, tool=tool, args={"target": "temporal-docs"})


async def test_execute_tool_heartbeats_with_progress_details():
    env = ActivityEnvironment()
    beats: list[tuple] = []
    env.on_heartbeat = lambda *details: beats.append(details)
    result = await env.run(execute_tool, call())
    assert result.ok
    assert beats == [("fetch:1/3",), ("fetch:2/3",), ("fetch:3/3",)]


async def test_execute_tool_is_idempotent_by_key():
    env = ActivityEnvironment()
    first = await env.run(execute_tool, call())
    second = await env.run(execute_tool, call())      # same key: a retry after a Worker death
    assert first == second
    assert workflows.effects() == ["fetch(temporal-docs)"]          # one effect, not two


async def test_a_different_key_is_a_different_effect():
    env = ActivityEnvironment()
    await env.run(execute_tool, call(key="agent-test:3:aaaa"))
    await env.run(execute_tool, call(key="agent-test:5:bbbb"))
    assert workflows.effects() == ["fetch(temporal-docs)", "fetch(temporal-docs)"]


async def test_cancellation_stops_the_tool_before_its_side_effect():
    env = ActivityEnvironment()                       # one env per test: cancel() is sticky
    beats: list[tuple] = []
    env.on_heartbeat = lambda *details: beats.append(details)
    task = asyncio.create_task(env.run(execute_tool, call()))
    await asyncio.sleep(0.07)                         # after the first heartbeat, before the last
    env.cancel()                                      # what the Service does on the next heartbeat
    with pytest.raises(asyncio.CancelledError):
        await task
    assert beats, "a long tool must heartbeat or it can never be cancelled"
    assert workflows.effects() == []                  # cancelled before the effect


async def test_unknown_tool_is_a_non_retryable_application_error():
    env = ActivityEnvironment()
    with pytest.raises(ApplicationError) as err:
        await env.run(execute_tool, call(tool="rm-rf"))
    assert err.value.type == "BadToolArguments" and err.value.non_retryable

"""8.4 — Failure injection: Worker crash, Activity timeout, duplicate execution, cancellation.

Each is a test, not a war story. Every one ends with the recurring question: what survived,
and what did the returning Worker do with it."""
from __future__ import annotations

import asyncio
import os
import signal
import subprocess
import sys
import uuid
from datetime import timedelta

import pytest
from temporalio import activity
from temporalio.api.enums.v1 import PendingActivityState, TimeoutType
from temporalio.client import WorkflowFailureError
from temporalio.exceptions import ApplicationError, CancelledError
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker

import fakes
import workflows
from workflows import AgentRun, ToolCall, ToolResult, execute_tool

HERE = os.path.dirname(os.path.abspath(__file__))


@pytest.fixture(autouse=True)
def _reset_fakes():
    fakes.reset()


async def pending_activity_started(handle) -> bool:
    """`describe` shows what the history does not: the attempt currently running."""
    desc = await handle.describe()
    return any(a.state == PendingActivityState.PENDING_ACTIVITY_STATE_STARTED
               for a in desc.raw_description.pending_activities)


# ---------------------------------------------------------------- 1. Worker crash
async def test_worker_crash_mid_activity_the_next_worker_finishes_the_run():
    """kill -9 a Worker while execute_tool is running; a fresh Worker finishes the run.

    Runs on `start_local()` — a real dev server in a subprocess — rather than the time-skipping
    server, because a crash exercises two Service timeouts the time-skipping server was observed
    not to fire for a dead Worker's sticky queue: the Activity's heartbeat timeout (5 s) and the
    sticky Workflow Task's schedule-to-start timeout (5 s). Expect ~10 s of real time; that is
    the honest cost of a crash, and the history below is what you saw in lab 1."""
    tq = f"lab08-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_local() as env:
        address = env.client.service_client.config.target_host
        proc = subprocess.Popen(
            [sys.executable, os.path.join(HERE, "crashable_worker.py")],
            env={**os.environ, "TEMPORAL_ADDRESS": address, "TEMPORAL_NAMESPACE": env.client.namespace,
                 "TASK_QUEUE": tq},
            stdout=subprocess.PIPE, text=True)
        try:
            assert proc.stdout.readline().strip() == "ready"
            handle = await env.client.start_workflow(
                AgentRun.run, "survive a crash", id=f"lab08-test-{uuid.uuid4()}", task_queue=tq)
            while not await pending_activity_started(handle):   # execute_tool is running in the subprocess
                await asyncio.sleep(0.1)
            proc.send_signal(signal.SIGKILL)                       # kill -9, not Ctrl-C
            proc.wait(timeout=10)
            # Nothing reports the death. The Service learns of it only when the heartbeat stops
            # arriving; heartbeat_timeout later the attempt is timed out and attempt 2 scheduled.
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                              activities=[fakes.fake_call_llm, fakes.fake_execute_tool]):
                assert await asyncio.wait_for(handle.result(), timeout=60) == "42"
            events = (await handle.fetch_history()).events
            started = [e.activity_task_started_event_attributes for e in events
                       if e.HasField("activity_task_started_event_attributes")]
            # The history records one ActivityTaskStarted per *closed* Activity, carrying the
            # attempt that closed it: attempt 2, on the new Worker, with attempt 1's failure —
            # a heartbeat timeout — as last_failure. Attempt 1 itself left no event.
            tool_start = next(a for a in started if a.attempt == 2)
            assert not tool_start.identity.startswith("crashable-")
            assert tool_start.last_failure.timeout_failure_info.timeout_type == TimeoutType.TIMEOUT_TYPE_HEARTBEAT
            # And the dead Worker's sticky queue: its Workflow Task timed out on schedule-to-start
            # before the Service re-offered it on the normal queue to the Worker that was alive.
            assert any(e.HasField("workflow_task_timed_out_event_attributes") and
                       e.workflow_task_timed_out_event_attributes.timeout_type == TimeoutType.TIMEOUT_TYPE_SCHEDULE_TO_START
                       for e in events)
        finally:
            if proc.poll() is None:
                proc.kill()


# ---------------------------------------------------------------- 2. Activity timeout
@activity.defn(name="call_llm")
async def call_llm_that_stops_on_error(prompt: str, key: str) -> str:
    if "RESULT: error" in prompt:
        return "FINAL: gave up, the tool timed out"
    return "TOOL: fetch the-answer"


async def test_activity_timeout_is_retried_then_surfaces_as_an_error_the_workflow_handles():
    tq = f"lab08-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                          activities=[call_llm_that_stops_on_error, fakes.hanging_execute_tool]):
            handle = await env.client.start_workflow(
                AgentRun.run, "wait forever", id=f"lab08-test-{uuid.uuid4()}", task_queue=tq)
            # A hung tool never heartbeats. Each attempt dies on heartbeat_timeout (5 s); the
            # RetryPolicy allows 3. Time is skipped past each one instead of waited for.
            for _ in range(3):
                await env.sleep(timedelta(seconds=7))
            result = await handle.result()
            assert result == "gave up, the tool timed out"
            status = await handle.query(AgentRun.status)
            assert "RESULT: error TimeoutError" in status["context_summary"]
            assert status["attempts"] == 1
            assert fakes.calls["execute_tool"] == 3          # three attempts ran, none finished


# ---------------------------------------------------------------- 3. Duplicate execution
@activity.defn(name="execute_tool")
async def execute_tool_reported_late(call: ToolCall) -> ToolResult:
    """The real tool, but attempt 1 dies *after* the side effect and before reporting — the
    duplicate-execution injection. Attempt 2 receives the same ToolCall, same key."""
    result = await execute_tool(call)
    if activity.info().attempt == 1:
        raise ApplicationError("worker lost after the effect", type="LostReport")
    return result


@activity.defn(name="execute_tool")
async def execute_tool_ignoring_its_key(call: ToolCall) -> ToolResult:
    """The bug this injection exists to find: a tool with no system of record."""
    workflows._EFFECT_LOG.append(f"{call.tool}({call.args.get('target', '')})")
    if activity.info().attempt == 1:
        raise ApplicationError("worker lost after the effect", type="LostReport")
    return ToolResult(key=call.key, ok=True, output="done")


async def _run_duplicate_injection(tool) -> list[str]:
    tq = f"lab08-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                          activities=[fakes.fake_call_llm, tool]):
            await env.client.execute_workflow(
                AgentRun.run, "do it once", id=f"lab08-test-{uuid.uuid4()}", task_queue=tq)
    return workflows.effects()


async def test_duplicate_delivery_with_the_same_key_has_one_effect():
    assert await _run_duplicate_injection(execute_tool_reported_late) == ["fetch(the-answer)"]


async def test_the_injection_catches_a_tool_that_ignores_its_key():
    # This is what the suite is for. A tool without a system of record double-fires under the
    # exact retry the happy path never exercises. (Asserting the bug so the lesson is visible;
    # in your own suite this test asserts == 1 and the tool that fails it gets fixed.)
    assert await _run_duplicate_injection(execute_tool_ignoring_its_key) == ["fetch(the-answer)", "fetch(the-answer)"]


# ---------------------------------------------------------------- 4. Cancellation
async def test_cancel_during_a_long_tool_stops_it_and_runs_the_compensation():
    tq = f"lab08-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                          activities=[fakes.fake_call_llm_two_tools, fakes.slow_execute_tool]):
            handle = await env.client.start_workflow(
                AgentRun.run, "take your time", id=f"lab08-test-{uuid.uuid4()}", task_queue=tq)
            # Let the first tool (fetch) finish so there is something to compensate, then
            # cancel while the second (summarize) is running.
            while fakes.calls["execute_tool"] < 2 or not await pending_activity_started(handle):
                await asyncio.sleep(0.05)
            await handle.cancel()
            with pytest.raises(WorkflowFailureError) as err:
                await handle.result()
            assert isinstance(err.value.cause, CancelledError)
            # The tool learns of it on its next *sent* heartbeat — the SDK throttles heartbeats
            # to 80% of heartbeat_timeout (5 s), so this can take up to ~4 s of real time.
            for _ in range(120):
                if fakes.cancelled["execute_tool"]:
                    break
                await asyncio.sleep(0.05)
            assert fakes.cancelled["execute_tool"] == 1           # the running tool was cancelled ...
            assert [k for k in fakes.tool_keys if k.endswith(":undo")], "... and undo_fetch ran"
            assert (await handle.query(AgentRun.status))["status"] == "cancelled"

import asyncio
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.worker import Worker  # noqa: E402
from workflows import AgentRun, call_llm, execute_tool  # noqa: E402


async def main() -> None:
    client = await connect()
    worker = Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=[AgentRun],
        activities=[call_llm, execute_tool],
    )
    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 8: AgentRun, stage 8 — the loop from labs 1–7, written to be tested.

Carried from earlier stages: `plan()` (a pure function of state, L01), `call_llm` and
`execute_tool` as Activities with timeouts, a RetryPolicy, heartbeats and an idempotency
key (L04), the `pause`/`resume`/`cancel_run`/`inject_context` Signals, the `status` Query
and the `change_goal` Update with `processed_command_ids` dedupe (L05), the `AgentState`
snapshot and `continue_as_new` (L06), and the `compensations` list run in reverse on
cancellation (L07). The `ResearchAgent` child and the four-Activity GPU lane stay in
labs 6 and 7 — this file is the surface the tests in `tests/` exercise.

Nothing here calls a real model: `call_llm` is a stand-in with fixed rules so the loop is
deterministic and free. The tests replace it anyway (`@activity.defn(name="call_llm")`).
"""
from __future__ import annotations

import asyncio
import hashlib
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


# The agent parks here when paused and waits for a human. Three days, not three seconds:
# the time-skipping test environment is what makes this testable.
CHECKPOINT_TIMEOUT = timedelta(hours=72)
MAX_TOOL_FAILURES = 3          # plan() gives up after this many failed tool steps
TOOL_HEARTBEATS = 3            # execute_tool heartbeats this many times per attempt
TOOL_TICK_SECONDS = 0.05       # ... this far apart


# ---------------------------------------------------------------- data types
@dataclass
class Step:
    kind: str                  # "llm" | "tool" | "finish"
    prompt: str = ""
    tool: str = ""
    args: dict = field(default_factory=dict)


@dataclass
class ToolCall:
    key: str                   # idempotency key: f"{workflow_id}:{step}:{sha256(tool+args)[:12]}"
    tool: str
    args: dict


@dataclass
class ToolResult:
    key: str
    ok: bool
    output: str


@dataclass
class AgentState:
    """The snapshot continue_as_new carries into the next run (L06)."""
    step: int
    goal: str
    context_summary: str
    attempts: int
    processed_command_ids: list[str]


# ---------------------------------------------------------------- activities
@activity.defn
async def call_llm(prompt: str, key: str) -> str:
    """The model call as an Activity. This stand-in follows three rules so the loop is
    deterministic: with no tool result in the prompt it asks for `fetch`; with one it asks
    for `summarize`; with two it answers. Tests replace it with their own fake."""
    activity.logger.info("call_llm key=%s", key)
    ok_results = prompt.count("RESULT: ok")
    if ok_results == 0:
        return "TOOL: fetch temporal-docs"
    if ok_results == 1:
        return "TOOL: summarize temporal-docs"
    return "FINAL: the agent fetched and summarized the docs"


# The tool's system of record: one effect per key. In production this is a table keyed by
# the idempotency key; here it is a dict in the Worker process, which is exactly enough to
# test "run it twice with the same key, assert one effect" (8.3) and exactly *not* enough to
# survive a Worker death (8.4 shows why the real one must be outside the process).
_EFFECTS: dict[str, ToolResult] = {}
_EFFECT_LOG: list[str] = []    # every side effect that actually happened, in order


def effects() -> list[str]:
    return list(_EFFECT_LOG)


def reset_effects() -> None:
    _EFFECTS.clear()
    _EFFECT_LOG.clear()


@activity.defn
async def execute_tool(call: ToolCall) -> ToolResult:
    """Run a tool. Heartbeats so a long tool can be cancelled and its progress seen from
    `describe`; idempotent by `call.key` so a retry after a Worker death does not repeat
    the side effect."""
    prior = _EFFECTS.get(call.key)
    if prior is not None:
        activity.logger.info("execute_tool key=%s already done; returning recorded result", call.key)
        return prior
    if call.tool not in ("fetch", "summarize", "undo_fetch", "undo_summarize"):
        raise ApplicationError(f"unknown tool {call.tool!r}", type="BadToolArguments", non_retryable=True)
    for i in range(TOOL_HEARTBEATS):
        # Heartbeat details are what a retry (or a forensic reader) sees as "how far it got".
        activity.heartbeat(f"{call.tool}:{i + 1}/{TOOL_HEARTBEATS}")
        await asyncio.sleep(TOOL_TICK_SECONDS)   # cancellation arrives here
    _EFFECT_LOG.append(f"{call.tool}({call.args.get('target', '')})")   # the side effect
    result = ToolResult(key=call.key, ok=True, output=f"{call.tool}: {call.args.get('target', '')} done")
    _EFFECTS[call.key] = result
    return result


# ---------------------------------------------------------------- the workflow
@workflow.defn
class AgentRun:
    def __init__(self) -> None:
        self.step = 0
        self.status = "running"
        self.goal = ""
        self.context_summary = ""
        self.attempts = 0
        self.paused = False
        self.processed_command_ids: set[str] = set()
        self.compensations: list[ToolCall] = []

    # -- plan(): Workflow code, a pure function of state. Never calls a model. --
    def plan(self) -> Step:
        if self.status == "cancelling":
            return Step("finish")
        if self.attempts >= MAX_TOOL_FAILURES:
            return Step("finish")
        last = self.context_summary.rsplit("\n", 1)[-1]
        if self.step == 0 or not last:
            return Step("llm", prompt=f"Goal: {self.goal}\nMake a plan.")
        if last.startswith("TOOL:"):
            tool, _, target = last[len("TOOL:"):].strip().partition(" ")
            return Step("tool", tool=tool, args={"target": target})
        if last.startswith("FINAL:"):
            return Step("finish")
        # A RESULT line (ok or error): ask the model what to do next.
        return Step("llm", prompt=f"Goal: {self.goal}\n{self.context_summary}\nWhat next?")

    def snapshot(self) -> AgentState:
        return AgentState(self.step, self.goal, self.context_summary, self.attempts,
                          sorted(self.processed_command_ids))

    # -- Signals / Query / Update (L05) --
    @workflow.signal
    def pause(self) -> None:
        self.paused = True

    @workflow.signal
    def resume(self) -> None:
        self.paused = False

    @workflow.signal
    def cancel_run(self) -> None:
        self.status = "cancelling"

    @workflow.signal
    def inject_context(self, text: str, command_id: str) -> None:
        if command_id in self.processed_command_ids:
            return                                   # a retried Signal is ignored
        self.processed_command_ids.add(command_id)
        self.context_summary += f"\nNOTE: {text}"

    @workflow.query
    def status(self) -> dict:
        return {"step": self.step, "status": self.status, "goal": self.goal,
                "attempts": self.attempts, "paused": self.paused,
                "context_summary": self.context_summary}

    @workflow.update
    def change_goal(self, goal: str, command_id: str) -> str:
        if command_id in self.processed_command_ids:
            return self.goal                         # a retried Update returns the same answer
        self.processed_command_ids.add(command_id)
        self.goal = goal
        return self.goal

    @change_goal.validator
    def _validate_change_goal(self, goal: str, command_id: str) -> None:
        if not goal.strip():
            raise ValueError("goal must not be empty")

    # -- the loop --
    @workflow.run
    async def run(self, goal: str, state: AgentState | None = None) -> str:
        self.goal = goal
        if state is not None:                        # continued-as-new: restore the snapshot
            self.step, self.goal = state.step, state.goal
            self.context_summary, self.attempts = state.context_summary, state.attempts
            self.processed_command_ids = set(state.processed_command_ids)
        try:
            while True:
                if self.paused:
                    self.status = "paused"
                    try:
                        await workflow.wait_condition(
                            lambda: not self.paused or self.status == "cancelling",
                            timeout=CHECKPOINT_TIMEOUT)
                    except asyncio.TimeoutError:
                        self.status = "expired"
                        return "expired: no resume within 72h"
                    if self.status != "cancelling":
                        self.status = "running"
                step = self.plan()
                if step.kind == "finish":
                    self.status = "cancelled" if self.status == "cancelling" else "done"
                    if self.status == "cancelled":
                        await self._compensate()
                    return self._final_answer()
                if step.kind == "llm":
                    await self._llm_step(step)
                elif step.kind == "tool":
                    await self._tool_step(step)
                self.step += 1
                if workflow.info().is_continue_as_new_suggested():
                    workflow.continue_as_new(args=[self.goal, self.snapshot()])
        except (asyncio.CancelledError, ActivityError) as err:
            # Workflow cancellation (handle.cancel()) is cooperative: it surfaces at the await
            # we are blocked on — as asyncio.CancelledError on a timer or wait_condition, as an
            # ActivityError whose cause is CancelledError on a running Activity. (Lab 8.4's
            # cancellation injection found the version of this code that treated the second
            # form as a failed tool step and kept going.) Undo what was done, then let the
            # cancellation stand so the execution closes as Cancelled.
            if isinstance(err, ActivityError) and not isinstance(err.cause, CancelledError):
                raise
            self.status = "cancelled"
            await self._compensate()
            raise

    async def _llm_step(self, step: Step) -> None:
        key = f"{workflow.info().workflow_id}:{self.step}:{_sha(step.prompt)}"
        answer = await workflow.execute_activity(
            call_llm, args=[step.prompt, key],
            start_to_close_timeout=timedelta(seconds=30),
            retry_policy=RetryPolicy(maximum_attempts=3))
        self.context_summary += f"\n{answer}"

    async def _tool_step(self, step: Step) -> None:
        call = ToolCall(
            key=f"{workflow.info().workflow_id}:{self.step}:{_sha(step.tool + repr(sorted(step.args.items())))}",
            tool=step.tool, args=step.args)
        try:
            result = await workflow.execute_activity(
                execute_tool, call,
                start_to_close_timeout=timedelta(seconds=30),
                heartbeat_timeout=timedelta(seconds=5),
                retry_policy=RetryPolicy(initial_interval=timedelta(seconds=1),
                                         backoff_coefficient=1.0, maximum_attempts=3))
        except ActivityError as err:
            if isinstance(err.cause, CancelledError):
                raise            # cancellation is not a failed tool step; run() handles it
            # All attempts spent (or a non-retryable error): record it and let plan() decide.
            self.attempts += 1
            self.context_summary += f"\nRESULT: error {_failure_label(err)}"
            return
        if result.ok:
            self.context_summary += f"\nRESULT: ok {result.output}"
            self.compensations.append(ToolCall(key=call.key + ":undo", tool=f"undo_{call.tool}", args=call.args))
        else:
            self.attempts += 1
            self.context_summary += f"\nRESULT: error {result.output}"

    async def _compensate(self) -> None:
        """L07: undo completed tool steps in reverse. Each undo is itself a keyed tool call."""
        while self.compensations:
            undo = self.compensations.pop()
            await workflow.execute_activity(
                execute_tool, undo, start_to_close_timeout=timedelta(seconds=30),
                retry_policy=RetryPolicy(maximum_attempts=3))

    def _final_answer(self) -> str:
        if self.status == "cancelled":
            return "cancelled"
        if self.attempts >= MAX_TOOL_FAILURES:
            return f"gave up after {self.attempts} failed tool steps"
        last = self.context_summary.rsplit("\n", 1)[-1]
        return last.removeprefix("FINAL:").strip()


def _failure_label(err: ActivityError) -> str:
    """'BadToolArguments' for an ApplicationError with a type, else the failure class name
    ('TimeoutError', 'CancelledError'): what plan() and the forensic reader get to see."""
    cause = err.cause
    if isinstance(cause, ApplicationError) and cause.type:
        return cause.type
    return type(cause).__name__


def _sha(text: str) -> str:
    return hashlib.sha256(text.encode()).hexdigest()[:12]