Lab 9: Deploying Changed Workflow Code

"""Which AgentRun to run: `AGENTRUN_CODE=v1|v2|v2_patched|v2_deprecated` (default v1)."""
from __future__ import annotations

import importlib
import os

CODE_MODULES = {
    "v1": "workflows",
    "v2": "workflows_v2",
    "v2_patched": "workflows_v2_patched",
    "v2_deprecated": "workflows_v2_deprecated",
}


def load(code: str | None = None):
    """Return (module, workflow class, activities) for a code version name."""
    code = code or os.environ.get("AGENTRUN_CODE", "v1")
    module = importlib.import_module(CODE_MODULES[code])
    activities = [module.call_llm, module.execute_tool]
    if hasattr(module, "evaluate"):
        activities.append(module.evaluate)          # v1 Workers do not know `evaluate`
    return module, module.AgentRun, activities

"""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()

[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

"""Deployment routing from Python, for a box without the Temporal CLI. Each subcommand names
the CLI form it mirrors (verified against CLI 1.8.3 / Server 1.31.2).

    python routing.py describe                      # temporal worker deployment describe --name agent-runs
    python routing.py set-current v2                # temporal worker deployment set-current-version --deployment-name agent-runs --build-id v2
    python routing.py set-ramping v2 10             # temporal worker deployment set-ramping-version ... --build-id v2 --percentage 10
    python routing.py show agent-42                 # temporal workflow describe -w agent-42   (the Versioning Info block)
    python routing.py resume agent-42               # temporal workflow signal -w agent-42 --name resume
    python routing.py pause agent-42                # temporal workflow signal -w agent-42 --name pause
    python routing.py resume agent-42 --pause-after 1.8
                                                    # resume, then pause 1.8 s later: parks the run at
                                                    # the checkpoint after the step now in flight (9.2b)

Env: DEPLOYMENT_NAME (agent-runs), TEMPORAL_ADDRESS, TEMPORAL_NAMESPACE.
"""
from __future__ import annotations

import asyncio
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import NAMESPACE, connect  # noqa: E402
from temporalio.api.enums.v1 import VersioningBehavior, WorkerDeploymentVersionStatus  # noqa: E402
from temporalio.api.workflowservice.v1 import (  # noqa: E402
    DescribeWorkerDeploymentRequest,
    SetWorkerDeploymentCurrentVersionRequest,
    SetWorkerDeploymentRampingVersionRequest,
)
from temporalio.client import Client  # noqa: E402

DEPLOYMENT_NAME = os.environ.get("DEPLOYMENT_NAME", "agent-runs")


async def describe_deployment(client: Client) -> str:
    resp = await client.workflow_service.describe_worker_deployment(
        DescribeWorkerDeploymentRequest(namespace=NAMESPACE, deployment_name=DEPLOYMENT_NAME))
    info = resp.worker_deployment_info
    rc = info.routing_config
    lines = [f"deployment {info.name}: current={rc.current_version or '__unversioned__'}"
             f" ramping={rc.ramping_version or '-'} ({rc.ramping_version_percentage:g}%)"]
    for v in info.version_summaries:
        status = WorkerDeploymentVersionStatus.Name(v.status).removeprefix("WORKER_DEPLOYMENT_VERSION_STATUS_")
        lines.append(f"  build {v.deployment_version.build_id:14s} {status}")
    return "\n".join(lines)


async def set_current(client: Client, build_id: str) -> None:
    await client.workflow_service.set_worker_deployment_current_version(
        SetWorkerDeploymentCurrentVersionRequest(
            namespace=NAMESPACE, deployment_name=DEPLOYMENT_NAME, build_id=build_id,
            identity="routing.py"))


async def set_ramping(client: Client, build_id: str, percentage: float) -> None:
    await client.workflow_service.set_worker_deployment_ramping_version(
        SetWorkerDeploymentRampingVersionRequest(
            namespace=NAMESPACE, deployment_name=DEPLOYMENT_NAME, build_id=build_id,
            percentage=percentage, identity="routing.py"))


async def describe_versioning(client: Client, workflow_id: str) -> str:
    """The Versioning Info block of `temporal workflow describe`: behavior, version, override."""
    desc = await client.get_workflow_handle(workflow_id).describe()
    info = desc.raw_description.workflow_execution_info
    vi = info.versioning_info
    behavior = VersioningBehavior.Name(vi.behavior).removeprefix("VERSIONING_BEHAVIOR_")
    version = f"{vi.deployment_version.deployment_name}.{vi.deployment_version.build_id}" \
        if vi.deployment_version.build_id else (vi.version or "unversioned")
    override = ""
    if vi.HasField("versioning_override"):
        o = vi.versioning_override
        if o.HasField("pinned"):
            override = f" override=PINNED:{o.pinned.version.deployment_name}.{o.pinned.version.build_id}"
        elif o.HasField("auto_upgrade"):
            override = " override=AUTO_UPGRADE"
    transition = ""
    if vi.HasField("version_transition"):
        t = vi.version_transition.deployment_version
        transition = f" transitioning-to={t.deployment_name}.{t.build_id}"
    return (f"{workflow_id} run={desc.run_id} status={desc.status.name} "
            f"behavior={behavior} version={version}{override}{transition}")


async def main(argv: list[str]) -> None:
    client = await connect()
    cmd, args = argv[0], argv[1:]
    if cmd == "describe":
        print(await describe_deployment(client))
    elif cmd == "set-current":
        await set_current(client, args[0]); print(await describe_deployment(client))
    elif cmd == "set-ramping":
        await set_ramping(client, args[0], float(args[1])); print(await describe_deployment(client))
    elif cmd == "show":
        print(await describe_versioning(client, args[0]))
    elif cmd == "resume":
        handle = client.get_workflow_handle(args[0])
        await handle.signal("resume"); print(f"resume -> {args[0]}")
        if args[1:2] == ["--pause-after"]:
            await asyncio.sleep(float(args[2]))
            await handle.signal("pause"); print(f"pause  -> {args[0]} ({args[2]} s later)")
    elif cmd == "pause":
        await client.get_workflow_handle(args[0]).signal("pause"); print(f"pause -> {args[0]}")
    else:
        raise SystemExit(__doc__)


if __name__ == "__main__":
    asyncio.run(main(sys.argv[1:] or ["describe"]))

"""Start an AgentRun. Options for the versioning exercises:

    python starter.py                                  # id agent-42, the class's behavior (PINNED)
    python starter.py --id agent-43 --pause            # start parked at the 72 h checkpoint
    python starter.py --id agent-44 --auto-upgrade     # versioning override: AUTO_UPGRADE
    python starter.py --id agent-42 --wait             # block for the result, print the history
"""
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, show_history, start_workflow  # noqa: E402
from temporalio.common import AutoUpgradeVersioningOverride  # noqa: E402
from workflows import AgentRun  # noqa: E402

from routing import describe_versioning  # noqa: E402


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--id", default=os.environ.get("WORKFLOW_ID", "agent-42"))
    ap.add_argument("--goal", default="summarize the Temporal docs")
    ap.add_argument("--pause", action="store_true", help="deliver `pause` as the first event: parks at the checkpoint")
    ap.add_argument("--auto-upgrade", action="store_true", help="start with an AUTO_UPGRADE versioning override")
    ap.add_argument("--wait", action="store_true", help="wait for the result and print the history")
    ns = ap.parse_args()
    client = await connect()
    handle = await start_workflow(client,
        AgentRun.run, ns.goal, id=ns.id, task_queue=TASK_QUEUE,
        start_signal="pause" if ns.pause else None,
        versioning_override=AutoUpgradeVersioningOverride() if ns.auto_upgrade else None)
    print(f"started {handle.id} run={handle.result_run_id}; open http://localhost:8233")
    await asyncio.sleep(1.0)                                  # let the first Workflow Task land
    print(await describe_versioning(client, handle.id))
    if ns.wait:
        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()

"""9.1 — The replay matrix: every recorded history against every code version.

Rows are histories in tests/histories/ (the Workflow ID is the filename); columns are the four
files of the patch lifecycle. A cell is "clean" or the NondeterminismError a Worker would raise
on that history's next Workflow Task. Read the matrix before the assertions: it is the whole
lesson of exercise 9.1 in one table."""
from __future__ import annotations

from pathlib import Path

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

import workflows
import workflows_v2
import workflows_v2_deprecated
import workflows_v2_patched

HISTORIES = Path(__file__).parent / "histories"
CODE = {"v1": workflows, "v2": workflows_v2, "v2_patched": workflows_v2_patched,
        "v2_deprecated": workflows_v2_deprecated}

# history -> {code: replays clean?}   (every cell observed with temporalio 1.32.0)
#   lab08-agent-42        recorded under v1: no marker, no evaluate
#   lab09-agent-43        recorded under v2_patched: marker + evaluate
#   lab09-agent-44        recorded under v2 (final): evaluate, no marker
#
# Read down the v2 column: the naive/final code fails BOTH the v1 history (it schedules
# `evaluate` where the history has `call_llm`) AND the patched history ("Non-deprecated patch
# marker encountered ... but there is no corresponding change command"). Removing the
# `patched()` call is deploy 3, and it waits until every execution that recorded the marker is
# closed and out of retention — not merely until the v1 ones are.
EXPECTED = {
    "lab08-agent-42": {"v1": True,  "v2": False, "v2_patched": True,  "v2_deprecated": False},
    "lab09-agent-43": {"v1": False, "v2": False, "v2_patched": True,  "v2_deprecated": True},
    "lab09-agent-44": {"v1": False, "v2": True,  "v2_patched": False, "v2_deprecated": True},
}


def load(path: Path) -> WorkflowHistory:
    return WorkflowHistory.from_json(path.stem, path.read_text())


async def replay(code: str, history: WorkflowHistory):
    return await Replayer(workflows=[CODE[code].AgentRun]).replay_workflow(
        history, raise_on_replay_failure=False)


@pytest.mark.parametrize("code", list(CODE))
@pytest.mark.parametrize("stem", list(EXPECTED))
async def test_replay_matrix(stem: str, code: str):
    result = await replay(code, load(HISTORIES / f"{stem}.json"))
    clean = result.replay_failure is None
    assert clean == EXPECTED[stem][code], (
        f"{stem} under {code}: expected {'clean' if EXPECTED[stem][code] else 'a failure'}, "
        f"got {result.replay_failure!r}")
    if not clean:
        assert isinstance(result.replay_failure, workflow.NondeterminismError)


# The four sentences of 9.1, as tests you can run one at a time.
async def test_the_naive_change_breaks_old_histories():
    result = await replay("v2", load(HISTORIES / "lab08-agent-42.json"))
    assert isinstance(result.replay_failure, workflow.NondeterminismError)
    # v1 scheduled call_llm after the tool; v2 schedules evaluate at the same point.
    assert ("Activity type of scheduled event 'call_llm' does not match activity type of "
            "activity command 'evaluate'") in str(result.replay_failure)


async def test_the_patch_makes_both_branches_replay_clean():
    for stem in ("lab08-agent-42", "lab09-agent-43"):
        result = await replay("v2_patched", load(HISTORIES / f"{stem}.json"))
        assert result.replay_failure is None, (stem, result.replay_failure)


async def test_deprecating_the_patch_needs_no_pre_patch_execution_left():
    ok = await replay("v2_deprecated", load(HISTORIES / "lab09-agent-43.json"))
    assert ok.replay_failure is None
    old = await replay("v2_deprecated", load(HISTORIES / "lab08-agent-42.json"))
    assert isinstance(old.replay_failure, workflow.NondeterminismError)


async def test_removing_the_patch_call_needs_no_marked_execution_left():
    result = await replay("v2", load(HISTORIES / "lab09-agent-43.json"))
    assert "Non-deprecated patch marker encountered for change insert-verify-step" in str(result.replay_failure)

"""The v2 loop in the time-skipping environment: the verify step runs, and a rejected
verdict counts as a failed attempt instead of being handed to the model as a result."""
from __future__ import annotations

import uuid

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

import workflows_v2
from workflows import ToolCall, ToolResult
from workflows_v2 import AgentRun, Verdict

calls = {"call_llm": 0, "execute_tool": 0, "evaluate": 0}


@activity.defn(name="call_llm")
async def fake_call_llm(prompt: str, key: str) -> str:
    calls["call_llm"] += 1
    return "FINAL: 42" if "RESULT: ok" in prompt else "TOOL: fetch the-answer"


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


@activity.defn(name="evaluate")
async def picky_evaluate(result: ToolResult) -> Verdict:
    calls["evaluate"] += 1
    return Verdict(ok=calls["evaluate"] >= 2, reason="first answer never trusted")


async def test_v2_verifies_every_tool_result():
    for k in calls:
        calls[k] = 0
    tq = f"lab09-{uuid.uuid4()}"
    async with await WorkflowEnvironment.start_time_skipping() as env:
        async with Worker(env.client, task_queue=tq, workflows=[AgentRun],
                          activities=[fake_call_llm, fake_execute_tool, picky_evaluate]):
            handle = await env.client.start_workflow(
                AgentRun.run, "find the answer", id=f"lab09-test-{uuid.uuid4()}", task_queue=tq)
            assert await handle.result() == "42"
            status = await handle.query(AgentRun.status)
    # tool -> rejected verdict (attempt 1) -> model asked again -> tool -> accepted -> answer
    assert calls == {"call_llm": 3, "execute_tool": 2, "evaluate": 2}
    assert status["attempts"] == 1
    assert "RESULT: error verify: first answer never trusted" in status["context_summary"]
    assert workflows_v2.BUILD_ID == "v2"

"""A Worker for one code version, optionally as one version of a Worker Deployment.

    AGENTRUN_CODE=v1 python worker.py                          # unversioned, v1 code (9.1)
    AGENTRUN_CODE=v1 BUILD_ID=v1 python worker.py              # versioned: deployment agent-runs, build v1 (9.2)
    AGENTRUN_CODE=v2 BUILD_ID=v2 python worker.py              # ... build v2, side by side

Env: TASK_QUEUE (agent-runs), DEPLOYMENT_NAME (agent-runs), AGENTRUN_STEPS_PER_RUN (0; 9.3 uses 2).
"""
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.common import VersioningBehavior  # noqa: E402
from temporalio.worker import Worker, WorkerDeploymentConfig, WorkerDeploymentVersion  # noqa: E402

from code_versions import load  # noqa: E402

DEPLOYMENT_NAME = os.environ.get("DEPLOYMENT_NAME", "agent-runs")
BUILD_ID = os.environ.get("BUILD_ID")               # unset -> an unversioned Worker


async def main() -> None:
    module, agent_run, activities = load()
    client = await connect()
    deployment_config = None
    if BUILD_ID:
        deployment_config = WorkerDeploymentConfig(   # as of 2026-09 (temporalio 1.32); see upstream page
            version=WorkerDeploymentVersion(deployment_name=DEPLOYMENT_NAME, build_id=BUILD_ID),
            use_worker_versioning=True,
            # Every AgentRun is pinned: it finishes its run on the build that started it. The
            # upstream shape is `@workflow.defn(versioning_behavior=...)` per Workflow Type; a
            # Worker-level default keeps the same code runnable by an unversioned Worker (9.1).
            default_versioning_behavior=VersioningBehavior.PINNED,
        )
    worker = Worker(client, task_queue=TASK_QUEUE, workflows=[agent_run], activities=activities,
                    deployment_config=deployment_config,
                    identity=f"{BUILD_ID or 'unversioned'}-{module.BUILD_ID}-{os.getpid()}")
    version = f"{DEPLOYMENT_NAME}.{BUILD_ID}" if BUILD_ID else "unversioned"
    print(f"worker pid={os.getpid()} code={module.BUILD_ID} version={version} polling {TASK_QUEUE!r}")
    await worker.run()


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

"""Lab 9: AgentRun v1 — the code that has been running for two months.

Byte-for-byte the loop from lab 8 (plan → call_llm → execute_tool, Signals, Query, Update,
72-hour checkpoint, continue-as-new, compensation) plus the three things this lab is about:

* Pinned by default. `worker.py` declares `default_versioning_behavior=PINNED` on its
  deployment config, so with a versioned Worker (`BUILD_ID=... python worker.py`) every
  AgentRun finishes its run on the build that started it. (Declaring it on the class —
  `@workflow.defn(versioning_behavior=...)` — is the upstream shape, but Server 1.31.2 then
  rejects every Workflow Task completion from an *unversioned* Worker, and 9.1 runs unversioned.)
* `AGENTRUN_STEPS_PER_RUN` — a lab knob. The course rule is `is_continue_as_new_suggested()`;
  setting the knob to 2 makes the agent continue-as-new after two steps so you can watch the
  boundary without 500 steps.
* Upgrade on continue-as-new (Public Preview as of 2026-09): when the Service reports a new
  Target Version, the next run starts with AUTO_UPGRADE and lands on it, then pins there.

v1 has no verify step. `workflows_v2*.py` are the change.
"""
from __future__ import annotations

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

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

BUILD_ID = "v1"                                 # what this file is, for log lines
# Lab knob: continue-as-new after this many steps per run (0 = only when the Service suggests it).
STEPS_PER_RUN = int(os.environ.get("AGENTRUN_STEPS_PER_RUN", "0"))

# 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
# ... this far apart. A knob, because 9.2b needs the tool slow enough to send `pause` while a
# step is in flight: AGENTRUN_TOOL_TICK_SECONDS=1 makes one tool call take ~3 s.
TOOL_TICK_SECONDS = float(os.environ.get("AGENTRUN_TOOL_TICK_SECONDS", "0.05"))


# ---------------------------------------------------------------- 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)
        run_start_step = self.step
        workflow.logger.info("AgentRun %s starting at step %d", BUILD_ID, self.step)
        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 self._should_continue_as_new(run_start_step):
                    self._continue_as_new()
        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

    def _should_continue_as_new(self, run_start_step: int) -> bool:
        if workflow.info().is_continue_as_new_suggested():
            return True
        return STEPS_PER_RUN > 0 and self.step - run_start_step >= STEPS_PER_RUN

    def _continue_as_new(self) -> None:
        """The upgrade point of a long-lived pinned agent (Public Preview as of 2026-09). When
        the Service has told this run that a newer Target Version exists, the next run starts
        with AUTO_UPGRADE — it lands on the current version, then pins there. Otherwise the
        next run inherits this run's pinned version."""
        behavior = None
        if workflow.info().is_target_worker_deployment_version_changed():
            workflow.logger.info("target version changed; next run will auto-upgrade")
            behavior = ContinueAsNewVersioningBehavior.AUTO_UPGRADE
        workflow.continue_as_new(args=[self.goal, self.snapshot()], initial_versioning_behavior=behavior)

    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, result = await self._run_tool(step)
        if result is None:
            return                                   # failure already recorded
        self._record_tool_result(call, result)

    async def _run_tool(self, step: Step) -> tuple[ToolCall, ToolResult | None]:
        """Schedule execute_tool; on a spent RetryPolicy record the failure and return 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
            self.attempts += 1
            self.context_summary += f"\nRESULT: error {_failure_label(err)}"
            return call, None
        return call, result

    def _record_tool_result(self, call: ToolCall, result: ToolResult) -> None:
        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]

"""Lab 9: AgentRun v2 — a `verify` step inserted after `execute_tool`.

The change is one method. After every tool result, v2 runs the `evaluate` Activity and feeds
the `Verdict` back into state: a rejected result counts as a failed attempt instead of being
handed to the model as if it were fine. Everything else is inherited from v1 unchanged.

This file is BOTH the naive change (deploy it over v1 histories and the replay test fails —
exercise 9.1) AND the final code of the patch lifecycle (deploy 3: no `patched()` call left).
"""
from __future__ import annotations

from dataclasses import dataclass
from datetime import timedelta

from temporalio import activity, workflow
from temporalio.common import RetryPolicy

from workflows import AgentRun as AgentRunV1
from workflows import AgentState, Step, ToolResult, call_llm, execute_tool  # noqa: F401  (re-exported for worker.py)

BUILD_ID = "v2"


@dataclass
class Verdict:
    ok: bool
    reason: str = ""


@activity.defn
async def evaluate(result: ToolResult) -> Verdict:
    """The verifier as an Activity. This stand-in accepts any result that reports completion."""
    activity.logger.info("evaluate key=%s", result.key)
    if result.ok and result.output.endswith("done"):
        return Verdict(ok=True)
    return Verdict(ok=False, reason=f"unverifiable output: {result.output[:40]!r}")


@workflow.defn(name="AgentRun")
class AgentRun(AgentRunV1):
    @workflow.run                                    # the SDK wants @workflow.run on the class it registers
    async def run(self, goal: str, state: AgentState | None = None) -> str:
        return await super().run(goal, state)

    async def _tool_step(self, step: Step) -> None:
        call, result = await self._run_tool(step)
        if result is None:
            return
        # NEW in v2: the verify step. A v1 history has no ActivityTaskScheduled for `evaluate`
        # here — where v1 went on to schedule call_llm, v2 schedules evaluate.
        verdict = await workflow.execute_activity(
            evaluate, result, start_to_close_timeout=timedelta(seconds=30),
            retry_policy=RetryPolicy(maximum_attempts=3))
        if not verdict.ok:
            self.attempts += 1
            self.context_summary += f"\nRESULT: error verify: {verdict.reason}"
            return
        self._record_tool_result(call, result)

"""Lab 9: AgentRun v2, patch deprecated — deploy 2 of the patch lifecycle.

Ship this once no open execution is still on the v1 branch. `deprecate_patch` keeps
recording the marker for executions that already have it but removes the branch: the code
reads as if verify had always been there. A pre-patch (v1) history can NOT replay under this
file any more — that is what "no open execution on the old branch" guards.
"""
from __future__ import annotations

from datetime import timedelta

from temporalio import workflow
from temporalio.common import RetryPolicy

from workflows import AgentRun as AgentRunV1
from workflows import AgentState, Step, call_llm, execute_tool  # noqa: F401
from workflows_v2 import evaluate  # noqa: F401
from workflows_v2_patched import PATCH

BUILD_ID = "v2-deprecated"


@workflow.defn(name="AgentRun")
class AgentRun(AgentRunV1):
    @workflow.run                                    # the SDK wants @workflow.run on the class it registers
    async def run(self, goal: str, state: AgentState | None = None) -> str:
        return await super().run(goal, state)

    async def _tool_step(self, step: Step) -> None:
        call, result = await self._run_tool(step)
        if result is None:
            return
        workflow.deprecate_patch(PATCH)              # no branch left; marker kept for old histories
        verdict = await workflow.execute_activity(
            evaluate, result, start_to_close_timeout=timedelta(seconds=30),
            retry_policy=RetryPolicy(maximum_attempts=3))
        if not verdict.ok:
            self.attempts += 1
            self.context_summary += f"\nRESULT: error verify: {verdict.reason}"
            return
        self._record_tool_result(call, result)

"""Lab 9: AgentRun v2, patched — deploy 1 of the patch lifecycle.

`workflow.patched("insert-verify-step")` is a branch tied to a marker. A new execution records
the marker and takes the verify branch; a v1 history has no marker, so replay takes the old
branch and the Commands still match. Both kinds of history replay clean under this file.
"""
from __future__ import annotations

from datetime import timedelta

from temporalio import workflow
from temporalio.common import RetryPolicy

from workflows import AgentRun as AgentRunV1
from workflows import AgentState, Step, call_llm, execute_tool  # noqa: F401
from workflows_v2 import evaluate  # noqa: F401

BUILD_ID = "v2-patched"
PATCH = "insert-verify-step"


@workflow.defn(name="AgentRun")
class AgentRun(AgentRunV1):
    @workflow.run                                    # the SDK wants @workflow.run on the class it registers
    async def run(self, goal: str, state: AgentState | None = None) -> str:
        return await super().run(goal, state)

    async def _tool_step(self, step: Step) -> None:
        call, result = await self._run_tool(step)
        if result is None:
            return
        if workflow.patched(PATCH):                  # new executions: marker recorded, verify runs
            verdict = await workflow.execute_activity(
                evaluate, result, start_to_close_timeout=timedelta(seconds=30),
                retry_policy=RetryPolicy(maximum_attempts=3))
            if not verdict.ok:
                self.attempts += 1
                self.context_summary += f"\nRESULT: error verify: {verdict.reason}"
                return
        else:                                        # v1 histories: no marker, no verify, as before
            pass
        self._record_tool_result(call, result)