vv_agent/runtime/state/

validation.rs

use std::collections::HashSet;

use crate::budget::BudgetEnforcementBoundary;
use crate::events::{ModelCallFailureOutcome, RunEventPayload};
use crate::types::{AgentResult, AgentStatus, ModelCallStatus};

use super::*;

pub fn validate_checkpoint(checkpoint: &Checkpoint) -> CheckpointResult<()> {
    if checkpoint.schema_version != CHECKPOINT_SCHEMA {
        return Err(CheckpointError::new(
            "checkpoint_schema_unsupported",
            "checkpoint schema_version is unsupported",
        ));
    }
    if checkpoint.run_definition_schema != RUN_DEFINITION_SCHEMA {
        return Err(CheckpointError::new(
            "checkpoint_definition_schema_unsupported",
            "run_definition_schema is missing or unsupported",
        ));
    }
    crate::checkpoint::validate_run_definition(&checkpoint.run_definition)?;
    validate_sha256(&checkpoint.run_definition_digest, "run_definition_digest").map_err(
        |error| CheckpointError::new("checkpoint_definition_digest_invalid", error.message()),
    )?;
    let digest = crate::checkpoint::run_definition_digest(&checkpoint.run_definition)?;
    if digest != checkpoint.run_definition_digest {
        return Err(CheckpointError::new(
            "checkpoint_definition_mismatch",
            "run_definition_digest does not match embedded run_definition",
        ));
    }
    validate_checkpoint_key(&checkpoint.checkpoint_key)?;
    for (value, field_name) in [
        (&checkpoint.task_id, "task_id"),
        (&checkpoint.root_run_id, "root_run_id"),
        (&checkpoint.trace_id, "trace_id"),
    ] {
        if value.trim().is_empty() {
            return Err(CheckpointError::new(
                "checkpoint_value_invalid",
                format!("{field_name} must be non-empty"),
            ));
        }
    }
    if checkpoint.resume_attempt == 0 || checkpoint.resume_attempt > MAX_WIRE_INTEGER {
        return Err(CheckpointError::new(
            "checkpoint_resume_attempt_invalid",
            "resume_attempt must be positive and JSON-safe",
        ));
    }
    if checkpoint.cycle_index > MAX_WIRE_INTEGER || checkpoint.revision > MAX_WIRE_INTEGER {
        return Err(CheckpointError::new(
            "checkpoint_integer_invalid",
            "checkpoint integer is outside the JSON-safe range",
        ));
    }
    let claim_values = [
        checkpoint.claim_token.is_some(),
        checkpoint.claimed_cycle.is_some(),
        checkpoint.lease_expires_at_ms.is_some(),
    ];
    if claim_values.iter().any(|value| *value) && claim_values.iter().any(|value| !*value) {
        return Err(CheckpointError::new(
            "checkpoint_claim_invalid",
            "claim fields must be all present or all null",
        ));
    }
    if let Some(claim_token) = &checkpoint.claim_token {
        if claim_token.trim().is_empty() {
            return Err(CheckpointError::new(
                "checkpoint_claim_invalid",
                "claim_token must be non-empty",
            ));
        }
        let claimed_cycle = checkpoint.claimed_cycle.expect("claim tuple checked");
        let expected = checkpoint.cycle_index.checked_add(1).ok_or_else(|| {
            CheckpointError::new("checkpoint_claim_invalid", "claimed cycle overflow")
        })?;
        if claimed_cycle != expected || claimed_cycle == 0 || claimed_cycle > MAX_WIRE_INTEGER {
            return Err(CheckpointError::new(
                "checkpoint_claim_invalid",
                "claimed_cycle must equal cycle_index + 1",
            ));
        }
        let lease = checkpoint.lease_expires_at_ms.expect("claim tuple checked");
        if lease > MAX_WIRE_INTEGER {
            return Err(CheckpointError::new(
                "checkpoint_claim_invalid",
                "lease expiry is outside the JSON-safe range",
            ));
        }
    }
    if checkpoint.terminal_result.is_some() && checkpoint.claim_token.is_some() {
        return Err(CheckpointError::new(
            "checkpoint_status_invalid",
            "terminal checkpoint cannot have an active claim",
        ));
    }
    if checkpoint.terminal_acknowledged && checkpoint.terminal_result.is_none() {
        return Err(CheckpointError::new(
            "checkpoint_status_invalid",
            "terminal acknowledgement requires a terminal result",
        ));
    }
    if checkpoint.terminal_result.is_none()
        && !matches!(
            checkpoint.status,
            CheckpointStatus::Running | CheckpointStatus::ReconciliationRequired
        )
    {
        return Err(CheckpointError::new(
            "checkpoint_status_invalid",
            "non-terminal checkpoint must be running or reconciliation_required",
        ));
    }
    if checkpoint.terminal_result.is_some() && !checkpoint.status.is_terminal() {
        return Err(CheckpointError::new(
            "checkpoint_status_invalid",
            "terminal_result requires a terminal checkpoint status",
        ));
    }
    let call_ids = checkpoint
        .model_calls
        .iter()
        .map(|record| record.call_id.as_str())
        .collect::<HashSet<_>>();
    if call_ids.len() != checkpoint.model_calls.len() {
        return Err(CheckpointError::new(
            "checkpoint_status_invalid",
            "checkpoint model_calls contains duplicate call ids",
        ));
    }

    let active_cycle = checkpoint.active_cycle()?;
    for entry in checkpoint
        .model_call_journal
        .iter()
        .chain(checkpoint.tool_journal.iter())
    {
        entry.validate()?;
        if entry.cycle_index != active_cycle {
            return Err(CheckpointError::new(
                "checkpoint_journal_cycle_invalid",
                "journal cycle_index must equal the active cycle",
            ));
        }
    }
    if checkpoint
        .model_call_journal
        .iter()
        .any(|entry| entry.kind != OperationKind::Model)
        || checkpoint
            .tool_journal
            .iter()
            .any(|entry| entry.kind != OperationKind::Tool)
    {
        return Err(CheckpointError::new(
            "checkpoint_journal_kind_invalid",
            "journal arrays contain an entry of the wrong kind",
        ));
    }
    for (namespace, entry) in &checkpoint.extension_state {
        validate_extension_namespace(namespace)?;
        entry.validate()?;
    }
    validate_extension_state_size(&checkpoint.extension_state, MAX_WIRE_INTEGER)?;
    if let Some(cursor) = &checkpoint.event_cursor {
        cursor.validate()?;
    }
    let mut event_ids = HashSet::new();
    for entry in &checkpoint.event_outbox {
        entry.verify_payload()?;
        if !event_ids.insert(entry.event_id.as_str()) {
            return Err(CheckpointError::new(
                "event_identity_conflict",
                "checkpoint event_outbox contains a duplicate event id",
            ));
        }
    }
    validate_model_journal_accounting(checkpoint)?;
    for value in checkpoint.shared_state.values() {
        validate_json(value, "shared_state")?;
    }
    if checkpoint.status == CheckpointStatus::ReconciliationRequired
        && (!checkpoint.has_ambiguous_operation() || checkpoint.claim_token.is_some())
    {
        return Err(CheckpointError::new(
            "checkpoint_status_invalid",
            "reconciliation_required needs an ambiguous journal and no claim",
        ));
    }
    if checkpoint.status == CheckpointStatus::Running
        && checkpoint.has_ambiguous_operation()
        && checkpoint.claim_token.is_none()
    {
        return Err(CheckpointError::new(
            "checkpoint_status_invalid",
            "running checkpoint with ambiguity needs an active recovery claim",
        ));
    }
    if checkpoint.terminal_result.is_some()
        && (!checkpoint.model_call_journal.is_empty() || !checkpoint.tool_journal.is_empty())
        && !checkpoint.is_operator_abort_terminal()
    {
        return Err(CheckpointError::new(
            "checkpoint_status_invalid",
            "terminal checkpoint cannot retain active journals",
        ));
    }
    if let Some(result) = &checkpoint.terminal_result {
        validate_json(result, "terminal_result")?;
        let result = AgentResult::from_dict(result).map_err(|error| {
            CheckpointError::new(
                "checkpoint_status_invalid",
                format!("terminal_result is not the current AgentResult shape: {error}"),
            )
        })?;
        if !agent_status_matches_checkpoint(result.status, checkpoint.status) {
            return Err(CheckpointError::new(
                "checkpoint_status_invalid",
                "terminal result status must match checkpoint status",
            ));
        }
        if result
            .checkpoint_key
            .as_deref()
            .is_some_and(|key| key != checkpoint.checkpoint_key)
        {
            return Err(CheckpointError::new(
                "checkpoint_status_invalid",
                "terminal result checkpoint_key must match checkpoint",
            ));
        }
        if result.token_usage.model_calls != checkpoint.model_calls {
            return Err(CheckpointError::new(
                "checkpoint_status_invalid",
                "terminal result model-call ledger does not match checkpoint",
            ));
        }
    }
    Ok(())
}

pub fn validate_model_journal_accounting(checkpoint: &Checkpoint) -> CheckpointResult<()> {
    for journal in &checkpoint.model_call_journal {
        if journal.kind != OperationKind::Model {
            return Err(CheckpointError::new(
                "operation_kind_fields_invalid",
                "model_call_journal contains a non-model entry",
            ));
        }
        validate_model_journal_entry_accounting(checkpoint, journal)?;
    }
    Ok(())
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct ModelAccountingIdentity {
    call_id: String,
    operation_id: String,
    attempt: u64,
    operation: ModelCallOperation,
    cycle_index: u64,
    backend: String,
    model: String,
}

fn validate_model_journal_entry_accounting(
    checkpoint: &Checkpoint,
    journal: &OperationJournalEntry,
) -> CheckpointResult<()> {
    journal.validate()?;
    let identity = model_journal_identity(journal)?;
    let record_candidates = checkpoint
        .model_calls
        .iter()
        .filter(|record| {
            record.call_id == identity.call_id
                || (record.operation_id == identity.operation_id
                    && u64::from(record.attempt) == identity.attempt)
        })
        .collect::<Vec<_>>();

    let mut event_candidates = Vec::new();
    for (index, entry) in checkpoint.event_outbox.iter().enumerate() {
        entry.verify_payload()?;
        let event: RunEvent = serde_json::from_value(entry.event.clone()).map_err(|error| {
            CheckpointError::new(
                "checkpoint_event_outbox_invalid",
                format!("checkpoint event payload is invalid: {error}"),
            )
        })?;
        let Some(event_identity) = model_event_identity(&event) else {
            continue;
        };
        if event_identity.call_id == identity.call_id
            || (event_identity.operation_id == identity.operation_id
                && event_identity.attempt == identity.attempt)
        {
            event_candidates.push((index, event, event_identity));
        }
    }
    let started_events = event_candidates
        .iter()
        .filter(|(_, event, _)| matches!(event.payload(), RunEventPayload::ModelCallStarted { .. }))
        .collect::<Vec<_>>();
    let terminal_events = event_candidates
        .iter()
        .filter(|(_, event, _)| {
            matches!(
                event.payload(),
                RunEventPayload::ModelCallCompleted { .. }
                    | RunEventPayload::ModelCallFailed { .. }
            )
        })
        .collect::<Vec<_>>();

    if record_candidates.len() > 1 || started_events.len() > 1 || terminal_events.len() > 1 {
        return Err(model_accounting_error(
            "model journal attempt has duplicate accounting evidence",
        ));
    }

    match journal.state {
        OperationState::Planned => {
            require_model_evidence_counts(
                &record_candidates,
                &started_events,
                &terminal_events,
                (0, 0, 0),
            )?;
            return Ok(());
        }
        OperationState::Started => {
            require_model_evidence_counts(
                &record_candidates,
                &started_events,
                &terminal_events,
                (0, 1, 0),
            )?;
            require_model_identity(&identity, &started_events[0].2)?;
            return Ok(());
        }
        OperationState::Failed
            if record_candidates.is_empty()
                && started_events.is_empty()
                && terminal_events.is_empty() =>
        {
            return Ok(());
        }
        OperationState::Succeeded | OperationState::Failed | OperationState::Ambiguous => {}
    }

    require_model_evidence_counts(
        &record_candidates,
        &started_events,
        &terminal_events,
        (1, 1, 1),
    )?;
    let record = record_candidates[0];
    let started_event = started_events[0];
    let terminal_event = terminal_events[0];
    require_model_identity(&identity, &model_record_identity(record))?;
    require_model_identity(&identity, &started_event.2)?;
    require_model_identity(&identity, &terminal_event.2)?;

    let status_matches = match journal.state {
        OperationState::Succeeded => matches!(
            record.status,
            ModelCallStatus::Completed | ModelCallStatus::Ambiguous
        ),
        OperationState::Failed => matches!(
            record.status,
            ModelCallStatus::Failed | ModelCallStatus::Ambiguous
        ),
        OperationState::Ambiguous => record.status == ModelCallStatus::Ambiguous,
        OperationState::Planned | OperationState::Started => false,
    };
    let event_type_matches = matches!(
        (record.status, terminal_event.1.payload()),
        (
            ModelCallStatus::Completed,
            RunEventPayload::ModelCallCompleted { .. }
        ) | (
            ModelCallStatus::Failed | ModelCallStatus::Ambiguous,
            RunEventPayload::ModelCallFailed { .. }
        )
    );
    if !status_matches || !event_type_matches {
        return Err(model_accounting_error(
            "model journal terminal state does not match its accounting evidence",
        ));
    }

    let event_usage = match terminal_event.1.payload() {
        RunEventPayload::ModelCallCompleted { usage, .. }
        | RunEventPayload::ModelCallFailed { usage, .. } => usage,
        _ => unreachable!("terminal event filtered above"),
    };
    if event_usage != &record.usage {
        return Err(model_accounting_error(
            "model terminal event usage does not match its ledger record",
        ));
    }
    if let RunEventPayload::ModelCallFailed {
        outcome,
        error_code,
        ..
    } = terminal_event.1.payload()
    {
        let expected_outcome = if record.status == ModelCallStatus::Ambiguous {
            ModelCallFailureOutcome::Ambiguous
        } else {
            ModelCallFailureOutcome::Definitive
        };
        if *outcome != expected_outcome || record.error_code.as_deref() != Some(error_code.as_str())
        {
            return Err(model_accounting_error(
                "model failed event does not match its ledger record",
            ));
        }
    }

    validate_terminal_budget_event_order(checkpoint, terminal_event.0)?;
    Ok(())
}

fn require_model_evidence_counts(
    records: &[&ModelCallRecord],
    started_events: &[&(usize, RunEvent, ModelAccountingIdentity)],
    terminal_events: &[&(usize, RunEvent, ModelAccountingIdentity)],
    expected: (usize, usize, usize),
) -> CheckpointResult<()> {
    if (records.len(), started_events.len(), terminal_events.len()) != expected {
        return Err(model_accounting_error(
            "model journal attempt is missing atomic accounting evidence",
        ));
    }
    Ok(())
}

fn model_journal_identity(
    journal: &OperationJournalEntry,
) -> CheckpointResult<ModelAccountingIdentity> {
    Ok(ModelAccountingIdentity {
        call_id: journal.call_id.clone().ok_or_else(|| {
            model_accounting_error("model journal call_id is missing from accounting identity")
        })?,
        operation_id: journal.operation_id.clone(),
        attempt: journal.attempt,
        operation: journal.model_operation.ok_or_else(|| {
            model_accounting_error("model journal operation is missing from accounting identity")
        })?,
        cycle_index: journal.cycle_index,
        backend: journal.backend.clone().ok_or_else(|| {
            model_accounting_error("model journal backend is missing from accounting identity")
        })?,
        model: journal.model.clone().ok_or_else(|| {
            model_accounting_error("model journal model is missing from accounting identity")
        })?,
    })
}

fn model_record_identity(record: &ModelCallRecord) -> ModelAccountingIdentity {
    ModelAccountingIdentity {
        call_id: record.call_id.clone(),
        operation_id: record.operation_id.clone(),
        attempt: u64::from(record.attempt),
        operation: record.operation,
        cycle_index: u64::from(record.cycle_index),
        backend: record.backend.clone(),
        model: record.model.clone(),
    }
}

fn model_event_identity(event: &RunEvent) -> Option<ModelAccountingIdentity> {
    let (call_id, operation_id, attempt, operation, backend, model) = match event.payload() {
        RunEventPayload::ModelCallStarted {
            call_id,
            operation_id,
            attempt,
            operation,
            backend,
            model,
        }
        | RunEventPayload::ModelCallCompleted {
            call_id,
            operation_id,
            attempt,
            operation,
            backend,
            model,
            ..
        }
        | RunEventPayload::ModelCallFailed {
            call_id,
            operation_id,
            attempt,
            operation,
            backend,
            model,
            ..
        } => (call_id, operation_id, attempt, operation, backend, model),
        _ => return None,
    };
    Some(ModelAccountingIdentity {
        call_id: call_id.clone(),
        operation_id: operation_id.clone(),
        attempt: u64::from(*attempt),
        operation: *operation,
        cycle_index: u64::from(event.cycle_index()?),
        backend: backend.clone(),
        model: model.clone(),
    })
}

fn require_model_identity(
    expected: &ModelAccountingIdentity,
    observed: &ModelAccountingIdentity,
) -> CheckpointResult<()> {
    if expected != observed {
        return Err(model_accounting_error(
            "model journal, event, and ledger identities do not match",
        ));
    }
    Ok(())
}

fn validate_terminal_budget_event_order(
    checkpoint: &Checkpoint,
    terminal_event_index: usize,
) -> CheckpointResult<()> {
    let budget_configured = checkpoint
        .run_definition
        .get("budget_limits")
        .is_some_and(|value| !value.is_null());
    let next_event = checkpoint
        .event_outbox
        .get(terminal_event_index + 1)
        .map(|entry| {
            serde_json::from_value::<RunEvent>(entry.event.clone()).map_err(|error| {
                CheckpointError::new(
                    "checkpoint_event_outbox_invalid",
                    format!("checkpoint event payload is invalid: {error}"),
                )
            })
        })
        .transpose()?;
    let next_is_model_budget = next_event.as_ref().is_some_and(|event| {
        matches!(
            event.payload(),
            RunEventPayload::BudgetSnapshot {
                enforcement_boundary: BudgetEnforcementBoundary::ModelCallComplete,
                ..
            } | RunEventPayload::BudgetExhausted {
                enforcement_boundary: BudgetEnforcementBoundary::ModelCallComplete,
                ..
            }
        )
    });
    if budget_configured != next_is_model_budget {
        return Err(model_accounting_error(
            "model terminal event must be followed immediately by its configured budget observation",
        ));
    }
    if next_is_model_budget {
        let duplicate_budget = checkpoint
            .event_outbox
            .get(terminal_event_index + 2)
            .map(|entry| {
                serde_json::from_value::<RunEvent>(entry.event.clone()).map_err(|error| {
                    CheckpointError::new(
                        "checkpoint_event_outbox_invalid",
                        format!("checkpoint event payload is invalid: {error}"),
                    )
                })
            })
            .transpose()?
            .as_ref()
            .is_some_and(|event| {
                matches!(
                    event.payload(),
                    RunEventPayload::BudgetSnapshot {
                        enforcement_boundary: BudgetEnforcementBoundary::ModelCallComplete,
                        ..
                    } | RunEventPayload::BudgetExhausted {
                        enforcement_boundary: BudgetEnforcementBoundary::ModelCallComplete,
                        ..
                    }
                )
            });
        if duplicate_budget {
            return Err(model_accounting_error(
                "budget_exhausted must replace budget_snapshot for a model-call boundary",
            ));
        }
    }
    Ok(())
}

fn model_accounting_error(message: impl Into<String>) -> CheckpointError {
    CheckpointError::new("checkpoint_status_invalid", message)
}

fn agent_status_matches_checkpoint(status: AgentStatus, checkpoint: CheckpointStatus) -> bool {
    matches!(
        (status, checkpoint),
        (AgentStatus::Pending, CheckpointStatus::Pending)
            | (AgentStatus::Running, CheckpointStatus::Running)
            | (AgentStatus::WaitUser, CheckpointStatus::WaitUser)
            | (AgentStatus::Completed, CheckpointStatus::Completed)
            | (AgentStatus::Failed, CheckpointStatus::Failed)
            | (AgentStatus::MaxCycles, CheckpointStatus::MaxCycles)
            | (
                AgentStatus::ReconciliationRequired,
                CheckpointStatus::ReconciliationRequired
            )
    )
}

pub fn validate_extension_state_size(
    extensions: &BTreeMap<String, ExtensionStateEntry>,
    max_total: u64,
) -> CheckpointResult<()> {
    let mut total = 0_u64;
    for (namespace, entry) in extensions {
        let bytes = canonical_json_bytes(&entry.to_value(), "extension state entry")?;
        if bytes.len() > MAX_EXTENSION_ENTRY_BYTES {
            return Err(CheckpointError::new(
                "checkpoint_extension_entry_too_large",
                format!("extension {namespace} exceeds {MAX_EXTENSION_ENTRY_BYTES} bytes"),
            ));
        }
        total = total.checked_add(bytes.len() as u64).ok_or_else(|| {
            CheckpointError::new(
                "checkpoint_extension_state_too_large",
                "extension state byte count overflow",
            )
        })?;
    }
    if total > max_total {
        return Err(CheckpointError::new(
            "checkpoint_extension_state_too_large",
            format!("extension state exceeds {max_total} bytes"),
        ));
    }
    Ok(())
}

pub(super) fn validate_json(value: &Value, field_name: &str) -> CheckpointResult<()> {
    crate::checkpoint::canonical_json_bytes(value, field_name).map(|_| ())
}

pub(super) fn required_string<'a>(
    object: &'a Map<String, Value>,
    field: &str,
    code: &str,
) -> CheckpointResult<&'a str> {
    object
        .get(field)
        .and_then(Value::as_str)
        .ok_or_else(|| CheckpointError::new(code, format!("{field} must be a string")))
}

pub(super) fn optional_string(
    object: &Map<String, Value>,
    field: &str,
) -> CheckpointResult<Option<String>> {
    match object.get(field) {
        None | Some(Value::Null) => Ok(None),
        Some(Value::String(value)) => Ok(Some(value.clone())),
        Some(_) => Err(CheckpointError::new(
            "operation_kind_fields_invalid",
            format!("{field} must be a string or null"),
        )),
    }
}

pub(super) fn required_u64(
    object: &Map<String, Value>,
    field: &str,
    code: &str,
) -> CheckpointResult<u64> {
    object
        .get(field)
        .and_then(Value::as_u64)
        .ok_or_else(|| CheckpointError::new(code, format!("{field} must be a JSON-safe integer")))
}