Lab 2: Event History and Replay

"""Print a Workflow's Event History with the facts you annotate in 2.1 — or save it as JSON.

    python history.py agent-42                  # one row per Event, with the attributes that matter
    python history.py agent-42 --json h.json    # the same history as data (Replayer input)
    python history.py agent-42 --first 13       # only Events 1..13 (2.4: cut the history mid-way)
    python history.py --file histories/x.json   # print a saved JSON history instead

The "caused by" column of the reading's table is deliberately NOT printed — writing it is the
exercise. Everything printed is read straight off the Event's attributes.
"""
from __future__ import annotations

import argparse
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, TaskQueueKind, TimeoutType, WorkflowTaskFailedCause  # noqa: E402
from temporalio.api.history.v1 import HistoryEvent  # noqa: E402
from temporalio.client import WorkflowHistory  # noqa: E402


def _payloads(payloads) -> str:
    out = []
    for p in payloads:
        try:
            out.append(p.data.decode("utf-8"))
        except UnicodeDecodeError:
            out.append(f"<{len(p.data)} bytes>")
    return ", ".join(out)


def _queue(tq) -> str:
    kind = TaskQueueKind.Name(tq.kind).removeprefix("TASK_QUEUE_KIND_").lower()
    return f"{tq.name} [{kind}]"


def describe(event: HistoryEvent) -> str:
    """The attributes that let you say what caused the Event."""
    t = event.event_type
    if t == EventType.EVENT_TYPE_WORKFLOW_EXECUTION_STARTED:
        a = event.workflow_execution_started_event_attributes
        return (f"type={a.workflow_type.name} input=[{_payloads(a.input.payloads)}] "
                f"task_queue={a.task_queue.name} client={a.identity} attempt={a.attempt}")
    if t == EventType.EVENT_TYPE_WORKFLOW_TASK_SCHEDULED:
        a = event.workflow_task_scheduled_event_attributes
        return f"queue={_queue(a.task_queue)} attempt={a.attempt}"
    if t == EventType.EVENT_TYPE_WORKFLOW_TASK_STARTED:
        return f"worker={event.workflow_task_started_event_attributes.identity}"
    if t == EventType.EVENT_TYPE_WORKFLOW_TASK_COMPLETED:
        return f"worker={event.workflow_task_completed_event_attributes.identity}  (Commands follow)"
    if t == EventType.EVENT_TYPE_WORKFLOW_TASK_TIMED_OUT:
        a = event.workflow_task_timed_out_event_attributes
        return f"timeout_type={TimeoutType.Name(a.timeout_type).removeprefix('TIMEOUT_TYPE_')}"
    if t == EventType.EVENT_TYPE_WORKFLOW_TASK_FAILED:
        a = event.workflow_task_failed_event_attributes
        cause = WorkflowTaskFailedCause.Name(a.cause).removeprefix("WORKFLOW_TASK_FAILED_CAUSE_")
        return f"cause={cause} worker={a.identity} message={a.failure.message!r}"
    if t == EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED:
        a = event.activity_task_scheduled_event_attributes
        return (f"activity={a.activity_type.name} activity_id={a.activity_id} "
                f"queue={a.task_queue.name} input=[{_payloads(a.input.payloads)}]")
    if t == EventType.EVENT_TYPE_ACTIVITY_TASK_STARTED:
        a = event.activity_task_started_event_attributes
        return f"worker={a.identity} attempt={a.attempt} scheduled_event_id={a.scheduled_event_id}"
    if t == EventType.EVENT_TYPE_ACTIVITY_TASK_COMPLETED:
        a = event.activity_task_completed_event_attributes
        return f"result=[{_payloads(a.result.payloads)}] worker={a.identity}"
    if t == EventType.EVENT_TYPE_ACTIVITY_TASK_FAILED:
        a = event.activity_task_failed_event_attributes
        return f"message={a.failure.message!r}"
    if t == EventType.EVENT_TYPE_ACTIVITY_TASK_TIMED_OUT:
        a = event.activity_task_timed_out_event_attributes
        tt = TimeoutType.Name(a.failure.timeout_failure_info.timeout_type).removeprefix("TIMEOUT_TYPE_")
        return f"timeout_type={tt}"
    if t == EventType.EVENT_TYPE_TIMER_STARTED:
        a = event.timer_started_event_attributes
        return f"timer_id={a.timer_id} fires_after={a.start_to_fire_timeout.ToTimedelta()}"
    if t == EventType.EVENT_TYPE_TIMER_FIRED:
        a = event.timer_fired_event_attributes
        return f"timer_id={a.timer_id} started_event_id={a.started_event_id}"
    if t == EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED:
        a = event.workflow_execution_completed_event_attributes
        return f"result=[{_payloads(a.result.payloads)}]"
    if t == EventType.EVENT_TYPE_WORKFLOW_EXECUTION_FAILED:
        return f"message={event.workflow_execution_failed_event_attributes.failure.message!r}"
    return ""


def print_history(history: WorkflowHistory, first: int | None = None) -> None:
    t0 = None
    for event in history.events:
        if first is not None and event.event_id > first:
            break
        ts = event.event_time.ToDatetime()
        t0 = t0 or ts
        name = EventType.Name(event.event_type).removeprefix("EVENT_TYPE_")
        print(f"{event.event_id:3d}  +{(ts - t0).total_seconds():6.2f}s  {name:30s} {describe(event)}")


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("workflow_id", nargs="?", default=os.environ.get("WORKFLOW_ID", "agent-42"))
    ap.add_argument("--run-id", default=None)
    ap.add_argument("--file", help="print a saved JSON history instead of fetching one")
    ap.add_argument("--json", metavar="OUT", help="also write the history as JSON to OUT")
    ap.add_argument("--first", type=int, help="show (and save) only Events 1..N")
    args = ap.parse_args()

    if args.file:
        with open(args.file) as f:
            history = WorkflowHistory.from_json(args.workflow_id, f.read())
    else:
        client = await connect()
        history = await client.get_workflow_handle(args.workflow_id, run_id=args.run_id).fetch_history()

    print_history(history, args.first)
    if args.json:
        events = history.events if args.first is None else history.events[: args.first]
        with open(args.json, "w") as f:
            f.write(WorkflowHistory(workflow_id=history.workflow_id, events=events).to_json())
        print(f"wrote {len(events)} events to {args.json}")


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

"""Replay AgentRun's code against a history, exactly as a fresh Worker would — locally,
without connecting a Worker to the Workflow's Task Queue and without running any Activity.

    python replay.py agent-42                          # fetch the live history and replay it
    python replay.py --file histories/agent-42.json    # replay a saved history
    python replay.py --file histories/agent-42.json --first 13   # 2.4: cut it, see what goes live

What to watch: the Workflow's "live:" log lines. While the SDK is replaying they are suppressed;
the first one printed is the first Command the code would issue for real. And
CALL_LLM_INVOCATIONS: the Activity function is never called during a replay.
"""
from __future__ import annotations

import argparse
import asyncio
import logging
import os
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.client import WorkflowHistory  # noqa: E402
from temporalio.worker import Replayer  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
logging.getLogger("temporalio.worker").setLevel(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("--file")
    ap.add_argument("--first", type=int, help="replay only Events 1..N (cut at a WorkflowTaskScheduled)")
    args = ap.parse_args()

    if args.file:
        with open(args.file) as f:
            history = WorkflowHistory.from_json(args.workflow_id, f.read())
    else:
        client = await connect()
        history = await client.get_workflow_handle(args.workflow_id).fetch_history()
    if args.first:
        history = WorkflowHistory(workflow_id=history.workflow_id, events=history.events[: args.first])

    # A saved file carries no Workflow ID, so name the file rather than a default ID it may not be.
    source = args.file if args.file else history.workflow_id
    print(f"replaying {len(history.events)} events of {source}; "
          f"CALL_LLM_INVOCATIONS before = {workflows.CALL_LLM_INVOCATIONS}")
    try:
        await Replayer(workflows=[workflows.AgentRun]).replay_workflow(history)
        print("REPLAY CLEAN: every Command the code issued matched the Event history")
    except Exception as e:  # temporalio.workflow.NondeterminismError, in lab 3
        print(f"REPLAY FAILED: {type(e).__name__}: {e}")
    print(f"CALL_LLM_INVOCATIONS after = {workflows.CALL_LLM_INVOCATIONS}")


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

"""Start AgentRun with a fixed Workflow ID and wait for it.

    python starter.py                       # agent-42, one step, 20 s timer
    python starter.py --steps 3 --sleep 2   # the counting exercise: three steps, short timers
    python starter.py --no-wait             # return as soon as it is started
"""
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=20, help="timer length in seconds (kill the Worker inside it)")
    ap.add_argument("--steps", type=int, default=1, help="how many call_llm steps before finish")
    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.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())

"""Check yourself: the history shape you annotated in 2.1, and that replay never re-runs call_llm.

Runs against Temporal's time-skipping test server (downloaded on first use), so the 20 s timer
costs nothing. Plain pytest; no plugin needed:  PYTHONPATH=.. pytest tests/
"""
from __future__ import annotations

import asyncio
import os
import sys

import pytest

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

import workflows  # noqa: E402
from temporalio.api.enums.v1 import EventType  # noqa: E402
from temporalio.client import WorkflowHistory  # noqa: E402
from temporalio.testing import WorkflowEnvironment  # noqa: E402
from temporalio.worker import Replayer, Worker  # noqa: E402

TASK_QUEUE = "lab-02-test"

# The reading's table, minus the Worker kill: one Worker, so the last Workflow Task lands on the
# sticky queue and there is no WorkflowTaskTimedOut / reschedule pair.
EXPECTED_ONE_STEP = [
    "WORKFLOW_EXECUTION_STARTED",
    "WORKFLOW_TASK_SCHEDULED", "WORKFLOW_TASK_STARTED", "WORKFLOW_TASK_COMPLETED",
    "ACTIVITY_TASK_SCHEDULED", "ACTIVITY_TASK_STARTED", "ACTIVITY_TASK_COMPLETED",
    "WORKFLOW_TASK_SCHEDULED", "WORKFLOW_TASK_STARTED", "WORKFLOW_TASK_COMPLETED",
    "TIMER_STARTED", "TIMER_FIRED",
    "WORKFLOW_TASK_SCHEDULED", "WORKFLOW_TASK_STARTED", "WORKFLOW_TASK_COMPLETED",
    "WORKFLOW_EXECUTION_COMPLETED",
]


def _types(history: WorkflowHistory) -> list[str]:
    return [EventType.Name(e.event_type).removeprefix("EVENT_TYPE_") for e in history.events]


async def _run(steps: int, wf_id: str) -> 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", 20, steps], id=wf_id, task_queue=TASK_QUEUE
            )
            result = await handle.result()
            assert result.startswith("[model answer to:")
            return await handle.fetch_history()


def test_one_step_history_is_sixteen_events():
    history = asyncio.run(_run(1, "lab02-test-one-step"))
    assert _types(history) == EXPECTED_ONE_STEP


def test_each_step_costs_eleven_events():
    """2.1, last part: an Activity is 3 Events, the Workflow Task that delivers its result 3 more,
    the timer 2, and the Workflow Task that wakes after the timer 3 — 11 per step."""
    history = asyncio.run(_run(3, "lab02-test-three-steps"))
    assert len(history.events) == 4 + 3 * 11 + 1
    assert _types(history).count("ACTIVITY_TASK_SCHEDULED") == 3


def test_replay_does_not_invoke_call_llm():
    """The Activity function ran once, in the Worker. Replaying the history — full, or cut at the
    Workflow Task after TimerFired — issues the same Commands and never calls it again."""
    history = asyncio.run(_run(1, "lab02-test-replay"))
    before = workflows.CALL_LLM_INVOCATIONS

    async def replay(h: WorkflowHistory) -> None:
        await Replayer(workflows=[workflows.AgentRun]).replay_workflow(h)

    asyncio.run(replay(history))
    cut = WorkflowHistory(workflow_id=history.workflow_id, events=history.events[:13])
    assert _types(cut)[-1] == "WORKFLOW_TASK_SCHEDULED"
    asyncio.run(replay(cut))
    assert workflows.CALL_LLM_INVOCATIONS == before


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

# Show the Workflow's "live:" lines and the Activity's invocation counter.
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],
    )
    print(
        f"worker pid={os.getpid()} identity={client.identity} polling {TASK_QUEUE!r}; "
        f"kill -9 {os.getpid()} to crash it",
        flush=True,
    )
    await worker.run()


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

"""Lab 2: AgentRun, stage 2 — the same program as lab 1, read through its Event History.

The code is lab 1's; the two additions are instruments, not behaviour:
  * `CALL_LLM_INVOCATIONS` counts how often the Activity function really ran in *this* process,
    so you can see that a replaying Worker never calls it.
  * `workflow.logger` lines before every await. Log lines are suppressed while the SDK is
    replaying and printed only once the code is "live" — so the first line a fresh Worker prints
    names the first Command it actually issued.
`max_steps` exists for the counting exercise (2.1, last part): every step costs one Activity
(three Events) plus the Workflow Task that delivers its result (three more).
"""
from __future__ import annotations

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

from temporalio import activity, workflow


# Lives in the Worker process, outside the Workflow sandbox: a plain global.
CALL_LLM_INVOCATIONS = 0


@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. Stage 2 still fakes the model; what matters here is that
    the function body runs once per attempt and never during replay."""
    global CALL_LLM_INVOCATIONS
    CALL_LLM_INVOCATIONS += 1
    activity.logger.info(
        "call_llm key=%s  (invocation #%d in pid %d)", key, CALL_LLM_INVOCATIONS, os.getpid()
    )
    return f"[model answer to: {prompt[:40]}]"


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

    def plan(self, max_steps: int) -> Step:
        """Workflow code: a pure function of state. It never calls a model."""
        if self.step < max_steps:
            return Step("llm", prompt=f"Step {self.step} of a plan for: {self.goal}")
        return Step("finish")

    @workflow.run
    async def run(self, goal: str, sleep_seconds: int = 20, max_steps: int = 1) -> str:
        self.goal = goal
        while True:
            step = self.plan(max_steps)
            if step.kind == "finish":
                self.status = "done"
                workflow.logger.info("live: returning -> Command CompleteWorkflowExecution")
                return self.context_summary
            key = f"{workflow.info().workflow_id}:{self.step}:{hashlib.sha256(step.prompt.encode()).hexdigest()[:12]}"
            workflow.logger.info("live: step %d -> Command ScheduleActivityTask(call_llm)", self.step)
            answer = await workflow.execute_activity(
                call_llm, args=[step.prompt, key], start_to_close_timeout=timedelta(seconds=10)
            )
            self.context_summary = answer
            self.step += 1
            workflow.logger.info("live: step %d -> Command StartTimer(%ds)", self.step, sleep_seconds)
            await asyncio.sleep(sleep_seconds)