vv_agent/runtime/sub_task_manager/

record.rs

use std::collections::BTreeMap;
use std::sync::Arc;
use std::thread::JoinHandle;

use serde_json::{json, Value};

use crate::runtime::sub_agent_sessions::SubAgentSession;
use crate::types::SubTaskOutcome;
use crate::workspace::WorkspaceBackend;

use super::helpers::{is_internal_workspace_file, preview_text, status_label};
use super::types::{ManagedSubTaskSnapshot, SubTaskLineage};

#[derive(Clone)]
pub(super) struct ManagedSubAgentSession {
    pub(super) session: Arc<dyn SubAgentSession>,
    pub(super) initial_lineage: SubTaskLineage,
}

#[derive(Clone)]
pub(super) struct ContinuationAdmissionState {
    task_title: String,
    outcome: Option<SubTaskOutcome>,
    recent_activity: Option<String>,
    parent_run_id: Option<String>,
    parent_tool_call_id: Option<String>,
    running: bool,
    worker_owned_run: bool,
    updated_at: String,
}

pub struct ManagedSubTask {
    pub(super) task_id: String,
    pub(super) session_id: String,
    pub(super) agent_name: String,
    pub(super) task_title: String,
    pub(super) workspace_backend: Option<Arc<dyn WorkspaceBackend>>,
    pub(super) session: Option<ManagedSubAgentSession>,
    pub(super) outcome: Option<SubTaskOutcome>,
    pub(super) resolved: BTreeMap<String, String>,
    pub(super) current_cycle_index: Option<u32>,
    pub(super) recent_activity: Option<String>,
    pub(super) latest_cycle: Option<Value>,
    pub(super) latest_tool_call: Option<Value>,
    pub(super) parent_run_id: Option<String>,
    pub(super) parent_tool_call_id: Option<String>,
    pub(super) running: bool,
    pub(super) worker_owned_run: bool,
    pub(super) handle: Option<JoinHandle<()>>,
    pub(super) updated_at: String,
    pub(super) session_generation: u64,
    pub(super) manager_listener_generation: Option<u64>,
}

impl ManagedSubTask {
    pub(super) fn is_running(&self) -> bool {
        self.running
    }

    pub(super) fn to_status_entry(&self, detail_level: &str, workspace_file_limit: usize) -> Value {
        let status = self.status_label();
        let mut entry = json!({
            "task_id": self.task_id,
            "session_id": self.session_id,
            "agent_name": self.agent_name,
            "status": status,
        });
        if !self.task_title.is_empty() {
            entry["task_description"] = Value::String(self.task_title.clone());
        }
        if !self.is_running() {
            if let Some(outcome) = &self.outcome {
                if let Some(final_answer) = &outcome.final_answer {
                    entry["final_answer"] = Value::String(final_answer.clone());
                }
                if let Some(wait_reason) = &outcome.wait_reason {
                    entry["wait_reason"] = Value::String(wait_reason.clone());
                }
                if let Some(error) = &outcome.error {
                    entry["error"] = Value::String(error.clone());
                }
                if let Some(error_code) = &outcome.error_code {
                    entry["error_code"] = Value::String(error_code.clone());
                }
                if let Some(completion_reason) = outcome.completion_reason {
                    entry["completion_reason"] =
                        Value::String(completion_reason.as_str().to_string());
                }
                if let Some(completion_tool_name) = &outcome.completion_tool_name {
                    entry["completion_tool_name"] = Value::String(completion_tool_name.clone());
                }
                if let Some(partial_output) = &outcome.partial_output {
                    entry["partial_output"] = Value::String(partial_output.clone());
                }
                if outcome.cycles > 0 {
                    entry["cycles"] = Value::Number(outcome.cycles.into());
                }
                if !outcome.todo_list.is_empty() {
                    entry["todo_list"] = Value::Array(outcome.todo_list.clone());
                }
                if !outcome.resolved.is_empty() {
                    entry["resolved"] = json!(outcome.resolved);
                }
            }
        }
        if let Some(parent_run_id) = &self.parent_run_id {
            entry["parent_run_id"] = Value::String(parent_run_id.clone());
        }
        if let Some(parent_tool_call_id) = &self.parent_tool_call_id {
            entry["parent_tool_call_id"] = Value::String(parent_tool_call_id.clone());
        }
        if detail_level == "snapshot" {
            let workspace_snapshot = self.workspace_snapshot(workspace_file_limit);
            entry["snapshot"] = json!({
                "current_cycle_index": self.current_cycle_index,
                "updated_at": self.updated_at,
                "workspace_files": workspace_snapshot.files,
                "workspace_file_count": workspace_snapshot.file_count,
                "workspace_files_truncated": workspace_snapshot.truncated,
            });
            if !self.task_title.is_empty() {
                entry["snapshot"]["task_title"] = Value::String(self.task_title.clone());
            }
            if let Some(recent_activity) = &self.recent_activity {
                entry["snapshot"]["recent_activity"] = Value::String(recent_activity.clone());
            }
            if let Some(latest_cycle) = &self.latest_cycle {
                entry["snapshot"]["latest_cycle"] = latest_cycle.clone();
            }
            if let Some(latest_tool_call) = &self.latest_tool_call {
                entry["snapshot"]["latest_tool_call"] = latest_tool_call.clone();
            }
        }
        entry
    }

    pub(super) fn snapshot(&self) -> ManagedSubTaskSnapshot {
        ManagedSubTaskSnapshot {
            task_id: self.task_id.clone(),
            session_id: self.session_id.clone(),
            agent_name: self.agent_name.clone(),
            task_title: self.task_title.clone(),
            status: self.status_label().to_string(),
            running: self.is_running(),
            outcome: (!self.is_running()).then(|| self.outcome.clone()).flatten(),
            resolved: self.resolved.clone(),
            current_cycle_index: self.current_cycle_index,
            recent_activity: self.recent_activity.clone(),
            latest_cycle: self.latest_cycle.clone(),
            latest_tool_call: self.latest_tool_call.clone(),
            parent_run_id: self.parent_run_id.clone(),
            parent_tool_call_id: self.parent_tool_call_id.clone(),
            updated_at: self.updated_at.clone(),
        }
    }

    pub(super) fn status_label(&self) -> &'static str {
        if self.is_running() {
            return "running";
        }
        self.outcome
            .as_ref()
            .map(|outcome| status_label(outcome.status))
            .unwrap_or("pending")
    }

    pub(super) fn update_from_outcome(&mut self, outcome: &SubTaskOutcome) {
        if outcome.cycles > 0 {
            self.current_cycle_index = Some(outcome.cycles);
        }
        self.set_latest_cycle_status(status_label(outcome.status));
        if let Some(detail) = outcome
            .final_answer
            .as_ref()
            .or(outcome.wait_reason.as_ref())
            .or(outcome.error.as_ref())
            .and_then(|value| preview_text(Some(&Value::String(value.clone()))))
        {
            self.recent_activity = Some(detail);
        }
    }

    pub(super) fn admit_continuation(
        &mut self,
        prompt: &str,
        lineage: &SubTaskLineage,
        updated_at: String,
    ) -> ContinuationAdmissionState {
        let previous = ContinuationAdmissionState {
            task_title: self.task_title.clone(),
            outcome: self.outcome.clone(),
            recent_activity: self.recent_activity.clone(),
            parent_run_id: self.parent_run_id.clone(),
            parent_tool_call_id: self.parent_tool_call_id.clone(),
            running: self.running,
            worker_owned_run: self.worker_owned_run,
            updated_at: self.updated_at.clone(),
        };
        self.task_title = prompt.to_string();
        self.outcome = None;
        self.recent_activity = Some(prompt.to_string());
        self.parent_run_id = lineage.parent_run_id.clone();
        self.parent_tool_call_id = lineage.parent_tool_call_id.clone();
        self.running = true;
        self.worker_owned_run = true;
        self.updated_at = updated_at;
        previous
    }

    pub(super) fn rollback_continuation(&mut self, previous: ContinuationAdmissionState) {
        self.task_title = previous.task_title;
        self.outcome = previous.outcome;
        self.recent_activity = previous.recent_activity;
        self.parent_run_id = previous.parent_run_id;
        self.parent_tool_call_id = previous.parent_tool_call_id;
        self.running = previous.running;
        self.worker_owned_run = previous.worker_owned_run;
        self.handle = None;
        self.updated_at = previous.updated_at;
    }

    pub(super) fn set_latest_cycle_status(&mut self, status: &str) {
        let mut latest_cycle = self
            .latest_cycle
            .as_ref()
            .and_then(Value::as_object)
            .cloned()
            .unwrap_or_default();
        latest_cycle.insert("status".to_string(), Value::String(status.to_string()));
        if let Some(current_cycle_index) = self.current_cycle_index {
            latest_cycle
                .entry("cycle_index".to_string())
                .or_insert_with(|| Value::from(current_cycle_index));
        }
        self.latest_cycle = Some(Value::Object(latest_cycle));
    }

    pub(super) fn mark_terminal_state(&mut self, status: &str, detail: Option<&str>) {
        self.set_latest_cycle_status(status);
        if let Some(detail) = detail {
            self.recent_activity = Some(detail.to_string());
        }
    }

    fn workspace_snapshot(&self, workspace_file_limit: usize) -> WorkspaceSnapshot {
        let Some(workspace_backend) = &self.workspace_backend else {
            return WorkspaceSnapshot::default();
        };
        let Ok(raw_files) = workspace_backend.list_files(".", "**/*") else {
            return WorkspaceSnapshot::default();
        };
        let visible_files = raw_files
            .into_iter()
            .filter(|path| !is_internal_workspace_file(path))
            .collect::<Vec<_>>();
        let file_count = visible_files.len();
        WorkspaceSnapshot {
            files: visible_files
                .into_iter()
                .take(workspace_file_limit)
                .collect::<Vec<_>>(),
            file_count,
            truncated: file_count > workspace_file_limit,
        }
    }
}

#[derive(Default)]
struct WorkspaceSnapshot {
    files: Vec<String>,
    file_count: usize,
    truncated: bool,
}