orchestral_runtime/

generic_agent_checkpoint.rs

//! Provider-private write-ahead journal for the in-process Generic Agent.
//!
//! This journal is deliberately separate from the Host Agent journal. The
//! Host journal proves public protocol state; this journal proves whether the
//! Generic Agent may safely reconstruct private model-loop work after restart.

use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use std::sync::RwLock;

use crate::session_context::observed_prefix::{ContextPlanningTrace, ObservedPrefixAnchor};
use orchestral_core::agent_protocol::wire::{
    AgentAdmission, AgentCommandEnvelope, AgentEvent, AgentEventDraft, AgentEventId,
    AgentExecutionRef, AgentStartRequest, CommandId, Digest, ProviderCommandOutcome, RunId,
};
use orchestral_core::agent_session::SessionSourceRange;
use orchestral_core::model_protocol::{
    ModelContent, ModelFinishReason, ModelRequestId, ModelToolCallId, ModelUsage,
};
use orchestral_core::tool_protocol::ApprovalCapability;
use serde::{Deserialize, Serialize};

macro_rules! string_id {
    ($name:ident) => {
        #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
        #[serde(transparent)]
        pub struct $name(String);

        impl $name {
            pub fn new(value: impl Into<String>) -> Self {
                Self(value.into())
            }

            pub fn as_str(&self) -> &str {
                &self.0
            }

            pub fn is_empty(&self) -> bool {
                self.0.trim().is_empty()
            }
        }

        impl fmt::Display for $name {
            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
                formatter.write_str(&self.0)
            }
        }
    };
}

string_id!(GenericCheckpointEventId);

/// Immutable Generic Agent identity bound before any model attempt begins.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GenericAgentRunRegistration {
    pub request: AgentStartRequest,
    pub execution: AgentExecutionRef,
    pub admission: AgentAdmission,
    /// Binds the model contract, system prompt, Host Tool authority and
    /// schemas, Skill catalog, and loop limits that determine reconstructed
    /// execution.
    pub config_digest: Digest,
}

/// One model Tool call reconstructed from the canonical streaming protocol.
/// Raw argument bytes are retained so recovery can apply the same JSON/schema
/// validation without asking the model to repeat the call.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GenericObservedToolCall {
    pub call_id: ModelToolCallId,
    pub name: String,
    #[serde(default)]
    pub arguments: String,
    #[serde(default)]
    pub extensions: BTreeMap<String, serde_json::Value>,
    pub ended: bool,
}

/// Durable aggregate produced only after a valid terminal Model stream event
/// has been observed. It is Provider-private execution state, not a public
/// Agent event and not a claim that any Tool effect has occurred.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GenericModelObservation {
    pub finish_reason: ModelFinishReason,
    #[serde(default)]
    pub response: String,
    /// Provider-owned message state, separate from visible text and Tool data.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub continuation: BTreeMap<String, serde_json::Value>,
    #[serde(default)]
    pub usage: Option<ModelUsage>,
    #[serde(default)]
    pub tool_calls: Vec<GenericObservedToolCall>,
}

/// Durable provenance for the exact Session Context projected before one
/// model attempt. The request digest proves bytes; this trace explains which
/// Journal ranges and Host limits produced those bytes.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GenericModelContextTrace {
    pub through_session_seq: u64,
    pub included_ranges: Vec<SessionSourceRange>,
    pub deferred_ranges: Vec<SessionSourceRange>,
    pub config_digest: Digest,
    pub history_limit: usize,
    pub used_input_tokens: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub context_estimate: Option<orchestral_core::model_protocol::ModelContextEstimate>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub planning: Option<ContextPlanningTrace>,
    pub input_budget_tokens: u64,
}

impl GenericModelContextTrace {
    pub(crate) fn validate(&self) -> Result<(), GenericCheckpointError> {
        if self.planning.as_ref().is_some_and(|planning| {
            !planning.input.validate()
                || planning.input.raw_estimate_tokens > self.used_input_tokens
                || self.context_estimate.is_none()
                || planning
                    .anchor
                    .as_ref()
                    .is_some_and(|anchor| !anchor.validate())
        }) {
            return Err(GenericCheckpointError::InvalidData(
                "invalid context planning provenance".to_owned(),
            ));
        }
        if let Some(estimate) = &self.context_estimate {
            if estimate.accounting
                != orchestral_core::model_protocol::ModelTokenAccounting::Estimated
                || estimate.tokens > self.used_input_tokens
            {
                return Err(GenericCheckpointError::InvalidData(
                    "model Context planning estimate must be marked estimated and within its input bound".to_owned(),
                ));
            }
        }
        if !self.config_digest.is_sha256()
            || self.history_limit == 0
            || self.input_budget_tokens == 0
            || self
                .context_estimate
                .as_ref()
                .map_or(self.used_input_tokens, |estimate| estimate.tokens)
                > self.input_budget_tokens
        {
            return Err(GenericCheckpointError::InvalidData(
                "model Context trace requires a config digest and valid Host limits".to_owned(),
            ));
        }
        let ranges = self
            .included_ranges
            .iter()
            .chain(self.deferred_ranges.iter())
            .collect::<Vec<_>>();
        for range in &ranges {
            range.validate().map_err(invalid_data)?;
            if range.last_session_seq > self.through_session_seq {
                return Err(GenericCheckpointError::InvalidData(
                    "model Context range exceeds its Session Journal cursor".to_owned(),
                ));
            }
        }
        for (index, left) in ranges.iter().enumerate() {
            if ranges.iter().skip(index + 1).any(|right| {
                left.first_session_seq <= right.last_session_seq
                    && right.first_session_seq <= left.last_session_seq
            }) {
                return Err(GenericCheckpointError::InvalidData(
                    "model Context trace ranges must not overlap".to_owned(),
                ));
            }
        }
        Ok(())
    }

    pub(crate) fn observed_prefix(
        &self,
        run_id: &RunId,
        request_id: &ModelRequestId,
        observation: &GenericModelObservation,
        max_output_tokens: Option<u64>,
    ) -> Option<ObservedPrefixAnchor> {
        let planning = self.planning.as_ref()?;
        let input_tokens = observation.usage.as_ref()?.input_tokens?;
        if input_tokens == 0
            || input_tokens > self.used_input_tokens
            || observation
                .usage
                .as_ref()
                .and_then(|usage| usage.output_tokens)
                .zip(max_output_tokens)
                .is_some_and(|(used, cap)| used > cap)
            || !matches!(
                observation.finish_reason,
                ModelFinishReason::Stop | ModelFinishReason::ToolCalls
            )
            || (observation.response.is_empty() && observation.tool_calls.is_empty())
            || observation.tool_calls.iter().any(|call| {
                !call.ended || serde_json::from_str::<serde_json::Value>(&call.arguments).is_err()
            })
        {
            return None;
        }
        Some(ObservedPrefixAnchor {
            run_id: run_id.clone(),
            config_digest: self.config_digest.clone(),
            source_request_id: request_id.clone(),
            input: planning.input.clone(),
            observed_input_tokens: input_tokens,
        })
    }
}

impl GenericModelObservation {
    pub(crate) fn assistant_content(&self) -> Vec<ModelContent> {
        let mut content = Vec::new();
        if !self.response.is_empty() {
            content.push(ModelContent::Text {
                text: self.response.clone(),
            });
        }
        content.extend(self.continuation.iter().map(|(namespace, value)| {
            ModelContent::Continuation {
                namespace: namespace.clone(),
                value: value.clone(),
            }
        }));
        content
    }

    fn validate(&self) -> Result<(), GenericCheckpointError> {
        for content in self.assistant_content() {
            content.validate().map_err(invalid_data)?;
        }
        let mut call_ids = BTreeSet::new();
        if self.tool_calls.iter().any(|call| {
            call.call_id.is_empty()
                || call.name.trim().is_empty()
                || !call_ids.insert(call.call_id.clone())
        }) {
            return Err(GenericCheckpointError::InvalidData(
                "model observation Tool calls require unique identities and names".to_owned(),
            ));
        }
        Ok(())
    }
}

impl GenericAgentRunRegistration {
    pub fn run_id(&self) -> &RunId {
        &self.execution.run_id
    }

    pub fn validate(&self) -> Result<(), GenericCheckpointError> {
        self.request
            .run
            .validate_integrity()
            .map_err(invalid_data)?;
        self.execution.validate_integrity().map_err(invalid_data)?;
        self.admission.validate_integrity().map_err(invalid_data)?;
        if !self.config_digest.is_sha256()
            || self.execution.run_id != self.request.run.spec.run_id
            || self.execution.session_id != self.request.run.spec.session_id
            || self.execution.spec_digest != self.request.run.spec_digest
            || self.execution.binding_ref != self.request.provider_binding
            || self.execution.descriptor_digest != self.request.expected_descriptor_digest
        {
            return Err(GenericCheckpointError::InvalidData(
                "Generic Agent registration identities or config digest do not agree".to_owned(),
            ));
        }
        Ok(())
    }
}

/// Private facts required to distinguish a stable loop boundary from an
/// uncertain model attempt. `ModelAttemptStarted` is written before polling
/// `ModelBackend::start`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
#[non_exhaustive]
// This is a versioned wire contract. Boxing individual fields would change the
// public Rust construction API without changing the serialized representation.
#[allow(clippy::large_enum_variant)]
pub enum GenericCheckpointEvent {
    LoopBoundaryCommitted {
        next_model_round: u64,
        #[serde(default)]
        usage: ModelUsage,
        tool_call_count: u64,
        #[serde(default)]
        last_response: String,
        #[serde(default)]
        supporting_event_ids: Vec<AgentEventId>,
    },
    ModelAttemptStarted {
        round: u64,
        request_id: ModelRequestId,
        request_digest: Digest,
        #[serde(default)]
        max_output_tokens: Option<u64>,
        context: GenericModelContextTrace,
    },
    ModelAttemptObserved {
        round: u64,
        request_id: ModelRequestId,
        observation: GenericModelObservation,
    },
    /// A definite capacity rejection, before any model usage or content.
    /// This closes the attempt without executing tools and durably reserves
    /// a smaller input or output budget, including after restart.
    ModelContextRejected {
        round: u64,
        request_id: ModelRequestId,
        retry_number: u32,
        input_budget_tokens: u64,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        output_budget_tokens: Option<u64>,
        error: orchestral_core::model_protocol::ModelError,
    },
    /// A retry within an open logical attempt, before any non-usage event was
    /// observed. Recovery still treats an open attempt as interrupted; this
    /// fact does not authorize replaying a model request after process loss.
    ModelRetryScheduled {
        round: u64,
        request_id: ModelRequestId,
        retry_number: u32,
        delay_ms: u64,
        error: orchestral_core::model_protocol::ModelError,
        /// Last reported usage snapshot of the failed attempt. Missing fields
        /// remain unknown; this does not certify the request's final cost.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        observed_usage: Option<ModelUsage>,
    },
    /// Written before entering the Workflow DAG executor. Once present, a
    /// missing durable Workflow output is intentionally outcome-unknown and
    /// must never be reconstructed by rerunning the DAG.
    WorkflowAttemptStarted {
        round: u64,
        request_id: ModelRequestId,
        call_id: ModelToolCallId,
        arguments_digest: Digest,
    },
    CommandCommitted {
        command: AgentCommandEnvelope,
        outcome: ProviderCommandOutcome,
        #[serde(default)]
        approval_capability: Option<ApprovalCapability>,
    },
    /// Provider drafts are persisted before they enter the live Provider
    /// stream. A delivery batch may contain OutputCommitted followed by the
    /// single terminal DeliveryCommitted event.
    ProviderEventsCommitted { events: Vec<AgentEventDraft> },
}

impl GenericCheckpointEvent {
    fn validate(&self, run_id: &RunId) -> Result<(), GenericCheckpointError> {
        match self {
            Self::ModelContextRejected {
                round,
                request_id,
                retry_number,
                input_budget_tokens,
                output_budget_tokens,
                error,
            } => {
                if *round == 0
                    || round.checked_add(1).is_none()
                    || request_id.is_empty()
                    || *retry_number == 0
                    || *input_budget_tokens == 0
                    || *output_budget_tokens == Some(0)
                    || error.code
                        != orchestral_core::model_protocol::ModelErrorCode::ContextLengthExceeded
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "context recovery requires a definite capacity rejection and positive budget".to_owned(),
                    ));
                }
            }
            Self::ModelRetryScheduled {
                round,
                request_id,
                retry_number,
                delay_ms,
                error,
                ..
            } => {
                if *round == 0
                    || request_id.is_empty()
                    || *retry_number == 0
                    || *delay_ms == 0
                    || !error.retryable
                    || !matches!(
                        error.code,
                        orchestral_core::model_protocol::ModelErrorCode::RateLimited
                            | orchestral_core::model_protocol::ModelErrorCode::Unavailable
                    )
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "model retry requires an attempt identity, delay, and transient error"
                            .to_owned(),
                    ));
                }
            }
            Self::LoopBoundaryCommitted {
                next_model_round,
                supporting_event_ids,
                ..
            } => {
                if *next_model_round == 0
                    || supporting_event_ids.iter().any(AgentEventId::is_empty)
                    || supporting_event_ids.iter().collect::<BTreeSet<_>>().len()
                        != supporting_event_ids.len()
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "loop boundary requires a positive round and unique event references"
                            .to_owned(),
                    ));
                }
            }
            Self::ModelAttemptStarted {
                round,
                request_id,
                request_digest,
                max_output_tokens,
                context,
            } => {
                if *round == 0
                    || request_id.is_empty()
                    || !request_digest.is_sha256()
                    || *max_output_tokens == Some(0)
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "model attempt requires a round, request identity, and digest".to_owned(),
                    ));
                }
                context.validate()?;
            }
            Self::ModelAttemptObserved {
                round,
                request_id,
                observation,
            } => {
                if *round == 0 || request_id.is_empty() {
                    return Err(GenericCheckpointError::InvalidData(
                        "model observation requires a round and request identity".to_owned(),
                    ));
                }
                observation.validate()?;
            }
            Self::WorkflowAttemptStarted {
                round,
                request_id,
                call_id,
                arguments_digest,
            } => {
                if *round == 0
                    || request_id.is_empty()
                    || call_id.is_empty()
                    || !arguments_digest.is_sha256()
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "workflow attempt requires model, call, and argument identities".to_owned(),
                    ));
                }
            }
            Self::CommandCommitted {
                command,
                outcome,
                approval_capability,
            } => {
                command.verify_digest().map_err(invalid_data)?;
                outcome.validate_shape().map_err(invalid_data)?;
                if command.run_id != *run_id {
                    return Err(GenericCheckpointError::InvalidData(
                        "checkpoint command crossed a Run boundary".to_owned(),
                    ));
                }
                let accepted_allow = matches!(
                    (&command.payload, outcome),
                    (
                        orchestral_core::agent_protocol::wire::AgentCommand::ResolveRequest {
                            response: orchestral_core::agent_protocol::wire::RequestResolution::Approval {
                                decision: orchestral_core::agent_protocol::wire::ApprovalDecision::Allow,
                                ..
                            }
                        },
                        ProviderCommandOutcome::Accepted
                    )
                );
                if accepted_allow != approval_capability.is_some()
                    || approval_capability.as_ref().is_some_and(|capability| {
                        capability.claims.binding.run_id != *run_id
                            || !capability.authenticator.is_sha256()
                    })
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "checkpoint approval capability does not match its accepted command"
                            .to_owned(),
                    ));
                }
            }
            Self::ProviderEventsCommitted { events } => {
                if events.is_empty() {
                    return Err(GenericCheckpointError::InvalidData(
                        "Provider event checkpoint batch must not be empty".to_owned(),
                    ));
                }
                let mut event_ids = BTreeSet::new();
                let mut terminal_seen = false;
                for (index, event) in events.iter().enumerate() {
                    event.validate_integrity().map_err(invalid_data)?;
                    if event.run_id != *run_id || !event_ids.insert(&event.event_id) {
                        return Err(GenericCheckpointError::InvalidData(
                            "Provider checkpoint events must be unique and Run-bound".to_owned(),
                        ));
                    }
                    let terminal = is_terminal_event(&event.payload);
                    if terminal_seen || (terminal && index + 1 != events.len()) {
                        return Err(GenericCheckpointError::InvalidData(
                            "terminal Provider event must be the final event in its batch"
                                .to_owned(),
                        ));
                    }
                    terminal_seen = terminal;
                }
            }
        }
        Ok(())
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GenericCheckpointDraft {
    pub event_id: GenericCheckpointEventId,
    pub run_id: RunId,
    pub payload: GenericCheckpointEvent,
}

impl GenericCheckpointDraft {
    pub fn validate(&self) -> Result<(), GenericCheckpointError> {
        if self.event_id.is_empty() || self.run_id.is_empty() {
            return Err(GenericCheckpointError::InvalidData(
                "Generic checkpoint identities must not be empty".to_owned(),
            ));
        }
        self.payload.validate(&self.run_id)
    }

    pub fn digest(&self) -> Result<Digest, GenericCheckpointError> {
        self.validate()?;
        canonical_digest(self)
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GenericCheckpointRecord {
    pub checkpoint_seq: u64,
    pub draft_digest: Digest,
    pub event_digest: Digest,
    pub event_id: GenericCheckpointEventId,
    pub run_id: RunId,
    pub payload: GenericCheckpointEvent,
}

#[derive(Serialize)]
struct GenericCheckpointRecordDigestView<'a> {
    checkpoint_seq: u64,
    draft_digest: &'a Digest,
    event_id: &'a GenericCheckpointEventId,
    run_id: &'a RunId,
    payload: &'a GenericCheckpointEvent,
}

impl GenericCheckpointRecord {
    /// Seals one validated private-WAL draft at the store-assigned sequence.
    /// Concrete checkpoint-store plugins use this after enforcing their
    /// compare-and-append boundary.
    pub fn seal(
        draft: GenericCheckpointDraft,
        checkpoint_seq: u64,
    ) -> Result<Self, GenericCheckpointError> {
        draft.validate()?;
        if checkpoint_seq == 0 {
            return Err(GenericCheckpointError::InvalidData(
                "checkpoint sequence must be positive".to_owned(),
            ));
        }
        let draft_digest = draft.digest()?;
        let mut record = Self {
            checkpoint_seq,
            draft_digest,
            event_digest: Digest::sha256([]),
            event_id: draft.event_id,
            run_id: draft.run_id,
            payload: draft.payload,
        };
        record.event_digest = record.computed_event_digest()?;
        Ok(record)
    }

    pub fn validate(&self) -> Result<(), GenericCheckpointError> {
        if self.checkpoint_seq == 0
            || self.event_id.is_empty()
            || self.run_id.is_empty()
            || !self.draft_digest.is_sha256()
            || self.computed_event_digest()? != self.event_digest
        {
            return Err(GenericCheckpointError::InvalidData(
                "Generic checkpoint record identity or digest is invalid".to_owned(),
            ));
        }
        self.payload.validate(&self.run_id)
    }

    fn computed_event_digest(&self) -> Result<Digest, GenericCheckpointError> {
        canonical_digest(&GenericCheckpointRecordDigestView {
            checkpoint_seq: self.checkpoint_seq,
            draft_digest: &self.draft_digest,
            event_id: &self.event_id,
            run_id: &self.run_id,
            payload: &self.payload,
        })
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct StoredGenericAgentRun {
    pub registration: GenericAgentRunRegistration,
    pub records: Vec<GenericCheckpointRecord>,
}

impl StoredGenericAgentRun {
    pub fn validate(&self) -> Result<GenericAgentCheckpointProjection, GenericCheckpointError> {
        self.registration.validate()?;
        replay_generic_agent_checkpoint(self)
    }

    pub fn last_checkpoint_seq(&self) -> u64 {
        self.records
            .last()
            .map(|record| record.checkpoint_seq)
            .unwrap_or(0)
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct GenericLoopBoundary {
    pub next_model_round: u64,
    pub usage: ModelUsage,
    pub tool_call_count: u64,
    pub last_response: String,
    pub supporting_event_ids: Vec<AgentEventId>,
}

#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub enum GenericCheckpointPhase {
    Prepared,
    Stable(GenericLoopBoundary),
    ModelAttemptOpen {
        boundary: GenericLoopBoundary,
        round: u64,
        request_id: ModelRequestId,
        request_digest: Digest,
    },
    ModelAttemptObserved {
        boundary: GenericLoopBoundary,
        round: u64,
        request_id: ModelRequestId,
        request_digest: Digest,
        observation: GenericModelObservation,
    },
    WorkflowAttemptOpen {
        boundary: GenericLoopBoundary,
        round: u64,
        request_id: ModelRequestId,
        request_digest: Digest,
        observation: GenericModelObservation,
        call_id: ModelToolCallId,
        arguments_digest: Digest,
    },
    Terminal,
}

impl GenericCheckpointPhase {
    pub fn is_stable(&self) -> bool {
        matches!(self, Self::Stable(_))
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct GenericAgentCheckpointProjection {
    pub phase: GenericCheckpointPhase,
    pub provider_events: Vec<AgentEventDraft>,
    pub commands: BTreeMap<CommandId, CommandCheckpoint>,
    pub last_checkpoint_seq: u64,
    /// Derived from the last completed request, never cumulative/output usage.
    pub observed_prefix: Option<ObservedPrefixAnchor>,
    pub context_recovery: Option<GenericContextRecovery>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GenericContextRecovery {
    /// Consecutive rejected attempts. Zero means generation succeeded; an
    /// input ceiling remains learned only when input reduction was required.
    pub retry_number: u32,
    pub input_budget_tokens: u64,
    pub output_budget_tokens: Option<u64>,
    pub compact_input: bool,
}

impl GenericContextRecovery {
    /// An output-only rejection does not establish a smaller input capacity.
    /// Keep its original input fixed during retry; after success, allow new
    /// observations to use the space released by the smaller output budget.
    pub(crate) fn input_capacity_tokens(&self) -> Option<u64> {
        (self.retry_number > 0 || self.compact_input).then_some(self.input_budget_tokens)
    }

    /// The ceiling is one token below the rejected input; rounding its half
    /// upward recovers half the original rejected input without overflowing.
    pub(crate) fn compaction_target_tokens(&self) -> Option<u64> {
        (self.retry_number > 0 && self.compact_input).then(|| self.input_budget_tokens.div_ceil(2))
    }

    pub(crate) fn generation_observed(&mut self) {
        self.retry_number = 0;
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct CommandCheckpoint {
    pub command: AgentCommandEnvelope,
    pub outcome: ProviderCommandOutcome,
    pub approval_capability: Option<ApprovalCapability>,
}

pub fn replay_generic_agent_checkpoint(
    run: &StoredGenericAgentRun,
) -> Result<GenericAgentCheckpointProjection, GenericCheckpointError> {
    run.registration.validate()?;
    let run_id = run.registration.run_id();
    let mut phase = GenericCheckpointPhase::Prepared;
    let mut provider_events = Vec::new();
    let mut provider_event_digests = BTreeMap::<AgentEventId, Digest>::new();
    let mut commands = BTreeMap::<CommandId, CommandCheckpoint>::new();
    let mut checkpoint_ids = BTreeMap::<GenericCheckpointEventId, Digest>::new();
    let mut last_retry_number = 0_u32;
    let mut observed_prefix = None;
    let mut context_recovery: Option<GenericContextRecovery> = None;
    let mut started_context: Option<(GenericModelContextTrace, Option<u64>)> = None;

    for (index, record) in run.records.iter().enumerate() {
        record.validate()?;
        let expected_seq = index as u64 + 1;
        if record.run_id != *run_id || record.checkpoint_seq != expected_seq {
            return Err(GenericCheckpointError::InvalidData(format!(
                "Generic checkpoint sequence mismatch at {expected_seq}"
            )));
        }
        if let Some(existing) =
            checkpoint_ids.insert(record.event_id.clone(), record.draft_digest.clone())
        {
            return Err(if existing == record.draft_digest {
                GenericCheckpointError::InvalidData(
                    "stored Generic checkpoint contains a duplicate record".to_owned(),
                )
            } else {
                GenericCheckpointError::EventConflict(record.event_id.clone())
            });
        }
        if matches!(phase, GenericCheckpointPhase::Terminal) {
            return Err(GenericCheckpointError::InvalidData(
                "Generic checkpoint contains facts after terminal".to_owned(),
            ));
        }

        match &record.payload {
            GenericCheckpointEvent::ModelContextRejected {
                round,
                request_id,
                retry_number,
                input_budget_tokens,
                output_budget_tokens,
                ..
            } => {
                let GenericCheckpointPhase::ModelAttemptOpen {
                    boundary,
                    round: open_round,
                    request_id: open_request_id,
                    ..
                } = &phase
                else {
                    return Err(GenericCheckpointError::InvalidData(
                        "context rejection must close an open model attempt".to_owned(),
                    ));
                };
                let (trace, rejected_output) = &started_context.as_ref().ok_or_else(|| {
                    GenericCheckpointError::InvalidData("missing rejected context trace".to_owned())
                })?;
                let planned_input = trace
                    .context_estimate
                    .as_ref()
                    .map_or(trace.used_input_tokens, |estimate| estimate.tokens);
                if round != open_round
                    || request_id != open_request_id
                    || context_recovery
                        .as_ref()
                        .map_or(Some(1), |prior| prior.retry_number.checked_add(1))
                        != Some(*retry_number)
                    || *input_budget_tokens > trace.input_budget_tokens
                    || output_budget_tokens
                        .is_some_and(|output| rejected_output.is_none_or(|prior| output > prior))
                    || !(*input_budget_tokens < planned_input
                        || output_budget_tokens.is_some_and(|output| {
                            rejected_output.is_some_and(|prior| output < prior)
                        }))
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "context recovery must advance its rejection count and reduce the rejected input or output budget".to_owned(),
                    ));
                }
                let mut next = boundary.clone();
                next.next_model_round = round.checked_add(1).ok_or_else(|| {
                    GenericCheckpointError::InvalidData(
                        "context recovery round overflow".to_owned(),
                    )
                })?;
                phase = GenericCheckpointPhase::Stable(next);
                context_recovery = Some(GenericContextRecovery {
                    retry_number: *retry_number,
                    input_budget_tokens: *input_budget_tokens,
                    output_budget_tokens: *output_budget_tokens,
                    compact_input: *input_budget_tokens < planned_input
                        || context_recovery
                            .as_ref()
                            .is_some_and(|prior| prior.compact_input),
                });
            }
            GenericCheckpointEvent::ModelRetryScheduled {
                round,
                request_id,
                retry_number,
                observed_usage,
                ..
            } => {
                if !matches!(&phase, GenericCheckpointPhase::ModelAttemptOpen {
                    round: open_round, request_id: open_request_id, ..
                } if round == open_round && request_id == open_request_id)
                    || last_retry_number.checked_add(1) != Some(*retry_number)
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "model retry must advance the retry sequence of its open attempt"
                            .to_owned(),
                    ));
                }
                last_retry_number = *retry_number;
                if let (Some(usage), GenericCheckpointPhase::ModelAttemptOpen { boundary, .. }) =
                    (observed_usage, &mut phase)
                {
                    for (total, observed) in [
                        (&mut boundary.usage.input_tokens, usage.input_tokens),
                        (&mut boundary.usage.output_tokens, usage.output_tokens),
                    ] {
                        if let Some(observed) = observed {
                            *total = Some(total.unwrap_or(0).saturating_add(observed));
                        }
                    }
                }
            }
            GenericCheckpointEvent::LoopBoundaryCommitted {
                next_model_round,
                usage,
                tool_call_count,
                last_response,
                supporting_event_ids,
            } => {
                match &phase {
                    GenericCheckpointPhase::Prepared if *next_model_round == 1 => {}
                    GenericCheckpointPhase::ModelAttemptOpen { round, .. }
                    | GenericCheckpointPhase::ModelAttemptObserved { round, .. }
                    | GenericCheckpointPhase::WorkflowAttemptOpen { round, .. }
                        if *next_model_round > *round => {}
                    _ => {
                        return Err(GenericCheckpointError::InvalidData(
                            "loop boundary does not close the current checkpoint phase".to_owned(),
                        ))
                    }
                }
                phase = GenericCheckpointPhase::Stable(GenericLoopBoundary {
                    next_model_round: *next_model_round,
                    usage: usage.clone(),
                    tool_call_count: *tool_call_count,
                    last_response: last_response.clone(),
                    supporting_event_ids: supporting_event_ids.clone(),
                });
            }
            GenericCheckpointEvent::ModelAttemptStarted {
                round,
                request_id,
                request_digest,
                max_output_tokens,
                context,
            } => {
                let GenericCheckpointPhase::Stable(boundary) = &phase else {
                    return Err(GenericCheckpointError::InvalidData(
                        "model attempt did not begin at a stable loop boundary".to_owned(),
                    ));
                };
                if *round != boundary.next_model_round {
                    return Err(GenericCheckpointError::InvalidData(
                        "model attempt round does not match the stable boundary".to_owned(),
                    ));
                }
                if context_recovery
                    .as_ref()
                    .and_then(GenericContextRecovery::input_capacity_tokens)
                    .is_some_and(|ceiling| context.input_budget_tokens > ceiling)
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "model retry exceeded its durable recovery input budget".to_owned(),
                    ));
                }
                if context.planning.as_ref().is_some_and(|planning| {
                    context.config_digest != run.registration.config_digest
                        || planning
                            .anchor
                            .as_ref()
                            .is_some_and(|anchor| Some(anchor) != observed_prefix.as_ref())
                }) {
                    return Err(GenericCheckpointError::InvalidData(
                        "context anchor does not match an earlier completed request in this Run"
                            .to_owned(),
                    ));
                }
                started_context = Some((context.clone(), *max_output_tokens));
                phase = GenericCheckpointPhase::ModelAttemptOpen {
                    boundary: boundary.clone(),
                    round: *round,
                    request_id: request_id.clone(),
                    request_digest: request_digest.clone(),
                };
                last_retry_number = 0;
            }
            GenericCheckpointEvent::ModelAttemptObserved {
                round,
                request_id,
                observation,
            } => {
                let GenericCheckpointPhase::ModelAttemptOpen {
                    boundary,
                    round: open_round,
                    request_id: open_request_id,
                    request_digest,
                } = &phase
                else {
                    return Err(GenericCheckpointError::InvalidData(
                        "model observation did not close an open attempt".to_owned(),
                    ));
                };
                if round != open_round || request_id != open_request_id {
                    return Err(GenericCheckpointError::InvalidData(
                        "model observation identity does not match its open attempt".to_owned(),
                    ));
                }
                observed_prefix = started_context.as_ref().and_then(|(context, cap)| {
                    context.observed_prefix(run_id, request_id, observation, *cap)
                });
                if let Some(recovery) = &mut context_recovery {
                    recovery.generation_observed();
                }
                phase = GenericCheckpointPhase::ModelAttemptObserved {
                    boundary: boundary.clone(),
                    round: *round,
                    request_id: request_id.clone(),
                    request_digest: request_digest.clone(),
                    observation: observation.clone(),
                };
            }
            GenericCheckpointEvent::WorkflowAttemptStarted {
                round,
                request_id,
                call_id,
                arguments_digest,
            } => {
                let GenericCheckpointPhase::ModelAttemptObserved {
                    boundary,
                    round: observed_round,
                    request_id: observed_request_id,
                    request_digest,
                    observation,
                } = &phase
                else {
                    return Err(GenericCheckpointError::InvalidData(
                        "workflow attempt did not begin from an observed model call".to_owned(),
                    ));
                };
                let matching_call = observation.tool_calls.iter().find(|call| {
                    call.call_id == *call_id
                        && call.name == "orchestral_workflow"
                        && call.ended
                        && Digest::sha256(call.arguments.as_bytes()) == *arguments_digest
                });
                if round != observed_round
                    || request_id != observed_request_id
                    || matching_call.is_none()
                {
                    return Err(GenericCheckpointError::InvalidData(
                        "workflow attempt identity does not match its observed model call"
                            .to_owned(),
                    ));
                }
                phase = GenericCheckpointPhase::WorkflowAttemptOpen {
                    boundary: boundary.clone(),
                    round: *round,
                    request_id: request_id.clone(),
                    request_digest: request_digest.clone(),
                    observation: observation.clone(),
                    call_id: call_id.clone(),
                    arguments_digest: arguments_digest.clone(),
                };
            }
            GenericCheckpointEvent::CommandCommitted {
                command,
                outcome,
                approval_capability,
            } => {
                if let Some(existing) = commands.get(&command.command_id) {
                    if existing.command != *command
                        || existing.outcome != *outcome
                        || existing.approval_capability != *approval_capability
                    {
                        return Err(GenericCheckpointError::InvalidData(
                            "command identity was reused with different checkpoint content"
                                .to_owned(),
                        ));
                    }
                    return Err(GenericCheckpointError::InvalidData(
                        "stored Generic checkpoint contains a duplicate command".to_owned(),
                    ));
                }
                commands.insert(
                    command.command_id.clone(),
                    CommandCheckpoint {
                        command: command.clone(),
                        outcome: outcome.clone(),
                        approval_capability: approval_capability.clone(),
                    },
                );
            }
            GenericCheckpointEvent::ProviderEventsCommitted { events } => {
                for event in events {
                    let digest = event.computed_digest().map_err(invalid_data)?;
                    if let Some(existing) = provider_event_digests.get(&event.event_id) {
                        if existing != &digest {
                            return Err(GenericCheckpointError::InvalidData(
                                "Provider event identity was reused with different content"
                                    .to_owned(),
                            ));
                        }
                        return Err(GenericCheckpointError::InvalidData(
                            "stored Generic checkpoint contains a duplicate Provider event"
                                .to_owned(),
                        ));
                    }
                    provider_event_digests.insert(event.event_id.clone(), digest);
                    provider_events.push(event.clone());
                    if is_terminal_event(&event.payload) {
                        phase = GenericCheckpointPhase::Terminal;
                    }
                }
            }
        }
    }

    Ok(GenericAgentCheckpointProjection {
        phase,
        provider_events,
        commands,
        last_checkpoint_seq: run.last_checkpoint_seq(),
        observed_prefix,
        context_recovery,
    })
}

fn is_terminal_event(event: &AgentEvent) -> bool {
    matches!(
        event,
        AgentEvent::DeliveryCommitted { .. }
            | AgentEvent::RunIncomplete { .. }
            | AgentEvent::RunFailed { .. }
            | AgentEvent::RunCancelled { .. }
    )
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CreateGenericRunOutcome {
    Created,
    ExactExisting,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AppendGenericCheckpointOutcome {
    Appended,
    ExactDuplicate,
}

#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum GenericCheckpointError {
    #[error("Generic Agent checkpoint storage is unavailable: {0}")]
    Unavailable(String),
    #[error("Generic Agent checkpoint Run does not exist: {0}")]
    RunNotFound(RunId),
    #[error("Generic Agent checkpoint Run conflicts with durable state: {0}")]
    RunConflict(RunId),
    #[error(
        "Generic Agent checkpoint sequence conflict for {run_id}: expected previous {expected_previous}, durable previous {actual_previous}"
    )]
    SequenceConflict {
        run_id: RunId,
        expected_previous: u64,
        actual_previous: u64,
    },
    #[error("Generic Agent checkpoint event identity conflict: {0}")]
    EventConflict(GenericCheckpointEventId),
    #[error("Generic Agent checkpoint data is invalid: {0}")]
    InvalidData(String),
}

/// Synchronous metadata WAL used on the Generic Agent's existing commit
/// critical sections. Implementations must keep records small and must not do
/// heavyweight parsing, embedding, or remote work in this call path.
pub trait GenericAgentCheckpointStore: Send + Sync {
    fn load_run(
        &self,
        run_id: &RunId,
    ) -> Result<Option<StoredGenericAgentRun>, GenericCheckpointError>;

    fn create_run(
        &self,
        registration: GenericAgentRunRegistration,
    ) -> Result<CreateGenericRunOutcome, GenericCheckpointError>;

    fn append(
        &self,
        run_id: &RunId,
        expected_previous: u64,
        draft: GenericCheckpointDraft,
    ) -> Result<AppendGenericCheckpointOutcome, GenericCheckpointError>;
}

#[derive(Default)]
pub struct InMemoryGenericAgentCheckpointStore {
    runs: RwLock<BTreeMap<RunId, StoredGenericAgentRun>>,
}

impl GenericAgentCheckpointStore for InMemoryGenericAgentCheckpointStore {
    fn load_run(
        &self,
        run_id: &RunId,
    ) -> Result<Option<StoredGenericAgentRun>, GenericCheckpointError> {
        let run = self
            .runs
            .read()
            .map_err(|_| GenericCheckpointError::Unavailable("reader lock poisoned".to_owned()))?
            .get(run_id)
            .cloned();
        if let Some(run) = &run {
            run.validate()?;
        }
        Ok(run)
    }

    fn create_run(
        &self,
        registration: GenericAgentRunRegistration,
    ) -> Result<CreateGenericRunOutcome, GenericCheckpointError> {
        registration.validate()?;
        let run_id = registration.run_id().clone();
        let mut runs = self
            .runs
            .write()
            .map_err(|_| GenericCheckpointError::Unavailable("writer lock poisoned".to_owned()))?;
        if let Some(existing) = runs.get(&run_id) {
            return if existing.registration == registration {
                Ok(CreateGenericRunOutcome::ExactExisting)
            } else {
                Err(GenericCheckpointError::RunConflict(run_id))
            };
        }
        runs.insert(
            run_id,
            StoredGenericAgentRun {
                registration,
                records: Vec::new(),
            },
        );
        Ok(CreateGenericRunOutcome::Created)
    }

    fn append(
        &self,
        run_id: &RunId,
        expected_previous: u64,
        draft: GenericCheckpointDraft,
    ) -> Result<AppendGenericCheckpointOutcome, GenericCheckpointError> {
        draft.validate()?;
        if draft.run_id != *run_id {
            return Err(GenericCheckpointError::InvalidData(
                "checkpoint append crossed a Run boundary".to_owned(),
            ));
        }
        let draft_digest = draft.digest()?;
        let mut runs = self
            .runs
            .write()
            .map_err(|_| GenericCheckpointError::Unavailable("writer lock poisoned".to_owned()))?;
        let run = runs
            .get_mut(run_id)
            .ok_or_else(|| GenericCheckpointError::RunNotFound(run_id.clone()))?;
        if let Some(existing) = run
            .records
            .iter()
            .find(|record| record.event_id == draft.event_id)
        {
            return if existing.draft_digest == draft_digest {
                Ok(AppendGenericCheckpointOutcome::ExactDuplicate)
            } else {
                Err(GenericCheckpointError::EventConflict(draft.event_id))
            };
        }
        let actual_previous = run.last_checkpoint_seq();
        if actual_previous != expected_previous {
            return Err(GenericCheckpointError::SequenceConflict {
                run_id: run_id.clone(),
                expected_previous,
                actual_previous,
            });
        }
        let record = GenericCheckpointRecord::seal(draft, actual_previous + 1)?;
        let mut candidate = run.clone();
        candidate.records.push(record.clone());
        candidate.validate()?;
        run.records.push(record);
        Ok(AppendGenericCheckpointOutcome::Appended)
    }
}

fn canonical_digest(value: &impl Serialize) -> Result<Digest, GenericCheckpointError> {
    serde_jcs::to_vec(value)
        .map(Digest::sha256)
        .map_err(|error| GenericCheckpointError::InvalidData(error.to_string()))
}

fn invalid_data(error: impl fmt::Display) -> GenericCheckpointError {
    GenericCheckpointError::InvalidData(error.to_string())
}

#[cfg(test)]
mod tests {
    use orchestral_core::agent_protocol::wire::{
        AgentDescriptor, AgentDescriptorEnvelope, AgentId, AgentProviderId, AgentRunEnvelope,
        AgentSessionId, Content, ProviderBindingRef,
    };
    use orchestral_core::agent_protocol::AGENT_PROTOCOL_V1;

    use super::*;

    fn registration() -> GenericAgentRunRegistration {
        let descriptor = AgentDescriptorEnvelope::seal(AgentDescriptor {
            provider_id: AgentProviderId::new("test/generic"),
            agent_id: AgentId::new("generic-v1"),
            supported_protocol_versions: vec![AGENT_PROTOCOL_V1],
            accepted_content_types: BTreeSet::from(["text/plain".to_owned()]),
            capabilities: Default::default(),
            extensions: Default::default(),
        })
        .unwrap();
        let run = AgentRunEnvelope::new(
            AGENT_PROTOCOL_V1,
            AgentSessionId::new("session-1"),
            RunId::new("run-1"),
            vec![Content::text("hello")],
        )
        .unwrap();
        let request =
            AgentStartRequest::new(run, ProviderBindingRef::new("binding-1"), &descriptor).unwrap();
        GenericAgentRunRegistration {
            execution: AgentExecutionRef::for_start(&request, &descriptor).unwrap(),
            request,
            admission: AgentAdmission::default(),
            config_digest: Digest::sha256("config-v1"),
        }
    }

    fn boundary(run_id: &RunId, next_model_round: u64) -> GenericCheckpointDraft {
        GenericCheckpointDraft {
            event_id: GenericCheckpointEventId::new(format!("boundary-{next_model_round}")),
            run_id: run_id.clone(),
            payload: GenericCheckpointEvent::LoopBoundaryCommitted {
                next_model_round,
                usage: ModelUsage::default(),
                tool_call_count: 0,
                last_response: String::new(),
                supporting_event_ids: Vec::new(),
            },
        }
    }

    fn context_trace() -> GenericModelContextTrace {
        GenericModelContextTrace {
            through_session_seq: 1,
            included_ranges: vec![SessionSourceRange {
                first_session_seq: 1,
                last_session_seq: 1,
            }],
            deferred_ranges: Vec::new(),
            config_digest: Digest::sha256("config-v1"),
            history_limit: 128,
            used_input_tokens: 10,
            context_estimate: None,
            planning: None,
            input_budget_tokens: 100,
        }
    }

    #[test]
    fn model_context_trace_rejects_overlapping_or_over_budget_provenance() {
        let mut trace = context_trace();
        trace.deferred_ranges = trace.included_ranges.clone();
        assert!(trace.validate().is_err());

        let mut trace = context_trace();
        trace.used_input_tokens = trace.input_budget_tokens + 1;
        assert!(trace.validate().is_err());
    }

    #[test]
    fn model_context_trace_keeps_legacy_bounds_and_explicit_estimates_distinct() {
        use orchestral_core::model_protocol::{ModelContextEstimate, ModelTokenAccounting};
        let legacy = context_trace();
        let serialized = serde_json::to_value(&legacy).unwrap();
        assert!(serialized.get("context_estimate").is_none());
        assert!(serialized.get("planning").is_none());
        let restored: GenericModelContextTrace =
            serde_json::from_value(serialized.clone()).unwrap();
        assert_eq!(
            serde_jcs::to_vec(&restored).unwrap(),
            serde_jcs::to_vec(&serialized).unwrap()
        );
        restored.validate().unwrap();
        assert_eq!(restored, legacy);

        let mut trace = legacy;
        trace.used_input_tokens = 900;
        trace.context_estimate = Some(ModelContextEstimate {
            tokens: 80,
            accounting: ModelTokenAccounting::Estimated,
        });
        trace.validate().unwrap();
        let restored: GenericModelContextTrace =
            serde_json::from_slice(&serde_json::to_vec(&trace).unwrap()).unwrap();
        assert_eq!(restored, trace);
        for (tokens, accounting) in [
            (101, ModelTokenAccounting::Estimated),
            (901, ModelTokenAccounting::Estimated),
            (80, ModelTokenAccounting::Exact),
            (80, ModelTokenAccounting::ConservativeUpperBound),
        ] {
            trace.context_estimate = Some(ModelContextEstimate { tokens, accounting });
            assert!(trace.validate().is_err());
        }
        trace.context_estimate = None;
        assert!(trace.validate().is_err());
    }

    #[test]
    fn observed_prefix_requires_positive_complete_in_bound_input_usage() {
        use crate::session_context::observed_prefix::ContextInputSignature;
        use orchestral_core::model_protocol::{ModelContextEstimate, ModelTokenAccounting};
        let mut trace = context_trace();
        trace.context_estimate = Some(ModelContextEstimate {
            tokens: 8,
            accounting: ModelTokenAccounting::Estimated,
        });
        trace.planning = Some(ContextPlanningTrace {
            input: ContextInputSignature {
                messages_len: 2,
                messages_digest: Digest::sha256("messages"),
                tools_digest: Digest::sha256("tools"),
                raw_estimate_tokens: 8,
            },
            anchor: None,
        });
        let observation = GenericModelObservation {
            finish_reason: ModelFinishReason::Stop,
            response: "complete".to_owned(),
            continuation: Default::default(),
            tool_calls: Vec::new(),
            usage: Some(ModelUsage {
                input_tokens: Some(5),
                output_tokens: Some(2),
            }),
        };
        let derive = |observation: &GenericModelObservation| {
            trace.observed_prefix(
                &RunId::new("run"),
                &ModelRequestId::new("request"),
                observation,
                Some(4),
            )
        };
        assert_eq!(derive(&observation).unwrap().observed_input_tokens, 5);
        for reason in [
            ModelFinishReason::Length,
            ModelFinishReason::Cancelled,
            ModelFinishReason::ContentFilter,
            ModelFinishReason::Other,
        ] {
            let mut rejected = observation.clone();
            rejected.finish_reason = reason;
            assert!(derive(&rejected).is_none());
        }
        for input in [None, Some(0), Some(11)] {
            let mut rejected = observation.clone();
            rejected.usage.as_mut().unwrap().input_tokens = input;
            assert!(derive(&rejected).is_none());
        }
        let mut rejected = observation.clone();
        rejected.usage = None;
        assert!(derive(&rejected).is_none());
        let mut rejected = observation.clone();
        rejected.usage.as_mut().unwrap().output_tokens = Some(5);
        assert!(derive(&rejected).is_none());
        let mut rejected = observation;
        rejected.tool_calls.push(GenericObservedToolCall {
            call_id: ModelToolCallId::new("half-call"),
            name: "inspect".to_owned(),
            arguments: "{}".to_owned(),
            extensions: Default::default(),
            ended: false,
        });
        assert!(derive(&rejected).is_none());
    }

    #[test]
    fn retry_checkpoints_must_match_an_open_attempt_and_advance_in_order() {
        use orchestral_core::model_protocol::{ModelError, ModelErrorCode};
        let store = InMemoryGenericAgentCheckpointStore::default();
        let registration = registration();
        let run_id = registration.run_id().clone();
        store.create_run(registration).unwrap();
        store.append(&run_id, 0, boundary(&run_id, 1)).unwrap();
        let retry = |number, request_id: &str| GenericCheckpointDraft {
            event_id: GenericCheckpointEventId::new(format!("retry-{number}-{request_id}")),
            run_id: run_id.clone(),
            payload: GenericCheckpointEvent::ModelRetryScheduled {
                round: 1,
                request_id: ModelRequestId::new(request_id),
                retry_number: number,
                delay_ms: 1,
                error: ModelError::new(ModelErrorCode::Unavailable, "temporary")
                    .with_retryable(true),
                observed_usage: None,
            },
        };
        // Omitting the new optional field preserves the serialized payload and
        // digest of checkpoints written before usage-only retries existed.
        let old_payload = serde_json::json!({
            "type": "model_retry_scheduled",
            "round": 1,
            "request_id": "model-1",
            "retry_number": 1,
            "delay_ms": 1,
            "error": {
                "code": "unavailable", "message": "temporary",
                "retryable": true, "details": null,
            },
        });
        let decoded: GenericCheckpointEvent = serde_json::from_value(old_payload.clone()).unwrap();
        assert_eq!(serde_json::to_value(&decoded).unwrap(), old_payload);
        assert_eq!(decoded, retry(1, "model-1").payload);
        assert!(store.append(&run_id, 1, retry(1, "model-1")).is_err());
        store
            .append(
                &run_id,
                1,
                GenericCheckpointDraft {
                    event_id: GenericCheckpointEventId::new("attempt-1"),
                    run_id: run_id.clone(),
                    payload: GenericCheckpointEvent::ModelAttemptStarted {
                        round: 1,
                        request_id: ModelRequestId::new("model-1"),
                        request_digest: Digest::sha256("request"),
                        max_output_tokens: None,
                        context: context_trace(),
                    },
                },
            )
            .unwrap();
        assert!(store.append(&run_id, 2, retry(2, "model-1")).is_err());
        assert!(store
            .append(&run_id, 2, retry(1, "different-model"))
            .is_err());
        store.append(&run_id, 2, retry(1, "model-1")).unwrap();
        assert!(matches!(
            store
                .load_run(&run_id)
                .unwrap()
                .unwrap()
                .validate()
                .unwrap()
                .phase,
            GenericCheckpointPhase::ModelAttemptOpen { .. }
        ));
        store
            .append(
                &run_id,
                3,
                GenericCheckpointDraft {
                    event_id: GenericCheckpointEventId::new("observed-1"),
                    run_id: run_id.clone(),
                    payload: GenericCheckpointEvent::ModelAttemptObserved {
                        round: 1,
                        request_id: ModelRequestId::new("model-1"),
                        observation: GenericModelObservation {
                            finish_reason: ModelFinishReason::Stop,
                            response: "done".to_owned(),
                            continuation: BTreeMap::new(),
                            usage: None,
                            tool_calls: vec![],
                        },
                    },
                },
            )
            .unwrap();
        assert!(store.append(&run_id, 4, retry(2, "model-1")).is_err());
    }

    #[test]
    fn stable_open_and_observed_model_boundaries_are_distinguishable_after_replay() {
        let store = InMemoryGenericAgentCheckpointStore::default();
        let registration = registration();
        let run_id = registration.run_id().clone();
        store.create_run(registration).unwrap();
        store.append(&run_id, 0, boundary(&run_id, 1)).unwrap();
        let stable = store.load_run(&run_id).unwrap().unwrap();
        assert!(stable.validate().unwrap().phase.is_stable());

        store
            .append(
                &run_id,
                1,
                GenericCheckpointDraft {
                    event_id: GenericCheckpointEventId::new("attempt-1"),
                    run_id: run_id.clone(),
                    payload: GenericCheckpointEvent::ModelAttemptStarted {
                        round: 1,
                        request_id: ModelRequestId::new("model-run-1-1"),
                        request_digest: Digest::sha256("request-1"),
                        max_output_tokens: None,
                        context: context_trace(),
                    },
                },
            )
            .unwrap();
        let uncertain = store.load_run(&run_id).unwrap().unwrap();
        assert!(matches!(
            uncertain.validate().unwrap().phase,
            GenericCheckpointPhase::ModelAttemptOpen { .. }
        ));

        store
            .append(
                &run_id,
                2,
                GenericCheckpointDraft {
                    event_id: GenericCheckpointEventId::new("observed-1"),
                    run_id: run_id.clone(),
                    payload: GenericCheckpointEvent::ModelAttemptObserved {
                        round: 1,
                        request_id: ModelRequestId::new("model-run-1-1"),
                        observation: GenericModelObservation {
                            finish_reason: ModelFinishReason::ToolCalls,
                            response: "calling a Tool".to_owned(),
                            continuation: BTreeMap::from([(
                                "fixture/native".to_owned(),
                                serde_json::json!({"opaque": "Ω\n"}),
                            )]),
                            usage: Some(ModelUsage {
                                input_tokens: Some(10),
                                output_tokens: Some(5),
                            }),
                            tool_calls: vec![GenericObservedToolCall {
                                call_id: ModelToolCallId::new("call-1"),
                                name: "echo".to_owned(),
                                arguments: r#"{"value":"hello"}"#.to_owned(),
                                extensions: Default::default(),
                                ended: true,
                            }],
                        },
                    },
                },
            )
            .unwrap();
        let observed = store.load_run(&run_id).unwrap().unwrap();
        let persisted = serde_json::to_vec(&observed).unwrap();
        let restored: StoredGenericAgentRun = serde_json::from_slice(&persisted).unwrap();
        let GenericCheckpointPhase::ModelAttemptObserved { observation, .. } =
            restored.validate().unwrap().phase
        else {
            panic!("restored terminal model observation");
        };
        assert_eq!(
            observation.continuation["fixture/native"],
            serde_json::json!({"opaque": "Ω\n"})
        );
        assert!(
            matches!(&observation.assistant_content()[1], ModelContent::Continuation { namespace, .. } if namespace == "fixture/native")
        );
        assert!(matches!(
            observed.validate().unwrap().phase,
            GenericCheckpointPhase::ModelAttemptObserved {
                round: 1,
                observation: GenericModelObservation { ref tool_calls, .. },
                ..
            } if tool_calls.len() == 1
        ));

        store.append(&run_id, 3, boundary(&run_id, 2)).unwrap();
        assert!(matches!(
            store
                .load_run(&run_id)
                .unwrap()
                .unwrap()
                .validate()
                .unwrap()
                .phase,
            GenericCheckpointPhase::Stable(GenericLoopBoundary {
                next_model_round: 2,
                ..
            })
        ));
    }

    #[test]
    fn invalid_transition_and_event_equivocation_never_advance_the_wal() {
        let store = InMemoryGenericAgentCheckpointStore::default();
        let registration = registration();
        let run_id = registration.run_id().clone();
        store.create_run(registration).unwrap();
        assert!(store.append(&run_id, 0, boundary(&run_id, 2)).is_err());
        assert_eq!(
            store
                .load_run(&run_id)
                .unwrap()
                .unwrap()
                .last_checkpoint_seq(),
            0
        );

        let original = boundary(&run_id, 1);
        store.append(&run_id, 0, original.clone()).unwrap();
        assert_eq!(
            store.append(&run_id, 1, original).unwrap(),
            AppendGenericCheckpointOutcome::ExactDuplicate
        );
        let mut conflict = boundary(&run_id, 2);
        conflict.event_id = GenericCheckpointEventId::new("boundary-1");
        assert!(matches!(
            store.append(&run_id, 1, conflict),
            Err(GenericCheckpointError::EventConflict(_))
        ));
        assert_eq!(
            store
                .load_run(&run_id)
                .unwrap()
                .unwrap()
                .last_checkpoint_seq(),
            1
        );
    }
}