vv_agent/runtime/backends/distributed/

driver.rs

use std::sync::Arc;

use crate::budget::RunBudgetLimits;
use crate::checkpoint::{CheckpointStatus, ClaimMode};
use crate::runtime::token_usage::summarize_task_token_usage;
use crate::runtime::CheckpointStore;
use crate::types::{last_assistant_output, AgentResult, AgentStatus, AgentTask, CompletionReason};

use super::super::RuntimeRecipe;
use super::backend::DistributedBackend;
use super::capabilities::DistributedCapabilityRegistry;
use super::contract::{now_unix_ms, DistributedCheckpointConfig, DistributedRunEnvelope};
use super::dispatch::CycleDispatchResult;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DistributedRunHandle {
    pub checkpoint_key: String,
    pub run_id: String,
    pub trace_id: String,
}

impl DistributedRunHandle {
    fn from_checkpoint(checkpoint: &crate::runtime::Checkpoint) -> Self {
        Self {
            checkpoint_key: checkpoint.checkpoint_key.clone(),
            run_id: checkpoint.root_run_id.clone(),
            trace_id: checkpoint.trace_id.clone(),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum DistributedDeliveryOutcome {
    Worker(Box<CycleDispatchResult>),
    TransportFailure(String),
}

impl DistributedDeliveryOutcome {
    pub fn worker(response: CycleDispatchResult) -> Self {
        Self::Worker(Box::new(response))
    }

    pub fn transport_failure(error: impl Into<String>) -> Result<Self, String> {
        let error = error.into();
        if error.trim().is_empty() {
            return Err("distributed transport error must be a non-empty string".to_string());
        }
        Ok(Self::TransportFailure(error))
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DistributedWaitReason {
    ActiveClaim,
    ReconciliationRequired,
    HostInteraction,
    SupersededDelivery,
}

#[derive(Debug, Clone, PartialEq)]
pub enum DistributedAdvanceDecision {
    Dispatch {
        handle: DistributedRunHandle,
        envelope: DistributedRunEnvelope,
    },
    RetryAt {
        handle: DistributedRunHandle,
        envelope: DistributedRunEnvelope,
        not_before_unix_ms: u64,
    },
    Wait {
        handle: DistributedRunHandle,
        reason: DistributedWaitReason,
    },
    FinalizeRequired {
        handle: DistributedRunHandle,
        checkpoint_revision: u64,
        result: AgentResult,
    },
    TerminalReplay {
        handle: DistributedRunHandle,
        checkpoint_revision: u64,
        result: AgentResult,
    },
}

pub trait CycleEnqueuer: Send + Sync {
    fn enqueue_envelope(
        &self,
        envelope: &DistributedRunEnvelope,
        not_before_unix_ms: Option<u64>,
    ) -> Result<(), String>;
}

type NonblockingComponents<'a> = (
    &'a RuntimeRecipe,
    Arc<dyn CheckpointStore>,
    &'a Arc<dyn CycleEnqueuer>,
);

impl DistributedBackend {
    pub fn nonblocking(
        runtime_recipe: RuntimeRecipe,
        capability_registry: DistributedCapabilityRegistry,
        cycle_enqueuer: Arc<dyn CycleEnqueuer>,
    ) -> Self {
        Self::inline_fallback().with_nonblocking_driver(
            runtime_recipe,
            capability_registry,
            cycle_enqueuer,
        )
    }

    pub fn with_nonblocking_driver(
        mut self,
        runtime_recipe: RuntimeRecipe,
        capability_registry: DistributedCapabilityRegistry,
        cycle_enqueuer: Arc<dyn CycleEnqueuer>,
    ) -> Self {
        self.runtime_recipe = Some(runtime_recipe);
        self.capability_registry = Some(capability_registry);
        self.cycle_enqueuer = Some(cycle_enqueuer);
        self
    }

    pub fn start(
        &self,
        task: AgentTask,
        checkpoint_config: DistributedCheckpointConfig,
        budget_limits: Option<RunBudgetLimits>,
    ) -> Result<DistributedRunHandle, String> {
        let (recipe, store, enqueuer) = self.nonblocking_components()?;
        checkpoint_config.validate()?;
        let checkpoint = load_checkpoint_once(store.as_ref(), &checkpoint_config.key)?;
        let handle = DistributedRunHandle::from_checkpoint(&checkpoint);
        if checkpoint.terminal_result.is_some() {
            return Ok(handle);
        }
        if checkpoint.status != CheckpointStatus::Running
            || checkpoint.claim_token.is_some()
            || checkpoint.cycle_index != 0
        {
            return Err(
                "distributed start requires an unclaimed running checkpoint before cycle 1"
                    .to_string(),
            );
        }
        if checkpoint.task_id != task.task_id {
            return Err(
                "distributed start task does not match the authoritative checkpoint".to_string(),
            );
        }
        let envelope = self.envelope_from_checkpoint(
            task,
            recipe.clone(),
            &checkpoint,
            checkpoint_config,
            1,
            ClaimMode::Continue,
            budget_limits,
            None,
        )?;
        enqueuer.enqueue_envelope(&envelope, None)?;
        Ok(handle)
    }

    pub fn advance(
        &self,
        previous_envelope: &DistributedRunEnvelope,
        outcome: DistributedDeliveryOutcome,
    ) -> Result<DistributedAdvanceDecision, String> {
        self.validate_nonblocking_recipe()?;
        previous_envelope.validate()?;
        if previous_envelope
            .recipe
            .capabilities
            .approval_provider_ref
            .is_some()
            || previous_envelope
                .recipe
                .capabilities
                .approval_broker_ref
                .is_some()
        {
            return Err(
                "nonblocking distributed runs do not support brokered approval waits".to_string(),
            );
        }
        if matches!(
            &outcome,
            DistributedDeliveryOutcome::TransportFailure(error) if error.trim().is_empty()
        ) {
            return Err("distributed transport error must be a non-empty string".to_string());
        }
        if let DistributedDeliveryOutcome::Worker(response) = &outcome {
            response.validate()?;
        }
        let registry = self.capability_registry.as_ref().ok_or_else(|| {
            "nonblocking DistributedBackend requires a DistributedCapabilityRegistry".to_string()
        })?;
        let store_reference = previous_envelope
            .recipe
            .capabilities
            .checkpoint_store_ref
            .as_ref()
            .ok_or_else(|| "distributed run requires checkpoint_store_ref".to_string())?;
        let store = registry
            .resolve_checkpoint_store_required(store_reference)
            .map_err(|error| error.to_string())?;

        // This is the one authoritative read performed by an advance invocation.
        let checkpoint =
            load_checkpoint_once(store.as_ref(), &previous_envelope.checkpoint_config.key)?;
        validate_checkpoint_identity(previous_envelope, &checkpoint)?;
        let handle = DistributedRunHandle::from_checkpoint(&checkpoint);
        let response = match &outcome {
            DistributedDeliveryOutcome::Worker(response) => Some(response.as_ref()),
            DistributedDeliveryOutcome::TransportFailure(_) => None,
        };

        if let Some(terminal_result) = checkpoint.terminal_result.as_ref() {
            let result = AgentResult::from_dict(terminal_result)
                .map_err(|error| format!("invalid durable terminal result: {error}"))?;
            if let Some(CycleDispatchResult::TerminalReplay {
                checkpoint_revision,
                result: observed,
            }) = response
            {
                if *checkpoint_revision != checkpoint.revision || observed != &result {
                    return Err(
                        "distributed terminal replay does not match the durable checkpoint"
                            .to_string(),
                    );
                }
            }
            return Ok(DistributedAdvanceDecision::TerminalReplay {
                handle,
                checkpoint_revision: checkpoint.revision,
                result,
            });
        }

        if matches!(response, Some(CycleDispatchResult::TerminalReplay { .. })) {
            return Err("distributed terminal replay has no matching durable terminal".to_string());
        }

        if let Some(CycleDispatchResult::TerminalCandidate {
            checkpoint_revision,
            result,
        }) = response
        {
            if *checkpoint_revision != checkpoint.revision {
                return Err(
                    "distributed terminal candidate revision does not match the checkpoint"
                        .to_string(),
                );
            }
            if result.status == AgentStatus::ReconciliationRequired {
                if checkpoint.status != CheckpointStatus::ReconciliationRequired
                    || checkpoint.claim_token.is_some()
                {
                    return Err(
                        "distributed reconciliation candidate does not match durable state"
                            .to_string(),
                    );
                }
                return Ok(DistributedAdvanceDecision::Wait {
                    handle,
                    reason: DistributedWaitReason::ReconciliationRequired,
                });
            }
            if checkpoint.claim_token.is_none()
                || checkpoint.claimed_cycle != Some(u64::from(previous_envelope.cycle_index))
            {
                return Err(
                    "distributed terminal candidate does not retain the dispatched cycle claim"
                        .to_string(),
                );
            }
            if result
                .cycles
                .last()
                .is_some_and(|cycle| cycle.index != previous_envelope.cycle_index)
            {
                return Err(
                    "distributed terminal candidate does not contain the dispatched cycle"
                        .to_string(),
                );
            }
            return Ok(DistributedAdvanceDecision::FinalizeRequired {
                handle,
                checkpoint_revision: checkpoint.revision,
                result: result.clone(),
            });
        }

        if checkpoint.status == CheckpointStatus::ReconciliationRequired {
            return Ok(DistributedAdvanceDecision::Wait {
                handle,
                reason: DistributedWaitReason::ReconciliationRequired,
            });
        }

        if let Some(CycleDispatchResult::Committed {
            checkpoint_revision: _,
            committed_cycle,
        }) = response
        {
            if *committed_cycle != u64::from(previous_envelope.cycle_index) {
                return Err(
                    "distributed committed response does not match the dispatched cycle"
                        .to_string(),
                );
            }
            if checkpoint.cycle_index < *committed_cycle {
                return Err("distributed committed response is ahead of the checkpoint".to_string());
            }
        }

        let previous_cycle = u64::from(previous_envelope.cycle_index);
        if checkpoint.cycle_index > previous_cycle
            || checkpoint
                .claimed_cycle
                .is_some_and(|claimed_cycle| claimed_cycle > previous_cycle)
        {
            return Ok(DistributedAdvanceDecision::Wait {
                handle,
                reason: DistributedWaitReason::SupersededDelivery,
            });
        }
        if let Some(CycleDispatchResult::Committed {
            checkpoint_revision,
            ..
        }) = response
        {
            if checkpoint.claim_token.is_none() && *checkpoint_revision != checkpoint.revision {
                return Err(
                    "distributed committed response revision does not match the checkpoint"
                        .to_string(),
                );
            }
        }

        let now_ms = now_unix_ms()?;
        let (cycle_index, claim_mode, not_before_unix_ms) = if checkpoint.claim_token.is_some() {
            let cycle_index = checkpoint
                .claimed_cycle
                .ok_or_else(|| "distributed checkpoint has a partial claim".to_string())?;
            let not_before = checkpoint
                .lease_expires_at_ms
                .filter(|lease_expires_at_ms| *lease_expires_at_ms > now_ms);
            (cycle_index, ClaimMode::Recovery, not_before)
        } else if checkpoint.cycle_index == previous_cycle {
            (
                checkpoint
                    .cycle_index
                    .checked_add(1)
                    .ok_or_else(|| "distributed cycle index overflow".to_string())?,
                ClaimMode::Continue,
                None,
            )
        } else if checkpoint.cycle_index.checked_add(1) == Some(previous_cycle) {
            (previous_cycle, ClaimMode::Recovery, None)
        } else {
            return Err("distributed delivery is out of order with the checkpoint".to_string());
        };

        if cycle_index > u64::from(previous_envelope.task.max_cycles) {
            let result = AgentResult {
                status: AgentStatus::MaxCycles,
                completion_reason: Some(CompletionReason::MaxCycles),
                partial_output: last_assistant_output(&checkpoint.cycles),
                token_usage: summarize_task_token_usage(&checkpoint.model_calls),
                messages: checkpoint.messages.clone(),
                cycles: checkpoint.cycles.clone(),
                budget_usage: checkpoint.budget_usage.clone(),
                final_answer: Some("Reached max cycles without finish signal.".to_string()),
                shared_state: checkpoint.shared_state.clone(),
                ..AgentResult::default()
            };
            return Ok(DistributedAdvanceDecision::FinalizeRequired {
                handle,
                checkpoint_revision: checkpoint.revision,
                result,
            });
        }

        let cycle_index = u32::try_from(cycle_index)
            .map_err(|_| "distributed cycle index exceeds u32".to_string())?;
        let envelope = self.envelope_from_checkpoint(
            previous_envelope.task.clone(),
            previous_envelope.recipe.clone(),
            &checkpoint,
            previous_envelope.checkpoint_config.clone(),
            cycle_index,
            claim_mode,
            previous_envelope.budget_limits.clone(),
            not_before_unix_ms,
        )?;
        let enqueuer = self
            .cycle_enqueuer
            .as_ref()
            .ok_or_else(|| "nonblocking DistributedBackend requires a CycleEnqueuer".to_string())?;
        if let Some(not_before_unix_ms) = not_before_unix_ms {
            enqueuer.enqueue_envelope(&envelope, Some(not_before_unix_ms))?;
            Ok(DistributedAdvanceDecision::RetryAt {
                handle,
                envelope,
                not_before_unix_ms,
            })
        } else {
            enqueuer.enqueue_envelope(&envelope, None)?;
            Ok(DistributedAdvanceDecision::Dispatch { handle, envelope })
        }
    }

    fn validate_nonblocking_recipe(&self) -> Result<(), String> {
        let recipe = self
            .runtime_recipe
            .as_ref()
            .ok_or_else(|| "nonblocking DistributedBackend requires a RuntimeRecipe".to_string())?;
        recipe.validate()?;
        if recipe.capabilities.approval_provider_ref.is_some()
            || recipe.capabilities.approval_broker_ref.is_some()
        {
            return Err(
                "nonblocking distributed runs do not support brokered approval waits".to_string(),
            );
        }
        Ok(())
    }

    fn nonblocking_components(&self) -> Result<NonblockingComponents<'_>, String> {
        self.validate_nonblocking_recipe()?;
        let recipe = self.runtime_recipe.as_ref().expect("validated above");
        let registry = self.capability_registry.as_ref().ok_or_else(|| {
            "nonblocking DistributedBackend requires a DistributedCapabilityRegistry".to_string()
        })?;
        registry
            .resolve(&recipe.capabilities)
            .map_err(|error| error.to_string())?;
        let reference = recipe
            .capabilities
            .checkpoint_store_ref
            .as_ref()
            .ok_or_else(|| "distributed run requires checkpoint_store_ref".to_string())?;
        let store = registry
            .resolve_checkpoint_store_required(reference)
            .map_err(|error| error.to_string())?;
        let enqueuer = self
            .cycle_enqueuer
            .as_ref()
            .ok_or_else(|| "nonblocking DistributedBackend requires a CycleEnqueuer".to_string())?;
        Ok((recipe, store, enqueuer))
    }

    #[allow(clippy::too_many_arguments)]
    fn envelope_from_checkpoint(
        &self,
        task: AgentTask,
        mut recipe: RuntimeRecipe,
        checkpoint: &crate::runtime::Checkpoint,
        checkpoint_config: DistributedCheckpointConfig,
        cycle_index: u32,
        claim_mode: ClaimMode,
        budget_limits: Option<RunBudgetLimits>,
        not_before_unix_ms: Option<u64>,
    ) -> Result<DistributedRunEnvelope, String> {
        let metadata_denials = crate::runtime::tool_planner::projected_metadata_denials(&task)?;
        recipe
            .capabilities
            .tool_policy
            .set_metadata_denials(&metadata_denials);
        let timeout_ms = u64::try_from(self.dispatch_timeout.as_millis())
            .map_err(|_| "distributed dispatch timeout exceeds u64 milliseconds".to_string())?;
        let deadline_base_ms = now_unix_ms()?.max(not_before_unix_ms.unwrap_or(0));
        let deadline_unix_ms = deadline_base_ms
            .checked_add(timeout_ms)
            .ok_or_else(|| "distributed dispatch deadline overflow".to_string())?;
        DistributedRunEnvelope::for_cycle(
            task,
            recipe,
            cycle_index,
            self.cycle_name.clone(),
            Some(checkpoint.root_run_id.clone()),
            Some(deadline_unix_ms),
            self.lease_duration_ms,
            budget_limits,
            checkpoint.root_run_id.clone(),
            checkpoint.trace_id.clone(),
            checkpoint.run_definition_digest.clone(),
            claim_mode,
            checkpoint.resume_attempt,
            checkpoint_config,
        )
    }
}

fn load_checkpoint_once(
    store: &dyn CheckpointStore,
    checkpoint_key: &str,
) -> Result<crate::runtime::Checkpoint, String> {
    store
        .load_checkpoint(checkpoint_key)
        .map_err(|error| error.to_string())?
        .ok_or_else(|| "checkpoint disappeared before distributed driver invocation".to_string())
}

fn validate_checkpoint_identity(
    envelope: &DistributedRunEnvelope,
    checkpoint: &crate::runtime::Checkpoint,
) -> Result<(), String> {
    if checkpoint.checkpoint_key != envelope.checkpoint_config.key
        || checkpoint.task_id != envelope.task.task_id
        || checkpoint.root_run_id != envelope.root_run_id
        || checkpoint.trace_id != envelope.trace_id
        || checkpoint.run_definition_digest != envelope.run_definition_digest
    {
        return Err(
            "distributed advance envelope does not match the authoritative checkpoint".to_string(),
        );
    }
    Ok(())
}