1use crate::agent::{AgentConfig, AgentEvent, AgentLoop};
18use crate::llm::structured::{
19 generate_blocking, parse_validated_output, StructuredMode, StructuredRequest,
20};
21use crate::llm::{LlmClient, ToolDefinition};
22use crate::mcp::manager::McpManager;
23use crate::orchestration::{AgentExecutor, AgentStepSpec, StepOutcome, ToolSourceAnchor};
24use crate::subagent::{AgentDefinition, AgentRegistry};
25use crate::tools::types::{Tool, ToolContext, ToolOutput};
26use anyhow::{Context, Result};
27use async_trait::async_trait;
28use futures::FutureExt;
29use serde::{Deserialize, Serialize};
30use std::any::Any;
31use std::panic::AssertUnwindSafe;
32use std::path::PathBuf;
33use std::sync::Arc;
34use tokio::sync::broadcast;
35use tokio::task::JoinSet;
36use tokio_util::sync::CancellationToken;
37
38const TASK_OUTPUT_CONTEXT_LIMIT: usize = 4_000;
39const TASK_OUTPUT_CONTEXT_HEAD: usize = 3_000;
40const TASK_OUTPUT_CONTEXT_TAIL: usize = 800;
41const MAX_TASK_SOURCE_ANCHORS: usize = 64;
42const MAX_TASK_SOURCE_CANDIDATES: usize = MAX_TASK_SOURCE_ANCHORS * 4;
43const MAX_TASK_SOURCE_TOOL_BYTES: usize = 64;
44const MAX_TASK_SOURCE_VALUE_BYTES: usize = 4 * 1024;
45const MAX_PARALLEL_TASK_SOURCE_ANCHORS: usize = MAX_TASK_SOURCE_ANCHORS;
46const TASK_TOOL_DESCRIPTION: &str = "Delegate one or more bounded tasks to specialized child runs. Pass one item for a focused child run or multiple INDEPENDENT items for concurrent fan-out. A single item may run in the background; multi-item calls are collected by the parent. By default any failed child makes a multi-item call fail; evidence-gathering callers may set allow_partial_failure=true. Choose canonical worker names from the live agent catalog. Custom agents from agent_dirs and .a3s/agents are supported; .claude/agents is read for compatibility.";
47const PARALLEL_TASK_TOOL_DESCRIPTION: &str = "Fan out 2 or more INDEPENDENT subtasks as delegated child runs that execute concurrently; results are returned when all complete. By default any failed child makes the tool fail; evidence-gathering callers may set allow_partial_failure=true to continue when at least one child succeeds. Child output never authorizes branch replay; provider and child runtimes own any typed retry policy below this boundary. Use this only when the work genuinely splits into branches that can be investigated or implemented separately. Do not use it for trivial, conversational, single-step, or dependent work. Choose canonical worker names from the live agent catalog.";
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
51#[serde(deny_unknown_fields)]
52pub struct TaskParams {
53 pub agent: String,
55 pub description: String,
57 pub prompt: String,
59 #[serde(default)]
61 pub background: bool,
62 #[serde(skip_serializing_if = "Option::is_none")]
64 pub max_steps: Option<usize>,
65 #[serde(default, skip_serializing_if = "Option::is_none")]
67 pub output_schema: Option<serde_json::Value>,
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct TaskResult {
73 pub output: String,
75 pub session_id: String,
77 pub agent: String,
79 pub success: bool,
81 pub task_id: String,
83 #[serde(default, skip_serializing_if = "Option::is_none")]
85 pub structured: Option<serde_json::Value>,
86 #[serde(default, skip_serializing_if = "Vec::is_empty")]
88 pub source_anchors: Vec<ToolSourceAnchor>,
89}
90
91mod result_projection;
92use result_projection::*;
93
94mod parallel_execution;
95
96const MAX_PARALLEL_TASKS_PER_CALL: usize = 32;
97
98#[derive(Clone)]
100pub struct TaskExecutor {
101 registry: Arc<AgentRegistry>,
103 llm_client: Arc<dyn LlmClient>,
105 workspace: String,
107 mcp_managers: Vec<Arc<McpManager>>,
109 parent_context: Option<crate::child_run::ChildRunContext>,
111 search_config: Option<Arc<crate::config::SearchConfig>>,
113 search_bulkhead: Option<a3s_search::Bulkhead>,
115 search_retry_budget: Option<a3s_search::RetryBudget>,
117 search_request_coalescer: Option<a3s_search::SearchCoalescer>,
119 parent_cancellation: Option<CancellationToken>,
123 max_parallel_tasks: usize,
124 parallel_permits: Arc<tokio::sync::Semaphore>,
128 subagent_tracker: Option<Arc<crate::subagent_task_tracker::InMemorySubagentTaskTracker>>,
131}
132
133impl TaskExecutor {
134 pub fn new(
136 registry: Arc<AgentRegistry>,
137 llm_client: Arc<dyn LlmClient>,
138 workspace: String,
139 ) -> Self {
140 Self {
141 registry,
142 llm_client,
143 workspace,
144 mcp_managers: Vec::new(),
145 parent_context: None,
146 search_config: None,
147 search_bulkhead: None,
148 search_retry_budget: None,
149 search_request_coalescer: None,
150 parent_cancellation: None,
151 max_parallel_tasks: crate::agent::DEFAULT_MAX_PARALLEL_TASKS,
152 parallel_permits: Arc::new(tokio::sync::Semaphore::new(
153 crate::agent::DEFAULT_MAX_PARALLEL_TASKS,
154 )),
155 subagent_tracker: None,
156 }
157 }
158
159 pub fn with_mcp(
161 registry: Arc<AgentRegistry>,
162 llm_client: Arc<dyn LlmClient>,
163 workspace: String,
164 mcp_manager: Arc<McpManager>,
165 ) -> Self {
166 Self::with_mcp_managers(registry, llm_client, workspace, vec![mcp_manager])
167 }
168
169 pub fn with_mcp_managers(
171 registry: Arc<AgentRegistry>,
172 llm_client: Arc<dyn LlmClient>,
173 workspace: String,
174 mcp_managers: Vec<Arc<McpManager>>,
175 ) -> Self {
176 Self {
177 registry,
178 llm_client,
179 workspace,
180 mcp_managers,
181 parent_context: None,
182 search_config: None,
183 search_bulkhead: None,
184 search_retry_budget: None,
185 search_request_coalescer: None,
186 parent_cancellation: None,
187 max_parallel_tasks: crate::agent::DEFAULT_MAX_PARALLEL_TASKS,
188 parallel_permits: Arc::new(tokio::sync::Semaphore::new(
189 crate::agent::DEFAULT_MAX_PARALLEL_TASKS,
190 )),
191 subagent_tracker: None,
192 }
193 }
194
195 pub fn with_parent_context(mut self, ctx: crate::child_run::ChildRunContext) -> Self {
197 if let Some(max_parallel_tasks) = ctx.max_parallel_tasks {
198 let max_parallel_tasks = max_parallel_tasks.max(1);
199 self.max_parallel_tasks = max_parallel_tasks;
200 self.parallel_permits = Arc::new(tokio::sync::Semaphore::new(max_parallel_tasks));
201 }
202 self.parent_context = Some(ctx);
203 self
204 }
205
206 fn scoped_for_invocation(self: &Arc<Self>, ctx: &ToolContext) -> Arc<Self> {
207 let mut scoped = self.as_ref().clone();
208 scoped.search_config = ctx.search_config.clone();
209 scoped.search_bulkhead = Some(ctx.search_bulkhead());
210 scoped.search_retry_budget = Some(ctx.search_retry_budget());
211 scoped.search_request_coalescer = Some(ctx.search_request_coalescer());
212 if ctx.has_run_governance() {
213 scoped.parent_context = scoped.parent_context.take().map(|parent| {
214 parent.with_run_governance(
215 ctx.run_permission_checker(),
216 ctx.run_confirmation_manager(),
217 )
218 });
219 }
220 Arc::new(scoped)
221 }
222
223 fn child_tool_context(
224 &self,
225 session_id: String,
226 cancellation: CancellationToken,
227 ) -> ToolContext {
228 let mut context = ToolContext::new(PathBuf::from(&self.workspace))
229 .with_session_id(session_id)
230 .with_cancellation(cancellation);
231 if let (Some(bulkhead), Some(retry_budget)) =
232 (&self.search_bulkhead, &self.search_retry_budget)
233 {
234 context = context.with_search_runtime(bulkhead.clone(), retry_budget.clone());
235 }
236 if let Some(search_config) = &self.search_config {
237 context = context.with_search_config(search_config.as_ref().clone());
238 }
239 if let Some(coalescer) = &self.search_request_coalescer {
240 context = context.with_search_request_coalescer(coalescer.clone());
241 }
242 context
243 }
244
245 pub fn with_parent_cancellation(mut self, cancellation: CancellationToken) -> Self {
252 self.parent_cancellation = Some(cancellation);
253 self
254 }
255
256 pub fn with_max_parallel_tasks(mut self, max_parallel_tasks: usize) -> Self {
257 let max_parallel_tasks = max_parallel_tasks.max(1);
258 self.max_parallel_tasks = max_parallel_tasks;
259 self.parallel_permits = Arc::new(tokio::sync::Semaphore::new(max_parallel_tasks));
260 self
261 }
262
263 pub fn with_subagent_tracker(
267 mut self,
268 tracker: Arc<crate::subagent_task_tracker::InMemorySubagentTaskTracker>,
269 ) -> Self {
270 self.subagent_tracker = Some(tracker);
271 self
272 }
273
274 fn visible_agents(&self) -> Vec<AgentDefinition> {
275 self.registry.list_visible()
276 }
277
278 pub async fn execute(
283 &self,
284 params: TaskParams,
285 event_tx: Option<broadcast::Sender<AgentEvent>>,
286 parent_session_id: Option<&str>,
287 ) -> Result<TaskResult> {
288 self.execute_with_parent_cancellation(
289 params,
290 event_tx,
291 parent_session_id,
292 self.parent_cancellation.as_ref(),
293 )
294 .await
295 }
296
297 async fn execute_with_parent_cancellation(
298 &self,
299 params: TaskParams,
300 event_tx: Option<broadcast::Sender<AgentEvent>>,
301 parent_session_id: Option<&str>,
302 parent_cancellation: Option<&CancellationToken>,
303 ) -> Result<TaskResult> {
304 let task_id = format!("task-{}", uuid::Uuid::new_v4());
305 self.execute_with_task_id_scoped(
306 task_id,
307 params,
308 event_tx,
309 parent_session_id,
310 true,
311 parent_cancellation,
312 )
313 .await
314 }
315
316 pub async fn execute_with_task_id(
321 &self,
322 task_id: String,
323 params: TaskParams,
324 event_tx: Option<broadcast::Sender<AgentEvent>>,
325 parent_session_id: Option<&str>,
326 emit_start: bool,
327 ) -> Result<TaskResult> {
328 self.execute_with_task_id_scoped(
329 task_id,
330 params,
331 event_tx,
332 parent_session_id,
333 emit_start,
334 self.parent_cancellation.as_ref(),
335 )
336 .await
337 }
338
339 async fn execute_with_task_id_scoped(
340 &self,
341 task_id: String,
342 params: TaskParams,
343 event_tx: Option<broadcast::Sender<AgentEvent>>,
344 parent_session_id: Option<&str>,
345 emit_start: bool,
346 parent_cancellation: Option<&CancellationToken>,
347 ) -> Result<TaskResult> {
348 if parent_cancellation.is_some_and(CancellationToken::is_cancelled) {
349 anyhow::bail!("Operation cancelled by parent session");
350 }
351
352 let session_id = format!("task-run-{}", task_id);
353 let started_ms = epoch_ms();
354 let output_schema = params.output_schema.clone();
355
356 let agent = self
357 .registry
358 .get(¶ms.agent)
359 .context(format!("Unknown agent type: '{}'", params.agent))?;
360 let tool_free = agent.tool_free;
361 let tool_free_system = agent.prompt.clone();
362 let inherited_security_provider = self
363 .parent_context
364 .as_ref()
365 .and_then(|context| context.security_provider.clone());
366
367 if emit_start {
368 let event = AgentEvent::SubagentStart {
369 task_id: task_id.clone(),
370 session_id: session_id.clone(),
371 parent_session_id: parent_session_id.unwrap_or_default().to_string(),
372 agent: params.agent.clone(),
373 description: params.description.clone(),
374 started_ms,
375 };
376 let event = inherited_security_provider
377 .as_deref()
378 .map(|provider| crate::security::sanitize_agent_event(provider, &event))
379 .unwrap_or(event);
380 if let Some(ref tracker) = self.subagent_tracker {
381 tracker.record_event(&event).await;
382 }
383 if let Some(ref tx) = event_tx {
384 let _ = tx.send(event);
385 }
386 }
387
388 let child_executor = if let Some(ref parent_ctx) = self.parent_context {
391 if let Some(ref services) = parent_ctx.workspace_services {
392 crate::tools::ToolExecutor::new_with_workspace_services_and_artifact_limits(
393 self.workspace.clone(),
394 Arc::clone(services),
395 crate::tools::ArtifactStoreLimits::default(),
396 )
397 } else {
398 crate::tools::ToolExecutor::new(self.workspace.clone())
399 }
400 } else {
401 crate::tools::ToolExecutor::new(self.workspace.clone())
402 };
403
404 for mcp in &self.mcp_managers {
406 let all_tools = match parent_cancellation {
407 Some(cancellation) => {
408 tokio::select! {
409 biased;
410 _ = cancellation.cancelled() => {
411 anyhow::bail!("Operation cancelled by parent session");
412 }
413 tools = mcp.get_all_tools() => tools,
414 }
415 }
416 None => mcp.get_all_tools().await,
417 };
418 let mut by_server: std::collections::HashMap<
419 String,
420 Vec<crate::mcp::protocol::McpTool>,
421 > = std::collections::HashMap::new();
422 for (server, tool) in all_tools {
423 by_server.entry(server).or_default().push(tool);
424 }
425 for (server_name, tools) in by_server {
426 let wrappers =
427 crate::mcp::tools::create_mcp_tools(&server_name, tools, Arc::clone(mcp));
428 for wrapper in wrappers {
429 child_executor.register_dynamic_tool(wrapper);
430 }
431 }
432 }
433
434 let child_executor = Arc::new(child_executor);
435
436 let mut child_config = AgentConfig {
437 tools: child_executor.definitions(),
438 ..AgentConfig::default()
439 };
440 agent.apply_to(&mut child_config);
441 if let Some(ref parent_ctx) = self.parent_context {
442 parent_ctx.apply_to(&mut child_config);
443 }
444 child_config.planning_mode = crate::prompts::PlanningMode::Disabled;
449 if let Some(max_steps) = params.max_steps {
450 child_config.max_tool_rounds = max_steps;
451 }
452 let child_security_provider = child_config.security_provider.clone();
453 let source_security_provider = child_security_provider.clone();
454
455 let cancel_token = parent_cancellation
456 .map(CancellationToken::child_token)
457 .unwrap_or_default();
458 let mut tool_context = self.child_tool_context(session_id.clone(), cancel_token.clone());
459 if let Some(ref parent_ctx) = self.parent_context {
460 if let Some(ref services) = parent_ctx.workspace_services {
461 tool_context = tool_context.with_workspace_services(Arc::clone(services));
462 }
463 if let Some(ref sandbox) = parent_ctx.sandbox_handle {
464 child_executor.registry().set_sandbox(Arc::clone(sandbox));
465 tool_context = tool_context.with_sandbox(Arc::clone(sandbox));
466 }
467 }
468
469 let source_context = tool_context.clone();
470 let agent_loop = AgentLoop::new(
471 Arc::clone(&self.llm_client),
472 child_executor,
473 tool_context,
474 child_config,
475 );
476
477 let (mpsc_tx, mut mpsc_rx) = tokio::sync::mpsc::channel(100);
481 let broadcast_tx = event_tx.clone();
482 let progress_task_id = task_id.clone();
483 let progress_session_id = session_id.clone();
484 let child_event_forwarder = tokio::spawn(async move {
485 let mut source_anchors = Vec::new();
486 let mut seen_source_anchors = std::collections::HashSet::new();
487 let mut scanned_source_candidates = 0usize;
488 while let Some(event) = mpsc_rx.recv().await {
489 let event = source_security_provider
490 .as_deref()
491 .map(|provider| crate::security::sanitize_agent_event(provider, &event))
492 .unwrap_or(event);
493 collect_tool_source_anchors(
494 &event,
495 &source_context,
496 &mut source_anchors,
497 &mut seen_source_anchors,
498 &mut scanned_source_candidates,
499 );
500 if let Some(ref broadcast_tx) = broadcast_tx {
501 if let Some(progress) = synthesize_subagent_progress(
502 &event,
503 &progress_task_id,
504 &progress_session_id,
505 ) {
506 let _ = broadcast_tx.send(progress);
507 }
508 let _ = broadcast_tx.send(event);
509 }
510 }
511 source_anchors
512 });
513 let child_event_tx = Some(mpsc_tx);
514 let child_llm_event_tx = child_event_tx.clone();
515
516 if let Some(ref tracker) = self.subagent_tracker {
519 tracker
520 .register_canceller(&task_id, cancel_token.clone())
521 .await;
522 }
523
524 let structured_prompt = output_schema
525 .as_ref()
526 .filter(|_| !tool_free)
527 .map(|schema| structured_task_prompt(¶ms.prompt, schema));
528 let execution_prompt = structured_prompt.as_deref().unwrap_or(¶ms.prompt);
529
530 let mut structured = None;
531 let (mut output, mut success) = if tool_free && output_schema.is_some() {
532 let llm_client = agent_loop.scoped_llm_client_for_parts(
533 Some(&session_id),
534 &child_llm_event_tx,
535 &cancel_token,
536 );
537 match Self::generate_structured_task(
538 &*llm_client,
539 ¶ms.prompt,
540 tool_free_system.as_deref(),
541 output_schema.clone().expect("schema checked above"),
542 &cancel_token,
543 )
544 .await
545 {
546 Ok(object) => {
547 let output = serde_json::to_string_pretty(&object)
548 .unwrap_or_else(|_| object.to_string());
549 structured = Some(object);
550 (output, true)
551 }
552 Err(error) if cancel_token.is_cancelled() => {
553 (format!("Task cancelled by caller: {error}"), false)
554 }
555 Err(error) => (format!("Task failed: {error}"), false),
556 }
557 } else {
558 match agent_loop
559 .execute_with_session(
560 &[],
561 execution_prompt,
562 Some(&session_id),
563 child_event_tx.clone(),
564 Some(&cancel_token),
565 )
566 .await
567 {
568 Ok(_) if cancel_token.is_cancelled() => {
569 ("Task cancelled by caller".to_string(), false)
570 }
571 Ok(result) if result.text.trim().is_empty() => (
572 "Task failed: child agent returned no final output".to_string(),
573 false,
574 ),
575 Ok(result) if AgentLoop::is_synthetic_failure_output(&result.text) => {
576 (format!("Task failed: {}", result.text), false)
577 }
578 Ok(result) => (result.text, true),
579 Err(e) if cancel_token.is_cancelled() => {
580 (format!("Task cancelled by caller: {}", e), false)
581 }
582 Err(e) => (format!("Task failed: {}", e), false),
583 }
584 };
585
586 if success && !tool_free {
587 if let Some(schema) = output_schema {
588 if let Some(object) = parse_validated_output(&output, &schema) {
589 structured = Some(object);
590 } else {
591 let llm_client = agent_loop.scoped_llm_client_for_parts(
592 Some(&session_id),
593 &child_llm_event_tx,
594 &cancel_token,
595 );
596 match Self::coerce_to_schema(&*llm_client, &output, schema, &cancel_token).await
597 {
598 Ok(object) => structured = Some(object),
599 Err(error) => {
600 success = false;
601 output = format!("{output}\n\n[structured output failed: {error}]");
602 }
603 }
604 }
605 }
606 }
607 if let Some(provider) = child_security_provider.as_deref() {
608 output = provider.sanitize_output(&output);
609 if let Some(value) = &mut structured {
610 *value = sanitize_task_json(provider, value);
611 }
612 }
613
614 drop(child_event_tx);
619 drop(child_llm_event_tx);
620 let source_anchors = match child_event_forwarder.await {
621 Ok(source_anchors) => source_anchors,
622 Err(error) => {
623 tracing::warn!(%error, task_id = %task_id, "subagent event bridge failed");
624 Vec::new()
625 }
626 };
627
628 let end_event = AgentEvent::SubagentEnd {
629 task_id: task_id.clone(),
630 session_id: session_id.clone(),
631 agent: params.agent.clone(),
632 output: output.clone(),
633 success,
634 finished_ms: epoch_ms(),
635 };
636 if let Some(ref tracker) = self.subagent_tracker {
637 if success {
640 tracker
641 .record_source_anchors(&task_id, &source_anchors)
642 .await;
643 }
644 tracker.record_event(&end_event).await;
645 tracker.clear_canceller(&task_id).await;
646 }
647 if let Some(ref tx) = event_tx {
648 let _ = tx.send(end_event);
649 }
650
651 Ok(TaskResult {
652 output,
653 session_id,
654 agent: params.agent,
655 success,
656 task_id,
657 structured,
658 source_anchors,
659 })
660 }
661
662 pub fn execute_background(
670 self: Arc<Self>,
671 params: TaskParams,
672 event_tx: Option<broadcast::Sender<AgentEvent>>,
673 parent_session_id: Option<String>,
674 ) -> String {
675 let parent_cancellation = self.parent_cancellation.clone();
676 self.execute_background_with_parent_cancellation(
677 params,
678 event_tx,
679 parent_session_id,
680 parent_cancellation,
681 )
682 }
683
684 fn execute_background_with_parent_cancellation(
685 self: Arc<Self>,
686 params: TaskParams,
687 event_tx: Option<broadcast::Sender<AgentEvent>>,
688 parent_session_id: Option<String>,
689 parent_cancellation: Option<CancellationToken>,
690 ) -> String {
691 let task_id = format!("task-{}", uuid::Uuid::new_v4());
692 let session_id = format!("task-run-{}", task_id);
693 let failure_session_id = session_id.clone();
694 let failure_agent = params.agent.clone();
695 let start_event = AgentEvent::SubagentStart {
696 task_id: task_id.clone(),
697 session_id,
698 parent_session_id: parent_session_id.clone().unwrap_or_default(),
699 agent: params.agent.clone(),
700 description: params.description.clone(),
701 started_ms: epoch_ms(),
702 };
703 let security_provider = self
704 .parent_context
705 .as_ref()
706 .and_then(|context| context.security_provider.clone());
707 let start_event = security_provider
708 .as_deref()
709 .map(|provider| crate::security::sanitize_agent_event(provider, &start_event))
710 .unwrap_or(start_event);
711
712 if let Some(ref tx) = event_tx {
713 let _ = tx.send(start_event.clone());
714 }
715
716 let task_id_for_spawn = task_id.clone();
717 let task_id_for_log = task_id.clone();
718 tokio::spawn(async move {
719 if let Some(ref tracker) = self.subagent_tracker {
720 tracker.record_event(&start_event).await;
721 }
722 let failure_event_tx = event_tx.clone();
723 if let Err(error) = self
724 .execute_with_task_id_scoped(
725 task_id_for_spawn,
726 params,
727 event_tx,
728 parent_session_id.as_deref(),
729 false,
730 parent_cancellation.as_ref(),
731 )
732 .await
733 {
734 let end_event = AgentEvent::SubagentEnd {
735 task_id: task_id_for_log.clone(),
736 session_id: failure_session_id,
737 agent: failure_agent,
738 output: format!("Task failed before child execution started: {error}"),
739 success: false,
740 finished_ms: epoch_ms(),
741 };
742 let end_event = security_provider
743 .as_deref()
744 .map(|provider| crate::security::sanitize_agent_event(provider, &end_event))
745 .unwrap_or(end_event);
746 if let Some(ref tracker) = self.subagent_tracker {
747 tracker.record_event(&end_event).await;
748 tracker.clear_canceller(&task_id_for_log).await;
749 }
750 if let Some(tx) = failure_event_tx {
751 let _ = tx.send(end_event);
752 }
753 tracing::error!("Background task {} failed: {}", task_id_for_log, error);
754 }
755 });
756
757 task_id
758 }
759}
760
761fn structured_task_prompt(prompt: &str, schema: &serde_json::Value) -> String {
762 let schema = serde_json::to_string_pretty(schema).unwrap_or_else(|_| schema.to_string());
763 format!(
764 "{prompt}\n\n\
765 FINAL OUTPUT CONTRACT\n\
766 Complete the requested investigation before answering. Your final response must contain \
767 exactly one JSON value matching the JSON Schema below, with no Markdown fence or prose \
768 outside the JSON. This contract applies to the final response only; use the available \
769 tools as needed before finalizing.\n\n\
770 {schema}"
771 )
772}
773
774#[derive(Debug, Clone)]
775struct AgentCatalogEntry {
776 name: String,
777 description: String,
778}
779
780fn agent_catalog_entries(agents: &[AgentDefinition]) -> Vec<AgentCatalogEntry> {
781 let mut entries = agents
782 .iter()
783 .map(|agent| AgentCatalogEntry {
784 name: agent.name.clone(),
785 description: agent
786 .description
787 .split_whitespace()
788 .collect::<Vec<_>>()
789 .join(" "),
790 })
791 .collect::<Vec<_>>();
792 entries.sort_by(|left, right| left.name.cmp(&right.name));
793 entries
794}
795
796fn agent_catalog_text(agents: &[AgentDefinition]) -> String {
797 agent_catalog_entries(agents)
798 .into_iter()
799 .map(|entry| format!("{}: {}", entry.name, entry.description))
800 .collect::<Vec<_>>()
801 .join("\n")
802}
803
804fn delegation_tool_description(base: &str, agents: &[AgentDefinition]) -> String {
805 format!(
806 "{base}\n\nAvailable agents (live catalog; use canonical names):\n{}",
807 agent_catalog_text(agents)
808 )
809}
810
811pub(super) fn task_agent_parameter_schema(agents: &[AgentDefinition]) -> serde_json::Value {
812 let entries = agent_catalog_entries(agents);
813 let examples = entries
814 .iter()
815 .map(|entry| serde_json::Value::String(entry.name.clone()))
816 .collect::<Vec<_>>();
817 let catalog = entries
818 .into_iter()
819 .map(|entry| format!("{}: {}", entry.name, entry.description))
820 .collect::<Vec<_>>()
821 .join("\n");
822 serde_json::json!({
823 "type": "string",
824 "description": format!(
825 "Required. Canonical agent type to use. Always provide this exact field name: 'agent'. Live agent catalog:\n{catalog}"
826 ),
827 "examples": examples
828 })
829}
830
831pub fn task_params_schema() -> serde_json::Value {
837 task_params_schema_for_agents(&AgentRegistry::new().list_visible())
838}
839
840fn task_params_schema_for_agents(agents: &[AgentDefinition]) -> serde_json::Value {
841 serde_json::json!({
842 "oneOf": [
843 legacy_task_params_schema_for_agents(agents),
844 task_model_params_schema_for_agents(agents)
845 ]
846 })
847}
848
849fn legacy_task_params_schema_for_agents(agents: &[AgentDefinition]) -> serde_json::Value {
850 let mut schema = task_item_params_schema_for_agents(agents, true);
851 schema["examples"] = serde_json::json!([
852 {
853 "agent": "explore",
854 "description": "Find Rust files",
855 "prompt": "Search the workspace for Rust files and summarize the layout."
856 },
857 {
858 "agent": "general",
859 "description": "Investigate test failure",
860 "prompt": "Inspect the failing tests and explain the root cause.",
861 "max_steps": 6
862 }
863 ]);
864 schema
865}
866
867fn task_model_params_schema_for_agents(agents: &[AgentDefinition]) -> serde_json::Value {
868 parallel_params::task_tool_params_schema_for_agents(agents)
869}
870
871pub(super) fn task_item_params_schema_for_agents(
872 agents: &[AgentDefinition],
873 include_background: bool,
874) -> serde_json::Value {
875 let mut properties = serde_json::Map::from_iter([
876 ("agent".to_string(), task_agent_parameter_schema(agents)),
877 (
878 "description".to_string(),
879 serde_json::json!({
880 "type": "string",
881 "description": "Required. Short task label for display and tracking. Always provide this exact field name: 'description'."
882 }),
883 ),
884 (
885 "prompt".to_string(),
886 serde_json::json!({
887 "type": "string",
888 "description": "Required. Detailed instruction for the delegated child run. Always provide this exact field name: 'prompt'."
889 }),
890 ),
891 (
892 "max_steps".to_string(),
893 serde_json::json!({
894 "type": "integer",
895 "description": "Optional. Maximum number of steps for this task."
896 }),
897 ),
898 (
899 "output_schema".to_string(),
900 serde_json::json!({
901 "type": "object",
902 "description": "Optional. JSON Schema object the delegated result must satisfy. When provided, the child output is coerced into a validated structured object and returned in metadata."
903 }),
904 ),
905 ]);
906 if include_background {
907 properties.insert(
908 "background".to_string(),
909 serde_json::json!({
910 "type": "boolean",
911 "description": "Optional. Run this task in the background. Only valid when the outer tasks array contains one item. Default: false.",
912 "default": false
913 }),
914 );
915 }
916
917 serde_json::json!({
918 "type": "object",
919 "additionalProperties": false,
920 "properties": properties,
921 "required": ["agent", "description", "prompt"]
922 })
923}
924
925pub struct TaskTool {
928 executor: Arc<TaskExecutor>,
929}
930
931impl TaskTool {
932 pub fn new(executor: Arc<TaskExecutor>) -> Self {
934 Self { executor }
935 }
936
937 async fn execute_single(&self, params: TaskParams, ctx: &ToolContext) -> Result<ToolOutput> {
938 let parent_cancellation = ctx.cancellation_token();
939 let executor = self.executor.scoped_for_invocation(ctx);
940
941 if params.background {
942 let task_id = executor.execute_background_with_parent_cancellation(
943 params,
944 ctx.agent_event_tx.clone(),
945 ctx.session_id.clone(),
946 Some(parent_cancellation),
947 );
948 return Ok(ToolOutput::success(format!(
949 "Task started in background. Task ID: {}",
950 task_id
951 )));
952 }
953
954 let result = executor
955 .execute_with_parent_cancellation(
956 params,
957 ctx.agent_event_tx.clone(),
958 ctx.session_id.as_deref(),
959 Some(&parent_cancellation),
960 )
961 .await?;
962 let (content, truncated) = format_task_result_for_context(&result);
963 let metadata = serde_json::json!({
964 "task_id": result.task_id,
965 "session_id": result.session_id,
966 "agent": result.agent,
967 "success": result.success,
968 "output_bytes": result.output.len(),
969 "truncated_for_context": truncated,
970 "artifact_id": task_artifact_id(&result),
971 "artifact_uri": task_artifact_uri(&result),
972 "structured": result.structured,
973 "source_anchors": result.source_anchors,
974 });
975
976 if result.success {
977 Ok(ToolOutput::success(content).with_metadata(metadata))
978 } else {
979 Ok(ToolOutput::error(content).with_metadata(metadata))
980 }
981 }
982}
983
984#[async_trait]
985impl Tool for TaskTool {
986 fn name(&self) -> &str {
987 "task"
988 }
989
990 fn description(&self) -> &str {
991 TASK_TOOL_DESCRIPTION
992 }
993
994 fn parameters(&self) -> serde_json::Value {
995 task_params_schema_for_agents(&self.executor.visible_agents())
996 }
997
998 fn definition(&self) -> ToolDefinition {
999 let agents = self.executor.visible_agents();
1000 ToolDefinition {
1001 name: self.name().to_string(),
1002 description: delegation_tool_description(self.description(), &agents),
1003 parameters: task_model_params_schema_for_agents(&agents),
1004 }
1005 }
1006
1007 async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
1008 if args.get("tasks").is_some() {
1009 let mut params: ParallelTaskParams = match serde_json::from_value(args.clone()) {
1010 Ok(params) => params,
1011 Err(error) => {
1012 return Ok(invalid_delegation_argument(format!(
1013 "Invalid task parameters: {error}"
1014 )));
1015 }
1016 };
1017 if params.tasks.is_empty() {
1018 return Ok(invalid_delegation_argument(
1019 "task requires at least 1 task".to_string(),
1020 ));
1021 }
1022
1023 let has_fanout_options = params.allow_partial_failure
1024 || params.timeout_ms.is_some()
1025 || params.min_success_count.is_some();
1026 if params.tasks.len() == 1 && !has_fanout_options {
1027 return self.execute_single(params.tasks.remove(0), ctx).await;
1028 }
1029
1030 return ParallelTaskTool::new(Arc::clone(&self.executor))
1031 .execute_params(params, ctx, "task", 1)
1032 .await;
1033 }
1034
1035 let params: TaskParams = match serde_json::from_value(args.clone()) {
1036 Ok(params) => params,
1037 Err(error) => {
1038 return Ok(invalid_delegation_argument(format!(
1039 "Invalid task parameters: {error}"
1040 )));
1041 }
1042 };
1043 self.execute_single(params, ctx).await
1044 }
1045}
1046
1047mod parallel_params;
1048pub use parallel_params::{parallel_task_params_schema, ParallelTaskParams};
1049
1050pub struct ParallelTaskTool {
1054 executor: Arc<TaskExecutor>,
1055}
1056
1057impl ParallelTaskTool {
1058 pub fn new(executor: Arc<TaskExecutor>) -> Self {
1060 Self { executor }
1061 }
1062
1063 async fn execute_params(
1064 &self,
1065 params: ParallelTaskParams,
1066 ctx: &ToolContext,
1067 tool_name: &str,
1068 min_tasks: usize,
1069 ) -> Result<ToolOutput> {
1070 let started_at = std::time::Instant::now();
1071 let parent_cancellation = ctx.cancellation_token();
1072 let executor = self.executor.scoped_for_invocation(ctx);
1073
1074 if params.tasks.len() < min_tasks {
1075 return Ok(invalid_delegation_argument(format!(
1076 "{tool_name} requires at least {min_tasks} task{}",
1077 if min_tasks == 1 { "" } else { "s" }
1078 )));
1079 }
1080 if params.tasks.len() > MAX_PARALLEL_TASKS_PER_CALL {
1081 return Ok(invalid_delegation_argument(format!(
1082 "{tool_name} accepts at most {MAX_PARALLEL_TASKS_PER_CALL} tasks"
1083 )));
1084 }
1085 if let Some((index, _)) = params
1086 .tasks
1087 .iter()
1088 .enumerate()
1089 .find(|(_, task)| task.background)
1090 {
1091 return Ok(invalid_delegation_argument(format!(
1092 "{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",
1093 index + 1
1094 )));
1095 }
1096 if params.timeout_ms == Some(0) {
1097 return Ok(invalid_delegation_argument(format!(
1098 "{tool_name} timeout_ms must be at least 1"
1099 )));
1100 }
1101 if let Some(min_success_count) = params.min_success_count {
1102 if !params.allow_partial_failure {
1103 return Ok(invalid_delegation_argument(format!(
1104 "{tool_name} min_success_count requires allow_partial_failure=true"
1105 )));
1106 }
1107 if min_success_count == 0 || min_success_count > params.tasks.len() {
1108 return Ok(invalid_delegation_argument(format!(
1109 "{tool_name} min_success_count must be between 1 and the task count ({})",
1110 params.tasks.len()
1111 )));
1112 }
1113 }
1114
1115 let task_count = params.tasks.len();
1116 let run = executor
1117 .execute_parallel_for_tool(
1118 params.tasks.clone(),
1119 ctx.agent_event_tx.clone(),
1120 parallel_execution::ParallelToolOptions {
1121 parent_session_id: ctx.session_id.as_deref(),
1122 timeout_ms: params.timeout_ms,
1123 min_success_count: params.min_success_count,
1124 allow_partial_failure: params.allow_partial_failure,
1125 parent_cancellation: Some(&parent_cancellation),
1126 },
1127 )
1128 .await;
1129 let results = run.results;
1130
1131 let mut output = format!("Executed {} tasks concurrently:\n\n", task_count);
1132 let mut metadata_results = Vec::new();
1133 let source_anchor_counts = parallel_source_anchor_counts(&results);
1134 for (i, result) in results.iter().enumerate() {
1135 let status = if result.success { "[OK]" } else { "[ERR]" };
1136 let (formatted, truncated) = format_task_result_for_context(result);
1137 let (output_excerpt, _) = compact_task_output(&result.output);
1138 let source_anchors = &result.source_anchors[..source_anchor_counts[i]];
1139 metadata_results.push(serde_json::json!({
1140 "task_id": result.task_id,
1141 "session_id": result.session_id,
1142 "agent": result.agent,
1143 "success": result.success,
1144 "error_message": (!result.success).then(|| {
1145 crate::text::truncate_utf8(&result.output, 1024).to_string()
1146 }),
1147 "output_excerpt": output_excerpt,
1148 "structured": result.structured,
1149 "source_anchors": source_anchors,
1150 "output_bytes": result.output.len(),
1151 "truncated_for_context": truncated,
1152 "artifact_id": task_artifact_id(result),
1153 "artifact_uri": task_artifact_uri(result),
1154 }));
1155 output.push_str(&format!(
1156 "--- Task {} ({}) {} ---\n{}\n\n",
1157 i + 1,
1158 result.agent,
1159 status,
1160 formatted
1161 ));
1162 }
1163
1164 let success_count = results.iter().filter(|result| result.success).count();
1165 let failed_count = results.len().saturating_sub(success_count);
1166 let all_success = failed_count == 0;
1167 let partial_failure = failed_count > 0 && success_count > 0;
1168 if params.allow_partial_failure && partial_failure {
1169 output.push_str(&format!(
1170 "Partial failure tolerated: {success_count} succeeded, {failed_count} failed.\n"
1171 ));
1172 }
1173 if run.timed_out {
1174 output.push_str(&format!(
1175 "Task fan-out timed out after {} ms; returned completed child results and marked unfinished children failed.\n",
1176 run.timeout_ms.unwrap_or_default()
1177 ));
1178 } else if run.returned_early {
1179 output.push_str(&format!(
1180 "Task fan-out returned after reaching min_success_count={}; unfinished children were marked failed.\n",
1181 run.min_success_count.unwrap_or_default()
1182 ));
1183 }
1184
1185 let tool_success = all_success || (params.allow_partial_failure && success_count > 0);
1186 let mut output = if tool_success {
1187 ToolOutput::success(output)
1188 } else {
1189 ToolOutput::error(output)
1190 };
1191 if !tool_success && failed_count > 0 {
1192 output.error_kind = Some(crate::tools::ToolErrorKind::PartialFailure {
1193 failed: failed_count,
1194 total: results.len(),
1195 });
1196 }
1197
1198 Ok(output.with_metadata(serde_json::json!({
1199 "task_count": task_count,
1200 "result_count": results.len(),
1201 "success_count": success_count,
1202 "failed_count": failed_count,
1203 "all_success": all_success,
1204 "partial_failure": partial_failure,
1205 "allow_partial_failure": params.allow_partial_failure,
1206 "timeout_ms": params.timeout_ms,
1207 "timed_out": run.timed_out,
1208 "min_success_count": params.min_success_count,
1209 "returned_early": run.returned_early,
1210 "duration_ms": started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64,
1211 "results": metadata_results,
1212 })))
1213 }
1214}
1215
1216#[async_trait]
1217impl Tool for ParallelTaskTool {
1218 fn name(&self) -> &str {
1219 "parallel_task"
1220 }
1221
1222 fn description(&self) -> &str {
1223 PARALLEL_TASK_TOOL_DESCRIPTION
1224 }
1225
1226 fn parameters(&self) -> serde_json::Value {
1227 parallel_params::parallel_task_params_schema_for_agents(&self.executor.visible_agents())
1228 }
1229
1230 fn definition(&self) -> ToolDefinition {
1231 let agents = self.executor.visible_agents();
1232 ToolDefinition {
1233 name: self.name().to_string(),
1234 description: delegation_tool_description(self.description(), &agents),
1235 parameters: parallel_params::parallel_task_params_schema_for_agents(&agents),
1236 }
1237 }
1238
1239 fn is_model_visible(&self) -> bool {
1240 false
1241 }
1242
1243 async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
1244 let params: ParallelTaskParams = match serde_json::from_value(args.clone()) {
1245 Ok(params) => params,
1246 Err(error) => {
1247 return Ok(invalid_delegation_argument(format!(
1248 "Invalid parallel_task parameters: {error}"
1249 )));
1250 }
1251 };
1252 self.execute_params(params, ctx, "parallel_task", 2).await
1253 }
1254}
1255
1256fn invalid_delegation_argument(message: String) -> ToolOutput {
1257 ToolOutput::error(&message)
1258 .with_error_kind(crate::tools::ToolErrorKind::InvalidArgument { message })
1259}
1260
1261#[cfg(test)]
1262mod tests;