areev-run-core 1.8.4

//! `SchedulerState` — everything a checkpoint carries (§5.4). Rebuilt-from
//! on resume with nothing inferred from wall-clock or environment; every
//! collection is a `Vec` or `BTreeMap` so iteration order is content order,
//! never hash order (§4's determinism rule).

use crate::types::{Ask, DecisionRecord, JournalKey, RunOutcome};
use serde_json::Value;
use std::collections::{BTreeMap, BTreeSet};

/// Per-node lifecycle at the node's current generation (§6.2's machine,
/// operationally realized with per-node re-entry generations).
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum NodeState {
    /// In-edges unresolved at the current generation.
    Waiting,
    /// Dispatchable: no in-edge Pending, ≥1 Fired.
    Ready,
    /// An effect is outstanding (Host dispatch or retry in flight).
    Dispatched,
    /// Parked on a Client ask (§6.6).
    AwaitingClient,
    DoneOk,
    DoneFailed,
    /// All in-edges Dead at the current generation.
    Dead,
}

/// Resolution of one in-edge at the destination's current generation.
/// Absent from the map = Pending.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum EdgeRes {
    Fired,
    Dead,
}

/// Budget accumulators — pure functions of the journal, carried here as the
/// running total and re-derived + asserted on replay (§6.7). USD in integer
/// micro-dollars: budgets never do float arithmetic.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct Spent {
    pub supersteps: u64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub usd_micros: u64,
    /// Active wall only — Σ(superstep close − open). Parked and CRASHED gaps
    /// accumulate in `elapsed_ms`, reported but never charged (§6.7). The
    /// crashed half of that promise needs `EventIn::Resumed`: without the
    /// marker a resume's gap is invisible to the scheduler, and a replay
    /// (which opens the next superstep at the previous close) charges it as
    /// wall while the live driver did not.
    pub wall_ms: u64,
    pub storage_bytes: u64,
    pub journal_grains: u64,
}

/// The run's phase within the BSP cycle.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum Phase {
    /// Between supersteps.
    Idle,
    /// A superstep is open: effects outstanding, results buffering, the
    /// decision record accumulating.
    Open {
        /// Keys still awaiting resolution (Host effects and Client asks).
        outstanding: BTreeSet<JournalKey>,
        /// Buffered results by node index — merged in canonical order at
        /// close, never in arrival order (§5.2).
        results: BTreeMap<usize, Value>,
        /// Buffered Send-task results by `task_path` (zero-padded ordinals,
        /// so lexicographic order IS spawn order) — merged after the static
        /// results, in path order.
        #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
        send_results: BTreeMap<String, Value>,
        record: DecisionRecord,
    },
    /// Terminal.
    Finished(RunOutcome),
}

/// One Send-spawned task (§5.1: `task_path = parent_path/spawn_ordinal`).
/// Task-scoped input; attempts continue across re-spawns of the same path
/// (a later generation re-activating a path keeps the journal keys unique).
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct SendTask {
    pub node_idx: usize,
    pub input: Value,
    pub attempt: u32,
    pub state: SendTaskState,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum SendTaskState {
    Queued,
    Running,
    Done,
}

/// One parked Client ask.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct PendingAsk {
    pub key: JournalKey,
    pub node_idx: usize,
    pub ask: Ask,
}

/// A tool call the model made, awaiting its result inside an abstract
/// node's LLM loop.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct PendingToolCall {
    /// The MODEL's call id (what the transcript's tool_result addresses).
    pub model_call_id: String,
    pub tool_name: String,
    pub arguments: Value,
}

/// What an abstract flow needs next — set by effect resolution, consumed by
/// the progress loop (resolution has no command sink; emission is the
/// loop's job, in canonical order).
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum FlowNeed {
    /// Dispatch the next LLM turn over the current transcript.
    NextTurn,
    /// Dispatch these model-issued tool calls (already appended to the
    /// transcript as the assistant entry).
    Tools(Vec<PendingToolCall>),
}

/// A summarizer turn in flight: which transcript entries its result will
/// stand in for, and the `effect_seq` it was dispatched under.
///
/// Held in state rather than re-derived, because the decision of WHAT to fold
/// is made before the model answers and must not be re-made afterwards: the
/// transcript has not changed in between, but re-deriving would couple the
/// splice to whatever `fold_range` returns at resolution time rather than to
/// what the journaled request actually summarized.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct FoldInFlight {
    /// First transcript index covered (inclusive).
    pub from: usize,
    /// One past the last covered index.
    pub to: usize,
    /// The summarizer turn's own effect ordinal.
    pub effect_seq: u32,
}

/// In-flight LLM-loop state for one abstract node at its current attempt
/// (§6.2's "a workflow with abstract nodes is an agent — same runtime").
/// Every entry is built from journaled outcomes, so the transcript is
/// replay-stable by construction.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AbstractFlow {
    /// Journal-shaped messages:
    /// `{"role":"user","content":…}` ·
    /// `{"role":"assistant","text":…,"tool_calls":[{id,name,arguments}]}` ·
    /// `{"role":"tool","tool_call_id":…,"content":…,"is_error":…}`.
    pub messages: Vec<Value>,
    /// The next effect ordinal within this attempt (LLM turns and tool
    /// calls share the sequence — §5.1).
    pub next_effect_seq: u32,
    /// The current tool round: effect_seq → the model call it serves.
    pub pending_tools: BTreeMap<u32, PendingToolCall>,
    /// Settled results of the current round, by effect_seq (canonical
    /// transcript order regardless of completion order).
    pub round_results: BTreeMap<u32, Value>,
    /// What the loop should do next (None while waiting on effects).
    pub need: Option<FlowNeed>,
    /// Unknown-tool corrections issued (§6.11: ONE re-prompt, then the
    /// node fails `ExecutorError`).
    pub unknown_strikes: u32,
    /// The provider's OWN reported prompt-token count for the last completed
    /// turn — the transcript's size at the model's tokenizer, not an estimate.
    /// 0 before the first turn settles, and reset to 0 by a fold so the next
    /// real turn measures the spliced transcript rather than the old one.
    ///
    /// Every one of the fold fields carries `skip_serializing_if` at its
    /// default: an `AbstractFlow` in a run that configured no ceiling must
    /// serialize to exactly the bytes it did before these existed, or every
    /// stored checkpoint diverges on the next `verify`.
    #[serde(default, skip_serializing_if = "is_zero")]
    pub last_prompt_tokens: u64,
    /// Set while a summarizer turn is outstanding; taken at its resolution.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub folding: Option<FoldInFlight>,
    /// The provider rejected the last turn as too long: fold before trying
    /// again, whatever the configured ceiling says.
    ///
    /// Separate from the ceiling because it answers a different question. The
    /// ceiling is a PREDICTION made from the previous turn's token count, and
    /// it can be unset, or set too high, or right but measuring a transcript
    /// that has since grown. This is the provider's own verdict on the
    /// transcript as actually sent, and it has to win over all three.
    #[serde(default, skip_serializing_if = "is_false")]
    pub fold_forced: bool,
    /// Folds performed on this attempt, and the number in the fold message the
    /// model reads. A completed node's flow is gone by the terminal checkpoint,
    /// so the durable record of a fold is its journaled intent (`input.fold`,
    /// visible in `run-trace`) rather than this counter.
    #[serde(default, skip_serializing_if = "is_zero_u32")]
    pub folds: u32,
}

fn is_zero(n: &u64) -> bool {
    *n == 0
}

fn is_zero_u32(n: &u32) -> bool {
    *n == 0
}

fn is_false(b: &bool) -> bool {
    !*b
}

/// The complete scheduler state. Serialized (as JSON) into each checkpoint's
/// `context_data`; byte-stability across resumes is what the
/// replay-equivalence gate compares.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct SchedulerState {
    pub run_id: String,
    pub superstep: u64,
    pub phase: Phase,
    /// Per node: re-entry generation (0 = first activation).
    pub node_gen: Vec<u32>,
    pub node_state: Vec<NodeState>,
    /// Per node: in-edge resolutions AT ITS CURRENT GENERATION, keyed by
    /// edge index. Absent = Pending.
    pub in_res: Vec<BTreeMap<usize, EdgeRes>>,
    /// Per edge: total Fired count across the run (`max_cycles` counts
    /// these — §6.3: "fires in at most n generations per run").
    pub edge_fired: Vec<u32>,
    /// Per node: attempt counter — MONOTONIC across re-entry generations
    /// (never reset), because the journal key is
    /// (run, task_path, node, attempt, effect_seq, kind) and a reset would
    /// mint generation 1 the same keys as generation 0 — the driver would
    /// then answer iteration 2 from iteration 1's journaled results and the
    /// loop body would execute once, ever. The per-generation retry budget
    /// measures from `attempt_base` instead.
    pub attempt: Vec<u32>,
    /// Per node: the attempt count when the CURRENT generation began —
    /// `attempt - attempt_base` is this generation's attempt number.
    pub attempt_base: Vec<u32>,
    /// The accumulated run state the reducers fold into.
    pub context: Value,
    pub spent: Spent,
    /// Un-charged calendar time (parks, crashes) — reported, never billed.
    pub elapsed_ms: u64,
    /// The latest journaled clock reading seen.
    pub clock_ms: u64,
    /// The active wall segment's start (None while parked/idle). Wall
    /// charges only active segments — open→park and resume→close — so a
    /// three-day approval pause costs nothing (§6.7).
    pub wall_open: Option<u64>,
    /// When the current park began (drives `elapsed_ms` on resume).
    pub paused_at: Option<u64>,
    /// Parked Client asks by `tool_call_id` (§6.6 — addressed by id, never
    /// by index).
    pub pending_asks: BTreeMap<String, PendingAsk>,
    /// In-flight abstract-node LLM loops, by node index.
    pub abstract_flows: BTreeMap<String, AbstractFlow>,
    /// Send-spawned tasks by `task_path`. A target node shows `Dispatched`
    /// while its batch is unresolved and completes when the batch drains.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub send_tasks: BTreeMap<String, SendTask>,
    /// Next spawn ordinal per parent path — monotonic for the whole run, so
    /// a re-entered spawner mints fresh task paths and journal keys never
    /// collide across generations.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub spawn_counter: BTreeMap<String, u32>,
    /// Ask ids already announced in an envelope — a parked run polled again
    /// must not re-emit its envelope every step call.
    pub announced_asks: BTreeSet<String>,
    /// Steering messages queued since the last superstep opened, INERT
    /// until one does: an open moves them into `context` under `$inbox`.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub inbox: Vec<Value>,
    /// How many steering messages the run has ever accepted.
    #[serde(default, skip_serializing_if = "is_zero")]
    pub inputs_seen: u64,
    /// A cancel marker observed (drain in progress).
    pub cancel: Option<(String, String)>,
    /// A budget axis exhausted mid-superstep (per-dispatch reservation,
    /// §6.7): no new work is emitted; the open superstep's effects drain,
    /// then the run finishes `BudgetExhausted` — checkpointed, resumable
    /// after a manifest supersession raises the ceiling.
    pub exhausted: Option<crate::error::BudgetAxis>,
    /// Set when fail-fast drain began: the lowest-index permanently-failed
    /// node and its detail (deterministic under permutation — chosen at
    /// superstep close, not at arrival).
    pub failed: Option<(usize, String)>,
    /// Transcript folds performed across the WHOLE run.
    ///
    /// Per-flow `AbstractFlow.folds` disappears with the flow when its node
    /// completes, so by the terminal checkpoint there is nothing left to count
    /// — and "did this workflow keep outgrowing the model's window?" is a
    /// question about the run, asked after it ended. Carried here so the
    /// run-outcome record can state it and the loop can act on it.
    #[serde(default, skip_serializing_if = "is_zero_u32")]
    pub folds: u32,
}

impl SchedulerState {
    /// Fresh state for a validated plan.
    pub fn new(run_id: &str, plan: &crate::plan::PlanGraph) -> Self {
        let n = plan.nodes.len();
        SchedulerState {
            run_id: run_id.to_string(),
            superstep: 0,
            phase: Phase::Idle,
            node_gen: vec![0; n],
            node_state: vec![NodeState::Waiting; n],
            in_res: vec![BTreeMap::new(); n],
            edge_fired: vec![0; plan.edges.len()],
            attempt: vec![0; n],
            attempt_base: vec![0; n],
            context: Value::Object(serde_json::Map::new()),
            spent: Spent::default(),
            elapsed_ms: 0,
            clock_ms: 0,
            wall_open: None,
            paused_at: None,
            pending_asks: BTreeMap::new(),
            abstract_flows: BTreeMap::new(),
            send_tasks: BTreeMap::new(),
            spawn_counter: BTreeMap::new(),
            announced_asks: BTreeSet::new(),
            inbox: Vec::new(),
            inputs_seen: 0,
            cancel: None,
            exhausted: None,
            failed: None,
            folds: 0,
        }
    }

    pub fn is_terminal(&self) -> bool {
        matches!(self.phase, Phase::Finished(_))
    }

    pub fn outcome(&self) -> Option<&RunOutcome> {
        match &self.phase {
            Phase::Finished(o) => Some(o),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn bare_flow() -> AbstractFlow {
        AbstractFlow {
            messages: vec![serde_json::json!({"role": "user", "content": "go"})],
            next_effect_seq: 1,
            pending_tools: BTreeMap::new(),
            round_results: BTreeMap::new(),
            need: None,
            unknown_strikes: 0,
            last_prompt_tokens: 0,
            folding: None,
            fold_forced: false,
            folds: 0,
        }
    }

    /// The replay-equivalence gate compares checkpoint BYTES, so a run that
    /// configured no context ceiling must serialize its flows exactly as it did
    /// before the fold existed. Every fold field is therefore `serde(default)`
    /// plus `skip_serializing_if` at its default — and this is the test that
    /// notices when the next field forgets one.
    #[test]
    fn a_flow_with_no_fold_state_serializes_as_it_did_before_folds_existed() {
        let json = serde_json::to_string(&bare_flow()).unwrap();
        assert_eq!(
            json,
            r#"{"messages":[{"content":"go","role":"user"}],"next_effect_seq":1,"pending_tools":{},"round_results":{},"need":null,"unknown_strikes":0}"#
        );
    }

    /// And the fold state round-trips byte-stably once it IS set — a checkpoint
    /// taken mid-fold has to resume into the same splice decision.
    #[test]
    fn fold_state_round_trips_byte_stably() {
        let flow = AbstractFlow {
            last_prompt_tokens: 120_000,
            folding: Some(FoldInFlight { from: 1, to: 9, effect_seq: 9 }),
            folds: 2,
            ..bare_flow()
        };
        let bytes = serde_json::to_string(&flow).unwrap();
        let back: AbstractFlow = serde_json::from_str(&bytes).unwrap();
        assert_eq!(back, flow);
        assert_eq!(serde_json::to_string(&back).unwrap(), bytes);
    }
}