vv-agent 0.6.2

VectorVein agent runtime, SDK, CLI, tools, and workspace backends
Documentation
use std::sync::Arc;

use crate::runtime::state::{Checkpoint, StateStore};
use crate::runtime::token_usage::summarize_task_token_usage;
use crate::runtime::CancellationToken;
use crate::types::{AgentResult, AgentStatus, AgentTask, Message, Metadata};

use super::super::{cancelled_backend_result, failed_backend_result, RuntimeRecipe};
use super::backend::DistributedBackend;
use super::checkpoint::checkpoint_snapshot;
use super::dispatch::CycleDispatcher;

pub(super) struct DistributedRunContext<'a> {
    pub task: &'a AgentTask,
    pub recipe: &'a RuntimeRecipe,
    pub state_store: &'a Arc<dyn StateStore>,
    pub cycle_dispatcher: &'a Arc<dyn CycleDispatcher>,
    pub cancellation_token: Option<&'a CancellationToken>,
    pub max_cycles: u32,
}

impl DistributedBackend {
    pub(super) fn execute_distributed(
        &self,
        initial_messages: Vec<Message>,
        shared_state: Metadata,
        context: DistributedRunContext<'_>,
    ) -> AgentResult {
        let checkpoint = Checkpoint {
            task_id: context.task.task_id.clone(),
            cycle_index: 0,
            status: AgentStatus::Running,
            messages: initial_messages.clone(),
            cycles: Vec::new(),
            shared_state: shared_state.clone(),
        };
        if let Err(error) = context.state_store.save_checkpoint(checkpoint) {
            return AgentResult {
                status: AgentStatus::Failed,
                messages: initial_messages,
                cycles: Vec::new(),
                final_answer: None,
                wait_reason: None,
                error: Some(format!("Failed to save initial checkpoint: {error}")),
                shared_state,
                token_usage: Default::default(),
            };
        }

        let result = self.distributed_loop(&context);
        let _ = context.state_store.delete_checkpoint(&context.task.task_id);
        result
    }

    fn distributed_loop(&self, context: &DistributedRunContext<'_>) -> AgentResult {
        for cycle_index in 1..=context.max_cycles {
            if context
                .cancellation_token
                .is_some_and(CancellationToken::is_cancelled)
            {
                let (messages, cycles, shared_state) =
                    checkpoint_snapshot(context.state_store, &context.task.task_id);
                return cancelled_backend_result(messages, cycles, shared_state);
            }

            match context.cycle_dispatcher.dispatch_cycle(
                context.task,
                context.recipe,
                &self.cycle_name,
                cycle_index,
            ) {
                Ok(dispatch_result) if dispatch_result.finished => {
                    return dispatch_result.result.unwrap_or_else(|| {
                        let (messages, cycles, shared_state) =
                            checkpoint_snapshot(context.state_store, &context.task.task_id);
                        failed_backend_result(
                            messages,
                            cycles,
                            shared_state,
                            format!(
                                "Distributed cycle {cycle_index} finished without result payload"
                            ),
                        )
                    });
                }
                Ok(_) => {}
                Err(error) => {
                    let (messages, cycles, shared_state) =
                        checkpoint_snapshot(context.state_store, &context.task.task_id);
                    return failed_backend_result(
                        messages,
                        cycles,
                        shared_state,
                        format!("Distributed cycle {cycle_index} failed: {error}"),
                    );
                }
            }
        }

        let (messages, cycles, shared_state) =
            checkpoint_snapshot(context.state_store, &context.task.task_id);
        let token_usage = summarize_task_token_usage(&cycles);
        AgentResult {
            status: AgentStatus::MaxCycles,
            messages,
            cycles,
            final_answer: Some("Reached max cycles without finish signal.".to_string()),
            wait_reason: None,
            error: None,
            shared_state,
            token_usage,
        }
    }
}