Lab 4: Activities and Failure Semantics

"""The GPU scheduler's callback, 45 minutes later — from any process that has a Client.

    python complete_async.py            # completes every token found under .state/async/
    python complete_async.py --fail     # or fails them

Nothing here knows about the Worker; the Activity is completed through the Service by task token.
"""
import argparse
import asyncio
import glob
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import connect  # noqa: E402
from temporalio.exceptions import ApplicationError  # noqa: E402
from workflows import STATE_DIR, ToolResult  # noqa: E402


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--fail", action="store_true")
    args = ap.parse_args()
    client = await connect()
    tokens = sorted(glob.glob(os.path.join(STATE_DIR, "async", "*.token")))
    if not tokens:
        print("no pending task tokens under", os.path.join(STATE_DIR, "async"))
    for path in tokens:
        with open(path, "rb") as f:
            token = f.read()
        handle = client.get_async_activity_handle(task_token=token)
        await handle.heartbeat("job finished")
        if args.fail:
            await handle.fail(ApplicationError("GPU job failed", type="GpuJobFailed"))
        else:
            await handle.complete(ToolResult(key="finetune", tool="finetune", output="artifact://agent-7b-ft"))
        os.remove(path)
        print(f"{'failed' if args.fail else 'completed'} the Activity behind {os.path.basename(path)}")


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

"""The Activity side of a history: every attempt the Service recorded, with timings.

    python history.py agent-42-start_to_close
"""
from __future__ import annotations

import asyncio
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import connect  # noqa: E402
from temporalio.api.enums.v1 import EventType, RetryState, TimeoutType  # noqa: E402


def _payloads(payloads) -> str:
    return ", ".join(p.data.decode("utf-8", "replace") for p in payloads)


def _failure(f) -> str:
    parts = []
    if f.HasField("timeout_failure_info"):
        ti = f.timeout_failure_info
        parts.append(f"timeout_type={TimeoutType.Name(ti.timeout_type).removeprefix('TIMEOUT_TYPE_')}")
        if ti.last_heartbeat_details.payloads:
            parts.append(f"last_heartbeat_details=[{_payloads(ti.last_heartbeat_details.payloads)}]")
    if f.HasField("application_failure_info"):
        ai = f.application_failure_info
        parts.append(f"type={ai.type!r} non_retryable={ai.non_retryable}")
    if f.message:
        parts.append(f"message={f.message!r}")
    if f.HasField("cause"):
        parts.append(f"cause=({_failure(f.cause)})")
    return " ".join(parts)


async def main() -> None:
    wf_id = sys.argv[1] if len(sys.argv) > 1 else "agent-42-default"
    client = await connect()
    handle = client.get_workflow_handle(wf_id)
    t0 = None
    async for ev in handle.fetch_history_events():
        ts = ev.event_time.ToDatetime()
        t0 = t0 or ts
        name = EventType.Name(ev.event_type).removeprefix("EVENT_TYPE_")
        d = ""
        if ev.HasField("activity_task_scheduled_event_attributes"):
            a = ev.activity_task_scheduled_event_attributes
            rp = a.retry_policy
            d = (f"activity={a.activity_type.name} queue={a.task_queue.name} "
                 f"s2s={a.schedule_to_start_timeout.ToTimedelta()} stc={a.start_to_close_timeout.ToTimedelta()} "
                 f"s2c={a.schedule_to_close_timeout.ToTimedelta()} hb={a.heartbeat_timeout.ToTimedelta()} "
                 f"max_attempts={rp.maximum_attempts}")
        elif ev.HasField("activity_task_started_event_attributes"):
            a = ev.activity_task_started_event_attributes
            d = f"worker={a.identity} attempt={a.attempt}"
            if a.HasField("last_failure"):
                d += f" last_failure=({_failure(a.last_failure)})"
        elif ev.HasField("activity_task_completed_event_attributes"):
            a = ev.activity_task_completed_event_attributes
            d = f"result=[{_payloads(a.result.payloads)[:120]}]"
        elif ev.HasField("activity_task_failed_event_attributes"):
            a = ev.activity_task_failed_event_attributes
            d = f"retry_state={RetryState.Name(a.retry_state).removeprefix('RETRY_STATE_')} {_failure(a.failure)}"
        elif ev.HasField("activity_task_timed_out_event_attributes"):
            a = ev.activity_task_timed_out_event_attributes
            d = f"retry_state={RetryState.Name(a.retry_state).removeprefix('RETRY_STATE_')} {_failure(a.failure)}"
        elif ev.HasField("workflow_execution_completed_event_attributes"):
            d = f"result=[{_payloads(ev.workflow_execution_completed_event_attributes.result.payloads)[:200]}]"
        elif ev.HasField("workflow_task_started_event_attributes"):
            d = f"worker={ev.workflow_task_started_event_attributes.identity}"
        print(f"{ev.event_id:3d}  +{(ts - t0).total_seconds():6.2f}s  {name:28s} {d}")
    desc = (await handle.describe()).raw_description
    for pa in desc.pending_activities:
        hb = _payloads(pa.heartbeat_details.payloads) if pa.HasField("heartbeat_details") else ""
        print(f"pending activity: {pa.activity_type.name} attempt={pa.attempt} state={pa.state} "
              f"heartbeat_details=[{hb}] last_failure=({_failure(pa.last_failure) if pa.HasField('last_failure') else ''})")


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

"""Start AgentRun in one of the lab's scenarios and print what the Workflow made of the outcome.

    python starter.py                          # default: call_llm, echo tool with the specified options, finish
    python starter.py schedule_to_start        # 4.1  (also: start_to_close, schedule_to_close, heartbeat)
    python starter.py shards                   # 4.2
    python starter.py charge | send            # 4.3
    python starter.py finetune                 # 4.5
    python starter.py heartbeat --id agent-42-hb --no-wait

The Workflow ID defaults to agent-42-<scenario>. The Workflow catches the ActivityError, so the
*result* names the exception class and timeout type; the history has the Event.
"""
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 workflows import SCENARIOS, AgentRun  # noqa: E402


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("scenario", nargs="?", default="default", choices=sorted(SCENARIOS))
    ap.add_argument("--id", default=None)
    ap.add_argument("--no-wait", action="store_true")
    args = ap.parse_args()
    wf_id = args.id or f"{os.environ.get('WORKFLOW_ID_PREFIX', 'agent-42')}-{args.scenario}"

    client = await connect()
    handle = await start_workflow(client,
        AgentRun.run, args=["summarize the Temporal docs", args.scenario], id=wf_id, task_queue=TASK_QUEUE
    )
    print(f"started {handle.id} scenario={args.scenario}; open http://localhost:8233/namespaces/default/workflows/{handle.id}")
    if args.no_wait:
        return
    print("result:", await handle.result())
    print(f"history: python history.py {handle.id}")


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

"""Check yourself: each timeout is observable, a checkpoint survives an attempt, an Activity is
at-least-once until you key it, and an Activity can complete from outside the Worker.

Runs on the time-skipping test server; Activity sleeps are real seconds, so this takes ~1 min.
    PYTHONPATH=.. pytest tests/
"""
from __future__ import annotations

import asyncio
import glob
import os
import sys

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, RetryState, TimeoutType  # noqa: E402
from temporalio.testing import WorkflowEnvironment  # noqa: E402
from temporalio.worker import Worker  # noqa: E402

TASK_QUEUE = "lab-04-test"


@pytest.fixture
def state_dir(tmp_path, monkeypatch):
    monkeypatch.setattr(workflows, "STATE_DIR", str(tmp_path))
    return tmp_path


async def _run(scenario: str, wf_id: str, before_result=None):
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=TASK_QUEUE, workflows=[workflows.AgentRun],
                          activities=[workflows.call_llm, workflows.execute_tool, workflows.charge]):
            handle = await env.client.start_workflow(
                workflows.AgentRun.run, args=["test goal", scenario], id=wf_id, task_queue=TASK_QUEUE
            )
            if before_result is not None:
                await before_result(env.client, handle)
            result = await handle.result()
            history = await handle.fetch_history()
            return result, history


def _terminal_activity_events(history) -> list[str]:
    """'TIMED_OUT:<type>' or 'FAILED:<retry_state>' for every closed execute_tool attempt chain."""
    out = []
    for e in history.events:
        if e.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_TIMED_OUT:
            tt = e.activity_task_timed_out_event_attributes.failure.timeout_failure_info.timeout_type
            out.append("TIMED_OUT:" + TimeoutType.Name(tt).removeprefix("TIMEOUT_TYPE_"))
        elif e.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_FAILED:
            rs = e.activity_task_failed_event_attributes.retry_state
            out.append("FAILED:" + RetryState.Name(rs).removeprefix("RETRY_STATE_"))
    return out


@pytest.mark.parametrize("scenario,expected", [
    ("schedule_to_start", "SCHEDULE_TO_START"),
    ("start_to_close", "START_TO_CLOSE"),
    ("heartbeat", "HEARTBEAT"),
])
def test_each_timeout_is_one_event_and_one_exception(state_dir, scenario, expected):
    result, history = asyncio.run(_run(scenario, f"lab04-test-{scenario}"))
    assert _terminal_activity_events(history) == [f"TIMED_OUT:{expected}"]
    assert f"FAILED: TimeoutError type={expected}" in result


def test_schedule_to_close_bounds_all_attempts(state_dir):
    """The budget for the whole retry chain. Its footprint depends on the server: the dev server
    writes ActivityTaskTimedOut(SCHEDULE_TO_CLOSE) carrying the last attempt's failure as cause;
    the time-skipping test server (which skips the backoff, so more attempts fit) writes
    ActivityTaskFailed with retry_state=TIMEOUT and the last ApplicationError. Either way: one
    terminal Event, several attempts, and the Workflow's decision afterwards."""
    result, history = asyncio.run(_run("schedule_to_close", "lab04-test-schedule_to_close"))
    terminal = _terminal_activity_events(history)
    assert terminal in (["TIMED_OUT:SCHEDULE_TO_CLOSE"], ["FAILED:TIMEOUT"])
    assert "FAILED: TimeoutError type=SCHEDULE_TO_CLOSE" in result or "FAILED: ApplicationError type=ToolCrash" in result
    started = [e.activity_task_started_event_attributes for e in history.events
               if e.HasField("activity_task_started_event_attributes")]
    if started[-1].scheduled_event_id > 5:          # a Started Event for the tool, if any, names attempt > 1
        assert started[-1].attempt > 1


def test_heartbeat_checkpoint_resumes_the_next_attempt(state_dir):
    """4.2: attempt 1 goes silent after shard 5; the heartbeat timeout fails it; attempt 2 reads
    heartbeat details [5] and runs shards 5..7 only. Shards 0..4 happened exactly once."""
    result, history = asyncio.run(_run("shards_hang", "lab04-test-shards"))
    assert result.endswith("| shards: 8 shards")
    lines = (state_dir / "effects.log").read_text().splitlines()
    shards = [int(line.split(" shard ")[1].split()[0]) for line in lines]
    assert shards == list(range(8))                               # no shard ran twice
    assert [line.split("attempt ")[1] for line in lines] == ["1"] * 5 + ["2"] * 3
    started = [e for e in history.events if e.HasField("activity_task_started_event_attributes")]
    assert started[-1].activity_task_started_event_attributes.attempt == 2


def test_tool_effect_happens_twice_until_keyed(state_dir, monkeypatch):
    """4.3: the send tool does its effect, then attempt 1 never reports. The Service cannot tell
    hung from dead, times the attempt out, retries, and the effect happens again. Keyed, attempt
    2 finds the key and returns without acting."""
    monkeypatch.setattr(workflows, "IDEMPOTENT", False)
    result, _ = asyncio.run(_run("send_hang", "lab04-test-send-dup"))
    sent = [l for l in (state_dir / "effects.log").read_text().splitlines() if " sent " in l]
    assert len(sent) == 2 and result.endswith("| send: sent")

    (state_dir / "effects.log").unlink()
    monkeypatch.setattr(workflows, "IDEMPOTENT", True)
    result, _ = asyncio.run(_run("send_hang", "lab04-test-send-keyed"))
    sent = [l for l in (state_dir / "effects.log").read_text().splitlines() if " sent " in l]
    assert len(sent) == 1 and result.endswith("| send: already sent (deduplicated by key)")


def test_charge_is_at_least_once_until_keyed(state_dir, monkeypatch):
    monkeypatch.setattr(workflows, "IDEMPOTENT", False)
    asyncio.run(_run("charge_hang", "lab04-test-charge-dup"))
    assert len((state_dir / "ledger.txt").read_text().splitlines()) == 2
    (state_dir / "ledger.txt").unlink()
    monkeypatch.setattr(workflows, "IDEMPOTENT", True)
    result, _ = asyncio.run(_run("charge_hang", "lab04-test-charge-keyed"))
    assert len((state_dir / "ledger.txt").read_text().splitlines()) == 1
    assert result.endswith("| charge: already charged order-1")


def test_async_completion_from_outside_the_worker(state_dir):
    """4.5: the Activity returns with raise_complete_async(); a Client completes it by task token."""
    async def complete_from_outside(client, handle):
        tokens = []
        for _ in range(50):
            tokens = glob.glob(str(state_dir / "async" / "*.token"))
            if tokens:
                break
            await asyncio.sleep(0.1)
        assert tokens, "the Activity never wrote its task token"
        with open(tokens[0], "rb") as f:
            token = f.read()
        await client.get_async_activity_handle(task_token=token).complete(
            workflows.ToolResult(key="k", tool="finetune", output="artifact://done")
        )

    result, history = asyncio.run(_run("finetune", "lab04-test-finetune", complete_from_outside))
    assert result.endswith("| finetune: artifact://done")
    assert any(e.HasField("activity_task_completed_event_attributes") for e in history.events)


if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v"]))

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

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")


async def main() -> None:
    client = await connect()
    worker = Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=[AgentRun],
        activities=[call_llm, execute_tool, charge],
    )
    print(
        f"worker pid={os.getpid()} polling {TASK_QUEUE!r} state={STATE_DIR} "
        f"idempotent={'on' if IDEMPOTENT else 'OFF'}; kill -9 {os.getpid()} to crash it",
        flush=True,
    )
    await worker.run()


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

"""Lab 4: AgentRun, stage 4 — `execute_tool` joins `call_llm`, with timeouts, a RetryPolicy,
heartbeats and an idempotency key. Everything in this file is at-least-once; the key is what
makes that safe.

AgentRun.run(goal, scenario) plans three steps: call_llm, one tool call, finish. `scenario`
selects the tool and how the Workflow invokes it (the timeouts and retry policy are Workflow
code — see SCENARIOS). The tools themselves live in `execute_tool` and write their side effects
to plain files under `.state/` so you can count how many times an effect really happened:

    .state/effects.log     one line per tool side effect  (4.2, 4.3)
    .state/ledger.txt      one line per charge            (4.3)
    .state/llm_cache/      call_llm's answers by key       (idempotent by construction)
    .state/async/          task tokens of Activities waiting for asynchronous completion (4.5)

LAB04_IDEMPOTENT=1 turns on the fix in 4.3: `charge` and the `send` tool look their key up
before acting.
"""
from __future__ import annotations

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

from common import TASK_QUEUE  # the NOBODY queue further down is named after it

STATE_DIR = os.environ.get("LAB04_STATE_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), ".state"))
IDEMPOTENT = os.environ.get("LAB04_IDEMPOTENT", "") == "1"


# ----------------------------------------------------------------------------- data types

@dataclass
class ToolCall:
    key: str            # idempotency key, computed once in Workflow code
    tool: str
    args: dict = field(default_factory=dict)


@dataclass
class ToolResult:
    key: str
    tool: str
    output: str


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


# ----------------------------------------------------------------------------- side-effect files

def _path(*parts: str) -> str:
    os.makedirs(os.path.join(STATE_DIR, *parts[:-1]), exist_ok=True)
    return os.path.join(STATE_DIR, *parts)


def _append(name: str, line: str) -> None:
    with open(_path(name), "a") as f:
        f.write(line + "\n")


def _lines(name: str) -> list[str]:
    try:
        with open(_path(name)) as f:
            return f.read().splitlines()
    except FileNotFoundError:
        return []


async def _hold(seconds: float) -> None:
    """The Worker 'dies' here in 4.3: the effect is done, the completion not yet reported."""
    await asyncio.sleep(seconds)


# ----------------------------------------------------------------------------- activities

@activity.defn
async def call_llm(prompt: str, key: str) -> str:
    """Keyed and cached: attempt 2 returns attempt 1's answer instead of paying again."""
    cache = _path("llm_cache", hashlib.sha256(key.encode()).hexdigest()[:16] + ".txt")
    if os.path.exists(cache):
        with open(cache) as f:
            activity.logger.info("call_llm key=%s cache HIT (attempt %d)", key, activity.info().attempt)
            return f.read()
    answer = f"[model answer to: {prompt[:40]}]"
    with open(cache, "w") as f:
        f.write(answer)
    activity.logger.info("call_llm key=%s billed (attempt %d)", key, activity.info().attempt)
    return answer


@activity.defn
async def charge(order_id: str, hold_seconds: float = 6.0, hang_first: bool = False) -> str:
    """The canonical illustration: the card is charged, then the Worker dies before reporting.
    `hang_first`: attempt 1 never reports (a hung Worker) — the Service cannot tell it from a dead one."""
    info = activity.info()
    if IDEMPOTENT and any(line.startswith(f"charged {order_id} ") for line in _lines("ledger.txt")):
        activity.logger.info("charge %s: already in the ledger, returning without charging (attempt %d)", order_id, info.attempt)
        return f"already charged {order_id}"
    _append("ledger.txt", f"charged {order_id} by pid {os.getpid()} attempt {info.attempt}")
    activity.logger.info("charged %s (attempt %d); holding %.0fs before reporting -- kill -9 %d now",
                         order_id, info.attempt, hold_seconds, os.getpid())
    await _hold(3600 if hang_first and info.attempt == 1 else hold_seconds)
    return f"charged {order_id}"


@activity.defn
async def execute_tool(call: ToolCall) -> ToolResult:
    info = activity.info()
    log = activity.logger

    if call.tool == "echo":
        return ToolResult(call.key, call.tool, f"echo: {call.args.get('text', '')}")

    if call.tool == "sleep":                           # no heartbeat: the Service hears nothing
        seconds = float(call.args.get("seconds", 5))
        log.info("sleep tool: %.0fs of silence (attempt %d)", seconds, info.attempt)
        await asyncio.sleep(seconds)
        return ToolResult(call.key, call.tool, f"slept {seconds}s")

    if call.tool == "fail":
        log.info("fail tool: raising (attempt %d)", info.attempt)
        raise ApplicationError(
            f"tool crashed on attempt {info.attempt}",
            type=call.args.get("type", "ToolCrash"),
            non_retryable=bool(call.args.get("non_retryable", False)),
        )

    if call.tool == "shards":                          # 4.2: heartbeat with a checkpoint
        total = int(call.args.get("shards", 10))
        done = info.heartbeat_details[0] if info.heartbeat_details else 0
        log.info("shards tool attempt %d: resuming at shard %d/%d (heartbeat details=%s)",
                 info.attempt, done, total, list(info.heartbeat_details))
        for i in range(done, total):
            await asyncio.sleep(1)                     # the work of one shard
            _append("effects.log", f"{call.key} shard {i} by pid {os.getpid()} attempt {info.attempt}")
            activity.heartbeat(i + 1)                  # checkpoint: the next attempt resumes here
            log.info("shard %d done -> heartbeat(%d)   [pid %d]", i, i + 1, os.getpid())
            if info.attempt == 1 and call.args.get("hang_after") == i + 1:
                log.info("attempt 1 hangs here: no more heartbeats (the Service cannot tell this from a dead Worker)")
                await asyncio.sleep(3600)
        return ToolResult(call.key, call.tool, f"{total} shards")

    if call.tool == "send":                            # 4.3: the effect, then the death
        if IDEMPOTENT and any(line.startswith(f"{call.key} sent ") for line in _lines("effects.log")):
            log.info("send: key %s already in effects.log, returning the recorded outcome (attempt %d)", call.key, info.attempt)
            return ToolResult(call.key, call.tool, "already sent (deduplicated by key)")
        _append("effects.log", f"{call.key} sent by pid {os.getpid()} attempt {info.attempt}")
        hold = float(call.args.get("hold_seconds", 6))
        log.info("sent (attempt %d); holding %.0fs before reporting -- kill -9 %d now", info.attempt, hold, os.getpid())
        await _hold(3600 if call.args.get("hang_first") and info.attempt == 1 else hold)
        return ToolResult(call.key, call.tool, "sent")

    if call.tool == "finetune":                        # 4.5: asynchronous completion
        token_file = _path("async", hashlib.sha256(call.key.encode()).hexdigest()[:16] + ".token")
        with open(token_file, "wb") as f:
            f.write(info.task_token)
        log.info("finetune: job submitted with task token; %s; the Activity stays open, the Worker is free",
                 token_file)
        activity.raise_complete_async()

    raise ApplicationError(f"unknown tool {call.tool!r}", type="InvalidToolArguments", non_retryable=True)


# ----------------------------------------------------------------------------- the scenarios (Workflow code)

NOBODY = TASK_QUEUE + "-gpu"   # a Task Queue no Worker polls

SCENARIOS: dict[str, dict] = {
    # the specified execute_tool call from the reading — 4.4 starts from this
    "default": dict(tool="echo", args={"text": "hello"}, options=dict(
        start_to_close_timeout=timedelta(minutes=2),
        schedule_to_close_timeout=timedelta(minutes=15),
        heartbeat_timeout=timedelta(seconds=10),
        retry_policy=RetryPolicy(maximum_attempts=3, backoff_coefficient=2.0,
                                 non_retryable_error_types=["InvalidToolArguments"]),
    )),
    # 4.1 — one timeout each
    "schedule_to_start": dict(tool="echo", args={"text": "gpu"}, options=dict(
        task_queue=NOBODY,                                  # queued where no Worker polls
        schedule_to_start_timeout=timedelta(seconds=2),
        start_to_close_timeout=timedelta(seconds=10),
    )),
    "start_to_close": dict(tool="sleep", args={"seconds": 5}, options=dict(
        start_to_close_timeout=timedelta(seconds=1),
        retry_policy=RetryPolicy(maximum_attempts=1),       # so the timeout is final, not retried
    )),
    "schedule_to_close": dict(tool="fail", args={}, options=dict(
        start_to_close_timeout=timedelta(seconds=10),
        schedule_to_close_timeout=timedelta(seconds=3),      # the budget for all attempts
        retry_policy=RetryPolicy(initial_interval=timedelta(milliseconds=500), backoff_coefficient=1.0),
    )),
    "heartbeat": dict(tool="sleep", args={"seconds": 5}, options=dict(
        start_to_close_timeout=timedelta(seconds=30),
        heartbeat_timeout=timedelta(seconds=1),              # 5 s of silence is 4 s too many
        retry_policy=RetryPolicy(maximum_attempts=1),
    )),
    # 4.2 — checkpoint resume
    "shards": dict(tool="shards", args={"shards": 10}, options=dict(
        start_to_close_timeout=timedelta(seconds=60),
        heartbeat_timeout=timedelta(seconds=3),
        retry_policy=RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3),
    )),
    "shards_hang": dict(tool="shards", args={"shards": 8, "hang_after": 5}, options=dict(
        start_to_close_timeout=timedelta(seconds=60),
        heartbeat_timeout=timedelta(seconds=3),          # silence is noticed in 3 s, not 60
        retry_policy=RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3),
    )),
    # 4.3 — duplicate execution under a Worker death
    "send": dict(tool="send", args={"hold_seconds": 6}, options=dict(
        start_to_close_timeout=timedelta(seconds=10),        # a live attempt reports in 6 s; a dead one is noticed in 10 s
        retry_policy=RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3),
    )),
    "charge": dict(tool="charge", args={"order_id": "order-1", "hold_seconds": 6}, options=dict(
        start_to_close_timeout=timedelta(seconds=10),
        retry_policy=RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3),
    )),
    # the same two without a kill: attempt 1 does the effect and then never reports
    "send_hang": dict(tool="send", args={"hold_seconds": 1, "hang_first": True}, options=dict(
        start_to_close_timeout=timedelta(seconds=3),
        retry_policy=RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3),
    )),
    "charge_hang": dict(tool="charge", args={"order_id": "order-1", "hold_seconds": 1, "hang_first": True}, options=dict(
        start_to_close_timeout=timedelta(seconds=3),
        retry_policy=RetryPolicy(initial_interval=timedelta(seconds=1), maximum_attempts=3),
    )),
    # 4.4 — a failure that must not retry: the default options, a tool that does not exist
    "bad_tool": dict(tool="no_such_tool", args={}, options=dict(
        start_to_close_timeout=timedelta(minutes=2),
        schedule_to_close_timeout=timedelta(minutes=15),
        heartbeat_timeout=timedelta(seconds=10),
        retry_policy=RetryPolicy(maximum_attempts=3, backoff_coefficient=2.0,
                                 non_retryable_error_types=["InvalidToolArguments"]),
    )),
    # 4.5 — asynchronous completion
    "finetune": dict(tool="finetune", args={"model": "agent-7b"}, options=dict(
        start_to_close_timeout=timedelta(minutes=5),
        retry_policy=RetryPolicy(maximum_attempts=1),
    )),
}


@workflow.defn
class AgentRun:
    def __init__(self) -> None:
        self.step = 0
        self.status = "running"
        self.goal = ""
        self.context_summary = ""
        self.attempts = 0
        self.scenario = "default"

    def plan(self) -> Step:
        """Workflow code: a pure function of state."""
        if self.step == 0:
            return Step("llm", prompt=f"Make a one-line plan for: {self.goal}")
        if self.step == 1:
            spec = SCENARIOS[self.scenario]
            return Step("charge" if spec["tool"] == "charge" else "tool", tool=spec["tool"], args=spec["args"])
        return Step("finish")

    @workflow.run
    async def run(self, goal: str, scenario: str = "default") -> str:
        self.goal = goal
        self.scenario = scenario
        wid = workflow.info().workflow_id
        while True:
            step = self.plan()
            if step.kind == "finish":
                self.status = "done"
                return self.context_summary

            if step.kind == "llm":
                key = f"{wid}:{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=30),
                    retry_policy=RetryPolicy(maximum_attempts=3),
                )
                self.step += 1
                continue

            options = SCENARIOS[self.scenario]["options"]
            try:
                if step.kind == "charge":
                    out = await workflow.execute_activity(
                        charge, args=[step.args["order_id"], step.args["hold_seconds"], step.args.get("hang_first", False)],
                        **options,
                    )
                else:
                    digest = hashlib.sha256(json.dumps(step.args, sort_keys=True).encode()).hexdigest()[:12]
                    call = ToolCall(key=f"{wid}:{self.step}:{digest}", tool=step.tool, args=step.args)
                    result = await workflow.execute_activity(execute_tool, call, **options)
                    out = result.output
                self.context_summary += f" | {step.tool}: {out}"
            except ActivityError as e:
                # The failure is history; what to do about it is Workflow code.
                self.attempts += 1
                cause = e.cause
                kind = getattr(getattr(cause, "type", None), "name", getattr(cause, "type", None))
                self.context_summary += f" | {step.tool} FAILED: {type(cause).__name__} type={kind}: {cause}"
            self.step += 1