Lab 11: Capstone: a durable agent runtime

# AgentRun — architecture defense (template, lab 11.4)

Fill this in from memory, then check against `DEFENSE_ANSWERS.md`. One line per cell. If a cell
takes more than one line, you have not found the mechanism yet — you have found the module.

Author: ______________   Date: __________   Server version you ran it on: ______________

## 1. The failure contract (one line per box)

Write this before the diagram. For every box, answer eight questions: What survives a Worker death?
Can this operation happen twice? Where is its durable identity? Who retries it? Who cancels it?
What happens across a deployment? What goes into Event History? How does it Continue-As-New? The
table holds the three that fit in a line; the others land in section 2 (who retries it, across a
deployment, continue-as-new) and in the ten questions (who cancels it).

| Box | A Worker dies here | On retry | What history records |
|---|---|---|---|
| `plan()` | | | |
| `call_llm` (LLM Activity) | | | |
| `execute_tool` (Tool Activity) | | | |
| `ResearchAgent` (child) | | | |
| pause wait (`wait_condition`) | | | |
| `continue_as_new` | | | |

## 2. Requirement → mechanism (12 rows)

For each requirement: the *one* Temporal mechanism that satisfies it, the name of the box in the
diagram (a method, an Activity, a handler, a call), and the lab where you saw it fail and recover.

| # | Requirement | Mechanism | Box in AgentRun | Where you saw it |
|---|---|---|---|---|
| 1 | Runs for up to 72 hours | | | |
| 2 | Makes 1,000+ model calls | | | |
| 3 | Tool calls can fail | | | |
| 4 | Expensive calls must not accidentally duplicate | | | |
| 5 | GPUs disappear | | | |
| 6 | Workers deploy twice a day; old executions survive | | | |
| 7 | User can pause and resume | | | |
| 8 | User can change instructions mid-run | | | |
| 9 | User can ask what it is doing | | | |
| 10 | May delegate to sub-agents | | | |
| 11 | History cannot grow forever | | | |
| 12 | Operators need exact execution diagnostics | | | |

## 3. Where Temporal ends

| Layer | Owns | For AgentRun, concretely |
|---|---|---|
| Temporal | | |
| Kubernetes | | |
| Kafka | | |
| Postgres | | |

## 4. The ten questions, and one from 11.5 (one sentence each)

1. What exactly survives when a Worker crashes?
2. How does replay reconstruct Workflow state?
3. Why must Workflow code be deterministic?
4. Why can an Activity execute twice?
5. Activity retry vs. Workflow replay?
6. Signal vs. Query vs. Update?
7. Activity vs. Child Workflow?
8. When should you continue-as-new?
9. How do you safely deploy changed Workflow code?
10. Where does Temporal end and Kubernetes / Kafka / Postgres begin?
11. (11.5) What does the user see, and from which store? (product projection, never the visibility list)

## 5. Evidence

Paste, for your own run: the Run ID list from `run_200_steps.py`, the `kill -9` handoff events, and
`pytest` output from `tests/`. Anything you cannot paste, you have not defended.

# AgentRun — architecture defense (answers)

Filled in from the course design notes and from what the lab 11 code and tests
actually do. Where a claim is about upstream behaviour it carries "as of 2026-09"; where it is
about this lab it names the file or test that shows it.

## 1. The failure contract

| Box | A Worker dies here | On retry | What history records |
|---|---|---|---|
| `plan()` | Replayed from history on the next Worker; nothing is lost because nothing happened outside history | Not applicable — it is a decision | Commands only (the Activity it scheduled, the timer it started) |
| `call_llm` | The attempt fails at Start-to-Close or heartbeat; the Activity is retried with the same key | The lab fake returns its cached answer. In production, a gateway ledger alone cannot rule out a provider charge that succeeded just before the gateway crashed; use provider-side idempotency or reconcile an ambiguous outcome before another paid call | `ActivityTaskScheduled` / `Started` / `Completed` with the result; the failed attempt is *not* an event |
| `execute_tool` | Heartbeat timeout fails the attempt quickly; the next attempt receives the last delivered heartbeat details | Activity code reads those details and decides what to repeat; a completed-but-unreported effect comes back from the ledger as `duplicate=True` | Same three events, plus heartbeat details in Describe while pending |
| `ResearchAgent` | The child replays on its own; the parent does not notice | The child's own retry policy | `StartChildWorkflowExecutionInitiated` / `ChildWorkflowExecutionStarted` / `Completed` in the parent |
| pause wait | The condition is Workflow state; nothing polls while paused; the next Worker replays into the same wait | Not applicable | `WorkflowExecutionSignaled` (pause, resume); `UpsertWorkflowSearchAttributes` for `paused` |
| `continue_as_new` | The new run starts atomically as the old one closes; state arrives as the next run's arguments | Not applicable | `WorkflowExecutionContinuedAsNew` with `new_execution_run_id`; the new run's `WorkflowExecutionStarted` carries the snapshot |

Observed on the dev server (server 1.31.2, SDK 1.32.0): a `kill -9` of the workflow Worker at step
6 of 12 produced `WORKFLOW_TASK_TIMED_OUT` (sticky Schedule-to-Start, ~10 s), a rescheduled Workflow
Task, and the new Worker's first task 23 ms later; the run finished with every Activity executed once.

## 2. Requirement → mechanism

| # | Requirement | Mechanism | Box in AgentRun | Where you saw it |
|---|---|---|---|---|
| 1 | Runs for up to 72 hours | A Workflow Execution outlives every process. Waiting is a durable timer or `wait_condition`; no Worker is occupied while waiting | `run()` loop; `wait_condition(not paused)` | L01 (`kill -9` during the sleep); 11.1 `kill -9` handoff |
| 2 | Makes 1,000+ model calls | Each call is an Activity whose result is recorded once; the loop is ordinary code over recorded results | `call_llm` on `cpu-tools` (production: `premium-models`) | 11.2 — 200 steps, 4 Run IDs, every result in history |
| 3 | Tool calls can fail | `RetryPolicy` (bad arguments non-retryable) + Start-to-Close + Schedule-to-Close + heartbeat timeout; a retry is server-side, invisible to the caller | `execute_tool` with `TOOL_RETRY`, `TOOL_START_TO_CLOSE`, `TOOL_HEARTBEAT` | `test_stalled_attempt_times_out_on_heartbeat_and_the_retry_resumes` — attempt 2 resumed from heartbeat tick 2; no timeout event |
| 4 | Expensive calls must not accidentally duplicate | Idempotency key `wf_id:step:hash(...)` chosen in Workflow code, so a retry presents the same key; the lab fake keeps a ledger by key. A real provider also needs to honor that key or support result lookup and reconciliation after an ambiguous outcome | `_llm()` / `_tool()` build the key; `LLM_CACHE`, `TOOL_LEDGER` in `activities.py` | `test_completion_lost_after_the_effect_is_not_a_second_effect` — one effect, `duplicate=True` on attempt 2 in the fake ledger |
| 5 | GPUs disappear | A `gpu-tools` lane with Schedule-to-Start (nobody took it) and heartbeat (somebody died holding it) timeouts; the L07 `compensations` list to release what a cancelled run holds | `_tool()` for `GPU_TOOL_NAMES`; `compensations`; `_compensate()` | Lab 10.1 stall (`ACTIVITY_TASK_TIMED_OUT ScheduleToStart`); `test_temporal_cancel_during_the_gpu_step_runs_the_compensation` |
| 6 | Workers deploy twice a day; old executions survive | Worker Versioning: executions pinned to the version that started them, upgrade at continue-as-new (GA 2026-03-30; upgrade-on-CaN Public Preview, as of 2026-09); `workflow.patched` for in-place changes | L09's `verify` step; not exercised in lab 11 code | L09 (v1 and v2 side by side) |
| 7 | User can pause and resume | Signals set a flag; the loop waits on `wait_condition`; a paused Workflow consumes no Worker compute while waiting | `pause` / `resume` / `cancel_run` handlers | `test_pause_update_validator_and_command_dedupe`; lab 10.2 fleet query `AgentStatus = 'paused'` |
| 8 | User can change instructions mid-run | Update: a tracked write with a validator (rejected = never in history) and a `command_id` so a retried Update is applied once | `change_goal` + `@change_goal.validator`; `processed_command_ids` | same test — empty goal → `WorkflowUpdateFailedError`; `c1` twice → `ignored`; history has `UPDATE_ACCEPTED ×2`, none for `c0` |
| 9 | User can ask what it is doing | Query: a read of Workflow state, answered by a Worker from replayed state, never recorded | `status()` | every lab; `test_worker_crash…` shows a Query needs a Worker (deadline exceeded with none alive) while `describe()` does not |
| 10 | May delegate to sub-agents | Child Workflow: its own history, replay and retry; the parent sees one result and a stable id | `ResearchAgent`, id `f"{wf_id}/research/{step}"` | 11.2 — 13 children under `agent-200/research/` |
| 11 | History cannot grow forever | `continue_as_new(AgentState)` when `is_continue_as_new_suggested()` or every `STEPS_PER_RUN` steps; memo and Search Attributes carry over | `_should_continue_as_new()`, `snapshot()`, `run(goal, state)` | 11.2 — longest run 519 events for 50 steps; `test_continue_as_new_keeps_every_run_under_the_bound` |
| 12 | Operators need exact execution diagnostics | History (what happened), Describe (what Temporal thinks, pending attempts, heartbeat details), Search Attributes for finding, `DescribeTaskQueue` for lanes | `_publish()` → `AgentStatus`/`CurrentStep`; `Owner` at start | Lab 10.2 five questions; `five_questions.py` |

## 3. Where Temporal ends

| Layer | Owns | For AgentRun, concretely |
|---|---|---|
| Temporal | Where execution logically is: the step, the outstanding attempt, the pending timer, the lane a Task waits on | `agent-42` is at step 137, attempt 2 of `embed` is outstanding on `gpu-tools`, paused since Tuesday |
| Kubernetes | Where processes run: one Deployment per lane, autoscaled on that lane's backlog age; pod identity in the history | `gpu-tools` pods on the GPU node pool; `Worker(identity=f"{pod}@{lane}")` |
| Kafka | Streams between systems — not execution, not history | Outside events reach the run as Signals (a consumer calling Signal-With-Start); the run emits through an Activity that produces |
| Postgres | The system of record for results: what the provider charged, what the customer owns, whether the answer was good | The ledgers in `activities.py` are its stand-in: rows keyed by the idempotency keys, written by Activities, outliving retention |

The engine that rendered this course draws the same line: `MiniLectureCourseWorkflow` keeps the
legs, retries, heartbeats, yields and the review checkpoint in Temporal, and leaves the job store,
`agent_state.json`, artifacts, `usage.json` and the verifier database outside it. History answers
what execution happened; it does not answer whether the lecture was good.

## 4. The ten questions, and one from 11.5

1. **What survives a Worker crash?** Everything in the Service: the Event History (every decision and every recorded result), pending timers, queued Tasks with their attempt counts, Signals, Search Attributes. Nothing that lived only in the process — and the program was written so nothing does.
2. **How does replay reconstruct state?** A fresh Worker runs the Workflow code from the top, and every time the code asks for something (an Activity result, a timer), the SDK answers from history instead of doing it again; the code arrives at the same state having executed no effect.
3. **Why deterministic?** Because replay must make the same decisions in the same order to match the history; a decision that depends on the clock, randomness or a network call would diverge and the SDK would (rightly) refuse to continue.
4. **Why can an Activity execute twice?** Because the Service cannot tell "the Worker died before starting" from "the Worker did the work and died before reporting"; it retries, so at-least-once is the guarantee and idempotency by key is your job.
5. **Activity retry vs. Workflow replay?** A retry re-executes an *effect* (server-side, invisible in history until the final outcome); a replay re-executes *decisions* with no effects at all. Retries cost; replays are free.
6. **Signal vs. Query vs. Update?** Signal: an asynchronous write, recorded, no reply. Query: a read of current state, answered by a Worker, not recorded. Update: a tracked synchronous write that can return a result; a validator is optional, and an accepted Update is recorded.
7. **Activity vs. Child Workflow?** An Activity is one effect with timeouts and retries and no state of its own; a Child Workflow is its own durable program with its own history, replay, Signals and children — use it when the work has decisions, not just an effect.
8. **When continue-as-new?** When the Service says so (`is_continue_as_new_suggested()`), and on a cadence you choose; carry a snapshot of what `plan()` needs, never the history; drain handlers first.
9. **How to deploy changed code?** Never change a decision an old history already recorded: pin running executions to the version that started them and upgrade at continue-as-new (Worker Versioning), or `patched()` for a change that must apply in place.
10. **Where does Temporal end?** Temporal owns where execution logically is; Kubernetes owns where processes run; Kafka owns streams; Postgres owns the results and the money. The boundary is the row your Activity writes, keyed by the idempotency key.
11. **What does the user see, and from which store?** The product's own read model: the `runs` and `run_events` rows the Workflow's projection Activity writes after every state change, served by the API and its event stream. Never the visibility list — it is eventually consistent, knows nothing about orgs, and is Temporal's operator view, not a product API; the Workflow is the source of truth and the projection is how the product reads it.

"""AgentRun, assembled: the one program this course built, as one runnable package.

    agentrun.workflows   AgentRun, ResearchAgent          (Workflow code: decisions)
    agentrun.activities  call_llm, execute_tool, evaluate, compensate   (effects, faked, keyed)
    agentrun.models      Step, ToolCall, ToolResult, Verdict, AgentState
    agentrun.lanes       agent-workflows / cpu-tools / gpu-tools
    agentrun.search_attributes   AgentStatus / CurrentStep / Owner
"""
from .activities import call_llm, compensate, evaluate, execute_tool  # noqa: F401
from .lanes import AGENT_WORKFLOWS, CPU_TOOLS, GPU_TOOLS, workflow_id  # noqa: F401
from .models import AgentState, Step, ToolCall, ToolResult, Verdict  # noqa: F401
from .search_attributes import AGENT_STATUS, CURRENT_STEP, OWNER  # noqa: F401
from .workflows import AgentRun, ResearchAgent  # noqa: F401

"""The effects, faked. Every one is idempotent by key, and every one records what it did in a
module-level ledger so a test (or you) can count real work against attempts.

In production the ledgers are a table keyed by the idempotency key — the rows Postgres owns (module
10's tenth question). Here they live in the Worker process, which is exactly enough to show the
mechanism: a retry with the same key finds the row and does not spend again.
"""
from __future__ import annotations

import asyncio
import os
from collections import defaultdict
from dataclasses import replace

from temporalio import activity

from .models import ToolCall, ToolResult, Verdict

LLM_SECONDS = float(os.environ.get("AGENTRUN_LLM_SECONDS", "0.05"))
TOOL_SECONDS = float(os.environ.get("AGENTRUN_TOOL_SECONDS", "0.2"))
HEARTBEAT_TICKS = 4  # heartbeats per tool call; details carry the tick reached

LLM_CACHE: dict[str, str] = {}              # key → answer
LLM_CALLS: list[str] = []                   # every *billed* call, by key
TOOL_LEDGER: dict[str, ToolResult] = {}     # key → the result of the one time the effect happened
TOOL_SIDE_EFFECTS: dict[str, int] = defaultdict(int)  # key → how many times the effect really ran
COMPENSATED: list[str] = []                 # compensations that ran, in order


def reset_ledgers() -> None:
    LLM_CACHE.clear(); LLM_CALLS.clear(); TOOL_LEDGER.clear(); TOOL_SIDE_EFFECTS.clear(); COMPENSATED.clear()


@activity.defn
async def call_llm(prompt: str, key: str) -> str:
    """One bill per key. A retry that presents the same key gets the cached answer."""
    if key in LLM_CACHE:
        activity.logger.info("call_llm cache hit key=%s attempt=%d", key, activity.info().attempt)
        return LLM_CACHE[key]
    await asyncio.sleep(LLM_SECONDS)
    LLM_CALLS.append(key)
    answer = f"[answer #{len(LLM_CALLS)} to: {prompt[:48]}]"
    LLM_CACHE[key] = answer
    return answer


@activity.defn
async def execute_tool(call: ToolCall) -> ToolResult:
    """Keyed, heartbeating, resumable. The effect is written to the ledger *last*, so an attempt
    that dies half-way leaves no row and the retry redoes the work; an attempt that finished the
    work but whose completion never reached the Service leaves a row and the retry returns it."""
    info = activity.info()
    if call.key in TOOL_LEDGER:
        activity.logger.info("execute_tool duplicate suppressed key=%s attempt=%d", call.key, info.attempt)
        return replace(TOOL_LEDGER[call.key], attempt=info.attempt, duplicate=True)
    start_tick = 0
    if info.heartbeat_details:  # a previous attempt of this same Activity got this far
        start_tick = int(info.heartbeat_details[0].get("tick", 0))
    for tick in range(start_tick, HEARTBEAT_TICKS):
        activity.heartbeat({"tick": tick + 1, "of": HEARTBEAT_TICKS})
        await asyncio.sleep(TOOL_SECONDS / HEARTBEAT_TICKS)
    TOOL_SIDE_EFFECTS[call.key] += 1
    result = ToolResult(call.key, call.tool, f"[{call.tool} output for {call.args}]", info.attempt, info.task_queue)
    TOOL_LEDGER[call.key] = result
    return result


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


@activity.defn
async def compensate(name: str) -> str:
    """Undo one saga step (module 7). Idempotent: undoing twice is the same as once."""
    if name not in COMPENSATED:
        COMPENSATED.append(name)
    return f"compensated {name}"

"""Task Queue lanes (module 10). Spelled once; imported by the Workflow and by every Worker.
Course names are the defaults; LAB_PREFIX=lab11- keeps a shared dev server tidy."""
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")
# `premium-models` and `browser-tools` are the same mechanism again; the lab keeps three lanes.

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:
    return LAB_PREFIX + base

"""The data that crosses the Workflow/Activity boundary. Plain dataclasses: the SDK serialises them
as JSON, so every field is something a history can hold."""
from __future__ import annotations

from dataclasses import dataclass, field


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


@dataclass
class ToolCall:
    key: str                       # idempotency key: wf_id:step:hash(tool, args)
    tool: str
    args: dict


@dataclass
class ToolResult:
    key: str
    tool: str
    output: str
    attempt: int                   # which attempt produced it (>1 means a retry happened)
    task_queue: str                # the lane that answered
    duplicate: bool = False        # True when served from the ledger: a retry after the effect happened


@dataclass
class Verdict:
    ok: bool
    score: float
    note: str


@dataclass
class AgentState:
    """The snapshot continue_as_new carries into the next run. Only what plan() needs to go on;
    never the history."""
    step: int
    goal: str
    context_summary: str
    attempts: int
    processed_command_ids: list[str]

"""Registered once per Namespace (setup_search_attributes.py); written by Workflow code."""
from temporalio.common import SearchAttributeKey

AGENT_STATUS = SearchAttributeKey.for_keyword("AgentStatus")
CURRENT_STEP = SearchAttributeKey.for_int("CurrentStep")
OWNER = SearchAttributeKey.for_keyword("Owner")

"""AgentRun, the full program. Every box of the capstone diagram is a method or a call here.

    plan()                       decision, Workflow code, pure function of state
    call_llm / execute_tool /    Activities on lanes, keyed, timed out, retried, heartbeating
    evaluate
    ResearchAgent                Child Workflow, id = f"{wf_id}/research/{step}"
    pause/resume/cancel_run/     Signals (async writes)
    inject_context
    change_goal                  Update with validator (tracked write); command_id dedupe
    status                       Query (read)
    continue_as_new(AgentState)  when the Service suggests it, or every STEPS_PER_RUN steps
    AgentStatus/CurrentStep      Search Attributes, upserted where state changes
    compensations                the module 7 saga list, run in reverse on cancel
"""
from __future__ import annotations

import asyncio
import hashlib
import os
from datetime import timedelta

from temporalio import workflow
from temporalio.common import Priority, RetryPolicy
from temporalio.exceptions import ActivityError, ApplicationError, TimeoutError, TimeoutType, is_cancelled_exception

from .activities import call_llm, compensate, evaluate, execute_tool
from .lanes import AGENT_WORKFLOWS, CPU_TOOLS, GPU_TOOL_NAMES, LANE_FOR_TOOL
from .models import AgentState, Step, ToolCall, ToolResult, Verdict
from .search_attributes import AGENT_STATUS, CURRENT_STEP, OWNER

# Knobs. Course values in comments; the lab shrinks them so a 200-step run takes seconds.
STEPS_PER_RUN = int(os.environ.get("AGENTRUN_STEPS_PER_RUN", "50"))      # continue-as-new cadence
RESEARCH_EVERY = int(os.environ.get("AGENTRUN_RESEARCH_EVERY", "4"))     # a ResearchAgent child every N cycles
DEFAULT_MAX_STEPS = 8
TOOL_START_TO_CLOSE = timedelta(seconds=float(os.environ.get("AGENTRUN_TOOL_START_TO_CLOSE_SECONDS", "60")))  # course: 30 min
TOOL_HEARTBEAT = timedelta(seconds=float(os.environ.get("AGENTRUN_TOOL_HEARTBEAT_SECONDS", "5")))            # course: 30 s
GPU_SCHEDULE_TO_START = timedelta(seconds=float(os.environ.get("AGENTRUN_GPU_SCHEDULE_TO_START_SECONDS", "120")))  # course: 5 min

LLM_RETRY = RetryPolicy(initial_interval=timedelta(milliseconds=200), backoff_coefficient=2.0,
                        maximum_interval=timedelta(seconds=5), maximum_attempts=5,
                        non_retryable_error_types=["BadArguments"])
TOOL_RETRY = RetryPolicy(initial_interval=timedelta(milliseconds=200), backoff_coefficient=2.0,
                         maximum_interval=timedelta(seconds=5), maximum_attempts=5,
                         non_retryable_error_types=["BadArguments"])
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


def _digest(*parts: str) -> str:
    return hashlib.sha256(":".join(parts).encode()).hexdigest()[:12]


@workflow.defn
class ResearchAgent:
    """The sub-agent (module 6). Its own history, its own replay; the parent sees one result."""

    @workflow.run
    async def run(self, question: str) -> str:
        key = f"{workflow.info().workflow_id}:0:{_digest(question)}"
        return await workflow.execute_activity(
            call_llm, args=[f"research: {question}", key], task_queue=CPU_TOOLS,
            schedule_to_close_timeout=timedelta(minutes=2), start_to_close_timeout=timedelta(seconds=30),
            retry_policy=LLM_RETRY)


@workflow.defn
class AgentRun:
    def __init__(self) -> None:
        self.step = 0
        self.status = "starting"
        self.goal = ""
        self.context_summary = ""
        self.attempts = 0                  # retries that happened (sum of attempt-1 over results)
        self.paused = False
        self.cancel_requested = False
        self.processed_command_ids: set[str] = set()
        self.compensations: list[str] = []
        self.owner = ""
        self.tier = "batch"
        self.max_steps = DEFAULT_MAX_STEPS
        self.run_started_at_step = 0
        self.duplicates_suppressed = 0
        self.children: list[str] = []

    # ------------------------------------------------------------------ decisions
    def plan(self) -> Step:
        """Pure function of state. Never calls a model (module 3)."""
        if self.step >= self.max_steps:
            return Step("finish")
        phase = self.step % 4
        if phase == 0:
            return Step("llm", prompt=f"step {self.step}: plan the next move toward: {self.goal}")
        if phase == 1:
            return Step("tool", tool="search", args={"q": self.goal, "step": self.step})
        if phase == 2:
            return Step("tool", tool="embed", args={"text": self.context_summary[:40], "step": self.step})
        if (self.step // 4) % RESEARCH_EVERY == 0:
            return Step("research", question=f"what is still missing for: {self.goal}")
        return Step("llm", prompt=f"step {self.step}: reflect on: {self.context_summary[:40]}")

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

    def _publish(self, status: str) -> None:
        self.status = status
        workflow.upsert_search_attributes([AGENT_STATUS.value_set(status), CURRENT_STEP.value_set(self.step)])

    def _should_continue_as_new(self) -> bool:
        if self.step == self.run_started_at_step:
            return False  # never on the first step of a run: that would loop forever
        return workflow.info().is_continue_as_new_suggested() or (self.step - self.run_started_at_step) >= STEPS_PER_RUN

    # ------------------------------------------------------------------ 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.signal
    def cancel_run(self) -> None:
        """Cooperative cancel: the loop notices at its next decision, compensates, returns cleanly.
        (handle.cancel() is the other kind — Temporal's — handled in run() below.)"""
        self.cancel_requested = True
        self.paused = False

    @workflow.signal
    def inject_context(self, text: str, command_id: str) -> None:
        if command_id in self.processed_command_ids:
            return  # a retried Signal is delivered again; the command id makes the second delivery a no-op
        self.processed_command_ids.add(command_id)
        self.context_summary = f"{self.context_summary}\n[injected: {text}]"

    @workflow.query
    def status(self) -> dict:
        return {"step": self.step, "status": self.status, "goal": self.goal, "paused": self.paused,
                "attempts": self.attempts, "duplicates_suppressed": self.duplicates_suppressed,
                "compensations": list(self.compensations), "children": list(self.children),
                "processed_command_ids": sorted(self.processed_command_ids), "max_steps": self.max_steps,
                "run_id": workflow.info().run_id, "run_started_at_step": self.run_started_at_step,
                "history_length": workflow.info().get_current_history_length()}

    @workflow.update
    def change_goal(self, goal: str, command_id: str) -> str:
        if command_id in self.processed_command_ids:
            return f"ignored: command {command_id} already applied"
        self.processed_command_ids.add(command_id)
        old, self.goal = self.goal, goal
        self.context_summary = f"{self.context_summary}\n[goal changed: {goal}]"
        return f"goal changed from {old!r} to {goal!r}"

    @change_goal.validator
    def _validate_change_goal(self, goal: str, command_id: str) -> None:
        # Rejected here = never in history. Accepted = a tracked write the client sees complete.
        if not goal.strip():
            raise ValueError("goal must not be empty")
        if len(goal) > 500:
            raise ValueError("goal longer than 500 characters")
        if not command_id:
            raise ValueError("command_id is required")

    # ------------------------------------------------------------------ effects
    async def _llm(self, step: Step) -> str:
        key = f"{workflow.info().workflow_id}:{self.step}:{_digest(step.prompt)}"
        return await workflow.execute_activity(
            call_llm, args=[step.prompt, key], task_queue=CPU_TOOLS,  # production: premium-models
            schedule_to_close_timeout=timedelta(minutes=2), start_to_close_timeout=timedelta(seconds=30),
            retry_policy=LLM_RETRY)

    async def _tool(self, step: Step) -> ToolResult:
        call = ToolCall(f"{workflow.info().workflow_id}:{self.step}:{_digest(step.tool, repr(sorted(step.args.items())))}",
                        step.tool, step.args)
        gpu = step.tool in GPU_TOOL_NAMES
        if gpu:
            # Module 7's saga in one line: what must be undone if this run is cancelled holding a GPU.
            self.compensations.append(f"release_gpu:{call.key}")
        result: ToolResult = await workflow.execute_activity(
            execute_tool, call, task_queue=LANE_FOR_TOOL[step.tool],
            schedule_to_start_timeout=GPU_SCHEDULE_TO_START if gpu else timedelta(minutes=2),
            start_to_close_timeout=TOOL_START_TO_CLOSE, heartbeat_timeout=TOOL_HEARTBEAT,
            retry_policy=TOOL_RETRY,
            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())
        if gpu:
            self.compensations.remove(f"release_gpu:{call.key}")
        self.attempts += result.attempt - 1
        if result.duplicate:
            self.duplicates_suppressed += 1
        verdict: Verdict = await workflow.execute_activity(
            evaluate, result, task_queue=CPU_TOOLS, start_to_close_timeout=timedelta(seconds=30), retry_policy=TOOL_RETRY)
        self.context_summary = f"{result.output} ({verdict.note})"
        return result

    async def _research(self, step: Step) -> str:
        child_id = f"{workflow.info().workflow_id}/research/{self.step}"
        self.children.append(child_id)
        answer = await workflow.execute_child_workflow(
            ResearchAgent.run, step.question, id=child_id, task_queue=AGENT_WORKFLOWS,
            execution_timeout=timedelta(minutes=10))
        self.context_summary = f"{self.context_summary}\n[research: {answer}]"
        return answer

    async def _compensate(self) -> None:
        for name in reversed(self.compensations):
            await workflow.execute_activity(compensate, name, task_queue=CPU_TOOLS,
                                            start_to_close_timeout=timedelta(seconds=30), retry_policy=TOOL_RETRY)
        self.compensations.clear()

    async def _execute(self, step: Step) -> None:
        try:
            if step.kind == "llm":
                self.context_summary = await self._llm(step)
            elif step.kind == "tool":
                await self._tool(step)
            elif step.kind == "research":
                await self._research(step)
        except ActivityError as err:
            if isinstance(err.cause, TimeoutError) and err.cause.type == TimeoutType.SCHEDULE_TO_START:
                self._publish("stalled")
                raise ApplicationError(f"lane stalled: nobody took {step.tool} in time", type="LaneStalled",
                                       non_retryable=True) from err
            raise

    # ------------------------------------------------------------------ the loop
    @workflow.run
    async def run(self, goal: str, state: AgentState | None = None) -> dict:
        info = workflow.info()
        self.owner = info.typed_search_attributes.get(OWNER) or ""
        self.tier = workflow.memo_value("tier", "batch", type_hint=str)
        self.max_steps = workflow.memo_value("max_steps", DEFAULT_MAX_STEPS, type_hint=int)
        if state is None:
            self.goal = goal
        else:  # a continued run: the snapshot is the whole past
            self.step, self.goal, self.context_summary, self.attempts = state.step, state.goal, state.context_summary, state.attempts
            self.processed_command_ids = set(state.processed_command_ids)
        self.run_started_at_step = self.step
        self._publish("running")
        try:
            while True:
                await workflow.wait_condition(lambda: not self.paused or self.cancel_requested)
                if self.cancel_requested:
                    await self._compensate()
                    self._publish("cancelled")
                    return self._summary()
                step = self.plan()
                if step.kind == "finish":
                    self._publish("done")
                    return self._summary()
                if self._should_continue_as_new():   # more work to do, and this run has done its share
                    await workflow.wait_condition(workflow.all_handlers_finished)  # never drop an Update mid-flight
                    self._publish("continuing")
                    workflow.continue_as_new(args=[self.goal, self.snapshot()])
                await self._execute(step)
                self.step += 1
                self._publish("paused" if self.paused else "running")
        except BaseException as err:
            # Temporal cancellation (handle.cancel()) arrives in one of two shapes: asyncio.CancelledError if we were
            # waiting on a timer/condition, or an ActivityError/ChildWorkflowError whose cause is Temporal's
            # CancelledError if we were waiting on an Activity or child (it was told to cancel first).
            # `is_cancelled_exception` recognises both. Cleanup Activities still run here: the SDK cleared the task's
            # cancellation when it sent the cancel command, so the Workflow is not yet complete.
            if not is_cancelled_exception(err):
                raise
            await self._compensate()
            self._publish("cancelled")
            raise

    def _summary(self) -> dict:
        return {"status": self.status, "steps": self.step, "attempts": self.attempts,
                "duplicates_suppressed": self.duplicates_suppressed, "children_this_run": list(self.children),
                "context_tail": self.context_summary[-160:]}

"""The two Activities. `run_case` is the fake GPU (gpu-eval lane); `record_run_event` is the projection
(console lane) — the only code that writes SQLite from inside a Workflow's life.

Both are idempotent by key: `run_case` by `run_id:case_id` (a ledger in the Worker process, the rows
Postgres would own in production), `record_run_event` by `(run_id, seq)` (the SQLite primary key).
"""
from __future__ import annotations

import asyncio
import random
from dataclasses import replace
from typing import Any

from temporalio import activity
from temporalio.exceptions import ApplicationError

import projection
from config import CASE_SECONDS, CASE_TICKS, DB_PATH, HICCUP_FRACTION
from models import CaseResult

CASE_LEDGER: dict[str, CaseResult] = {}      # key → the one time the case really ran
CASE_EFFECTS: dict[str, int] = {}            # key → how many times the GPU work really happened
CANCELLED_AT_TICK: dict[str, int] = {}       # key → tick at which a cancel was observed


def reset_ledgers() -> None:
    CASE_LEDGER.clear(); CASE_EFFECTS.clear(); CANCELLED_AT_TICK.clear()


@activity.defn
async def run_case(run_id: str, case_id: str, fail_fraction: float = 0.25) -> CaseResult:
    """Evaluate one benchmark case on the (fake) GPU.

    - sleeps CASE_SECONDS in CASE_TICKS heartbeating ticks; a cancel is observed at the next heartbeat
    - a *benchmark failure* (fraction `fail_fraction`, seeded by run_id:case_id so a run is reproducible and a
      retry run — new run_id — rolls again) is `ApplicationError(type="CaseFailed", non_retryable=True)`:
      the Workflow records it, the RetryPolicy does not fight it
    - a *transient GPU hiccup* (HICCUP_FRACTION, seeded per attempt) is a plain retryable error: the
      RetryPolicy retries it and the history shows the attempt count, never a failed case
    - the ledger makes a retry after a lost completion return the recorded result (`duplicate=True`)
    """
    info = activity.info()
    key = f"{run_id}:{case_id}"
    if key in CASE_LEDGER:
        activity.logger.info("run_case duplicate suppressed key=%s attempt=%d", key, info.attempt)
        return replace(CASE_LEDGER[key], attempt=info.attempt, duplicate=True)
    rng = random.Random(f"{key}:attempt{info.attempt}")
    if rng.random() < HICCUP_FRACTION:
        raise ApplicationError(f"GPU hiccup on {key} attempt {info.attempt}", type="GpuHiccup")
    start_tick = int(info.heartbeat_details[0].get("tick", 0)) if info.heartbeat_details else 0
    try:
        for tick in range(start_tick, CASE_TICKS):
            activity.heartbeat({"tick": tick + 1, "of": CASE_TICKS})
            await asyncio.sleep(CASE_SECONDS / CASE_TICKS)
    except asyncio.CancelledError:
        CANCELLED_AT_TICK[key] = tick + 1
        activity.logger.info("run_case cancelled key=%s at tick %d/%d", key, tick + 1, CASE_TICKS)
        raise
    CASE_EFFECTS[key] = CASE_EFFECTS.get(key, 0) + 1
    ok = random.Random(key).random() >= fail_fraction
    result = CaseResult(run_id, case_id, ok, round(random.Random(key + ":score").uniform(0.5, 1.0), 3) if ok else 0.0, info.attempt)
    CASE_LEDGER[key] = result
    if not ok:
        raise ApplicationError(f"case {case_id} failed the benchmark", type="CaseFailed", non_retryable=True)
    return result


@activity.defn
async def record_run_event(run_id: str, seq: int, kind: str, payload: dict[str, Any]) -> bool:
    """Project one Workflow state change into SQLite. Returns True if the event was new, False if this
    (run_id, seq) had already been recorded — a retried execution changes nothing."""
    conn = projection.connect(DB_PATH)
    try:
        return projection.apply_event(conn, run_id, int(seq), kind, payload)
    finally:
        conn.close()

"""The evaluation console's HTTP surface. Reads come from SQLite (the projection); writes go to
Temporal (start, cancel, Update). Never Temporal visibility for a list — the projection is the
product's read model, visibility is the operator's.

    POST /runs                      start an EvalRun (202); Workflow ID == run_id
    GET  /runs?org=&status=&model=&cursor=&limit=     keyset pagination on (created_at, run_id)
    GET  /runs/{id}                 row + last events
    GET  /runs/{id}/events?after=   SSE, resumes from the cursor; polls SQLite every 0.5 s
    POST /runs/{id}/cancel          handle.cancel() + an immediate `cancel_requested` projection event
    POST /runs/{id}/retry-failed    new EvalRun over the failed case ids; id `{id}-retry-{n}`; memo parent_run_id
    POST /runs/{id}/approve         execute_update(EvalRun.approve)

Every request carries `X-Org`. A run that belongs to another org does not exist (404).
"""
from __future__ import annotations

import asyncio
import uuid
from contextlib import asynccontextmanager
from typing import Any

from fastapi import Depends, FastAPI, Header, HTTPException, Query, Request
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field
from temporalio.api.operatorservice.v1 import ListSearchAttributesRequest
from temporalio.client import Client, WorkflowUpdateFailedError
from temporalio.common import SearchAttributePair, TypedSearchAttributes
from temporalio.exceptions import ApplicationError
from temporalio.service import RPCError, RPCStatusCode

import projection
from config import (ALL_STATUSES, CONSOLE_QUEUE, DB_PATH, DEFAULT_FAIL_FRACTION, DEFAULT_REVIEWERS, NAMESPACE,
                    TEMPORAL_ADDRESS, TERMINAL_STATUSES, workflow_id)
from models import EvalRunInput
from workflows import MODEL, ORG, PARENT_RUN_ID, RUN_STATUS, EvalRun

SSE_POLL_SECONDS = 0.5


class StartRun(BaseModel):
    run_id: str | None = None
    model: str = "fake-model-7b"
    suite: str = "smoke"
    case_ids: list[str] | None = None
    cases: int = Field(12, ge=1, le=500)
    fail_fraction: float = Field(DEFAULT_FAIL_FRACTION, ge=0.0, le=1.0)
    reviewers: list[str] = Field(default_factory=lambda: list(DEFAULT_REVIEWERS))


class Approve(BaseModel):
    reviewer: str
    command_id: str


# ------------------------------------------------------------------ pure helpers (tested directly)
def retry_failed_input(parent: dict[str, Any], state: dict[str, Any], n: int, *, reviewers: list[str] | None = None,
                       fail_fraction: float = DEFAULT_FAIL_FRACTION) -> EvalRunInput:
    """The retry run: only the failed case ids, id `{parent}-retry-{n}`, same org/model/suite."""
    failed = list(state.get("failed_case_ids") or [])
    if not failed:
        raise ValueError("nothing to retry: the run has no failed cases")
    return EvalRunInput(run_id=f"{parent['run_id']}-retry-{n}", model=parent["model"], suite=parent["suite"],
                        case_ids=failed, org=parent["org"], fail_fraction=fail_fraction,
                        reviewers=reviewers or list(DEFAULT_REVIEWERS))


def start_search_attributes(inp: EvalRunInput, parent_run_id: str | None, parent_sa_registered: bool) -> TypedSearchAttributes:
    pairs = [SearchAttributePair(ORG, inp.org), SearchAttributePair(MODEL, inp.model), SearchAttributePair(RUN_STATUS, "queued")]
    if parent_run_id and parent_sa_registered:
        pairs.append(SearchAttributePair(PARENT_RUN_ID, parent_run_id))
    return TypedSearchAttributes(pairs)


# ------------------------------------------------------------------ app
@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.client = await Client.connect(TEMPORAL_ADDRESS, namespace=NAMESPACE)
    listed = await app.state.client.operator_service.list_search_attributes(ListSearchAttributesRequest(namespace=NAMESPACE))
    app.state.parent_sa_registered = "EvalParentRunId" in listed.custom_attributes
    app.state.db = projection.connect(DB_PATH)
    print(f"console api: temporal={TEMPORAL_ADDRESS} db={DB_PATH} EvalParentRunId registered={app.state.parent_sa_registered}")
    yield
    app.state.db.close()


app = FastAPI(title="EvalRun console", lifespan=lifespan)


def org_of(x_org: str | None = Header(default=None, alias="X-Org")) -> str:
    if not x_org:
        raise HTTPException(400, "X-Org header is required")
    return x_org


def owned_run(request: Request, run_id: str, org: str) -> dict[str, Any]:
    row = projection.get_run(request.app.state.db, run_id)
    if row is None or row["org"] != org:
        raise HTTPException(404, f"run {run_id!r} not found")   # another org's run is indistinguishable from none
    return row


async def _start(request: Request, inp: EvalRunInput, parent_run_id: str | None) -> dict[str, Any]:
    client: Client = request.app.state.client
    try:
        handle = await client.start_workflow(
            EvalRun.run, inp, id=workflow_id(inp.run_id), task_queue=CONSOLE_QUEUE,
            search_attributes=start_search_attributes(inp, parent_run_id, request.app.state.parent_sa_registered),
            memo={"parent_run_id": parent_run_id or ""})
    except Exception as err:  # WorkflowAlreadyStartedError is the one you will see
        if type(err).__name__ == "WorkflowAlreadyStartedError":
            raise HTTPException(409, f"run {inp.run_id!r} already exists")
        raise
    projection.insert_queued_run(request.app.state.db, run_id=inp.run_id, org=inp.org, model=inp.model,
                                 suite=inp.suite, total=len(inp.case_ids), parent_run_id=parent_run_id)
    return {"run_id": inp.run_id, "workflow_id": handle.id, "run_id_temporal": handle.result_run_id,
            "status": "queued", "total": len(inp.case_ids), "parent_run_id": parent_run_id}


@app.post("/runs", status_code=202)
async def create_run(body: StartRun, request: Request, org: str = Depends(org_of)):
    run_id = workflow_id(body.run_id or f"run-{uuid.uuid4().hex[:8]}")
    case_ids = body.case_ids or [f"case-{i:03d}" for i in range(1, body.cases + 1)]
    inp = EvalRunInput(run_id=run_id, model=body.model, suite=body.suite, case_ids=case_ids, org=org,
                       fail_fraction=body.fail_fraction, reviewers=body.reviewers)
    return JSONResponse(await _start(request, inp, None), status_code=202)


@app.get("/runs")
def list_runs(request: Request, org: str = Depends(org_of), status: str | None = None, model: str | None = None,
              cursor: str | None = None, limit: int = Query(20, ge=1, le=200)):
    if status and status not in ALL_STATUSES:
        raise HTTPException(400, f"status must be one of {ALL_STATUSES}")
    try:
        page, next_cursor = projection.list_runs(request.app.state.db, org=org, status=status, model=model, cursor=cursor, limit=limit)
    except ValueError as err:
        raise HTTPException(400, str(err))
    return {"runs": page, "next_cursor": next_cursor}


@app.get("/runs/{run_id}")
def get_run(run_id: str, request: Request, org: str = Depends(org_of)):
    row = owned_run(request, run_id, org)
    events = projection.last_events(request.app.state.db, run_id, 10)
    return {**row, "state": projection.latest_state(request.app.state.db, run_id),
            "last_events": [e.as_dict() for e in events], "last_seq": events[-1].seq if events else 0}


@app.get("/runs/{run_id}/events")
async def run_events(run_id: str, request: Request, org: str = Depends(org_of), after: int = Query(0, ge=0),
                     last_event_id: str | None = Header(default=None, alias="Last-Event-ID")):
    """SSE. The cursor is `?after=<seq>` or the browser's automatic `Last-Event-ID`; whichever is larger wins.
    Frames carry `id: <seq>` so nothing with seq <= cursor is ever sent, across any number of reconnects."""
    owned_run(request, run_id, org)
    cursor = max(after, int(last_event_id) if last_event_id and last_event_id.isdigit() else 0)
    conn = projection.connect(DB_PATH)   # this stream's own connection; the app one is shared by other requests

    async def gen():
        try:
            yield f": resuming after seq {cursor}\n\n"
            pos = cursor
            while not await request.is_disconnected():
                frames, pos, terminal = projection.poll_frames(conn, run_id, pos, TERMINAL_STATUSES)
                for frame in frames:
                    yield frame
                if terminal:
                    return
                if not frames:
                    await asyncio.sleep(SSE_POLL_SECONDS)
        finally:
            conn.close()

    return StreamingResponse(gen(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


@app.post("/runs/{run_id}/cancel", status_code=202)
async def cancel_run(run_id: str, request: Request, org: str = Depends(org_of)):
    """handle.cancel() is a request: the Workflow decides when it is cancelled (cases stop within one heartbeat).
    The console does not wait for that — it writes `cancel_requested` now, so the UI shows `cancelling` at once."""
    row = owned_run(request, run_id, org)
    if row["status"] in TERMINAL_STATUSES:
        raise HTTPException(409, f"run is already {row['status']}")
    client: Client = request.app.state.client
    try:
        await client.get_workflow_handle(workflow_id(run_id)).cancel()
    except RPCError as err:
        if err.status == RPCStatusCode.NOT_FOUND:
            raise HTTPException(404, f"run {run_id!r} not found")
        raise
    ev = projection.console_event(request.app.state.db, run_id, "cancel_requested", {"by": org}, status="cancelling")
    return {"run_id": run_id, "status": "cancelling", "event": ev.as_dict()}


@app.post("/runs/{run_id}/retry-failed", status_code=202)
async def retry_failed(run_id: str, request: Request, org: str = Depends(org_of)):
    db = request.app.state.db
    row = owned_run(request, run_id, org)
    if row["status"] not in TERMINAL_STATUSES | {"awaiting_review"}:
        raise HTTPException(409, f"run is {row['status']}; retry once it has stopped")
    state = projection.latest_state(db, run_id)
    n = projection.count_retries(db, run_id) + 1
    try:
        inp = retry_failed_input(row, state, n, fail_fraction=DEFAULT_FAIL_FRACTION)
    except ValueError as err:
        raise HTTPException(409, str(err))
    started = await _start(request, inp, parent_run_id=run_id)
    projection.console_event(db, run_id, "retry_requested", {"retry_run_id": inp.run_id, "case_ids": inp.case_ids})
    return {**started, "case_ids": inp.case_ids}


@app.post("/runs/{run_id}/approve")
async def approve_run(run_id: str, body: Approve, request: Request, org: str = Depends(org_of)):
    row = owned_run(request, run_id, org)
    if row["status"] in TERMINAL_STATUSES:       # a replayed click after publish: the run is closed, not missing
        raise HTTPException(409, f"run is already {row['status']}")
    client: Client = request.app.state.client
    try:
        result = await client.get_workflow_handle(workflow_id(run_id)).execute_update(
            EvalRun.approve, args=[body.reviewer, body.command_id])
    except WorkflowUpdateFailedError as err:
        cause = err.cause
        msg = cause.message if isinstance(cause, ApplicationError) else str(cause)
        raise HTTPException(422, f"approve rejected: {msg}")
    except RPCError as err:
        if err.status == RPCStatusCode.NOT_FOUND:   # the row exists, so the Workflow closed under us
            raise HTTPException(409, f"run {run_id!r} is closed")
        raise
    return {"run_id": run_id, "result": result}


@app.get("/healthz")
def healthz():
    return {"ok": True}

"""Names and knobs for the evaluation console, spelled once. Every process (API, both Workers, the
setup script, the tests) imports these; nothing else reads the environment.

Shared dev server: LAB_PREFIX (default `lab115-`) prefixes both task queues and every Workflow ID."""
from __future__ import annotations

import os

TEMPORAL_ADDRESS = os.environ.get("TEMPORAL_ADDRESS", "localhost:7233")
NAMESPACE = os.environ.get("TEMPORAL_NAMESPACE", "default")

LAB_PREFIX = os.environ.get("LAB_PREFIX", "lab115-")
CONSOLE_QUEUE = os.environ.get("CONSOLE_QUEUE", LAB_PREFIX + "console")      # EvalRun + record_run_event
GPU_EVAL_QUEUE = os.environ.get("GPU_EVAL_QUEUE", LAB_PREFIX + "gpu-eval")   # run_case

DB_PATH = os.environ.get("EVAL_CONSOLE_DB", os.path.join(os.path.dirname(os.path.abspath(__file__)), "eval_console.sqlite3"))
API_HOST = os.environ.get("EVAL_CONSOLE_HOST", "127.0.0.1")
API_PORT = int(os.environ.get("EVAL_CONSOLE_PORT", "8115"))

# The fake GPU: one case sleeps CASE_SECONDS in CASE_TICKS heartbeating ticks.
CASE_SECONDS = float(os.environ.get("EVAL_CASE_SECONDS", "0.8"))
CASE_TICKS = int(os.environ.get("EVAL_CASE_TICKS", "4"))
DEFAULT_FAIL_FRACTION = float(os.environ.get("EVAL_FAIL_FRACTION", "0.25"))     # benchmark failures (recorded, not retried)
HICCUP_FRACTION = float(os.environ.get("EVAL_HICCUP_FRACTION", "0.1"))           # transient GPU errors (retried by policy)
CASE_PARALLELISM = int(os.environ.get("EVAL_CASE_PARALLELISM", "4"))             # cases in flight per run
GPU_SLOTS = int(os.environ.get("GPU_SLOTS", "4"))                                # slots on the gpu-eval Worker
DEFAULT_REVIEWERS = tuple(x for x in os.environ.get("EVAL_REVIEWERS", "alice,bob").split(",") if x)

TERMINAL_STATUSES = frozenset({"published", "cancelled", "failed"})
ALL_STATUSES = ("queued", "running", "awaiting_review", "cancelling", "cancelled", "failed", "published")


def workflow_id(run_id: str) -> str:
    """Workflow ID == run_id; the prefix is applied once, never twice."""
    return run_id if run_id.startswith(LAB_PREFIX) else LAB_PREFIX + run_id

"""Dataclasses crossing the Workflow boundary (Temporal's default JSON converter handles them)."""
from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class EvalRunInput:
    run_id: str
    model: str
    suite: str
    case_ids: list[str]
    org: str
    fail_fraction: float = 0.25
    reviewers: list[str] = field(default_factory=lambda: ["alice", "bob"])


@dataclass
class CaseResult:
    run_id: str
    case_id: str
    ok: bool
    score: float
    attempt: int
    duplicate: bool = False


@dataclass
class EvalRunResult:
    run_id: str
    status: str
    done: int
    total: int
    failed_case_ids: list[str]
    approved_by: str | None = None

"""The read model. SQLite owns what the UI reads; Temporal owns what is true.

    runs(run_id, org, model, suite, status, done, total, created_at, parent_run_id)
    run_events(run_id, seq, kind, payload, ts, PRIMARY KEY (run_id, seq))

`apply_event` is the one write path: it is *idempotent by (run_id, seq)* — the projection Activity
that calls it may run twice (a retry after a lost completion), and the second execution must change
nothing. The `runs` row is rebuilt from the event's `state` payload, but only when that event is the
newest one for the run, so an out-of-order duplicate can never move a run backwards.

Sequence numbers: the Workflow owns the EVEN numbers (2, 4, 6, …: `2 * self.seq`). The console
itself writes ODD numbers (`cancel_requested`, `retry_requested`) in the gap after the newest event,
so a console-originated event can never collide with — or be overwritten by — a Workflow event.
"""
from __future__ import annotations

import json
import sqlite3
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Iterator

SCHEMA = """
CREATE TABLE IF NOT EXISTS runs (
    run_id        TEXT PRIMARY KEY,
    org           TEXT NOT NULL,
    model         TEXT NOT NULL,
    suite         TEXT NOT NULL,
    status        TEXT NOT NULL,
    done          INTEGER NOT NULL DEFAULT 0,
    total         INTEGER NOT NULL DEFAULT 0,
    created_at    TEXT NOT NULL,
    parent_run_id TEXT
);
CREATE INDEX IF NOT EXISTS runs_org_created ON runs(org, created_at DESC, run_id DESC);
CREATE TABLE IF NOT EXISTS run_events (
    run_id  TEXT NOT NULL,
    seq     INTEGER NOT NULL,
    kind    TEXT NOT NULL,
    payload TEXT NOT NULL,
    ts      TEXT NOT NULL,
    PRIMARY KEY (run_id, seq)
);
"""


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds")


def connect(path: str) -> sqlite3.Connection:
    conn = sqlite3.connect(path, check_same_thread=False, isolation_level=None, timeout=5.0)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=5000")
    conn.executescript(SCHEMA)
    return conn


@dataclass
class Event:
    run_id: str
    seq: int
    kind: str
    payload: dict[str, Any]
    ts: str

    def as_dict(self) -> dict[str, Any]:
        return {"run_id": self.run_id, "seq": self.seq, "kind": self.kind, "payload": self.payload, "ts": self.ts}


def _row_event(r: sqlite3.Row) -> Event:
    return Event(r["run_id"], r["seq"], r["kind"], json.loads(r["payload"]), r["ts"])


# ------------------------------------------------------------------ writes
def upsert_event(conn: sqlite3.Connection, run_id: str, seq: int, kind: str, payload: dict[str, Any], ts: str) -> bool:
    """INSERT the event; a second call with the same (run_id, seq) is a no-op. Returns True if inserted."""
    cur = conn.execute("INSERT OR IGNORE INTO run_events(run_id, seq, kind, payload, ts) VALUES (?,?,?,?,?)",
                       (run_id, seq, kind, json.dumps(payload, sort_keys=True), ts))
    return cur.rowcount == 1


def upsert_run(conn: sqlite3.Connection, *, run_id: str, org: str, model: str, suite: str, status: str,
               done: int, total: int, created_at: str, parent_run_id: str | None) -> None:
    conn.execute("""INSERT INTO runs(run_id, org, model, suite, status, done, total, created_at, parent_run_id)
                    VALUES (?,?,?,?,?,?,?,?,?)
                    ON CONFLICT(run_id) DO UPDATE SET
                      status=excluded.status, done=excluded.done, total=excluded.total""",
                 (run_id, org, model, suite, status, done, total, created_at, parent_run_id or None))


def insert_queued_run(conn: sqlite3.Connection, *, run_id: str, org: str, model: str, suite: str, total: int,
                      parent_run_id: str | None) -> None:
    """The row the API writes when it starts a run, so GET /runs/{id} and its event stream answer at once
    instead of 404ing until the Workflow's first projection lands. Insert-or-ignore: if that projection
    already landed, it is newer, and this must not move the run back to `queued`."""
    with conn:
        conn.execute("""INSERT OR IGNORE INTO runs(run_id, org, model, suite, status, done, total, created_at, parent_run_id)
                        VALUES (?,?,?,?,'queued',0,?,?,?)""",
                     (run_id, org, model, suite, total, now_iso(), parent_run_id or None))


def apply_event(conn: sqlite3.Connection, run_id: str, seq: int, kind: str, payload: dict[str, Any],
                ts: str | None = None) -> bool:
    """The projection: record the event, then rebuild the run row from its `state` iff it is the newest.
    Idempotent: replaying (run_id, seq) inserts nothing and moves nothing. Returns True if the event was new."""
    ts = ts or now_iso()
    with conn:  # one transaction
        inserted = upsert_event(conn, run_id, seq, kind, payload, ts)
        newest = conn.execute("SELECT MAX(seq) FROM run_events WHERE run_id=?", (run_id,)).fetchone()[0]
        state = payload.get("state")
        if inserted and state and seq == newest:
            row = conn.execute("SELECT * FROM runs WHERE run_id=?", (run_id,)).fetchone()
            upsert_run(conn,
                       run_id=run_id,
                       org=state.get("org", row["org"] if row else ""),
                       model=state.get("model", row["model"] if row else ""),
                       suite=state.get("suite", row["suite"] if row else ""),
                       status=state["status"], done=int(state.get("done", 0)), total=int(state.get("total", 0)),
                       created_at=row["created_at"] if row else state.get("created_at", ts),
                       parent_run_id=row["parent_run_id"] if row else state.get("parent_run_id"))
    return inserted


def console_event(conn: sqlite3.Connection, run_id: str, kind: str, payload: dict[str, Any],
                  status: str | None = None) -> Event:
    """A console-originated event (`cancel_requested`, `retry_requested`): takes the next ODD seq after
    the newest event, so it sorts after what the Workflow has said so far and can never collide with what
    it will say next (Workflow seqs are even). Optionally moves the run row's status (e.g. `cancelling`)."""
    ts = now_iso()
    with conn:
        newest = conn.execute("SELECT COALESCE(MAX(seq), 0) FROM run_events WHERE run_id=?", (run_id,)).fetchone()[0]
        seq = newest + 1 if newest % 2 == 0 else newest + 2
        payload = dict(payload, source="console")
        if status:
            if get_run(conn, run_id) is None:
                raise KeyError(run_id)
            payload["state"] = {**_state_of(conn, run_id), "status": status}
            conn.execute("UPDATE runs SET status=? WHERE run_id=?", (status, run_id))
        upsert_event(conn, run_id, seq, kind, payload, ts)
    return Event(run_id, seq, kind, payload, ts)


def _state_of(conn: sqlite3.Connection, run_id: str) -> dict[str, Any]:
    """The newest `state` payload the Workflow projected, or {}."""
    r = conn.execute("SELECT payload FROM run_events WHERE run_id=? AND payload LIKE '%\"state\"%' ORDER BY seq DESC LIMIT 1",
                     (run_id,)).fetchone()
    return json.loads(r["payload"]).get("state", {}) if r else {}


def latest_state(conn: sqlite3.Connection, run_id: str) -> dict[str, Any]:
    return _state_of(conn, run_id)


# ------------------------------------------------------------------ reads
def get_run(conn: sqlite3.Connection, run_id: str) -> dict[str, Any] | None:
    r = conn.execute("SELECT * FROM runs WHERE run_id=?", (run_id,)).fetchone()
    return dict(r) if r else None


def encode_cursor(created_at: str, run_id: str) -> str:
    return f"{created_at}|{run_id}"


def decode_cursor(cursor: str) -> tuple[str, str]:
    created_at, _, run_id = cursor.partition("|")
    if not created_at or not run_id:
        raise ValueError("malformed cursor")
    return created_at, run_id


def list_runs(conn: sqlite3.Connection, *, org: str, status: str | None = None, model: str | None = None,
              cursor: str | None = None, limit: int = 20) -> tuple[list[dict[str, Any]], str | None]:
    """Newest first; keyset pagination on (created_at, run_id). Every query is scoped to one org."""
    limit = max(1, min(int(limit), 200))
    where, args = ["org=?"], [org]
    if status:
        where.append("status=?"); args.append(status)
    if model:
        where.append("model=?"); args.append(model)
    if cursor:
        c_created, c_id = decode_cursor(cursor)
        where.append("(created_at, run_id) < (?, ?)"); args += [c_created, c_id]
    rows = conn.execute(f"SELECT * FROM runs WHERE {' AND '.join(where)} ORDER BY created_at DESC, run_id DESC LIMIT ?",
                        (*args, limit + 1)).fetchall()
    page = [dict(r) for r in rows[:limit]]
    next_cursor = encode_cursor(page[-1]["created_at"], page[-1]["run_id"]) if len(rows) > limit else None
    return page, next_cursor


def events_after(conn: sqlite3.Connection, run_id: str, after: int, limit: int = 500) -> list[Event]:
    """Every event with seq > after, ascending. This is the whole SSE cursor contract: a client that
    reconnects with the last seq it saw gets nothing twice."""
    rows = conn.execute("SELECT * FROM run_events WHERE run_id=? AND seq>? ORDER BY seq ASC LIMIT ?",
                        (run_id, int(after), limit)).fetchall()
    return [_row_event(r) for r in rows]


def last_events(conn: sqlite3.Connection, run_id: str, n: int = 10) -> list[Event]:
    rows = conn.execute("SELECT * FROM run_events WHERE run_id=? ORDER BY seq DESC LIMIT ?", (run_id, n)).fetchall()
    return [_row_event(r) for r in reversed(rows)]


def count_retries(conn: sqlite3.Connection, run_id: str) -> int:
    return conn.execute("SELECT COUNT(*) FROM runs WHERE run_id LIKE ?", (f"{run_id}-retry-%",)).fetchone()[0]


# ------------------------------------------------------------------ SSE framing (pure; tested directly)
def sse_frame(event: Event) -> str:
    """One text/event-stream record. `id:` carries the seq so a browser EventSource resumes with Last-Event-ID."""
    return f"id: {event.seq}\nevent: {event.kind}\ndata: {json.dumps(event.as_dict(), sort_keys=True)}\n\n"


def poll_frames(conn: sqlite3.Connection, run_id: str, after: int, terminal_statuses=frozenset()) -> tuple[list[str], int, bool]:
    """One poll: frames for every event with seq > after, the new cursor, and whether a terminal state was
    reached (after which the stream ends with an `end` frame). Pure; this is what the SSE test drives."""
    frames, cursor, terminal = [], int(after), False
    for ev in events_after(conn, run_id, after):
        frames.append(sse_frame(ev))
        cursor = ev.seq
        if (ev.payload.get("state") or {}).get("status") in terminal_statuses:
            frames.append("event: end\ndata: {}\n\n")
            terminal = True
            break
    return frames, cursor, terminal


def stream_events(conn: sqlite3.Connection, run_id: str, after: int, *, sleep, poll_seconds: float = 0.5,
                  should_stop=lambda: False, terminal_statuses=frozenset()) -> Iterator[str]:
    """Poll SQLite every `poll_seconds`, yield frames for seq > cursor, advance the cursor; end after a terminal
    state. Synchronous, so a test can drive it with a fake `sleep`; `api.py` runs the same `poll_frames` loop
    with `asyncio.sleep`."""
    cursor = int(after)
    while not should_stop():
        frames, cursor, terminal = poll_frames(conn, run_id, cursor, terminal_statuses)
        yield from frames
        if terminal:
            return
        if not frames:
            sleep(poll_seconds)