Skip to main content

a3s_code_core/tools/
task.rs

1//! Task tools for delegated child runs.
2//!
3//! The Task tool allows the main agent to delegate one or more specialized
4//! tasks to focused child runs. Each child run gets bounded context and the
5//! permissions declared by its agent definition.
6//!
7//! ## Usage
8//!
9//! ```json
10//! {"tasks": [{
11//!   "agent": "explore",
12//!   "description": "Find authentication code",
13//!   "prompt": "Search for files related to user authentication..."
14//! }]}
15//! ```
16
17use 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/// Task tool parameters
50#[derive(Debug, Clone, Serialize, Deserialize)]
51#[serde(deny_unknown_fields)]
52pub struct TaskParams {
53    /// Agent type to use (explore, general, plan, verification, review, etc.)
54    pub agent: String,
55    /// Short description of the task (for display)
56    pub description: String,
57    /// Detailed prompt for the agent
58    pub prompt: String,
59    /// Optional: run in background (default: false)
60    #[serde(default)]
61    pub background: bool,
62    /// Optional: maximum steps for this task
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub max_steps: Option<usize>,
65    /// Optional: JSON schema the child result must satisfy.
66    #[serde(default, skip_serializing_if = "Option::is_none")]
67    pub output_schema: Option<serde_json::Value>,
68}
69
70/// Task tool result
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct TaskResult {
73    /// Task output from the delegated child run.
74    pub output: String,
75    /// Child session ID
76    pub session_id: String,
77    /// Agent type used
78    pub agent: String,
79    /// Whether the task succeeded
80    pub success: bool,
81    /// Task ID for tracking
82    pub task_id: String,
83    /// Structured child output validated against an optional output schema.
84    #[serde(default, skip_serializing_if = "Option::is_none")]
85    pub structured: Option<serde_json::Value>,
86    /// Source locations observed by successful built-in child tool calls.
87    #[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/// Task executor for delegated child runs.
99#[derive(Clone)]
100pub struct TaskExecutor {
101    /// Agent registry for looking up agent definitions
102    registry: Arc<AgentRegistry>,
103    /// LLM client used to power child agent loops
104    llm_client: Arc<dyn LlmClient>,
105    /// Workspace path shared with child agents
106    workspace: String,
107    /// Ordered MCP managers for registering inherited tools in child sessions.
108    mcp_managers: Vec<Arc<McpManager>>,
109    /// Parent capabilities to inherit into child runs.
110    parent_context: Option<crate::child_run::ChildRunContext>,
111    /// Search configuration captured from the invoking parent context.
112    search_config: Option<Arc<crate::config::SearchConfig>>,
113    /// Agent-scoped search admission shared with delegated child runs.
114    search_bulkhead: Option<a3s_search::Bulkhead>,
115    /// Agent-scoped headless retry allowance shared with delegated child runs.
116    search_retry_budget: Option<a3s_search::RetryBudget>,
117    /// Parent-session request flights shared with delegated child runs.
118    search_request_coalescer: Option<a3s_search::SearchCoalescer>,
119    /// Optional lifetime boundary inherited from the session that created this
120    /// executor. Keeping it on the executor prevents cached workflow/executor
121    /// handles from starting new child runs after their session is closed.
122    parent_cancellation: Option<CancellationToken>,
123    max_parallel_tasks: usize,
124    /// Shared across every fan-out started by this executor. Per-call wave
125    /// limits alone are insufficient when a dynamic workflow launches several
126    /// `parallel_task` host steps concurrently.
127    parallel_permits: Arc<tokio::sync::Semaphore>,
128    /// Optional shared tracker — when present each task registers a
129    /// `CancellationToken` so callers can cancel by `task_id`.
130    subagent_tracker: Option<Arc<crate::subagent_task_tracker::InMemorySubagentTaskTracker>>,
131}
132
133impl TaskExecutor {
134    /// Create a new task executor
135    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    /// Create a new task executor with MCP manager for tool inheritance
160    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    /// Create a task executor with ordered MCP capability sources.
170    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    /// Set parent session capabilities to inherit into child runs.
196    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    /// Bind every run started by this executor to a parent lifetime.
246    ///
247    /// A token that is already cancelled makes execution fail before emitting
248    /// `SubagentStart` or performing MCP/LLM work. In-flight children derive
249    /// their own token so cancellation still cascades without granting them the
250    /// ability to cancel the parent.
251    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    /// Share a tracker with this executor. When set, each task registers
264    /// a `CancellationToken` against the tracker so the parent session
265    /// can cancel by `task_id`.
266    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    /// Execute a task by spawning an isolated child AgentLoop.
279    ///
280    /// `parent_session_id` flows into the emitted `SubagentStart`/`SubagentEnd`
281    /// events so dashboards can associate child runs with the parent session.
282    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    /// Execute a task using a caller-supplied task id. Used by `execute_background`
317    /// so the synchronously-returned task id matches the one in lifecycle events.
318    /// When `emit_start` is `false` the caller is responsible for emitting
319    /// `SubagentStart` themselves (e.g. to avoid a race against a tracker query).
320    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(&params.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        // Build a child ToolExecutor. Task tools are intentionally omitted
389        // here to prevent unlimited delegation nesting.
390        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        // Register MCP tools so child agents can access MCP servers.
405        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        // A delegated task is already the output of a parent planning
445        // decision. Running the generic pre-analysis/planning classifier again
446        // adds an unrelated LLM round to every child and can consume the whole
447        // fan-out deadline before any task tool runs.
448        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        // Always observe the child event stream so successful source tool calls
478        // survive in TaskResult metadata even when nobody subscribed to live
479        // progress. Forward the same events when a parent broadcast exists.
480        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        // Register a CancellationToken with the tracker (if shared) so the
517        // parent session's `cancel_subagent_task` can interrupt this run.
518        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(&params.prompt, schema));
528        let execution_prompt = structured_prompt.as_deref().unwrap_or(&params.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                &params.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        // The child loop and optional structured-output pass are the only
615        // producers. Close their sender and drain the bridge before emitting
616        // SubagentEnd so callers never observe a terminal event followed by
617        // stale child deltas or progress events.
618        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            // The tracker is authoritative even when a background child
638            // finishes after the parent run's event forwarder has closed.
639            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    /// Execute a task in the background.
663    ///
664    /// Returns immediately with the task ID; the same id is used in the emitted
665    /// `SubagentStart`/`SubagentEnd` events so callers can correlate. Pre-emits
666    /// `SubagentStart` synchronously when an event channel is available so a
667    /// caller that queries the subagent task tracker right after this call
668    /// observes the task in `Running` state without a race window.
669    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
831/// Get the compatibility JSON schema accepted by the `task` executor.
832///
833/// The model sees the more compact array-only schema. This public schema also
834/// accepts the pre-6.8 single-task object so persisted and host-direct calls do
835/// not break during the tool-surface migration.
836pub 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
925/// TaskTool wraps TaskExecutor as a Tool for registration in ToolExecutor.
926/// This allows the LLM to delegate tasks through the standard tool interface.
927pub struct TaskTool {
928    executor: Arc<TaskExecutor>,
929}
930
931impl TaskTool {
932    /// Create a new TaskTool
933    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
1050/// ParallelTaskTool allows the LLM to fan out multiple delegated tasks concurrently.
1051///
1052/// All tasks execute in parallel and the tool returns when all complete.
1053pub struct ParallelTaskTool {
1054    executor: Arc<TaskExecutor>,
1055}
1056
1057impl ParallelTaskTool {
1058    /// Create a new ParallelTaskTool
1059    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;