Skip to main content

a3s_code_core/tools/task/
parallel_task.rs

1use super::*;
2
3const PARALLEL_TASK_TOOL_DESCRIPTION: &str = "REMOVED from the model-visible registry (`HARNESS-CONV4`). Prefer `task` with multiple `tasks[]` items. This type remains for focused unit tests that construct ParallelTaskTool directly.";
4
5/// ParallelTaskTool allows the LLM to fan out multiple delegated tasks concurrently.
6///
7/// All tasks execute in parallel and the tool returns when all complete.
8pub struct ParallelTaskTool {
9    executor: Arc<TaskExecutor>,
10}
11
12impl ParallelTaskTool {
13    /// Create a new ParallelTaskTool
14    pub fn new(executor: Arc<TaskExecutor>) -> Self {
15        Self { executor }
16    }
17
18    pub(super) async fn execute_params(
19        &self,
20        params: ParallelTaskParams,
21        ctx: &ToolContext,
22        tool_name: &str,
23        min_tasks: usize,
24    ) -> Result<ToolOutput> {
25        let started_at = std::time::Instant::now();
26        let parent_cancellation = ctx.cancellation_token();
27        let executor = self.executor.scoped_for_invocation(ctx);
28
29        if params.tasks.len() < min_tasks {
30            return Ok(invalid_delegation_argument(format!(
31                "{tool_name} requires at least {min_tasks} task{}",
32                if min_tasks == 1 { "" } else { "s" }
33            )));
34        }
35        if params.tasks.len() > MAX_PARALLEL_TASKS_PER_CALL {
36            return Ok(invalid_delegation_argument(format!(
37                "{tool_name} accepts at most {MAX_PARALLEL_TASKS_PER_CALL} tasks"
38            )));
39        }
40        if let Some((index, _)) = params
41            .tasks
42            .iter()
43            .enumerate()
44            .find(|(_, task)| task.background)
45        {
46            return Ok(invalid_delegation_argument(format!(
47                "{tool_name} task {} cannot set background=true when fan-out options are used or multiple tasks are submitted; every branch is collected by the parent call",
48                index + 1
49            )));
50        }
51        if params.timeout_ms == Some(0) {
52            return Ok(invalid_delegation_argument(format!(
53                "{tool_name} timeout_ms must be at least 1"
54            )));
55        }
56        if let Some(min_success_count) = params.min_success_count {
57            if !params.allow_partial_failure {
58                return Ok(invalid_delegation_argument(format!(
59                    "{tool_name} min_success_count requires allow_partial_failure=true"
60                )));
61            }
62            if min_success_count == 0 || min_success_count > params.tasks.len() {
63                return Ok(invalid_delegation_argument(format!(
64                    "{tool_name} min_success_count must be between 1 and the task count ({})",
65                    params.tasks.len()
66                )));
67            }
68        }
69
70        let task_count = params.tasks.len();
71        let run = executor
72            .execute_parallel_for_tool(
73                params.tasks.clone(),
74                ctx.agent_event_tx.clone(),
75                parallel_execution::ParallelToolOptions {
76                    parent_session_id: ctx.session_id.as_deref(),
77                    timeout_ms: params.timeout_ms,
78                    min_success_count: params.min_success_count,
79                    allow_partial_failure: params.allow_partial_failure,
80                    parent_cancellation: Some(&parent_cancellation),
81                },
82            )
83            .await;
84        let results = run.results;
85
86        let mut output = format!("Executed {} tasks concurrently:\n\n", task_count);
87        let mut metadata_results = Vec::new();
88        let source_anchor_counts = parallel_source_anchor_counts(&results);
89        for (i, result) in results.iter().enumerate() {
90            let status = if result.success { "[OK]" } else { "[ERR]" };
91            let (formatted, truncated) = format_task_result_for_context(result);
92            let (output_excerpt, _) = compact_task_output(&result.output);
93            let source_anchors = &result.source_anchors[..source_anchor_counts[i]];
94            metadata_results.push(serde_json::json!({
95                "task_id": result.task_id,
96                "session_id": result.session_id,
97                "agent": result.agent,
98                "success": result.success,
99                "error_message": (!result.success).then(|| {
100                    crate::text::truncate_utf8(&result.output, 1024).to_string()
101                }),
102                "output_excerpt": output_excerpt,
103                "structured": result.structured,
104                "source_anchors": source_anchors,
105                "output_bytes": result.output.len(),
106                "truncated_for_context": truncated,
107                "artifact_id": task_artifact_id(result),
108                "artifact_uri": task_artifact_uri(result),
109            }));
110            output.push_str(&format!(
111                "--- Task {} ({}) {} ---\n{}\n\n",
112                i + 1,
113                result.agent,
114                status,
115                formatted
116            ));
117        }
118
119        let success_count = results.iter().filter(|result| result.success).count();
120        let failed_count = results.len().saturating_sub(success_count);
121        let all_success = failed_count == 0;
122        let partial_failure = failed_count > 0 && success_count > 0;
123        if params.allow_partial_failure && partial_failure {
124            output.push_str(&format!(
125                "Partial failure tolerated: {success_count} succeeded, {failed_count} failed.\n"
126            ));
127        }
128        if run.timed_out {
129            output.push_str(&format!(
130                "Task fan-out timed out after {} ms; returned completed child results and marked unfinished children failed.\n",
131                run.timeout_ms.unwrap_or_default()
132            ));
133        } else if run.returned_early {
134            output.push_str(&format!(
135                "Task fan-out returned after reaching min_success_count={}; unfinished children were marked failed.\n",
136                run.min_success_count.unwrap_or_default()
137            ));
138        }
139
140        let tool_success = all_success || (params.allow_partial_failure && success_count > 0);
141        let mut output = if tool_success {
142            ToolOutput::success(output)
143        } else {
144            ToolOutput::error(output)
145        };
146        if !tool_success && failed_count > 0 {
147            output.error_kind = Some(crate::tools::ToolErrorKind::PartialFailure {
148                failed: failed_count,
149                total: results.len(),
150            });
151        }
152
153        Ok(output.with_metadata(serde_json::json!({
154            "task_count": task_count,
155            "result_count": results.len(),
156            "success_count": success_count,
157            "failed_count": failed_count,
158            "all_success": all_success,
159            "partial_failure": partial_failure,
160            "allow_partial_failure": params.allow_partial_failure,
161            "timeout_ms": params.timeout_ms,
162            "timed_out": run.timed_out,
163            "min_success_count": params.min_success_count,
164            "returned_early": run.returned_early,
165            "duration_ms": started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64,
166            "results": metadata_results,
167        })))
168    }
169}
170
171#[async_trait]
172impl Tool for ParallelTaskTool {
173    fn name(&self) -> &str {
174        "parallel_task"
175    }
176
177    fn description(&self) -> &str {
178        PARALLEL_TASK_TOOL_DESCRIPTION
179    }
180
181    fn parameters(&self) -> serde_json::Value {
182        parallel_params::parallel_task_params_schema_for_agents(&self.executor.visible_agents())
183    }
184
185    fn definition(&self) -> ToolDefinition {
186        let agents = self.executor.visible_agents();
187        ToolDefinition {
188            name: self.name().to_string(),
189            description: delegation_tool_description(self.description(), &agents),
190            parameters: parallel_params::parallel_task_params_schema_for_agents(&agents),
191        }
192    }
193
194    fn is_model_visible(&self) -> bool {
195        false
196    }
197
198    async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
199        let params: ParallelTaskParams = match serde_json::from_value(args.clone()) {
200            Ok(params) => params,
201            Err(error) => {
202                return Ok(invalid_delegation_argument(format!(
203                    "Invalid parallel_task parameters: {error}"
204                )));
205            }
206        };
207        self.execute_params(params, ctx, "parallel_task", 2).await
208    }
209}