109 lines
4.1 KiB
Rust
109 lines
4.1 KiB
Rust
//! Linearization utilities for task messages.
|
|
//!
|
|
//! This module provides pure functions for linearizing task messages in a conversation,
|
|
//! following a DFS traversal that interleaves subtask messages at subagent tool calls.
|
|
|
|
use std::collections::{HashMap, HashSet};
|
|
|
|
use warp_multi_agent_api as api;
|
|
|
|
use crate::ai::agent::task::helper::TaskExt as _;
|
|
|
|
/// Computes the set of "active" task IDs in a task tree.
|
|
///
|
|
/// An active task is one that is still in progress. The algorithm:
|
|
/// 1. Start with a queue containing the root task ID
|
|
/// 2. For each task in the queue, walk through its messages:
|
|
/// - When encountering a subagent ToolCall, add the subtask to the queue
|
|
/// - When encountering a ToolCallResult matching a subagent call, remove from queue
|
|
/// 3. After processing all messages, add the task to the active set
|
|
/// 4. Repeat until the queue is empty
|
|
pub fn compute_active_task_ids<'a>(
|
|
root_task_id: &str,
|
|
tasks: &HashMap<&str, &'a api::Task>,
|
|
) -> HashSet<&'a str> {
|
|
let mut active_tasks = HashSet::new();
|
|
let mut visited = HashSet::new();
|
|
let mut queue = vec![root_task_id];
|
|
|
|
while let Some(task_id) = queue.pop() {
|
|
// Cycle protection: skip tasks we've already processed.
|
|
if !visited.insert(task_id) {
|
|
log::error!("Cycle detected in active task computation at task {task_id}");
|
|
continue;
|
|
}
|
|
|
|
let Some(task) = tasks.get(task_id) else {
|
|
// Task not found - skip it.
|
|
continue;
|
|
};
|
|
|
|
// Track subagent tool calls: tool_call_id -> subtask_id.
|
|
let mut pending_subagents: HashMap<&str, &str> = HashMap::new();
|
|
|
|
for message in &task.messages {
|
|
match &message.message {
|
|
Some(api::message::Message::ToolCall(tool_call)) => {
|
|
// Check if this is a subagent call.
|
|
if let Some(api::message::tool_call::Tool::Subagent(subagent)) = &tool_call.tool
|
|
{
|
|
if !subagent.task_id.is_empty() {
|
|
// Add subtask to the queue.
|
|
queue.push(subagent.task_id.as_str());
|
|
// Track this subagent call so we can remove it when we see the result.
|
|
pending_subagents
|
|
.insert(tool_call.tool_call_id.as_str(), subagent.task_id.as_str());
|
|
}
|
|
}
|
|
}
|
|
Some(api::message::Message::ToolCallResult(result)) => {
|
|
// If this result matches a pending subagent call, remove from queue.
|
|
if let Some(subtask_id) = pending_subagents.remove(result.tool_call_id.as_str())
|
|
{
|
|
queue.retain(|id| *id != subtask_id);
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
// After processing all messages, add this task to the active set.
|
|
active_tasks.insert(task.id.as_str());
|
|
}
|
|
|
|
active_tasks
|
|
}
|
|
|
|
/// Computes the depth (distance from root) for each task in the map.
|
|
///
|
|
/// Tasks with no parent have depth 0. Tasks whose parent chain contains a cycle or leads to a
|
|
/// missing task are assigned depth 0.
|
|
pub fn compute_task_depths(tasks: &HashMap<String, api::Task>) -> HashMap<&str, usize> {
|
|
let mut depths = HashMap::new();
|
|
for (task_id, _) in tasks.iter() {
|
|
let mut depth = 0;
|
|
let mut current_id: &str = task_id;
|
|
let mut visited = HashSet::new();
|
|
while let Some(task) = tasks.get(current_id) {
|
|
if !visited.insert(current_id) {
|
|
// Cycle detected; treat as depth 0.
|
|
log::error!("Cycle detected in task parent chain starting from task {task_id}");
|
|
depth = 0;
|
|
break;
|
|
}
|
|
if let Some(parent_id) = task.parent_id() {
|
|
depth += 1;
|
|
current_id = parent_id;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
depths.insert(task_id.as_str(), depth);
|
|
}
|
|
depths
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "linearization_tests.rs"]
|
|
mod tests;
|