Lab 10: Operating Temporal in production

"""The operator's five questions, from the Python client (lab 10.2).

    python five_questions.py five <workflow-id>      # all five, naming the surface that answered each
    python five_questions.py history <workflow-id>   # Q1  what happened            (fetch_history)
    python five_questions.py describe <workflow-id>  # Q2-4 what Temporal thinks     (describe)
    python five_questions.py list "<list filter>"    # fleet: which runs             (list_workflows)
    python five_questions.py lane <task-queue>       # Q5  who is polling, backlog   (DescribeTaskQueue)

CLI equivalents: `temporal workflow show|describe -w <id>`, `temporal workflow list --query`,
`temporal task-queue describe --task-queue <lane> --task-queue-type activity`.
"""
import asyncio
import os
import sys
from datetime import datetime, timezone

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import NAMESPACE, connect  # noqa: E402
from temporalio.api.enums.v1 import DescribeTaskQueueMode, EventType, TaskQueueType  # noqa: E402
from temporalio.api.taskqueue.v1 import TaskQueue  # noqa: E402
from temporalio.api.workflowservice.v1 import DescribeTaskQueueRequest  # noqa: E402
from temporalio.client import Client  # noqa: E402

from lanes import AGENT_WORKFLOWS, CPU_TOOLS, GPU_TOOLS  # noqa: E402

TERMINAL = {"ACTIVITY_TASK_COMPLETED", "ACTIVITY_TASK_FAILED", "ACTIVITY_TASK_TIMED_OUT", "ACTIVITY_TASK_CANCELED"}


def _name(event) -> str:
    return EventType.Name(event.event_type).removeprefix("EVENT_TYPE_")


async def q1_history(client: Client, wf_id: str) -> None:
    """Q1 What happened? — the Event History, in order (strongly consistent)."""
    history = await client.get_workflow_handle(wf_id).fetch_history()
    print(f"Q1 what happened: {len(history.events)} events (fetch_history)")
    scheduled = {}
    for ev in history.events:
        detail = ""
        if ev.HasField("activity_task_scheduled_event_attributes"):
            a = ev.activity_task_scheduled_event_attributes
            scheduled[ev.event_id] = a.activity_type.name
            detail = f"{a.activity_type.name} on {a.task_queue.name}"
            if a.priority.priority_key:
                detail += f" priority_key={a.priority.priority_key}"
        elif ev.HasField("activity_task_started_event_attributes"):
            a = ev.activity_task_started_event_attributes
            detail = f"{scheduled.get(a.scheduled_event_id, '?')} attempt={a.attempt} by {a.identity}"
        elif ev.HasField("activity_task_timed_out_event_attributes"):
            a = ev.activity_task_timed_out_event_attributes
            detail = f"{a.failure.message} retry_state={a.retry_state}"
        elif ev.HasField("workflow_task_started_event_attributes"):
            detail = ev.workflow_task_started_event_attributes.identity
        elif ev.HasField("workflow_execution_failed_event_attributes"):
            detail = ev.workflow_execution_failed_event_attributes.failure.message
        elif ev.HasField("upsert_workflow_search_attributes_event_attributes"):
            detail = ", ".join(sorted(ev.upsert_workflow_search_attributes_event_attributes.search_attributes.indexed_fields))
        elif ev.HasField("timer_started_event_attributes"):
            detail = f"fires in {ev.timer_started_event_attributes.start_to_fire_timeout.ToTimedelta()}"
        print(f"  {ev.event_id:4d}  {_name(ev):36s} {detail}")


async def q2_q4_describe(client: Client, wf_id: str) -> None:
    """Q2 What did Temporal think happened? Q3 What side effects may have occurred? Q4 What next?
    — all from DescribeWorkflowExecution (the current view, strongly consistent)."""
    d = await client.get_workflow_handle(wf_id).describe()
    sa = {k.name: v for k, v in d.typed_search_attributes}
    print(f"Q2 what Temporal thinks: status={d.status.name} task_queue={d.task_queue} history_length={d.history_length}")
    print(f"   search attributes: {sa}")
    raw = d.raw_description
    if not raw.pending_activities:
        print("   pending activities: none")
    for pa in raw.pending_activities:
        details = [p.data.decode(errors='replace') for p in pa.heartbeat_details.payloads] if pa.HasField("heartbeat_details") else []
        print(f"   pending {pa.activity_type.name} state={pa.state} attempt={pa.attempt} max={pa.maximum_attempts} "
              f"last_worker={pa.last_worker_identity!r} last_failure={pa.last_failure.message!r} heartbeat={details}")
    # Q3: any ActivityTaskStarted without a terminal event is an attempt that ran somewhere.
    history = await client.get_workflow_handle(wf_id).fetch_history()
    open_attempts = {}
    for ev in history.events:
        if ev.HasField("activity_task_started_event_attributes"):
            open_attempts[ev.activity_task_started_event_attributes.scheduled_event_id] = ev.event_id
        elif _name(ev) in TERMINAL:
            attrs = getattr(ev, _name(ev).lower() + "_event_attributes")
            open_attempts.pop(attrs.scheduled_event_id, None)
    print(f"Q3 side effects: {len(open_attempts)} attempt(s) started without a terminal event; "
          f"attempt counts above say how many times an idempotency key was presented")
    if raw.pending_activities:
        nxt = [f"{pa.activity_type.name} next attempt {pa.scheduled_time.ToDatetime().replace(tzinfo=timezone.utc).isoformat()}"
               for pa in raw.pending_activities if pa.HasField("scheduled_time")]
        print(f"Q4 what next: {nxt}")
    else:
        print("Q4 what next: no pending Activity; pending timers/children:",
              [f"timer {t.timer_id}" for t in raw.pending_timers] if hasattr(raw, "pending_timers") else "n/a",
              [c.workflow_id for c in raw.pending_children])
    print(f"   (now {datetime.now(timezone.utc).isoformat(timespec='seconds')})")


async def q5_lane(client: Client, lane: str) -> None:
    """Q5 What if every Worker disappears now? — the lane's pollers and backlog (DescribeTaskQueue)."""
    r = await client.workflow_service.describe_task_queue(DescribeTaskQueueRequest(
        namespace=NAMESPACE, task_queue=TaskQueue(name=lane),
        task_queue_types=[TaskQueueType.TASK_QUEUE_TYPE_ACTIVITY, TaskQueueType.TASK_QUEUE_TYPE_WORKFLOW],
        api_mode=DescribeTaskQueueMode.DESCRIBE_TASK_QUEUE_MODE_ENHANCED, report_stats=True, report_pollers=True))
    now = datetime.now(timezone.utc)
    for _build, vinfo in r.versions_info.items():
        for ttype, tinfo in vinfo.types_info.items():
            kind = TaskQueueType.Name(ttype).removeprefix("TASK_QUEUE_TYPE_").lower()
            age = tinfo.stats.approximate_backlog_age.ToTimedelta()
            print(f"Q5 lane {lane} [{kind}]: backlog={tinfo.stats.approximate_backlog_count} age={age} "
                  f"add_rate={tinfo.stats.tasks_add_rate:.3f}/s dispatch_rate={tinfo.stats.tasks_dispatch_rate:.3f}/s "
                  f"increase_rate={tinfo.stats.tasks_add_rate - tinfo.stats.tasks_dispatch_rate:+.3f}/s")
            for p in tinfo.pollers:
                seen = now - p.last_access_time.ToDatetime().replace(tzinfo=timezone.utc)
                print(f"   poller {p.identity} last poll {seen.total_seconds():.0f}s ago"
                      + ("  <- gone (a Worker polls every ~60 s; the list forgets it after 5 min)" if seen.total_seconds() > 90 else ""))
            if not tinfo.pollers:
                print("   no pollers")


async def live_pollers(client: Client, lane: str, kind: int = TaskQueueType.TASK_QUEUE_TYPE_ACTIVITY) -> list[str]:
    """Identities that polled `lane` in the last 90 s: Workers that are really there."""
    r = await client.workflow_service.describe_task_queue(DescribeTaskQueueRequest(
        namespace=NAMESPACE, task_queue=TaskQueue(name=lane), task_queue_types=[kind],
        api_mode=DescribeTaskQueueMode.DESCRIBE_TASK_QUEUE_MODE_ENHANCED, report_pollers=True))
    now = datetime.now(timezone.utc)
    return [p.identity for vinfo in r.versions_info.values() for tinfo in vinfo.types_info.values()
            for p in tinfo.pollers
            if (now - p.last_access_time.ToDatetime().replace(tzinfo=timezone.utc)).total_seconds() <= 90]


async def fleet(client: Client, query: str) -> None:
    print(f"list_workflows({query!r}):")
    n = 0
    async for wf in client.list_workflows(query):
        n += 1
        sa = {k.name: v for k, v in wf.typed_search_attributes}
        print(f"  {wf.id:32s} {wf.status.name:10s} AgentStatus={sa.get('AgentStatus')} CurrentStep={sa.get('CurrentStep')} Owner={sa.get('Owner')}")
    print(f"  ({n} executions; visibility is eventually consistent)")


async def main(argv: list[str]) -> None:
    client = await connect()
    cmd, arg = (argv + [None, None])[:2]
    if cmd == "history":
        await q1_history(client, arg)
    elif cmd == "describe":
        await q2_q4_describe(client, arg)
    elif cmd == "list":
        await fleet(client, arg)
    elif cmd == "lane":
        await q5_lane(client, arg)
    elif cmd == "five":
        await q1_history(client, arg)
        await q2_q4_describe(client, arg)
        for lane in (AGENT_WORKFLOWS, CPU_TOOLS, GPU_TOOLS):
            await q5_lane(client, lane)
    else:
        print(__doc__)


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

"""Lane names for lab 10. Spell every Task Queue name once, here, and import it on both sides.

A name mismatch between Client and Worker is not an error in Temporal: it creates two different
Task Queues, the Worker receives nothing, and the Workflow does not progress.

The course names are the defaults. When several people share one dev server, set LAB_PREFIX so your
queues and Workflow IDs do not collide with theirs:  LAB_PREFIX=lab10- python worker_cpu.py
"""
from __future__ import annotations

import os

LAB_PREFIX = os.environ.get("LAB_PREFIX", "")

AGENT_WORKFLOWS = os.environ.get("AGENT_WORKFLOWS_QUEUE", LAB_PREFIX + "agent-workflows")
CPU_TOOLS = os.environ.get("CPU_TOOLS_QUEUE", LAB_PREFIX + "cpu-tools")
GPU_TOOLS = os.environ.get("GPU_TOOLS_QUEUE", LAB_PREFIX + "gpu-tools")

# Which lane each tool needs. Read in Workflow code, so the lane a step went to is replayed like
# any other decision (moving a tool between lanes mid-run is a module 9 change).
CPU_TOOL_NAMES = frozenset({"search", "parse", "fetch"})
GPU_TOOL_NAMES = frozenset({"embed", "vision", "rerank"})
LANE_FOR_TOOL = {**{t: CPU_TOOLS for t in CPU_TOOL_NAMES}, **{t: GPU_TOOLS for t in GPU_TOOL_NAMES}}


def workflow_id(base: str) -> str:
    """`agent-42` in the course; `lab10-agent-42` on a shared server."""
    return LAB_PREFIX + base

"""Lab 10.3 — does the lane honour Priority / Fairness? Build a backlog on `gpu-tools` with no GPU
Worker, then bring up ONE single-slot GPU Worker and read the dispatch order off the histories.

    python priority_race.py                     # N batch runs (priority 5) queued first, then 1 enterprise run (priority 1)
    python priority_race.py --mode fairness     # 2 tenants, same priority, enqueued a a a a b b b b

Everything runs in this process (workflow + cpu Workers first, gpu Worker last) so the order in
which Tasks enter the backlog is exactly the order printed. Prediction (upstream rules, as of
2026-09): priority mode → the enterprise run's GPU Task starts before every batch Task that was
queued ahead of it; fairness mode → a b a b … if `matching.enableFairness` is on, a a a a b b b b
(FIFO) if it is not. Whatever you see, write it down with the server version.
"""
import argparse
import asyncio
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import connect  # noqa: E402
from temporalio.api.enums.v1 import EventType  # noqa: E402
from temporalio.client import Client, WorkflowHandle  # noqa: E402
from temporalio.common import SearchAttributePair, TypedSearchAttributes  # noqa: E402
from temporalio.worker import Worker  # noqa: E402

from lanes import AGENT_WORKFLOWS, CPU_TOOLS, GPU_TOOLS, workflow_id  # noqa: E402
from five_questions import live_pollers, q5_lane  # noqa: E402
from worker_gpu import execute_tool_on_gpu  # noqa: E402
from workflows import OWNER, AgentRun, call_llm, evaluate, execute_tool  # noqa: E402


async def start(client: Client, base: str, owner: str, tier: str) -> WorkflowHandle:
    return await client.start_workflow(
        AgentRun.run, f"race {base}", id=workflow_id(base), task_queue=AGENT_WORKFLOWS,
        search_attributes=TypedSearchAttributes([SearchAttributePair(OWNER, owner)]), memo={"tier": tier})


async def wait_gpu_pending(handle: WorkflowHandle) -> None:
    """Until the run's next Activity is the GPU tool, sitting in the backlog with nobody to take it."""
    while True:
        d = await handle.describe()
        for pa in d.raw_description.pending_activities:
            if pa.activity_type.name == "execute_tool" and (await handle.query(AgentRun.status))["step"] == 2:
                return
        await asyncio.sleep(0.2)


async def gpu_start_time(handle: WorkflowHandle) -> tuple[float, int, str]:
    """(ActivityTaskStarted time of the GPU Task, its priority_key, identity of the Worker) from history."""
    history = await handle.fetch_history()
    gpu_scheduled = {}
    for ev in history.events:
        if ev.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED:
            a = ev.activity_task_scheduled_event_attributes
            if a.task_queue.name == GPU_TOOLS:
                gpu_scheduled[ev.event_id] = a.priority.priority_key
        elif ev.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_STARTED:
            a = ev.activity_task_started_event_attributes
            if a.scheduled_event_id in gpu_scheduled:
                return ev.event_time.ToDatetime().timestamp(), gpu_scheduled[a.scheduled_event_id], a.identity
    raise RuntimeError(f"{handle.id}: GPU Task never started")


async def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--mode", choices=["priority", "fairness"], default="priority")
    p.add_argument("--batch", type=int, default=4, help="runs queued ahead (priority mode) / per tenant (fairness mode)")
    p.add_argument("--tag", default=time.strftime("%H%M%S"), help="suffix so Workflow IDs are fresh on every run")
    a = p.parse_args()
    os.environ.setdefault("GPU_SCHEDULE_TO_START_SECONDS", "180")

    client = await connect()
    # The race needs a backlog, and a backlog needs the GPU lane to be empty: a worker_gpu.py left
    # running from 10.1 takes every Task as it arrives, and the order printed is plain FIFO.
    # (A previous race's own GPU Worker is named race-gpu-* and exited with its process; skip it.)
    others = [i for i in await live_pollers(client, GPU_TOOLS) if not i.startswith("race-gpu-")]
    if others:
        raise SystemExit(f"A GPU Worker is still polling {GPU_TOOLS}: {', '.join(others)}.\n"
                         f"Stop worker_gpu.py (Ctrl-C) and wait ~90 s for the poller to age out, then rerun.")
    wf_worker = Worker(client, task_queue=AGENT_WORKFLOWS, workflows=[AgentRun], identity=f"wf-{os.getpid()}@{AGENT_WORKFLOWS}")
    cpu_worker = Worker(client, task_queue=CPU_TOOLS, activities=[call_llm, execute_tool, evaluate], identity=f"cpu-{os.getpid()}@{CPU_TOOLS}")
    gpu_worker = Worker(client, task_queue=GPU_TOOLS, activities=[execute_tool_on_gpu], max_concurrent_activities=1,
                        identity=f"race-gpu-{os.getpid()}@{GPU_TOOLS}")
    async with wf_worker, cpu_worker:
        queued: list[tuple[str, WorkflowHandle]] = []
        if a.mode == "priority":
            plan = [(f"race-{a.tag}-batch-{i}", "acct_batch", "batch") for i in range(1, a.batch + 1)]
            plan.append((f"race-{a.tag}-enterprise", "acct_ent", "enterprise"))
        else:
            plan = [(f"race-{a.tag}-a-{i}", "acct_a", "batch") for i in range(1, a.batch + 1)]
            plan += [(f"race-{a.tag}-b-{i}", "acct_b", "batch") for i in range(1, a.batch + 1)]
        for base, owner, tier in plan:
            # Strictly one at a time: each run must be IN the GPU backlog before the next is started,
            # so enqueue order is the order of this list.
            h = await start(client, base, owner, tier)
            await wait_gpu_pending(h)
            queued.append((owner, h))
            print(f"queued  {h.id:40s} owner={owner:10s} tier={tier}")
        print("--- backlog before the GPU Worker exists:")
        await q5_lane(client, GPU_TOOLS)
        print(f"--- starting one GPU Worker with 1 slot on {GPU_TOOLS!r}")
        async with gpu_worker:
            await asyncio.gather(*(h.result() for _, h in queued))
        rows = []
        for owner, h in queued:
            t, prio, ident = await gpu_start_time(h)
            rows.append((t, h.id, owner, prio))
        t0 = min(r[0] for r in rows)
        print("--- dispatch order observed (ActivityTaskStarted on the GPU lane):")
        for i, (t, wid, owner, prio) in enumerate(sorted(rows), 1):
            print(f"  {i:2d}. +{t - t0:5.2f}s  {wid:40s} owner={owner:10s} priority_key={prio}")
        enqueue_order = [h.id for _, h in queued]
        observed = [wid for _, wid, _, _ in sorted(rows)]
        print("--- enqueue order was:", " > ".join(x.rsplit('-', 2)[-2] + "-" + x.rsplit('-', 2)[-1] for x in enqueue_order))
        print("--- observed order is:", " > ".join(x.rsplit('-', 2)[-2] + "-" + x.rsplit('-', 2)[-1] for x in observed))
        print("FIFO" if observed == enqueue_order else "NOT FIFO: the lane reordered the backlog")


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

"""Register the custom Search Attributes AgentRun writes: AgentStatus, CurrentStep, Owner.

A Search Attribute must exist in the Namespace before a Workflow upserts it (an upsert of an
unregistered key fails the Workflow Task, forever, until someone registers it). Registration is
an operator action, once per Namespace, through the Operator Service — the same call as
`temporal operator search-attribute create --name AgentStatus --type Keyword`.

Equivalent at server start (dev server only):
    temporal server start-dev --search-attribute AgentStatus=Keyword --search-attribute CurrentStep=Int --search-attribute Owner=Keyword
"""
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 IndexedValueType  # noqa: E402
from temporalio.api.operatorservice.v1 import AddSearchAttributesRequest, ListSearchAttributesRequest  # noqa: E402

WANTED = {
    "AgentStatus": IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
    "CurrentStep": IndexedValueType.INDEXED_VALUE_TYPE_INT,
    "Owner": IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
}


async def main() -> None:
    client = await connect()
    listed = await client.operator_service.list_search_attributes(ListSearchAttributesRequest(namespace=NAMESPACE))
    existing = dict(listed.custom_attributes)
    missing = {k: v for k, v in WANTED.items() if k not in existing}
    for k in WANTED:
        if k in existing:
            print(f"exists  {k:12s} {IndexedValueType.Name(existing[k]).removeprefix('INDEXED_VALUE_TYPE_')}")
    if missing:
        await client.operator_service.add_search_attributes(
            AddSearchAttributesRequest(namespace=NAMESPACE, search_attributes=missing))
        for k, v in missing.items():
            print(f"created {k:12s} {IndexedValueType.Name(v).removeprefix('INDEXED_VALUE_TYPE_')}")
    listed = await client.operator_service.list_search_attributes(ListSearchAttributesRequest(namespace=NAMESPACE))
    assert all(k in listed.custom_attributes for k in WANTED), listed.custom_attributes
    print("namespace", NAMESPACE, "custom attributes:", sorted(listed.custom_attributes))


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

"""Start one AgentRun. `--tier enterprise` puts its GPU step at priority 1; `--tier batch` at 5.
`Owner` is a Search Attribute from the first event; tier rides in memo."""
import argparse
import asyncio
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from common import connect, show_history, start_workflow  # noqa: E402
from temporalio.common import SearchAttributePair, TypedSearchAttributes  # noqa: E402

from lanes import AGENT_WORKFLOWS, workflow_id  # noqa: E402
from workflows import OWNER, AgentRun  # noqa: E402


async def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--goal", default="summarize the Temporal docs")
    p.add_argument("--id", default="agent-42", help="Workflow ID = agent id (LAB_PREFIX is prepended)")
    p.add_argument("--owner", default="acct_7f3a", help="opaque account id; never an email (Search Attributes are unencrypted)")
    p.add_argument("--tier", default="batch", choices=["enterprise", "team", "batch"])
    p.add_argument("--no-wait", action="store_true", help="start and return; do not wait for the result")
    p.add_argument("--gpu-schedule-to-start", type=int, default=0, metavar="SECONDS",
                   help="this run's Schedule-to-Start on the GPU lane (10.1's stall); default: the Worker's 60 s")
    a = p.parse_args()

    client = await connect()
    handle = await start_workflow(client,
        AgentRun.run, a.goal, id=workflow_id(a.id), task_queue=AGENT_WORKFLOWS,
        search_attributes=TypedSearchAttributes([SearchAttributePair(OWNER, a.owner)]),
        memo={"tier": a.tier, **({"gpu_schedule_to_start_seconds": a.gpu_schedule_to_start}
                                 if a.gpu_schedule_to_start else {})},
    )
    print(f"started {handle.id} run={handle.result_run_id} owner={a.owner} tier={a.tier}")
    print(f"UI: http://localhost:8233/namespaces/default/workflows/{handle.id}/{handle.result_run_id}/history")
    if a.no_wait:
        return
    print("result:", await handle.result())
    await show_history(client, handle.id)


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

"""Test knobs must be in the environment before `workflows` is imported — the Worker's sandbox
re-imports the module and reads them again, so os.environ is the one channel both sides see."""
import os
import sys

LAB_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
os.environ.setdefault("LAB_PREFIX", "test10-")
os.environ["GPU_SCHEDULE_TO_START_SECONDS"] = "2"
os.environ["TOOL_SECONDS"] = "0.2"
sys.path.insert(0, LAB_DIR)
sys.path.insert(0, os.path.dirname(LAB_DIR))

"""Lab 10 self-check. Runs on the SDK's time-skipping test server (downloaded once; no Docker).
No pytest-asyncio: each test drives its coroutine on a module-scoped event loop."""
import asyncio

import pytest
from temporalio.api.enums.v1 import EventType, IndexedValueType
from temporalio.api.operatorservice.v1 import AddSearchAttributesRequest
from temporalio.client import WorkflowFailureError
from temporalio.common import SearchAttributePair, TypedSearchAttributes
from temporalio.exceptions import ApplicationError
from temporalio.testing import ActivityEnvironment, WorkflowEnvironment
from temporalio.worker import Worker

from lanes import AGENT_WORKFLOWS, CPU_TOOLS, GPU_TOOLS
from worker_gpu import execute_tool_on_gpu
from workflows import OWNER, AgentRun, ToolCall, call_llm, evaluate, execute_tool


@pytest.fixture(scope="module")
def loop():
    loop = asyncio.new_event_loop()
    yield loop
    loop.close()


@pytest.fixture(scope="module")
def env(loop):
    async def start():
        env = await WorkflowEnvironment.start_time_skipping()
        # The test server, like the dev server, rejects an unregistered Search Attribute
        # ("search attribute Owner is not defined"); register them the way setup_search_attributes.py does.
        await env.client.operator_service.add_search_attributes(AddSearchAttributesRequest(
            namespace="default", search_attributes={
                "AgentStatus": IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
                "CurrentStep": IndexedValueType.INDEXED_VALUE_TYPE_INT,
                "Owner": IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD}))
        return env
    env = loop.run_until_complete(start())
    yield env
    loop.run_until_complete(env.shutdown())


def workflow_worker(env):
    return Worker(env.client, task_queue=AGENT_WORKFLOWS, workflows=[AgentRun])


def cpu_worker(env):
    return Worker(env.client, task_queue=CPU_TOOLS, activities=[call_llm, execute_tool, evaluate])


def gpu_worker(env):
    return Worker(env.client, task_queue=GPU_TOOLS, activities=[execute_tool_on_gpu], max_concurrent_activities=1)


async def start(env, wf_id: str, tier: str):
    return await env.client.start_workflow(
        AgentRun.run, "summarize the Temporal docs", id=wf_id, task_queue=AGENT_WORKFLOWS,
        search_attributes=TypedSearchAttributes([SearchAttributePair(OWNER, "acct_test")]), memo={"tier": tier})


def gpu_scheduled_events(history):
    return [ev.activity_task_scheduled_event_attributes for ev in history.events
            if ev.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED
            and ev.activity_task_scheduled_event_attributes.task_queue.name == GPU_TOOLS]


def test_each_activity_runs_on_the_lane_the_workflow_named(loop, env):
    async def body():
        async with workflow_worker(env), cpu_worker(env), gpu_worker(env):
            handle = await start(env, "agent-42", "enterprise")
            result = await handle.result()
            assert result["lanes_used"] == [CPU_TOOLS, CPU_TOOLS, GPU_TOOLS]
            desc = await handle.describe()
            sa = {k.name: v for k, v in desc.typed_search_attributes}
            assert sa["AgentStatus"] == "done" and sa["CurrentStep"] == 3 and sa["Owner"] == "acct_test"
            history = await handle.fetch_history()
            [gpu] = gpu_scheduled_events(history)
            assert gpu.priority.priority_key == 1, "enterprise tier → priority 1 on the GPU lane"
    loop.run_until_complete(body())


def test_batch_tier_schedules_gpu_work_at_priority_5(loop, env):
    async def body():
        async with workflow_worker(env), cpu_worker(env), gpu_worker(env):
            handle = await start(env, "agent-batch", "batch")
            await handle.result()
            [gpu] = gpu_scheduled_events(await handle.fetch_history())
            assert gpu.priority.priority_key == 5
    loop.run_until_complete(body())


def test_gpu_lane_without_workers_fails_at_schedule_to_start(loop, env):
    async def body():
        async with workflow_worker(env), cpu_worker(env):  # no GPU pool
            handle = await start(env, "agent-stalled", "batch")
            with pytest.raises(WorkflowFailureError) as err:
                await handle.result()
            assert isinstance(err.value.cause, ApplicationError) and err.value.cause.type == "LaneStalled"
            history = await handle.fetch_history()
            timed_out = [ev for ev in history.events if ev.event_type == EventType.EVENT_TYPE_ACTIVITY_TASK_TIMED_OUT]
            assert len(timed_out) == 1, "Schedule-to-Start is not retried: one timeout event, no second attempt"
            assert timed_out[0].activity_task_timed_out_event_attributes.failure.timeout_failure_info.timeout_type == 2  # SCHEDULE_TO_START
            sa = {k.name: v for k, v in (await handle.describe()).typed_search_attributes}
            assert sa["AgentStatus"] == "stalled" and sa["CurrentStep"] == 2
    loop.run_until_complete(body())


def test_gpu_worker_refuses_a_cpu_tool(loop):
    async def body():
        with pytest.raises(ApplicationError) as err:
            await ActivityEnvironment().run(execute_tool_on_gpu, ToolCall("k", "search", {"q": "x"}))
        assert err.value.type == "BadArguments" and err.value.non_retryable
    loop.run_until_complete(body())

"""The `cpu-tools` pool: cheap, autoscaled. Registers exactly the Activities that need no GPU."""
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.worker import Worker  # noqa: E402

from lanes import CPU_TOOLS  # noqa: E402
from workflows import call_llm, evaluate, execute_tool  # noqa: E402


async def main() -> None:
    client = await connect()
    worker = Worker(client, task_queue=CPU_TOOLS, activities=[call_llm, execute_tool, evaluate],
                    identity=f"cpu-{os.getpid()}@{CPU_TOOLS}")
    print(f"cpu worker pid={os.getpid()} polling {CPU_TOOLS!r}: call_llm, execute_tool, evaluate")
    await worker.run()


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

"""The `gpu-tools` pool: scarce, preemptible. One Activity slot by default (GPU_SLOTS) so the
dispatch order on the lane is observable one Task at a time (lab 10.3).

The set of Activities registered here is exactly the set that needs a GPU. Any Worker polling
`gpu-tools` must register all of them, so a CPU-only tool arriving here is a routing bug — the
guard below turns it into a non-retryable failure instead of a silently wrong answer."""
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 import activity  # noqa: E402
from temporalio.exceptions import ApplicationError  # noqa: E402
from temporalio.worker import Worker  # noqa: E402

from lanes import GPU_TOOL_NAMES, GPU_TOOLS  # noqa: E402
from workflows import ToolCall, ToolResult, execute_tool  # noqa: E402

GPU_SLOTS = int(os.environ.get("GPU_SLOTS", "1"))


@activity.defn(name="execute_tool")
async def execute_tool_on_gpu(call: ToolCall) -> ToolResult:
    if call.tool not in GPU_TOOL_NAMES:
        raise ApplicationError(f"{call.tool} does not need a GPU; it was routed to {GPU_TOOLS}",
                               type="BadArguments", non_retryable=True)
    return await execute_tool(call)


async def main() -> None:
    client = await connect()
    worker = Worker(client, task_queue=GPU_TOOLS, activities=[execute_tool_on_gpu],
                    max_concurrent_activities=GPU_SLOTS, identity=f"gpu-{os.getpid()}@{GPU_TOOLS}")
    print(f"gpu worker pid={os.getpid()} polling {GPU_TOOLS!r} with {GPU_SLOTS} slot(s): execute_tool (GPU tools only)")
    await worker.run()


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

"""The `agent-workflows` pool: Workflow code only. Deliberately boring, so you can run many of
these and redeploy them twice a day without touching a GPU."""
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.worker import Worker  # noqa: E402

from lanes import AGENT_WORKFLOWS  # noqa: E402
from workflows import AgentRun  # noqa: E402


async def main() -> None:
    client = await connect()
    worker = Worker(client, task_queue=AGENT_WORKFLOWS, workflows=[AgentRun],
                    identity=f"wf-{os.getpid()}@{AGENT_WORKFLOWS}")
    print(f"workflow worker pid={os.getpid()} polling {AGENT_WORKFLOWS!r} (Workflow Tasks only)")
    await worker.run()


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

"""Lab 10: AgentRun, stage 10 — the same program, split across compute lanes.

Nothing about the agent loop changes. What changes is that every Activity call names the lane it
needs (`task_queue=`), carries a Schedule-to-Start timeout so an empty lane becomes visible, and
the run publishes `AgentStatus` / `CurrentStep` Search Attributes so an operator can find it among
thousands without opening a history.
"""
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 Priority, RetryPolicy, SearchAttributeKey
from temporalio.exceptions import ActivityError, ApplicationError, TimeoutError, TimeoutType

from lanes import CPU_TOOLS, GPU_TOOL_NAMES, GPU_TOOLS, LANE_FOR_TOOL

# ---------------------------------------------------------------- Search Attributes (registered
# once per Namespace by setup_search_attributes.py; the Workflow only writes them)
AGENT_STATUS = SearchAttributeKey.for_keyword("AgentStatus")
CURRENT_STEP = SearchAttributeKey.for_int("CurrentStep")
OWNER = SearchAttributeKey.for_keyword("Owner")

# Lab knobs. The course values are in the comments; the lab shrinks them so you can watch.
GPU_SCHEDULE_TO_START = timedelta(seconds=int(os.environ.get("GPU_SCHEDULE_TO_START_SECONDS", "60")))  # course: 5 min
TOOL_SECONDS = float(os.environ.get("TOOL_SECONDS", "3"))  # how long the fake tool "computes"

# Priority on the shared GPU lane: 1 is served first, 5 last, default 3 (upstream, as of 2026-09).
PRIORITY_FOR_TIER = {"enterprise": 1, "team": 3, "batch": 5}
# One weight per fairness key, and it does not move. Upstream is explicit: "There should only be
# one fairness weight assigned to each fairness key within a Task Queue. Having multiple fairness
# weights on a fairness key will result in unspecific behavior." The key here is the tenant, so the
# weight has to belong to the tenant -- deriving it from the tier breaks that rule the moment one
# tenant runs at two tiers. The tier is already expressed, once, as the priority key: priority for
# the SLA tier, fairness for equity between tenants inside it. If weights genuinely must differ by
# tier, make the key f"{tier}:{owner}" rather than moving the weight.
# Empty means every tenant is equal. In production these are set without a deploy:
# `temporal task-queue config set --fairness-key-weight <tenant>=<weight>`.
FAIRNESS_WEIGHT_FOR_OWNER = {}
DEFAULT_FAIRNESS_WEIGHT = 1.0


@dataclass
class Step:
    kind: str  # "llm" | "tool" | "finish"
    prompt: str = ""
    tool: str = ""
    args: dict = field(default_factory=dict)


@dataclass
class ToolCall:
    key: str
    tool: str
    args: dict


@dataclass
class ToolResult:
    key: str
    tool: str
    output: str
    attempt: int
    task_queue: str  # which lane answered — the Workflow never chose a host, only a lane
    worker_pid: int


@dataclass
class Verdict:
    ok: bool
    note: str


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


@activity.defn
async def execute_tool(call: ToolCall) -> ToolResult:
    """One registration for both lanes; the Worker that polls `gpu-tools` must be the one with a GPU.
    The fake heartbeats every 0.5 s so a dead Worker is noticed at the heartbeat timeout, not
    at Start-to-Close."""
    info = activity.info()
    activity.logger.info("execute_tool %s key=%s attempt=%d on %s", call.tool, call.key, info.attempt, info.task_queue)
    ticks = max(1, int(TOOL_SECONDS / 0.5))
    for i in range(ticks):
        activity.heartbeat({"tick": i, "of": ticks})
        await asyncio.sleep(TOOL_SECONDS / ticks)
    return ToolResult(call.key, call.tool, f"[{call.tool} output for {call.args}]", info.attempt, info.task_queue, os.getpid())


@activity.defn
async def evaluate(result: ToolResult) -> Verdict:
    return Verdict(ok=bool(result.output), note=f"{result.tool} answered from {result.task_queue}")


# ---------------------------------------------------------------- The Workflow
@workflow.defn
class AgentRun:
    def __init__(self) -> None:
        self.step = 0
        self.status = "starting"
        self.goal = ""
        self.context_summary = ""
        self.attempts = 0
        self.paused = False
        self.tier = "batch"
        self.owner = ""
        self.gpu_schedule_to_start = GPU_SCHEDULE_TO_START
        self.lanes_used: list[str] = []

    # -- decisions: pure functions of state ---------------------------------------------------
    def plan(self) -> Step:
        if self.step == 0:
            return Step("llm", prompt=f"Make a three-step plan for: {self.goal}")
        if self.step == 1:
            return Step("tool", tool="search", args={"q": self.goal})       # cpu-tools
        if self.step == 2:
            return Step("tool", tool="embed", args={"text": self.context_summary[:60]})  # gpu-tools
        return Step("finish")

    def _key(self, *parts: str) -> str:
        digest = hashlib.sha256(":".join(parts).encode()).hexdigest()[:12]
        return f"{workflow.info().workflow_id}:{self.step}:{digest}"

    def _publish(self, status: str) -> None:
        """The Query reads self.status (authoritative); the fleet view reads the Search Attributes
        (eventually consistent). The upsert is a Command recorded in History and replays."""
        self.status = status
        workflow.upsert_search_attributes([AGENT_STATUS.value_set(status), CURRENT_STEP.value_set(self.step)])

    # -- control surface ---------------------------------------------------------------------
    @workflow.signal
    def pause(self) -> None:
        self.paused = True
        self._publish("paused")

    @workflow.signal
    def resume(self) -> None:
        self.paused = False
        self._publish("running")

    @workflow.query
    def status(self) -> dict:
        return {"step": self.step, "status": self.status, "goal": self.goal, "paused": self.paused,
                "tier": self.tier, "owner": self.owner, "lanes_used": self.lanes_used}

    # -- effects: each names its lane ---------------------------------------------------------
    async def _run_tool(self, step: Step) -> ToolResult:
        call = ToolCall(self._key(step.tool, repr(sorted(step.args.items()))), step.tool, step.args)
        lane = LANE_FOR_TOOL[step.tool]
        gpu = step.tool in GPU_TOOL_NAMES
        result: ToolResult = await workflow.execute_activity(
            execute_tool, call,
            task_queue=lane,
            # Queued with no Worker on the lane → this is the timeout that fires (lab 10.1).
            schedule_to_start_timeout=self.gpu_schedule_to_start if gpu else timedelta(minutes=1),
            start_to_close_timeout=timedelta(minutes=2),
            heartbeat_timeout=timedelta(seconds=10),
            retry_policy=RetryPolicy(maximum_attempts=3, non_retryable_error_types=["BadArguments"]),
            # Lab 10.3: a Task Queue property, set per call. Same lane, different service order.
            priority=Priority(
                priority_key=PRIORITY_FOR_TIER[self.tier],
                fairness_key=self.owner or None,
                fairness_weight=FAIRNESS_WEIGHT_FOR_OWNER.get(self.owner, DEFAULT_FAIRNESS_WEIGHT),
            ) if gpu else Priority(),
        )
        self.lanes_used.append(result.task_queue)
        verdict: Verdict = await workflow.execute_activity(
            evaluate, result, task_queue=CPU_TOOLS, start_to_close_timeout=timedelta(seconds=30),
        )
        self.context_summary = f"{result.output} ({verdict.note})"
        return result

    @workflow.run
    async def run(self, goal: str) -> dict:
        self.goal = goal
        # Who owns this run and how it is prioritised come in from the starter: Owner as a Search
        # Attribute (so it is queryable), tier as memo (metadata the Workflow reads, nobody queries).
        self.owner = workflow.info().typed_search_attributes.get(OWNER) or ""
        self.tier = workflow.memo_value("tier", "batch", type_hint=str)
        # The stall experiment (10.1) shortens the GPU lane's Schedule-to-Start per run, from the
        # starter: memo is fixed at start, so this is as deterministic as a constant. Without it
        # the Worker's GPU_SCHEDULE_TO_START_SECONDS applies.
        seconds = workflow.memo_value("gpu_schedule_to_start_seconds", 0, type_hint=int)
        self.gpu_schedule_to_start = timedelta(seconds=seconds) if seconds else GPU_SCHEDULE_TO_START
        self._publish("running")
        while True:
            await workflow.wait_condition(lambda: not self.paused)
            step = self.plan()
            if step.kind == "finish":
                self._publish("done")
                return {"summary": self.context_summary, "lanes_used": self.lanes_used, "steps": self.step}
            try:
                if step.kind == "llm":
                    self.context_summary = await workflow.execute_activity(
                        call_llm, args=[step.prompt, self._key(step.prompt)],
                        task_queue=CPU_TOOLS,  # production: the `premium-models` lane, rate-limited
                        schedule_to_start_timeout=timedelta(minutes=1),
                        start_to_close_timeout=timedelta(seconds=30),
                    )
                    self.lanes_used.append(CPU_TOOLS)
                else:
                    await self._run_tool(step)
            except ActivityError as err:
                cause = err.cause
                if isinstance(cause, TimeoutError) and cause.type == TimeoutType.SCHEDULE_TO_START:
                    # No Worker on that lane took the Task before the deadline. Record it where an
                    # operator will look (Search Attribute), then fail loudly.
                    self._publish("stalled")
                    raise ApplicationError(
                        f"lane stalled: no Worker took {step.tool or 'call_llm'} within {cause.type.name}",
                        type="LaneStalled", non_retryable=True) from err
                raise
            self.step += 1
            # A pause Signal may have landed while the step ran: republish *its* status, not ours.
            self._publish("paused" if self.paused else "running")