agentyk-core 0.1.0

//! Turn state machine — the execution abstraction.
//!
//! Adopted from everruns-core's `TurnStateMachine` / `RuntimeTurnState`, with
//! one sharpening: the machine here is **sans-IO**. Transitions are pure —
//! they mutate bookkeeping and return the [`EventData`] to record — and the
//! only side-effecting operations are the two atoms ([`crate::atoms::reason`]
//! and [`crate::atoms::act`]) that a host runs between transitions.
//!
//! ```text
//! start ──▶ PendingReason ──reason──▶ PendingAct ──act──▶ PendingReason ──▶ …
//!                │                                                │
//!                └────────────▶ Completed(TurnOutcome) ◀──────────┘
//! ```
//!
//! The driving contract, shared by every host:
//!
//! ```text
//! let (mut state, effects) = TurnState::start(session_id, max_iterations, &input);
//! record(effects);
//! loop {
//!     match state.next_action() {                    // pure peek
//!         TurnAction::Reason => {
//!             record(state.on_reason_started(model_name));
//!             // Streaming deltas (output.message.delta) are ephemeral: the
//!             // driver reports them straight to `record` as it generates,
//!             // outside the state machine — see `atoms::reason_streaming`.
//!             let response = atoms::reason(…).await;  // effectful atom
//!             record(state.on_reason_completed(&response));
//!         }
//!         TurnAction::ExecuteTool { call } => {
//!             record(state.on_tool_started(&call.id));   // idempotent
//!             let output = atoms::act(…, &call, …).await;
//!             record(state.on_tool_completed(&call.id, &output));
//!         }
//!         TurnAction::Complete(outcome) => break,
//!     }
//! }
//! ```
//!
//! `next_action` walks the current batch in order, one call at a time — a
//! sequential executor (like agentyk's `InProcessExecutor`) can ignore
//! everything below and just follow that loop. A **parallel-capable**
//! executor instead calls [`TurnState::pending_tool_actions`] to fan out the
//! whole not-yet-started batch at once, starting/completing each call by id
//! (batch order is not completion order); the durable-execution contract
//! still holds one `TurnAction` == one activity, they just may run
//! concurrently.
//!
//! ## Why this supports durable execution
//!
//! - [`TurnState`] is `Serialize`/`Deserialize` and carries **bookkeeping
//!   only** — no credentials, no tool objects, no message history (mirroring
//!   everruns' `RuntimeTurnState`). A durable host checkpoints it between
//!   activities and rebuilds the environment (assembled tools, history) from
//!   the agent value and the event log on each step.
//! - Each [`TurnAction`] is one activity. Re-issuing the current action after
//!   a crash is safe: `next_action` is a pure read, and `on_tool_started` is
//!   idempotent per call id.
//! - Effects are data, so a host can persist the new state and append its
//!   events in one transaction — event emission never races execution.
//! - Message history is a fold over the event log
//!   ([`crate::replay::messages_from_events`]), so replay is sufficient to
//!   resume mid-turn.

use serde::{Deserialize, Serialize};

use crate::driver::{ChatResponse, Usage};
use crate::event::EventData;
use crate::id::{SessionId, TurnId};
use crate::message::{Message, ToolCall};
use crate::tool::ToolOutput;

/// Why a turn was deliberately sealed rather than left to run to
/// `MaxIterations` or fail. Mirrors everruns' `SealReason`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SealReason {
    /// A durable turn was reclaimed repeatedly without making forward
    /// progress. Host concern — agentyk's in-process executor can't
    /// crash-loop, so it never seals for this reason; a durable host would.
    NoProgress,
    /// A [`crate::budget::BudgetChecker`] returned
    /// [`crate::budget::BudgetDecision::Seal`].
    BudgetExhausted,
}

/// How a turn ended. Mirrors everruns' `TurnOutcome`
/// (Success/Failed/MaxIterations/Sealed) — `Cancelled` is the agentyk
/// analogue of a caller-driven seal via
/// [`crate::cancellation::CancellationToken`]; `Sealed` is a host- or
/// budget-driven one.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum TurnOutcome {
    Success {
        response: String,
    },
    Failed {
        error: String,
    },
    MaxIterations,
    /// Stopped cooperatively via a
    /// [`crate::cancellation::CancellationToken`] rather than failing.
    Cancelled,
    /// Stopped deliberately to prevent wasted work — see [`SealReason`].
    Sealed(SealReason),
}

/// One tool call within a reason step's batch, tracked independently so a
/// parallel executor can start/complete calls out of order.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PendingCall {
    pub call: ToolCall,
    /// Whether `tool.started` was already recorded for this call — makes
    /// `on_tool_started` idempotent per call id across crash/retry.
    pub started: bool,
    /// Set once the call resolves, whatever ran it.
    pub output: Option<ToolOutput>,
}

/// Current phase of the turn. `PendingAct` carries the whole batch from the
/// triggering reason step so a resumed host — sequential or parallel —
/// knows exactly what remains.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "phase", rename_all = "snake_case")]
pub enum TurnPhase {
    PendingReason,
    PendingAct { calls: Vec<PendingCall> },
    Completed(TurnOutcome),
}

/// The next thing a host must do. One `TurnAction` maps to one durable
/// activity.
#[derive(Debug, Clone, PartialEq)]
pub enum TurnAction {
    /// Run the LLM (the reason atom) over the current history.
    Reason,
    /// Execute one tool call (the act atom).
    ExecuteTool { call: ToolCall },
    /// The turn is finished.
    Complete(TurnOutcome),
}

/// Serializable turn bookkeeping. Everything else a step needs — model,
/// credentials, assembled tools, message history — is environment the host
/// provides per step; it is deliberately NOT part of this state.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TurnState {
    pub session_id: SessionId,
    pub turn_id: TurnId,
    pub max_iterations: usize,
    pub phase: TurnPhase,
    /// Completed reason (LLM) calls.
    pub iterations: usize,
    /// Executed tool calls.
    pub tool_calls_executed: usize,
    pub usage: Usage,
}

impl TurnState {
    /// Begin a turn. Returns the state (in `PendingReason`) and the effects
    /// to record: `turn.started` and `input.message`.
    pub fn start(
        session_id: SessionId,
        max_iterations: usize,
        input: &Message,
    ) -> (Self, Vec<EventData>) {
        let state = Self {
            session_id,
            turn_id: TurnId::new(),
            max_iterations,
            phase: TurnPhase::PendingReason,
            iterations: 0,
            tool_calls_executed: 0,
            usage: Usage::default(),
        };
        let effects = vec![
            EventData::TurnStarted,
            EventData::InputMessage {
                message: input.clone(),
            },
        ];
        (state, effects)
    }

    /// Pure peek at the next action: the first call in the batch that
    /// hasn't resolved yet, in order. A sequential executor can drive the
    /// whole turn from this alone — see [`Self::pending_tool_actions`] for
    /// the parallel-capable alternative. Idempotent — safe to call again
    /// after a crash without re-recording anything.
    pub fn next_action(&self) -> TurnAction {
        match &self.phase {
            TurnPhase::PendingReason => TurnAction::Reason,
            TurnPhase::PendingAct { calls } => {
                match calls.iter().find(|pending| pending.output.is_none()) {
                    Some(pending) => TurnAction::ExecuteTool {
                        call: pending.call.clone(),
                    },
                    // Unreachable by construction; degrade to another reason step.
                    None => TurnAction::Reason,
                }
            }
            TurnPhase::Completed(outcome) => TurnAction::Complete(outcome.clone()),
        }
    }

    /// Every call in the current batch that hasn't been started yet, all at
    /// once — what a parallel-capable executor fans out concurrently
    /// instead of following `next_action` one call at a time. Empty outside
    /// `PendingAct` or once every call has been started.
    pub fn pending_tool_actions(&self) -> Vec<TurnAction> {
        match &self.phase {
            TurnPhase::PendingAct { calls } => calls
                .iter()
                .filter(|pending| !pending.started)
                .map(|pending| TurnAction::ExecuteTool {
                    call: pending.call.clone(),
                })
                .collect(),
            _ => Vec::new(),
        }
    }

    /// About to run the reason atom for the upcoming iteration. Emits
    /// `output.message.started`. Doesn't mutate `phase` — purely
    /// informational, so it's safe to call again after a crash without
    /// double-counting iterations (unlike `on_reason_completed`, which does
    /// advance `self.iterations`).
    pub fn on_reason_started(&self, model: Option<&str>) -> Vec<EventData> {
        vec![EventData::OutputMessageStarted {
            model: model.map(str::to_string),
            iteration: Some(self.iterations as u32 + 1),
        }]
    }

    /// Apply an LLM response. Emits `output.message.completed`, then either
    /// finishes the turn (`turn.completed` / `turn.failed` on
    /// max-iterations) or moves to `PendingAct` with the requested tool
    /// calls.
    pub fn on_reason_completed(&mut self, response: &ChatResponse) -> Vec<EventData> {
        self.iterations += 1;
        self.usage.add(response.usage);
        let mut effects = vec![EventData::OutputMessageCompleted {
            message: response.message.clone(),
        }];

        if response.message.tool_calls.is_empty() {
            let outcome = TurnOutcome::Success {
                response: response.message.text(),
            };
            effects.push(EventData::TurnCompleted {
                iterations: self.iterations,
                tool_calls: self.tool_calls_executed,
            });
            self.phase = TurnPhase::Completed(outcome);
        } else if self.iterations >= self.max_iterations {
            // Sealing rather than executing tools that can never be reasoned
            // over again.
            let outcome = TurnOutcome::MaxIterations;
            effects.push(EventData::TurnFailed {
                error: format!("turn exceeded max iterations ({})", self.max_iterations),
            });
            self.phase = TurnPhase::Completed(outcome);
        } else {
            let calls = response
                .message
                .tool_calls
                .iter()
                .map(|call| PendingCall {
                    call: call.clone(),
                    started: false,
                    output: None,
                })
                .collect();
            self.phase = TurnPhase::PendingAct { calls };
        }
        effects
    }

    /// Record that the call with this id is starting. Idempotent per call
    /// id: re-running after a crash records nothing twice.
    pub fn on_tool_started(&mut self, call_id: &str) -> Vec<EventData> {
        let TurnPhase::PendingAct { calls } = &mut self.phase else {
            return Vec::new();
        };
        let Some(pending) = calls.iter_mut().find(|p| p.call.id == call_id) else {
            return Vec::new();
        };
        if pending.started {
            return Vec::new();
        }
        pending.started = true;
        vec![EventData::ToolStarted {
            call: pending.call.clone(),
        }]
    }

    /// Apply a call's result by id. Emits `tool.completed`; when every call
    /// in the batch has resolved, moves back to `PendingReason` — the batch
    /// can complete in any order.
    pub fn on_tool_completed(&mut self, call_id: &str, output: &ToolOutput) -> Vec<EventData> {
        let TurnPhase::PendingAct { calls } = &mut self.phase else {
            return Vec::new();
        };
        let Some(pending) = calls.iter_mut().find(|p| p.call.id == call_id) else {
            return Vec::new();
        };
        if pending.output.is_some() {
            return Vec::new();
        }
        pending.output = Some(output.clone());
        self.tool_calls_executed += 1;
        let effects = vec![EventData::ToolCompleted {
            call_id: pending.call.id.clone(),
            name: pending.call.name.clone(),
            output: output.content.clone(),
            is_error: output.is_error,
        }];
        if calls.iter().all(|p| p.output.is_some()) {
            self.phase = TurnPhase::PendingReason;
        }
        effects
    }

    /// Fail the turn (driver error, host abort). Emits `turn.failed`.
    pub fn on_failure(&mut self, error: impl Into<String>) -> Vec<EventData> {
        let error = error.into();
        self.phase = TurnPhase::Completed(TurnOutcome::Failed {
            error: error.clone(),
        });
        vec![EventData::TurnFailed { error }]
    }

    /// Stop the turn cooperatively (see
    /// [`crate::cancellation::CancellationToken`]), abandoning whatever is
    /// pending — including the rest of a not-yet-drained tool-call batch.
    /// Emits `turn.cancelled`.
    pub fn on_cancel(&mut self) -> Vec<EventData> {
        self.phase = TurnPhase::Completed(TurnOutcome::Cancelled);
        vec![EventData::TurnCancelled]
    }

    /// Stop the turn deliberately (see [`SealReason`]), abandoning whatever
    /// is pending. Emits `turn.sealed`.
    pub fn on_seal(&mut self, reason: SealReason) -> Vec<EventData> {
        self.phase = TurnPhase::Completed(TurnOutcome::Sealed(reason));
        vec![EventData::TurnSealed { reason }]
    }

    pub fn is_complete(&self) -> bool {
        matches!(self.phase, TurnPhase::Completed(_))
    }

    pub fn outcome(&self) -> Option<&TurnOutcome> {
        match &self.phase {
            TurnPhase::Completed(outcome) => Some(outcome),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::message::Message;
    use serde_json::json;

    fn tool_call_response(names: &[&str]) -> ChatResponse {
        let calls = names
            .iter()
            .enumerate()
            .map(|(index, name)| ToolCall {
                id: format!("call_{index}"),
                name: (*name).to_string(),
                arguments: json!({}),
            })
            .collect();
        ChatResponse {
            message: Message::assistant_with_calls("", calls),
            usage: Usage::default(),
        }
    }

    fn text_response(text: &str) -> ChatResponse {
        ChatResponse {
            message: Message::assistant(text),
            usage: Usage::default(),
        }
    }

    #[test]
    fn happy_path_transitions() {
        let input = Message::user("hi");
        let (mut state, effects) = TurnState::start(SessionId::new(), 4, &input);
        assert_eq!(effects.len(), 2);
        assert_eq!(state.next_action(), TurnAction::Reason);

        state.on_reason_completed(&tool_call_response(&["a", "b"]));
        let TurnAction::ExecuteTool { call } = state.next_action() else {
            panic!("expected tool action");
        };
        assert_eq!(call.name, "a");

        assert_eq!(state.on_tool_started(&call.id).len(), 1);
        assert_eq!(state.on_tool_started(&call.id).len(), 0); // idempotent
        state.on_tool_completed(&call.id, &ToolOutput::text("ok"));

        let TurnAction::ExecuteTool { call } = state.next_action() else {
            panic!("expected second tool action");
        };
        assert_eq!(call.name, "b");
        state.on_tool_started(&call.id);
        state.on_tool_completed(&call.id, &ToolOutput::text("ok"));

        assert_eq!(state.next_action(), TurnAction::Reason);
        let effects = state.on_reason_completed(&text_response("done"));
        assert!(matches!(
            effects.last(),
            Some(EventData::TurnCompleted { .. })
        ));
        assert_eq!(
            state.outcome(),
            Some(&TurnOutcome::Success {
                response: "done".into()
            })
        );
        assert_eq!(state.iterations, 2);
        assert_eq!(state.tool_calls_executed, 2);
    }

    #[test]
    fn max_iterations_seals_instead_of_acting() {
        let input = Message::user("go");
        let (mut state, _) = TurnState::start(SessionId::new(), 1, &input);
        let effects = state.on_reason_completed(&tool_call_response(&["a"]));
        assert!(matches!(effects.last(), Some(EventData::TurnFailed { .. })));
        assert_eq!(state.outcome(), Some(&TurnOutcome::MaxIterations));
    }

    #[test]
    fn state_survives_serialization_mid_turn() {
        let input = Message::user("hi");
        let (mut state, _) = TurnState::start(SessionId::new(), 4, &input);
        state.on_reason_completed(&tool_call_response(&["a"]));
        let TurnAction::ExecuteTool { call } = state.next_action() else {
            panic!("expected tool action");
        };
        state.on_tool_started(&call.id);

        let json = serde_json::to_string(&state).unwrap();
        let mut restored: TurnState = serde_json::from_str(&json).unwrap();
        assert_eq!(restored, state);

        // The resumed host re-issues the same action, records nothing twice,
        // and continues.
        assert!(matches!(
            restored.next_action(),
            TurnAction::ExecuteTool { .. }
        ));
        assert_eq!(restored.on_tool_started(&call.id).len(), 0);
        restored.on_tool_completed(&call.id, &ToolOutput::text("ok"));
        assert_eq!(restored.next_action(), TurnAction::Reason);
    }

    #[test]
    fn pending_tool_actions_returns_the_whole_not_started_batch() {
        let input = Message::user("hi");
        let (mut state, _) = TurnState::start(SessionId::new(), 4, &input);
        state.on_reason_completed(&tool_call_response(&["a", "b", "c"]));

        let actions = state.pending_tool_actions();
        assert_eq!(actions.len(), 3);

        // Starting one call removes only that call from the ready batch.
        let TurnAction::ExecuteTool { call } = &actions[0] else {
            panic!("expected tool action");
        };
        state.on_tool_started(&call.id);
        assert_eq!(state.pending_tool_actions().len(), 2);
    }

    #[test]
    fn batch_completes_out_of_order() {
        let input = Message::user("hi");
        let (mut state, _) = TurnState::start(SessionId::new(), 4, &input);
        state.on_reason_completed(&tool_call_response(&["a", "b"]));

        let ids: Vec<String> = state
            .pending_tool_actions()
            .into_iter()
            .map(|action| {
                let TurnAction::ExecuteTool { call } = action else {
                    unreachable!()
                };
                call.id
            })
            .collect();
        for id in &ids {
            state.on_tool_started(id);
        }

        // Complete "b" (the second call) before "a" — still ends up back at
        // PendingReason with both counted, regardless of completion order.
        state.on_tool_completed(&ids[1], &ToolOutput::text("b done"));
        let TurnAction::ExecuteTool { call: remaining } = state.next_action() else {
            panic!("expected the still-incomplete call");
        };
        assert_eq!(remaining.id, ids[0]);
        state.on_tool_completed(&ids[0], &ToolOutput::text("a done"));

        assert_eq!(state.next_action(), TurnAction::Reason);
        assert_eq!(state.tool_calls_executed, 2);
    }
}