Lab 5: Signals, Queries, and Updates

"""Send messages to a running AgentRun from the command line.

    python control.py pause|resume|cancel|status|history [--id agent-42] [--run-id RUN]
    python control.py inject "Prefer sources after 2025" --command-id cmd-7f3a [--twice]
    python control.py change-goal "write the tests first" --command-id cmd-9f1e [--twice]
    python control.py interleave [--no-drain]          # 5.5: two slow Updates + pause + continue-as-new

`--twice` simulates a client whose connection broke after the first send: it sends the
same message again with the same command_id, exactly as a retry loop would.
"""
from __future__ import annotations

import argparse
import asyncio
import os
import sys
from datetime import timedelta

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import connect  # noqa: E402
from temporalio.client import Client, WorkflowUpdateFailedError, WorkflowUpdateStage  # noqa: E402
from temporalio.api.enums.v1 import EventType  # noqa: E402
from temporalio.service import RPCError  # noqa: E402
from workflows import AgentRun  # noqa: E402


async def history(client: Client, workflow_id: str, run_id: str | None = None) -> None:
    """Like common.show_history, plus the detail the lab reads: which Signal, which Update, which Activity.
    Without run_id: the latest run of the ID (after continue-as-new, that is the new one)."""
    handle = client.get_workflow_handle(workflow_id, run_id=run_id)
    async for ev in handle.fetch_history_events():
        name = EventType.Name(ev.event_type).removeprefix("EVENT_TYPE_")
        detail = ""
        if ev.HasField("workflow_execution_signaled_event_attributes"):
            detail = ev.workflow_execution_signaled_event_attributes.signal_name
        elif ev.HasField("workflow_execution_update_accepted_event_attributes"):
            a = ev.workflow_execution_update_accepted_event_attributes
            detail = f"{a.accepted_request.input.name} update_id={a.protocol_instance_id}"
        elif ev.HasField("workflow_execution_update_completed_event_attributes"):
            detail = f"update_id={ev.workflow_execution_update_completed_event_attributes.meta.update_id}"
        elif ev.HasField("activity_task_scheduled_event_attributes"):
            detail = ev.activity_task_scheduled_event_attributes.activity_type.name
        elif ev.HasField("workflow_task_started_event_attributes"):
            detail = f"identity={ev.workflow_task_started_event_attributes.identity}"
        print(f"{ev.event_id:4d}  {name:38s} {detail}")


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("verb", choices=["pause", "resume", "cancel", "status", "history", "inject", "change-goal",
                                     "interleave"])
    ap.add_argument("text", nargs="?", default="", help="the note (inject) or the new goal (change-goal)")
    ap.add_argument("--id", default=os.environ.get("WORKFLOW_ID", "agent-42"))
    ap.add_argument("--run-id", default=None, help="history only: a specific run (e.g. the old run after continue-as-new)")
    ap.add_argument("--command-id", default=None, help="idempotency key chosen once per command; reused on retry")
    ap.add_argument("--twice", action="store_true", help="send the same message twice (simulated retry)")
    ap.add_argument("--reuse-update-id", action="store_true",
                    help="change-goal only: also reuse command_id as the Temporal Update ID, so the Service dedupes it")
    ap.add_argument("--no-drain", action="store_true",
                    help="interleave only: continue-as-new without waiting for handlers (reproduces the lost handler)")
    args = ap.parse_args()

    client = await connect()
    handle = client.get_workflow_handle_for(AgentRun.run, args.id)
    sends = 2 if args.twice else 1

    try:
        await dispatch(ap, args, client, handle, sends)
    except RPCError as e:
        if e.status.name == "NOT_FOUND":
            print(f"{e.message}: {args.id!r} is not running — a Signal or Update needs a live execution")
        else:
            raise


async def dispatch(ap, args, client: Client, handle, sends: int) -> None:
    if args.verb == "pause":
        await handle.signal(AgentRun.pause)                     # returns once the Service recorded it
        print("pause signalled (recorded in history; no Worker needed to accept it)")
    elif args.verb == "resume":
        await handle.signal(AgentRun.resume)
        print("resume signalled")
    elif args.verb == "cancel":
        await handle.signal(AgentRun.cancel_run)
        print("cancel_run signalled — the loop exits at its next check; the execution closes as Completed")
    elif args.verb == "status":
        try:
            print(await handle.query(AgentRun.status, rpc_timeout=timedelta(seconds=10)))
        except RPCError as e:
            print(f"query failed: {e.status.name}: {e.message} — a Query is answered by a Worker; is one running?")
    elif args.verb == "history":
        await history(client, args.id, args.run_id)
    elif args.verb == "inject":
        if not args.command_id:
            ap.error("inject needs --command-id")
        for i in range(sends):
            await handle.signal(AgentRun.inject_context, args=[args.text, args.command_id])
            print(f"send {i + 1}: inject_context({args.text!r}, {args.command_id!r}) recorded")
        print("status:", await handle.query(AgentRun.status))
    elif args.verb == "change-goal":
        if not args.command_id:
            ap.error("change-goal needs --command-id")
        for i in range(sends):
            try:
                new_goal = await handle.execute_update(
                    AgentRun.change_goal, args=[args.text, args.command_id],
                    id=args.command_id if args.reuse_update_id else None,
                )
                print(f"send {i + 1}: accepted; handler returned {new_goal!r}")
            except WorkflowUpdateFailedError as e:
                print(f"send {i + 1}: rejected: {type(e.cause).__name__}: {e.cause}")
        print("status:", await handle.query(AgentRun.status))
    elif args.verb == "interleave":
        await interleave(client, handle, drain=not args.no_drain)


async def interleave(client: Client, handle, drain: bool) -> None:
    """5.5: two `enrich_context` Updates (each awaits a slow Activity), a `pause`, and a
    continue-as-new request, sent back to back. Then wait for the Updates and read the new run."""
    old_run = (await handle.describe()).run_id
    tag = old_run[:4]
    updates = []
    for n, text in enumerate(["Prefer primary sources", "Cite line numbers"], start=1):
        u = await handle.start_update(                          # accepted = the Service has it; not finished
            AgentRun.enrich_context, args=[text, f"cmd-{tag}-{n}"], wait_for_stage=WorkflowUpdateStage.ACCEPTED,
        )
        updates.append((text, u))
        print(f"update {n}: enrich_context({text!r}) accepted; handler now awaiting expand_note")
    await handle.signal(AgentRun.pause)
    print("pause signalled")
    await handle.signal(AgentRun.request_continue_as_new, drain)
    print(f"request_continue_as_new(drain={drain}) signalled")
    for text, u in updates:
        try:
            print(f"update {text!r}: handler returned {await u.result()!r}")
        except Exception as e:                                  # noqa: BLE001 — print what the caller sees
            cause = getattr(e, "cause", None)
            print(f"update {text!r}: LOST — {type(e).__name__}: {cause or e}")
    old = await client.get_workflow_handle(handle.id, run_id=old_run).describe()
    latest = await handle.describe()                                # the handle has no run_id: latest run
    print(f"old run {old_run[:8]} closed as {old.status.name}; latest run {latest.run_id[:8]} is {latest.status.name}")
    print("status of the new run:", await handle.query(AgentRun.status))
    print(f"hint: `python control.py history --run-id {old_run}` reads the old run; plain `history` the new one.")
    print("      `python control.py resume` when you are done reading.")


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

"""Start AgentRun and wait for its result. Talk to it from another terminal with control.py."""
import argparse
import asyncio
import os
import sys

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


async def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--id", default=os.environ.get("WORKFLOW_ID", "agent-42"), help="Workflow ID (= agent ID)")
    ap.add_argument("--goal", default="summarize the Temporal docs")
    ap.add_argument("--no-wait", action="store_true", help="start and return; do not wait for the result")
    args = ap.parse_args()

    client = await connect()
    handle = await start_workflow(client, AgentRun.run, args.goal, 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:")
    print(await handle.result())


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

"""Lab 5 self-check. Runs against the time-skipping test server (no dev server needed):

    python -m pytest tests/ -q
"""
from __future__ import annotations

import asyncio
import os
import sys
import uuid

import pytest

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

from temporalio.client import WorkflowUpdateFailedError, WorkflowUpdateStage  # noqa: E402
from temporalio.service import RPCError  # noqa: E402
from temporalio.testing import WorkflowEnvironment  # noqa: E402
from temporalio.worker import Worker  # noqa: E402
from workflows import AgentRun, call_llm, expand_note  # noqa: E402


async def _update_events(handle) -> list[str]:
    """Update IDs of every accepted Update in history. A rejected Update writes no Update Event
    (the dev server writes nothing at all; the time-skipping test server records the Workflow
    Task that ran the validator, which is why this counts Update Events and not all Events)."""
    return [
        ev.workflow_execution_update_accepted_event_attributes.protocol_instance_id
        async for ev in handle.fetch_history_events()
        if ev.HasField("workflow_execution_update_accepted_event_attributes")
    ]


async def _signal_names(handle) -> list[str]:
    return [
        ev.workflow_execution_signaled_event_attributes.signal_name
        async for ev in handle.fetch_history_events()
        if ev.HasField("workflow_execution_signaled_event_attributes")
    ]


async def _start(env: WorkflowEnvironment, task_queue: str):
    return await env.client.start_workflow(
        AgentRun.run, "summarize the Temporal docs", id=f"lab05-test-{uuid.uuid4()}", task_queue=task_queue
    )


def test_pause_blocks_the_loop_and_resume_releases_it() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            tq = f"lab-05-test-{uuid.uuid4()}"
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=[call_llm]):
                handle = await _start(env, tq)
                await handle.signal(AgentRun.pause)
                status = await handle.query(AgentRun.status)
                assert status["paused"] is True and status["status"] == "running"
                await env.sleep(60)                      # the step in flight finishes; the loop blocks
                step_when_paused = (await handle.query(AgentRun.status))["step"]
                assert step_when_paused <= 1
                await env.sleep(600)                     # ten minutes pass; nothing moves while paused
                assert (await handle.query(AgentRun.status))["step"] == step_when_paused
                await handle.signal(AgentRun.resume)
                await handle.result()                    # time-skipping runs the 20 timers instantly
                final = await handle.query(AgentRun.status)
                assert final["status"] == "done" and final["paused"] is False
                assert (await _signal_names(handle)) == ["pause", "resume"]

    asyncio.run(go())


def test_rejected_update_leaves_no_event_and_accepted_update_returns_the_goal() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            tq = f"lab-05-test-{uuid.uuid4()}"
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=[call_llm]):
                handle = await _start(env, tq)
                await handle.signal(AgentRun.pause)
                await env.sleep(60)                      # the step in flight finishes; the loop blocks
                assert (await handle.query(AgentRun.status))["paused"] is True
                assert await _update_events(handle) == []
                with pytest.raises(WorkflowUpdateFailedError) as exc:
                    await handle.execute_update(AgentRun.change_goal, args=["   ", "cmd-9f20"])
                assert "goal must not be empty" in str(exc.value.cause)
                assert await _update_events(handle) == []         # rejected: no Update Event, ever

                goal = await handle.execute_update(AgentRun.change_goal, args=["write the tests first", "cmd-9f1e"])
                assert goal == "write the tests first"
                again = await handle.execute_update(AgentRun.change_goal, args=["write the tests first", "cmd-9f1e"])
                assert again == "write the tests first"           # retried: same answer, no second write
                assert len(await _update_events(handle)) == 2     # both deliveries accepted; one write
                status = await handle.query(AgentRun.status)
                assert status["goal"] == "write the tests first"
                assert status["notes"] == ["Operator: goal changed to: write the tests first"]
                await handle.signal(AgentRun.cancel_run)
                await handle.result()

    asyncio.run(go())


def test_retried_signal_is_recorded_twice_and_applied_once() -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            tq = f"lab-05-test-{uuid.uuid4()}"
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=[call_llm]):
                handle = await _start(env, tq)
                await handle.signal(AgentRun.pause)
                for _ in range(2):                       # the client's connection "broke"; it re-sends
                    await handle.signal(AgentRun.inject_context, args=["Prefer sources after 2025", "cmd-7f3a"])
                status = await handle.query(AgentRun.status)
                assert status["notes"] == ["Operator: Prefer sources after 2025"]
                assert status["processed_command_ids"] == ["cmd-7f3a"]
                assert (await _signal_names(handle)).count("inject_context") == 2
                await handle.signal(AgentRun.cancel_run)
                result = await handle.result()
                assert "Operator: Prefer sources after 2025" in result.splitlines()
                assert (await handle.query(AgentRun.status))["status"] == "cancelled"

    asyncio.run(go())


@pytest.mark.parametrize("drain", [False, True])
def test_continue_as_new_under_a_live_handler_loses_it_unless_drained(drain: bool) -> None:
    async def go() -> None:
        async with await WorkflowEnvironment.start_time_skipping() as env:
            tq = f"lab-05-test-{uuid.uuid4()}"
            async with Worker(env.client, task_queue=tq, workflows=[AgentRun], activities=[call_llm, expand_note]):
                handle = await _start(env, tq)
                first_run = (await handle.describe()).run_id
                updates = [
                    await handle.start_update(AgentRun.enrich_context, args=[text, f"cmd-{n}"],
                                              wait_for_stage=WorkflowUpdateStage.ACCEPTED)
                    for n, text in enumerate(["Prefer primary sources", "Cite line numbers"])
                ]
                await handle.signal(AgentRun.pause)
                await handle.signal(AgentRun.request_continue_as_new, drain)
                outcomes = []
                for u in updates:
                    try:
                        outcomes.append(await u.result())
                    except (WorkflowUpdateFailedError, RPCError) as e:
                        # dev server 1.31: WorkflowUpdateFailedError (AcceptedUpdateCompletedWorkflow);
                        # the time-skipping test server: RPCError "workflow execution already completed"
                        outcomes.append(f"lost: {type(e).__name__}")
                status = await handle.query(AgentRun.status)           # the new run
                assert (await handle.describe()).run_id != first_run
                assert status["paused"] is True                         # the pause rode along
                if drain:
                    assert outcomes == ["Prefer primary sources (expanded)", "Cite line numbers (expanded)"]
                    assert status["processed_command_ids"] == ["cmd-0", "cmd-1"]
                    assert len(status["notes"]) == 2
                else:
                    assert all(o.startswith("lost:") for o in outcomes)                # both lost
                    assert status["notes"] == [] and status["processed_command_ids"] == []
                await handle.signal(AgentRun.cancel_run)
                await handle.result()

    asyncio.run(go())

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


async def main() -> None:
    client = await connect()
    worker = Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=[AgentRun],
        activities=[call_llm, expand_note],
    )
    print(f"worker pid={os.getpid()} polling {TASK_QUEUE!r}; kill -9 {os.getpid()} to crash it")
    await worker.run()


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

"""Lab 5: AgentRun, stage 5 — the control surface.

Stage 4 gave AgentRun timeouts, retries and idempotency keys. Stage 5 lets a
client talk to it while it runs: Signals `pause`, `resume`, `cancel_run`,
`inject_context(text, command_id)`; the Query `status()`; the Update
`change_goal(goal, command_id)` with a validator. Handlers are Workflow code —
same determinism rules, and they interleave with the main loop.

The loop itself is unchanged in spirit: plan a step, call the model, fold the
answer into `context_summary`, wait a durable timer, repeat. The one new line
that matters is the `wait_condition` at the top of the loop: that is where a
paused agent waits, and while it waits no Worker needs to exist.

Exercise 5.5 adds the interleaving case: `enrich_context` is an Update whose
handler awaits an Activity, so it runs concurrently with the main loop and with
other handlers; `request_continue_as_new` makes the loop checkpoint into a fresh
run (module 6 does this properly) so you can see what happens to a handler that
is still awaiting when the run ends.
"""
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

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import TASK_QUEUE  # noqa: E402,F401  (re-exported for worker.py / starter.py)

MAX_STEPS = 20      # the loop finishes on its own after this many steps (~45 s)
STEP_SECONDS = 2    # durable timer between steps — long enough to send messages by hand
NOTE_SECONDS = 6    # how long `expand_note` takes: long enough to continue-as-new underneath it (5.5)


@dataclass
class AgentState:
    """The snapshot continue-as-new carries into the next run (5.5). Module 6 makes this the
    real thing; here it exists so the new run keeps its step, its notes and its dedupe set."""
    step: int
    context_summary: str = ""
    paused: bool = False
    processed_command_ids: list[str] = field(default_factory=list)


@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 1 fakes the model; stage 4 put a
    cache behind `key`). Kept trivial so the lab is about the control surface."""
    activity.logger.info("call_llm key=%s", key)
    return f"[model answer to: {prompt[:48]}]"


@activity.defn
async def expand_note(text: str) -> str:
    """A slow Activity an Update handler awaits (5.5): the model expands an operator's note."""
    activity.logger.info("expand_note %r", text)
    await asyncio.sleep(NOTE_SECONDS)
    return f"{text} (expanded)"


def fold(summary: str, answer: str) -> str:
    """Workflow code: pure string work. Keep every operator note, replace the
    model's last answer with the new one. (A real agent would summarize.)"""
    notes = [line for line in summary.splitlines() if line.startswith("Operator:")]
    return "\n".join(notes + [answer])


@workflow.defn
class AgentRun:
    def __init__(self) -> None:
        self.step = 0
        self.status = "running"          # running | done | cancelled
        self.goal = ""
        self.context_summary = ""
        self.attempts = 0
        self.paused = False
        self.processed_command_ids: set[str] = set()   # Workflow state: rebuilt by replay
        self.lock = asyncio.Lock()                     # 5.5: one mutating handler at a time
        self.continue_requested = False                # 5.5: set by a Signal; acted on by the loop
        self.drain_before_continue = True              # 5.5: the fix; the lab turns it off first

    def snapshot(self) -> AgentState:
        return AgentState(
            step=self.step, context_summary=self.context_summary, paused=self.paused,
            processed_command_ids=sorted(self.processed_command_ids),
        )

    def plan(self) -> Step:
        """Workflow code: a pure function of state. It never calls a model."""
        if self.step >= MAX_STEPS:
            return Step("finish")
        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 is not None:                           # 5.5: the run after a continue-as-new
            self.step = state.step
            self.context_summary = state.context_summary
            self.paused = state.paused
            self.processed_command_ids = set(state.processed_command_ids)
        while self.status == "running":
            # pause: the main loop blocks here. Nothing polls; no Worker is required.
            await workflow.wait_condition(
                lambda: not self.paused or self.status != "running" or self.continue_requested
            )
            if self.continue_requested:                 # 5.5: checkpoint into a fresh run
                if self.drain_before_continue:
                    # The fix: a handler still awaiting an Activity finishes before the run ends.
                    await workflow.wait_condition(workflow.all_handlers_finished)
                workflow.continue_as_new(args=[self.goal, self.snapshot()])   # nothing after this runs
            if self.status != "running":
                break                                   # cancel_run: exit cleanly (true Cancellation is lab 7)
            step = self.plan()
            if step.kind == "finish":
                self.status = "done"
                break
            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 = fold(self.context_summary, answer)
            self.step += 1
            await asyncio.sleep(STEP_SECONDS)           # TimerStarted / TimerFired
        # Drain: an `async` handler still awaiting something must not be cut off by our return.
        await workflow.wait_condition(workflow.all_handlers_finished)
        return self.context_summary

    # ---- Signals: asynchronous writes, recorded as WorkflowExecutionSignaled ----

    @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"                       # the loop exits at its next check

    @workflow.signal
    def request_continue_as_new(self, drain: bool = True) -> None:
        """5.5: ask the main loop to continue-as-new. Never call continue_as_new from a handler;
        the loop does it. `drain=False` reproduces the bug (the run ends under a live handler)."""
        self.drain_before_continue = drain
        self.continue_requested = True

    @workflow.signal
    def inject_context(self, text: str, command_id: str) -> None:
        if command_id in self.processed_command_ids:    # the client retried after a broken connection
            return                                      # the first delivery already did the work
        self.processed_command_ids.add(command_id)
        self.context_summary += f"\nOperator: {text}"

    # ---- Query: a synchronous read; answered by a live Worker; never in history ----

    @workflow.query
    def status(self) -> dict:
        return {
            "step": self.step,
            "status": self.status,
            "paused": self.paused,
            "goal": self.goal,
            "notes": [l for l in self.context_summary.splitlines() if l.startswith("Operator:")],
            "processed_command_ids": sorted(self.processed_command_ids),
        }

    # ---- Update: a synchronous, tracked write; validated first ----

    @workflow.update
    async def change_goal(self, goal: str, command_id: str) -> str:
        if command_id in self.processed_command_ids:    # a retried Update: same answer, no second write
            return self.goal
        self.processed_command_ids.add(command_id)
        self.goal = goal
        self.context_summary += f"\nOperator: goal changed to: {goal}"
        return self.goal

    @change_goal.validator
    def _check_goal(self, goal: str, command_id: str) -> None:
        if not goal.strip():
            raise ValueError("goal must not be empty")  # rejected: never enters history

    # ---- Update whose handler awaits an Activity (5.5): it interleaves with everything else ----

    @workflow.update
    async def enrich_context(self, text: str, command_id: str) -> str:
        async with self.lock:                           # a second enrich_context waits here
            if command_id in self.processed_command_ids:
                return "ignored"                        # a retry after the first one finished
            expanded = await workflow.execute_activity(  # the handler yields: the loop, a pause, a
                expand_note, text, start_to_close_timeout=timedelta(seconds=30)   # continue-as-new can all happen here
            )
            self.processed_command_ids.add(command_id)  # after the await, under the lock: a retry
            self.context_summary += f"\nOperator: {expanded}"   # that arrived meanwhile is deduped
            return expanded

    @enrich_context.validator
    def _check_note(self, text: str, command_id: str) -> None:
        if not text.strip():
            raise ValueError("note must not be empty")