Lab 6: Continue-As-New and Child Workflows

"""Start AgentRun, send one operator note mid-run, wait for the result, then print the run chain."""
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 temporalio.api.enums.v1 import EventType  # noqa: E402
from temporalio.client import Client  # noqa: E402
from workflows import AgentRun, AgentState  # noqa: E402


async def run_chain(client: Client, workflow_id: str, first_run_id: str) -> list[dict]:
    """Follow the continue-as-new chain from the first run: each run's last Event names the next Run ID."""
    chain, run_id = [], first_run_id
    while run_id:
        handle = client.get_workflow_handle(workflow_id, run_id=run_id)
        desc = await handle.describe()
        events = [ev async for ev in handle.fetch_history_events()]
        last = events[-1]
        nxt = None
        if last.HasField("workflow_execution_continued_as_new_event_attributes"):
            nxt = last.workflow_execution_continued_as_new_event_attributes.new_execution_run_id
        chain.append({"run_id": run_id, "status": desc.status.name, "events": len(events),
                      "last_event": EventType.Name(last.event_type).removeprefix("EVENT_TYPE_")})
        run_id = nxt
    return chain


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",
                    help='a goal containing "unreliable" makes the ResearchAgent child fail (6.2)')
    ap.add_argument("--steps", type=int, default=0, help="total steps before the agent finishes (default 20; 6.1b uses 60)")
    ap.add_argument("--inject-after", type=float, default=1.0,
                    help="seconds after start to send one inject_context Signal (0 = none)")
    args = ap.parse_args()

    client = await connect()
    # --steps starts the chain from an explicit step-0 checkpoint that carries the length; every
    # continue-as-new passes it on in AgentState. Without it, run 1 starts with no state at all.
    first = [args.goal, AgentState(step=0, goal=args.goal, max_steps=args.steps)] if args.steps else [args.goal]
    handle = await start_workflow(client, AgentRun.run, args=first, id=args.id, task_queue=TASK_QUEUE)
    first_run_id = handle.result_run_id
    print(f"started {handle.id} first run={first_run_id}; open http://localhost:8233/namespaces/default/workflows/{handle.id}")

    if args.inject_after > 0:
        await asyncio.sleep(args.inject_after)
        await handle.signal(AgentRun.inject_context, args=["Prefer sources after 2025", "cmd-7f3a"])
        status = await handle.query(AgentRun.status)
        print(f"inject_context sent; it landed in run={status['run_id']} at step {status['step']}")

    result = await handle.result()          # follows the chain: the handle is the Workflow ID, not a run
    print("result:")
    print(result)
    print(f"\nfinal status(): {await handle.query(AgentRun.status)}")

    print("\nrun chain (one Workflow ID, several Run IDs):")
    for i, run in enumerate(await run_chain(client, handle.id, first_run_id), 1):
        print(f"  run {i}: {run['run_id']}  {run['status']:17s} {run['events']:3d} events  last={run['last_event']}")


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

"""Lab 6 self-check against the time-skipping test server:  python -m pytest tests/ -q"""
from __future__ import annotations

import asyncio
import os
import sys
import uuid

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

from temporalio.api.enums.v1 import EventType  # noqa: E402
from temporalio.client import Client  # noqa: E402
from temporalio.testing import WorkflowEnvironment  # noqa: E402
from temporalio.worker import Worker  # noqa: E402
from workflows import (  # noqa: E402
    MAX_STEPS,
    RESEARCH_TASK_QUEUE,
    STEPS_PER_RUN,
    AgentRun,
    ResearchAgent,
    call_llm,
    fetch_source,
)


async def _events(client: Client, workflow_id: str, run_id: str) -> list:
    return [ev async for ev in client.get_workflow_handle(workflow_id, run_id=run_id).fetch_history_events()]


async def _run_chain(client: Client, workflow_id: str, first_run_id: str) -> list[list]:
    """Histories of every run, following WorkflowExecutionContinuedAsNew.new_execution_run_id."""
    chain, run_id = [], first_run_id
    while run_id:
        events = await _events(client, workflow_id, run_id)
        chain.append(events)
        last = events[-1]
        run_id = (last.workflow_execution_continued_as_new_event_attributes.new_execution_run_id
                  if last.HasField("workflow_execution_continued_as_new_event_attributes") else None)
    return chain


def _names(events) -> list[str]:
    return [EventType.Name(ev.event_type).removeprefix("EVENT_TYPE_") for ev in events]


async def _with_workers(env: WorkflowEnvironment, goal: str, inject: bool):
    tq = f"lab-06-test-{uuid.uuid4()}"
    agent = Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=[call_llm])
    research = Worker(env.client, task_queue=RESEARCH_TASK_QUEUE, workflows=[ResearchAgent], activities=[fetch_source])
    async with agent, research:
        handle = await env.client.start_workflow(AgentRun.run, goal, id=f"lab06-test-{uuid.uuid4()}", task_queue=tq)
        if inject:
            await handle.signal(AgentRun.inject_context, args=["Prefer sources after 2025", "cmd-7f3a"])
            first = await handle.query(AgentRun.status)          # answered by run 1
            assert first["run_id"] == handle.result_run_id and first["processed_command_ids"] == ["cmd-7f3a"]
            await env.sleep(6)                                   # skip ahead: 4 steps of 0.5 s per run -> run 3 or so
            later = await handle.query(AgentRun.status)          # same handle, a later run answers
            assert later["run_id"] != first["run_id"] and later["step"] > first["step"]
            assert later["processed_command_ids"] == ["cmd-7f3a"]   # the dedupe set travelled in AgentState
        result = await handle.result()                 # follows the chain; time-skipping runs every timer
        return handle, result


def test_one_workflow_id_several_run_ids_each_history_short() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            handle, result = await _with_workers(env, "summarize the Temporal docs", inject=True)
            chain = await _run_chain(env.client, handle.id, handle.result_run_id)
            assert len(chain) == MAX_STEPS // STEPS_PER_RUN == 5
            for events in chain[:-1]:
                assert _names(events)[-1] == "WORKFLOW_EXECUTION_CONTINUED_AS_NEW"
                assert len(events) < 80                # each run's history is short; the state travels as AgentState
            assert _names(chain[-1])[-1] == "WORKFLOW_EXECUTION_COMPLETED"
            started = chain[1][0].workflow_execution_started_event_attributes
            assert started.continued_execution_run_id == handle.result_run_id   # run 2 names run 1 as its origin
            # the operator note was signalled to run 1 and is in the final result: AgentState carried it
            assert result.count("Operator: Prefer sources after 2025") == 1
            assert (await env.client.get_workflow_handle(handle.id, run_id=handle.result_run_id).describe()).status.name == "CONTINUED_AS_NEW"

    asyncio.run(go())


def test_research_child_has_its_own_history_and_one_event_pair_in_the_parent() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            handle, result = await _with_workers(env, "summarize the Temporal docs", inject=False)
            assert "3 sources on sources for: summarize the Temporal docs" in result
            child_id = f"{handle.id}/research/1"
            child = env.client.get_workflow_handle(child_id)
            child_names = _names([ev async for ev in child.fetch_history_events()])
            assert child_names.count("ACTIVITY_TASK_COMPLETED") == 3     # the child's steps live in the child's history
            assert (await child.describe()).status.name == "COMPLETED"
            parent_names = [n for run in await _run_chain(env.client, handle.id, handle.result_run_id) for n in _names(run)]
            assert parent_names.count("START_CHILD_WORKFLOW_EXECUTION_INITIATED") == 1
            assert parent_names.count("CHILD_WORKFLOW_EXECUTION_STARTED") == 1
            assert parent_names.count("CHILD_WORKFLOW_EXECUTION_COMPLETED") == 1
            assert "fetch_source" not in str(parent_names)

    asyncio.run(go())


def test_child_failure_reaches_the_parent_as_a_catchable_error() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            handle, result = await _with_workers(env, "summarize the unreliable Temporal docs", inject=False)
            assert "research 1 failed: source registry rejected the task" in result
            assert result.count("[model answer") == MAX_STEPS - 1     # every other step still ran
            parent_names = [n for run in await _run_chain(env.client, handle.id, handle.result_run_id) for n in _names(run)]
            assert parent_names.count("CHILD_WORKFLOW_EXECUTION_FAILED") == 1
            child = env.client.get_workflow_handle(f"{handle.id}/research/1")
            assert (await child.describe()).status.name == "FAILED"

    asyncio.run(go())

"""Worker for lab 6.

    python worker.py                   # AgentRun + ResearchAgent, both queues (one process)
    python worker.py --role agent      # AgentRun and call_llm only, on TASK_QUEUE
    python worker.py --role research   # ResearchAgent and fetch_source only, on RESEARCH_TASK_QUEUE

6.2 wants the child on a Worker you can kill without touching the parent: run one of each role.
"""
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  # noqa: E402
from temporalio.worker import Worker  # noqa: E402
from workflows import RESEARCH_TASK_QUEUE, AgentRun, ResearchAgent, call_llm, fetch_source  # noqa: E402


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--role", choices=["both", "agent", "research"], default="both")
    args = ap.parse_args()

    client = await connect()
    workers = []
    if args.role in ("both", "agent"):
        workers.append(Worker(client, task_queue=TASK_QUEUE, workflows=[AgentRun], activities=[call_llm]))
    if args.role in ("both", "research"):
        workers.append(Worker(client, task_queue=RESEARCH_TASK_QUEUE, workflows=[ResearchAgent], activities=[fetch_source]))
    queues = [w.task_queue for w in workers]
    print(f"worker pid={os.getpid()} role={args.role} polling {queues}; kill -9 {os.getpid()} to crash it")
    await asyncio.gather(*(w.run() for w in workers))


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

"""Lab 6: AgentRun, stage 6 — continue-as-new and the ResearchAgent child.

Two additions to stage 5. (1) The loop checks, at the top of every step, whether it
should checkpoint into a fresh run: `workflow.info().is_continue_as_new_suggested()`
is the rule; `STEPS_PER_RUN` is a lab knob so the transition happens in a 20-step
run instead of a 20,000-step one. The checkpoint is `AgentState`, passed as the
next run's `run()` arguments. (2) `plan()` can now return a `research` step, which
AgentRun delegates to `ResearchAgent` — a Child Workflow with its own history,
its own timers and its own Worker.
"""
from __future__ import annotations

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

from temporalio import activity, workflow
from temporalio.exceptions import ApplicationError, ChildWorkflowError
from temporalio.workflow import ParentClosePolicy

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

# The child runs on its own queue only so that "the child's Worker" is a separate process you
# can kill in 6.2 (`python worker.py --role research`). By default a child inherits its parent's
# Task Queue; splitting queues on purpose is module 10.
RESEARCH_TASK_QUEUE = f"{TASK_QUEUE}-research"

MAX_STEPS = 20          # the loop finishes on its own after this many steps
STEPS_PER_RUN = 4       # lab knob: continue-as-new after this many steps in one run -> 5 runs
STEP_SECONDS = 0.5      # durable timer between steps
RESEARCH_STEP = 6       # plan() delegates to ResearchAgent at this step (run 2 of 5)
SOURCES = 3             # ResearchAgent reads this many sources ...
SOURCE_SECONDS = 3      # ... with a durable timer after each: ~10 s to kill its Worker


@dataclass
class AgentState:
    """The checkpoint: everything the next run needs to decide. A summary, never the transcript."""
    step: int
    goal: str
    context_summary: str = ""
    attempts: int = 0
    processed_command_ids: list[str] = field(default_factory=list)   # module 5's dedupe set rides along
    research_n: int = 0
    max_steps: int = MAX_STEPS                                         # the run's length, set once at start (6.1b)


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


@dataclass
class Report:
    task: str
    summary: str
    sources: int


@activity.defn
async def call_llm(prompt: str, key: str) -> str:
    activity.logger.info("call_llm key=%s", key)
    return f"[model answer to: {prompt[:48]}]"


@activity.defn
async def fetch_source(task: str, n: int) -> str:
    activity.logger.info("fetch_source %s #%d", task, n)
    return f"source {n}: notes on {task}"


@workflow.defn
class ResearchAgent:
    """A delegated unit with a lifecycle: several steps, timers, a Query of its own, and its own
    history in the UI. As one Activity this would be one giant timeout and all-or-nothing retry."""

    def __init__(self) -> None:
        self.notes: list[str] = []

    @workflow.run
    async def run(self, task: str) -> Report:
        if "unreliable" in task:                      # 6.2, second half: a child that fails
            raise ApplicationError("source registry rejected the task", non_retryable=True)
        for n in range(1, SOURCES + 1):
            note = await workflow.execute_activity(
                fetch_source, args=[task, n], start_to_close_timeout=timedelta(seconds=10)
            )
            self.notes.append(note)
            await asyncio.sleep(SOURCE_SECONDS)       # a durable timer per source: kill this Worker here
        return Report(task=task, summary=f"{len(self.notes)} sources on {task}", sources=len(self.notes))

    @workflow.query
    def progress(self) -> int:
        return len(self.notes)


@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.pending_context = ""       # accepted by a handler, not yet folded in: part of the snapshot
        self.research_n = 0
        self.steps_this_run = 0         # not in AgentState on purpose: it is about this run
        self.max_steps = MAX_STEPS

    def restore(self, state: AgentState) -> None:
        self.step = state.step
        self.goal = state.goal
        self.context_summary = state.context_summary
        self.attempts = state.attempts
        self.processed_command_ids = set(state.processed_command_ids)
        self.research_n = state.research_n
        self.max_steps = state.max_steps

    def snapshot(self) -> AgentState:
        return AgentState(
            step=self.step,
            goal=self.goal,
            context_summary=self.context_summary + self.pending_context,   # accepted, unprocessed: carry it
            attempts=self.attempts,
            processed_command_ids=sorted(self.processed_command_ids),
            research_n=self.research_n,
            max_steps=self.max_steps,
        )

    def plan(self) -> Step:
        """Workflow code: a pure function of state."""
        if self.step >= self.max_steps:
            return Step("finish")
        if self.step == RESEARCH_STEP:
            return Step("research", task=f"sources for: {self.goal}")
        return Step("llm", prompt=f"Step {self.step} toward: {self.goal}")

    @workflow.run
    async def run(self, goal: str, state: AgentState | None = None) -> str:
        self.goal = goal
        if state:
            self.restore(state)                        # a new run picks up where the old one left off
        while self.status == "running":
            if self.step < self.max_steps and (
                workflow.info().is_continue_as_new_suggested() or self.steps_this_run >= STEPS_PER_RUN
            ):
                await workflow.wait_condition(workflow.all_handlers_finished)   # drain Signals/Updates first
                workflow.continue_as_new(args=[self.goal, self.snapshot()])      # nothing after this line runs
            await workflow.wait_condition(lambda: not self.paused or self.status != "running")
            if self.status != "running":
                break
            step = self.plan()
            if step.kind == "finish":
                self.status = "done"
                break
            if step.kind == "research":
                self.research_n += 1
                try:
                    report = await workflow.execute_child_workflow(
                        ResearchAgent.run, step.task,
                        id=f"{workflow.info().workflow_id}/research/{self.research_n}",
                        task_queue=RESEARCH_TASK_QUEUE,
                        parent_close_policy=ParentClosePolicy.TERMINATE,
                    )
                    self.context_summary += f"\n{report.summary}"
                except ChildWorkflowError as e:        # the child failed: decide, as with an ActivityError
                    self.attempts += 1
                    self.context_summary += f"\nresearch {self.research_n} failed: {e.cause}"
            else:
                key = f"{workflow.info().workflow_id}:{self.step}:{hashlib.sha256(step.prompt.encode()).hexdigest()[:12]}"
                answer = await workflow.execute_activity(
                    call_llm, args=[step.prompt, key], start_to_close_timeout=timedelta(seconds=10)
                )
                self.context_summary = answer if not self.context_summary else f"{self.context_summary}\n{answer}"
            if self.pending_context:                   # fold operator notes at a step boundary
                self.context_summary += self.pending_context
                self.pending_context = ""
            self.step += 1
            self.steps_this_run += 1
            await asyncio.sleep(STEP_SECONDS)
        await workflow.wait_condition(workflow.all_handlers_finished)
        return self.context_summary

    @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 = "cancelled"

    @workflow.signal
    def inject_context(self, text: str, command_id: str) -> None:
        if command_id in self.processed_command_ids:
            return
        self.processed_command_ids.add(command_id)
        self.pending_context += f"\nOperator: {text}"

    @workflow.query
    def status(self) -> dict:
        return {
            "run_id": workflow.info().run_id,          # changes at every continue-as-new; the handle does not
            "step": self.step,
            "steps_this_run": self.steps_this_run,
            "status": self.status,
            "paused": self.paused,
            "goal": self.goal,
            "research_n": self.research_n,
            "attempts": self.attempts,
            "processed_command_ids": sorted(self.processed_command_ids),
        }