Lab 3: Determinism

"""What the Service recorded when a Worker's Workflow Task failed — and what it is still retrying.

    python failures.py agent-42
"""
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, PendingWorkflowTaskState, TimeoutType, WorkflowTaskFailedCause  # noqa: E402


async def main() -> None:
    wf_id = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("WORKFLOW_ID", "agent-42")
    client = await connect()
    handle = client.get_workflow_handle(wf_id)
    async for ev in handle.fetch_history_events():
        name = EventType.Name(ev.event_type).removeprefix("EVENT_TYPE_")
        if ev.HasField("workflow_task_failed_event_attributes"):
            a = ev.workflow_task_failed_event_attributes
            cause = WorkflowTaskFailedCause.Name(a.cause).removeprefix("WORKFLOW_TASK_FAILED_CAUSE_")
            print(f"{ev.event_id:3d}  {name}  cause={cause}  worker={a.identity}\n     {a.failure.message}")
        elif ev.HasField("workflow_task_timed_out_event_attributes"):
            a = ev.workflow_task_timed_out_event_attributes
            print(f"{ev.event_id:3d}  {name}  timeout_type={TimeoutType.Name(a.timeout_type).removeprefix('TIMEOUT_TYPE_')}")
        elif ev.HasField("workflow_task_scheduled_event_attributes"):
            a = ev.workflow_task_scheduled_event_attributes
            if a.attempt > 1:
                print(f"{ev.event_id:3d}  {name}  attempt={a.attempt}  (earlier attempts failed without an Event)")
        elif name in ("WORKFLOW_EXECUTION_COMPLETED", "WORKFLOW_EXECUTION_FAILED", "WORKFLOW_EXECUTION_TERMINATED"):
            print(f"{ev.event_id:3d}  {name}")
    desc = await handle.describe()
    raw = desc.raw_description
    print(f"status: {desc.status.name}")
    if raw.HasField("pending_workflow_task"):
        t = raw.pending_workflow_task
        print(f"pending Workflow Task: attempt={t.attempt} "
              f"state={PendingWorkflowTaskState.Name(t.state).removeprefix('PENDING_WORKFLOW_TASK_STATE_')}")


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

"""A model with state: the fake LLM behind lab 3.3.

    python model_server.py            # http://127.0.0.1:8765
    python model_server.py --reset    # tell the running server to start its count over

GET /decide?prompt=...  answers "next: continue" the first time it is asked and "next: finish"
every time after that, and prints one line per request — so you can count how often the model
was really called. GET /reset starts the count over.

That is all a real model shares with this one for our purposes: asking it twice does not give
the same answer twice, and every ask is billed.
"""
from __future__ import annotations

import argparse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlparse

REQUESTS = 0


class Handler(BaseHTTPRequestHandler):
    def do_GET(self) -> None:  # noqa: N802
        global REQUESTS
        url = urlparse(self.path)
        if url.path == "/reset":
            REQUESTS = 0
            body = "reset"
        elif url.path == "/decide":
            REQUESTS += 1
            body = "next: continue" if REQUESTS == 1 else "next: finish"
            prompt = parse_qs(url.query).get("prompt", [""])[0]
            print(f"model call #{REQUESTS}: {prompt[:50]!r} -> {body!r}", flush=True)
        else:
            self.send_response(404)
            self.end_headers()
            return
        data = body.encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "text/plain; charset=utf-8")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def log_message(self, *_args) -> None:  # quiet: only our own lines
        pass


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8765)
    ap.add_argument("--reset", action="store_true", help="reset the running server's count, then exit")
    args = ap.parse_args()
    if args.reset:
        from urllib.request import urlopen
        with urlopen(f"http://127.0.0.1:{args.port}/reset", timeout=5) as r:
            print(f"model server: {r.read().decode()} (the next ask answers 'next: continue')")
        raise SystemExit(0)
    print(f"model server on http://127.0.0.1:{args.port}  (GET /decide, GET /reset)", flush=True)
    ThreadingHTTPServer(("127.0.0.1", args.port), Handler).serve_forever()

"""Force a replay without killing anything: run the Worker's current code against a history.

    python replay.py agent-42              # replay 5 times, stop at the first mismatch
    python replay.py agent-42 --times 20   # the coin flip: count how many replays pass by luck

The code that replays is whatever LAB03_BREAK selects in *this* process — so
`LAB03_BREAK=clock python replay.py x` replays the broken build and
`python replay.py x` replays the fixed build against the same history.
"""
from __future__ import annotations

import argparse
import asyncio
import logging
import os
import re
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
import workflows  # noqa: E402
from common import connect  # noqa: E402
from temporalio.runtime import LoggingConfig, Runtime, TelemetryConfig, TelemetryFilter  # noqa: E402
from temporalio.worker import Replayer  # noqa: E402

# The SDK core logs a WARN for every failed replay, a long line that lands above this script's own
# summary. The summary below says the same thing more readably, so only errors get through.
QUIET = Runtime(telemetry=TelemetryConfig(logging=LoggingConfig(
    filter=TelemetryFilter(core_level="ERROR", other_level="ERROR"))))

logging.basicConfig(level=logging.WARNING)


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("workflow_id", nargs="?", default=os.environ.get("WORKFLOW_ID", "agent-42"))
    ap.add_argument("--times", type=int, default=5)
    args = ap.parse_args()

    client = await connect()
    history = await client.get_workflow_handle(args.workflow_id).fetch_history()
    print(f"code: LAB03_BREAK={workflows.BREAK!r}; history: {len(history.events)} events of {args.workflow_id}")
    clean = 0
    for i in range(1, args.times + 1):
        try:
            await Replayer(workflows=[workflows.AgentRun], runtime=QUIET).replay_workflow(history)
            clean += 1
            print(f"replay {i}: clean")
        except Exception as e:
            msg = str(e)
            m = re.search(r'message: "([^"]*)"', msg)
            print(f"replay {i}: {type(e).__name__}")
            print(f"  {m.group(1) if m else msg}")
            ev = re.search(r"HistoryEvent\(id: (\d+), (\w+)\)", msg)
            if ev:
                print(f"  -> the code's Command did not match Event {ev.group(1)} ({ev.group(2)}). "
                      f"Open that Event in the UI and find the `await` that produced the Command.")
            break
    else:
        print(f"{clean}/{args.times} replays clean")


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

"""Start AgentRun with a fixed Workflow ID.

    python starter.py                           # agent-42: 15 s timer, 5 s time budget, up to 5 steps
    python starter.py --id coinflip-1 --no-wait # return immediately; kill the Worker during the timer
    python starter.py --budget 300              # 3.3: give the model-driven run room for two steps
"""
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 workflows import AgentRun  # noqa: E402


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--id", default=os.environ.get("WORKFLOW_ID", "agent-42"))
    ap.add_argument("--sleep", type=int, default=15, help="timer after each step (kill the Worker inside it)")
    ap.add_argument("--budget", type=int, default=5, help="plan() finishes once the run is older than this")
    ap.add_argument("--steps", type=int, default=5, help="hard cap on steps")
    ap.add_argument("--no-wait", action="store_true")
    args = ap.parse_args()

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


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

"""End runs this lab leaves stuck, so the next Worker you start does not adopt them.

    python terminate.py coinflip-naive clock-naive net-naive

A run whose Workflow Task keeps failing stays RUNNING. That costs no compute, but it is not
harmless: every Worker on the same Task Queue polls for it, so the next build you start picks it
up and runs it (in 3.3, spending the model's one `next: continue` on the wrong run). Terminating is
what an operator does with an execution that has no future; the Web UI's Terminate button is the
same call.
"""
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.service import RPCError  # noqa: E402


async def main(ids: list[str]) -> None:
    client = await connect()
    for wid in ids:
        try:
            await client.get_workflow_handle(wid).terminate(reason="lab 3: stuck on a broken build")
            print(f"terminated {wid}")
        except RPCError as e:
            print(f"{wid}: {e.message}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        raise SystemExit(__doc__)
    asyncio.run(main(sys.argv[1:]))

"""Check yourself: the fixed build replays clean; the broken builds do not — and you can see why.

Runs on Temporal's time-skipping test server. The build under test is selected the way the
Worker selects it, through LAB03_BREAK, set before each run (the sandbox re-imports
workflows.py per execution, so the variable is read fresh).  PYTHONPATH=.. pytest tests/
"""
from __future__ import annotations

import asyncio
import os
import sys
import threading
import time
from http.server import ThreadingHTTPServer

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 model_server  # noqa: E402
import workflows  # noqa: E402
from temporalio.client import WorkflowHistory  # noqa: E402
from temporalio.testing import WorkflowEnvironment  # noqa: E402
from temporalio.worker import Replayer, Worker  # noqa: E402
from temporalio.workflow import NondeterminismError  # noqa: E402

TASK_QUEUE = "lab-03-test"


async def _run(wf_id: str, budget: int = 5) -> WorkflowHistory:
    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]):
            handle = await env.client.start_workflow(
                workflows.AgentRun.run, args=["test goal", 15, budget, 5], id=wf_id, task_queue=TASK_QUEUE
            )
            await handle.result()
            return await handle.fetch_history()


def _replay(history: WorkflowHistory) -> Exception | None:
    async def go():
        await Replayer(workflows=[workflows.AgentRun]).replay_workflow(history)
    try:
        asyncio.run(go())
        return None
    except Exception as e:  # noqa: BLE001
        return e


@pytest.fixture
def build(monkeypatch):
    """Select the build: build("clock") is `LAB03_BREAK=clock python worker.py`."""
    def select(mode: str | None) -> None:
        if mode is None:
            monkeypatch.delenv("LAB03_BREAK", raising=False)
        else:
            monkeypatch.setenv("LAB03_BREAK", mode)
    return select


def test_fixed_build_replays_clean(build):
    build(None)
    history = asyncio.run(_run("lab03-test-fixed"))
    assert _replay(history) is None


def test_clock_break_fails_replay_once_the_budget_has_passed(build):
    """3.2: the run passed step 0 with the wall clock inside the 1 s budget; any replay after
    that reads a later wall clock, plans `finish`, and proposes CompleteWorkflowExecution where
    history has the Activity (or the jitter timer) at event 5."""
    build("clock")
    history = asyncio.run(_run("lab03-test-clock", budget=1))
    time.sleep(1.5)
    err = _replay(history)
    assert isinstance(err, NondeterminismError)
    assert "HistoryEvent(id: 5" in str(err)
    # The fix replays the same history clean: workflow.now() is history time.
    build(None)
    assert _replay(history) is None


def test_model_in_plan_fails_replay_and_pays_again(build, monkeypatch):
    """3.3: with the model asked inside plan(), the replay asks it again (request #3) and gets a
    different answer. With the model behind call_llm, the replay asks nothing."""
    server = ThreadingHTTPServer(("127.0.0.1", 0), model_server.Handler)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    url = f"http://127.0.0.1:{server.server_port}"
    try:
        model_server.REQUESTS = 0
        monkeypatch.setenv("LAB03_MODEL_URL", url)
        build("network")
        history = asyncio.run(_run("lab03-test-network", budget=300))
        assert model_server.REQUESTS == 2                      # step 0: continue; step 1: finish
        assert isinstance(_replay(history), NondeterminismError)
        assert model_server.REQUESTS == 3                      # the replay paid for one more

        model_server.REQUESTS = 0
        build(None)
        monkeypatch.setattr(workflows, "MODEL_URL", url)       # call_llm runs outside the sandbox
        history = asyncio.run(_run("lab03-test-network-fixed", budget=300))
        assert model_server.REQUESTS == 2
        assert history.events and sum(
            1 for e in history.events if e.HasField("activity_task_scheduled_event_attributes")
        ) == 2                                                  # two call_llm steps, model-driven
        assert _replay(history) is None
        assert model_server.REQUESTS == 2                      # replay asked nothing
    finally:
        server.shutdown()


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 BREAK, MODEL_URL, AgentRun, call_llm  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
logging.getLogger("temporalio.worker._workflow_instance").setLevel(logging.ERROR)  # the traceback is in history


async def main() -> None:
    client = await connect()
    worker = Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=[AgentRun],
        activities=[call_llm],
    )
    print(
        f"worker pid={os.getpid()} polling {TASK_QUEUE!r} with LAB03_BREAK={BREAK!r}"
        f"{' model=' + MODEL_URL if MODEL_URL else ''}; kill -9 {os.getpid()} to crash it",
        flush=True,
    )
    await worker.run()


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

"""Lab 3: AgentRun, stage 3 — break replay three ways, read the error, then keep the rule.

The program is lab 2's, plus the two things a real agent is tempted to do in Workflow code:
consult the clock in `plan()` (a time budget) and consult the model in `plan()` (the next step).
`LAB03_BREAK` selects which *version of the code* a Worker runs — think "which build is
deployed", which is why it is an environment variable and not a Workflow argument: the fix
is a redeploy, and the same history must replay under the fixed code.

  LAB03_BREAK=          what the Worker runs                                  what you get
  none (default)        workflow.random(), workflow.now(), model via call_llm  clean replays
  random-naive          random.random() in run()                              sandbox: RestrictedWorkflowAccessError
  random                random.Random().random() in run()  (unseeded)         NondeterminismError, on some replays
  clock-naive           datetime.now() in plan()                              sandbox
  clock                 datetime.now() under sandbox_unrestricted() in plan()  NondeterminismError, on every replay after the budget
  network-naive         urllib.request.urlopen() in plan()                    sandbox
  network               the same under sandbox_unrestricted()                 NondeterminismError, and the replay paid for a model call

The model is `model_server.py` (3.3): it answers "next: continue" once, then "next: finish".
"""
from __future__ import annotations

import asyncio
import hashlib
import os
import random
import urllib.parse
import urllib.request
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from temporalio import activity, workflow

BREAK = os.environ.get("LAB03_BREAK", "none")
# Set in 3.3 only. Empty => call_llm fakes the model (one step, then finish).
MODEL_URL = os.environ.get("LAB03_MODEL_URL", "")
DEFAULT_MODEL_URL = "http://127.0.0.1:8765"


def ask_model(prompt: str) -> str:
    """One HTTP GET to the model server. An effect: the answer depends on how often it was asked."""
    url = f"{MODEL_URL or DEFAULT_MODEL_URL}/decide?prompt={urllib.parse.quote(prompt)}"
    with urllib.request.urlopen(url, timeout=5) as resp:
        return resp.read().decode("utf-8").strip()


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


@activity.defn
async def call_llm(prompt: str, key: str) -> str:
    """The model call as an Activity. Its result is recorded in ActivityTaskCompleted, so on
    replay the Workflow gets the same answer without the model being asked again."""
    activity.logger.info("call_llm key=%s", key)
    if MODEL_URL:
        return ask_model(prompt)
    return f"[stub model answer to: {prompt[:40]}] next: finish"


@workflow.defn
class AgentRun:
    def __init__(self) -> None:
        self.step = 0
        self.status = "running"
        self.goal = ""
        self.context_summary = ""
        self.started_at: datetime | None = None
        self.budget_seconds = 0
        self.max_steps = 0

    def plan(self) -> Step:
        """Workflow code: a pure function of state. It never calls a model — except in the
        versions below that are here to be broken."""
        # --- the time budget: has the run been going for too long? -------------------------
        if BREAK == "clock-naive":
            now = datetime.now(timezone.utc)                      # sandbox stops this
        elif BREAK == "clock":
            with workflow.unsafe.sandbox_unrestricted():
                now = datetime.now(timezone.utc)                  # wall clock: different on replay
        else:
            now = workflow.now()                                  # history time: same on replay
        if now - self.started_at > timedelta(seconds=self.budget_seconds):
            return Step("finish")
        if self.step >= self.max_steps:
            return Step("finish")
        # --- the next action: ask the model, or decide from the recorded answer? -----------
        prompt = f"Step {self.step} for: {self.goal}. Context: {self.context_summary[-60:]}"
        if BREAK == "network-naive":
            decision = ask_model(prompt)                          # sandbox stops this
            return Step("llm", prompt) if "continue" in decision else Step("finish")
        if BREAK == "network":
            with workflow.unsafe.sandbox_unrestricted():
                decision = ask_model(prompt)                      # a different answer on replay
            return Step("llm", prompt) if "continue" in decision else Step("finish")
        if "next: finish" in self.context_summary:                # the recorded answer decides
            return Step("finish")
        return Step("llm", prompt)

    @workflow.run
    async def run(self, goal: str, sleep_seconds: int = 15, budget_seconds: int = 5, max_steps: int = 5) -> str:
        self.goal = goal
        self.budget_seconds = budget_seconds
        self.max_steps = max_steps
        self.started_at = workflow.now()
        while True:
            step = self.plan()
            workflow.logger.info("live: step %d plan -> %s", self.step, step.kind)
            if step.kind == "finish":
                self.status = "done"
                return self.context_summary
            # Spread load: sometimes wait a second before calling the model.
            if BREAK == "random-naive":
                jitter = random.random() < 0.5                    # sandbox stops this
            elif BREAK == "random":
                jitter = random.Random().random() < 0.5           # unseeded: a fresh coin on every replay
            else:
                jitter = workflow.random().random() < 0.5         # seeded from history: the same coin
            if jitter:
                await asyncio.sleep(1)
            key = f"{workflow.info().workflow_id}:{self.step}:{hashlib.sha256(step.prompt.encode()).hexdigest()[:12]}"
            self.context_summary = await workflow.execute_activity(
                call_llm, args=[step.prompt, key], start_to_close_timeout=timedelta(seconds=10)
            )
            self.step += 1
            await asyncio.sleep(sleep_seconds)