vv_agent/runtime/sub_task_manager/

submission.rs

use std::collections::BTreeMap;
use std::fmt;
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::sync::Arc;
use std::thread;

use crate::types::{AgentStatus, SubTaskOutcome};
use crate::workspace::WorkspaceBackend;

use super::helpers::{normalize_failed_outcome, now_iso, panic_payload_to_string};
use super::manager::SubTaskManager;
use super::record::ManagedSubTask;
use super::types::{SubTaskLineage, SubTaskSubmissionContext};

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum SubTaskSubmitError {
    AlreadyRunning { task_id: String },
    SpawnFailed { task_id: String, error: String },
}

impl SubTaskSubmitError {
    pub(crate) fn error_code(&self) -> &'static str {
        match self {
            Self::AlreadyRunning { .. } => "sub_task_already_running",
            Self::SpawnFailed { .. } => "sub_task_submit_failed",
        }
    }
}

impl fmt::Display for SubTaskSubmitError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::AlreadyRunning { task_id } => {
                write!(formatter, "Sub-task {task_id} is already running.")
            }
            Self::SpawnFailed { task_id, error } => {
                write!(
                    formatter,
                    "Sub-task {task_id} thread failed to spawn: {error}"
                )
            }
        }
    }
}

impl SubTaskManager {
    pub fn submit(
        &self,
        task_id: impl Into<String>,
        session_id: impl Into<String>,
        agent_name: impl Into<String>,
        task_title: impl Into<String>,
        runner: impl FnOnce() -> SubTaskOutcome + Send + 'static,
    ) -> Result<(), String> {
        self.submit_with_workspace(task_id, session_id, agent_name, task_title, None, runner)
    }

    pub fn submit_with_workspace(
        &self,
        task_id: impl Into<String>,
        session_id: impl Into<String>,
        agent_name: impl Into<String>,
        task_title: impl Into<String>,
        workspace_backend: Option<Arc<dyn WorkspaceBackend>>,
        runner: impl FnOnce() -> SubTaskOutcome + Send + 'static,
    ) -> Result<(), String> {
        self.submit_with_context(
            task_id,
            session_id,
            agent_name,
            task_title,
            SubTaskSubmissionContext {
                workspace_backend,
                lineage: SubTaskLineage::default(),
            },
            runner,
        )
    }

    pub fn submit_with_context(
        &self,
        task_id: impl Into<String>,
        session_id: impl Into<String>,
        agent_name: impl Into<String>,
        task_title: impl Into<String>,
        context: SubTaskSubmissionContext,
        runner: impl FnOnce() -> SubTaskOutcome + Send + 'static,
    ) -> Result<(), String> {
        self.submit_with_context_detailed(
            task_id, session_id, agent_name, task_title, context, runner,
        )
        .map_err(|error| error.to_string())
    }

    pub(crate) fn submit_with_context_detailed(
        &self,
        task_id: impl Into<String>,
        session_id: impl Into<String>,
        agent_name: impl Into<String>,
        task_title: impl Into<String>,
        context: SubTaskSubmissionContext,
        runner: impl FnOnce() -> SubTaskOutcome + Send + 'static,
    ) -> Result<(), SubTaskSubmitError> {
        let SubTaskSubmissionContext {
            workspace_backend,
            lineage,
        } = context;
        let task_id = task_id.into();
        let session_id = session_id.into();
        let agent_name = agent_name.into();
        let task_title = task_title.into();
        let mut task_records = self.tasks.lock().expect("sub-task manager poisoned");
        if task_records
            .get(&task_id)
            .is_some_and(ManagedSubTask::is_running)
        {
            return Err(SubTaskSubmitError::AlreadyRunning { task_id });
        }
        let previous_record = task_records.insert(
            task_id.clone(),
            ManagedSubTask {
                task_id: task_id.clone(),
                session_id: session_id.clone(),
                agent_name: agent_name.clone(),
                task_title,
                workspace_backend,
                session: None,
                outcome: None,
                resolved: BTreeMap::new(),
                current_cycle_index: None,
                recent_activity: None,
                latest_cycle: None,
                latest_tool_call: None,
                parent_run_id: lineage.parent_run_id,
                parent_tool_call_id: lineage.parent_tool_call_id,
                running: true,
                worker_owned_run: true,
                handle: None,
                updated_at: now_iso(),
                session_generation: 0,
                manager_listener_generation: None,
            },
        );

        let tasks = self.tasks.clone();
        let task_id_for_thread = task_id.clone();
        let spawn_result = catch_unwind(AssertUnwindSafe(|| {
            thread::Builder::new()
                .name(format!("vv-agent-sub-task-{session_id}"))
                .spawn(move || {
                    let outcome = catch_unwind(AssertUnwindSafe(runner));
                    let mut tasks = tasks.lock().expect("sub-task manager poisoned");
                    if let Some(record) = tasks.get_mut(&task_id_for_thread) {
                        let outcome = normalize_failed_outcome(match outcome {
                            Ok(outcome) => outcome,
                            Err(payload) => SubTaskOutcome {
                                task_id: record.task_id.clone(),
                                agent_name: record.agent_name.clone(),
                                status: AgentStatus::Failed,
                                session_id: Some(record.session_id.clone()),
                                final_answer: None,
                                wait_reason: None,
                                error: Some(panic_payload_to_string(payload.as_ref())),
                                error_code: Some("sub_task_failed".to_string()),
                                completion_reason: Some(crate::types::CompletionReason::Failed),
                                completion_tool_name: None,
                                partial_output: None,
                                cycles: 0,
                                todo_list: Vec::new(),
                                resolved: record.resolved.clone(),
                            },
                        });
                        if !outcome.resolved.is_empty() {
                            record.resolved = outcome.resolved.clone();
                        }
                        record.update_from_outcome(&outcome);
                        record.outcome = Some(outcome);
                        record.updated_at = now_iso();
                        record.running = false;
                        record.worker_owned_run = false;
                    }
                })
        }));
        let handle = match spawn_result {
            Ok(Ok(handle)) => handle,
            Ok(Err(error)) => {
                restore_after_spawn_failure(&mut task_records, &task_id, previous_record);
                return Err(SubTaskSubmitError::SpawnFailed {
                    task_id,
                    error: error.to_string(),
                });
            }
            Err(payload) => {
                let error = panic_payload_to_string(payload.as_ref());
                restore_after_spawn_failure(&mut task_records, &task_id, previous_record);
                return Err(SubTaskSubmitError::SpawnFailed { task_id, error });
            }
        };

        if let Some(record) = task_records.get_mut(&task_id) {
            record.handle = Some(handle);
            record.updated_at = now_iso();
        }
        Ok(())
    }

    pub fn record_outcome(&self, task_id: &str, outcome: SubTaskOutcome) {
        self.record_outcome_with_context(task_id, outcome, None, SubTaskLineage::default());
    }

    pub fn record_outcome_with_context(
        &self,
        task_id: &str,
        outcome: SubTaskOutcome,
        workspace_backend: Option<Arc<dyn WorkspaceBackend>>,
        lineage: SubTaskLineage,
    ) {
        let outcome = normalize_failed_outcome(outcome);
        let mut tasks = self.tasks.lock().expect("sub-task manager poisoned");
        let task_id = task_id.to_string();
        match tasks.get_mut(&task_id) {
            Some(record) => {
                if workspace_backend.is_some() {
                    record.workspace_backend = workspace_backend;
                }
                if lineage.parent_run_id.is_some() {
                    record.parent_run_id = lineage.parent_run_id;
                }
                if lineage.parent_tool_call_id.is_some() {
                    record.parent_tool_call_id = lineage.parent_tool_call_id;
                }
                record.session_id = outcome
                    .session_id
                    .clone()
                    .unwrap_or_else(|| record.session_id.clone());
                record.agent_name = outcome.agent_name.clone();
                if !outcome.resolved.is_empty() {
                    record.resolved = outcome.resolved.clone();
                }
                if record.running && record.worker_owned_run {
                    record.updated_at = now_iso();
                    return;
                }
                record.running = false;
                record.worker_owned_run = false;
                record.update_from_outcome(&outcome);
                record.outcome = Some(outcome);
                record.updated_at = now_iso();
            }
            None => {
                let mut record = ManagedSubTask {
                    session_id: outcome.session_id.clone().unwrap_or_default(),
                    agent_name: outcome.agent_name.clone(),
                    task_title: String::new(),
                    workspace_backend,
                    session: None,
                    outcome: None,
                    resolved: outcome.resolved.clone(),
                    current_cycle_index: None,
                    recent_activity: None,
                    latest_cycle: None,
                    latest_tool_call: None,
                    parent_run_id: lineage.parent_run_id,
                    parent_tool_call_id: lineage.parent_tool_call_id,
                    task_id: task_id.clone(),
                    running: false,
                    worker_owned_run: false,
                    handle: None,
                    updated_at: now_iso(),
                    session_generation: 0,
                    manager_listener_generation: None,
                };
                record.update_from_outcome(&outcome);
                record.outcome = Some(outcome);
                tasks.insert(task_id.clone(), record);
            }
        }
    }
}

fn restore_after_spawn_failure(
    task_records: &mut BTreeMap<String, ManagedSubTask>,
    task_id: &str,
    previous_record: Option<ManagedSubTask>,
) {
    task_records.remove(task_id);
    if let Some(previous_record) = previous_record {
        task_records.insert(task_id.to_string(), previous_record);
    }
}