Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b

app/src/ai/agent/linearization.rs

@@ -0,0 +1,108 @@
+//! 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;