1use crate::agent::{AgentConfig, AgentEvent, AgentLoop};
18use crate::llm::structured::{generate_blocking, StructuredMode, StructuredRequest};
19use crate::llm::LlmClient;
20use crate::mcp::manager::McpManager;
21use crate::orchestration::{AgentExecutor, AgentStepSpec, StepOutcome};
22use crate::subagent::AgentRegistry;
23use crate::tools::types::{Tool, ToolContext, ToolOutput};
24use anyhow::{Context, Result};
25use async_trait::async_trait;
26use serde::{Deserialize, Serialize};
27use std::path::PathBuf;
28use std::sync::Arc;
29use tokio::sync::broadcast;
30
31const TASK_OUTPUT_CONTEXT_LIMIT: usize = 4_000;
32const TASK_OUTPUT_CONTEXT_HEAD: usize = 3_000;
33const TASK_OUTPUT_CONTEXT_TAIL: usize = 800;
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
37#[serde(deny_unknown_fields)]
38pub struct TaskParams {
39 pub agent: String,
41 pub description: String,
43 pub prompt: String,
45 #[serde(default)]
47 pub background: bool,
48 #[serde(skip_serializing_if = "Option::is_none")]
50 pub max_steps: Option<usize>,
51 #[serde(default, skip_serializing_if = "Option::is_none")]
53 pub output_schema: Option<serde_json::Value>,
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct TaskResult {
59 pub output: String,
61 pub session_id: String,
63 pub agent: String,
65 pub success: bool,
67 pub task_id: String,
69 #[serde(default, skip_serializing_if = "Option::is_none")]
71 pub structured: Option<serde_json::Value>,
72}
73
74fn compact_task_output(output: &str) -> (String, bool) {
75 if output.len() <= TASK_OUTPUT_CONTEXT_LIMIT {
76 return (output.to_string(), false);
77 }
78
79 let head = crate::text::truncate_utf8(output, TASK_OUTPUT_CONTEXT_HEAD);
80 let tail_start = output
81 .char_indices()
82 .find_map(|(idx, _)| {
83 if output.len().saturating_sub(idx) <= TASK_OUTPUT_CONTEXT_TAIL {
84 Some(idx)
85 } else {
86 None
87 }
88 })
89 .unwrap_or(output.len());
90 let tail = &output[tail_start..];
91
92 (
93 format!(
94 "{}\n\n[{} bytes omitted from delegated task output]\n\n{}",
95 head,
96 output.len().saturating_sub(head.len() + tail.len()),
97 tail
98 ),
99 true,
100 )
101}
102
103fn synthesize_subagent_progress(
114 event: &AgentEvent,
115 task_id: &str,
116 session_id: &str,
117) -> Option<AgentEvent> {
118 match event {
119 AgentEvent::ToolEnd {
120 name,
121 output,
122 exit_code,
123 error_kind,
124 ..
125 } => {
126 let mut metadata = serde_json::json!({
127 "tool": name,
128 "exit_code": exit_code,
129 "output_bytes": output.len(),
130 });
131 if let Some(kind) = error_kind {
132 metadata["error_kind"] =
133 serde_json::to_value(kind).unwrap_or(serde_json::Value::Null);
134 }
135 Some(AgentEvent::SubagentProgress {
136 task_id: task_id.to_string(),
137 session_id: session_id.to_string(),
138 status: "tool_completed".to_string(),
139 metadata,
140 })
141 }
142 AgentEvent::TurnEnd { turn, usage } => Some(AgentEvent::SubagentProgress {
143 task_id: task_id.to_string(),
144 session_id: session_id.to_string(),
145 status: "turn_completed".to_string(),
146 metadata: serde_json::json!({
147 "turn": turn,
148 "total_tokens": usage.total_tokens,
149 "prompt_tokens": usage.prompt_tokens,
150 "completion_tokens": usage.completion_tokens,
151 }),
152 }),
153 _ => None,
154 }
155}
156
157fn task_artifact_id(result: &TaskResult) -> String {
158 format!("task-output:{}", result.task_id)
159}
160
161fn task_artifact_uri(result: &TaskResult) -> String {
162 format!(
163 "a3s://tasks/{}/runs/{}/output",
164 result.session_id, result.task_id
165 )
166}
167
168fn epoch_ms() -> u64 {
169 use std::time::{SystemTime, UNIX_EPOCH};
170 SystemTime::now()
171 .duration_since(UNIX_EPOCH)
172 .map(|duration| duration.as_millis() as u64)
173 .unwrap_or(0)
174}
175
176fn format_task_result_for_context(result: &TaskResult) -> (String, bool) {
177 let (output, truncated) = compact_task_output(&result.output);
178 let status = if result.success {
179 "completed"
180 } else {
181 "failed"
182 };
183 let artifact_id = task_artifact_id(result);
184 let artifact_uri = task_artifact_uri(result);
185 let mut formatted = format!(
186 "Task {status}: {}\nAgent: {}\nSession: {}\nTask ID: {}\nArtifact ID: {}\nArtifact URI: {}\n",
187 result.task_id, result.agent, result.session_id, result.task_id, artifact_id, artifact_uri
188 );
189 if truncated {
190 formatted.push_str(
191 "Output excerpt: truncated for parent context. Use the artifact URI or child run session/events if exact omitted content is needed.\n",
192 );
193 } else {
194 formatted.push_str("Output:\n");
195 }
196 formatted.push_str(&output);
197 (formatted, truncated)
198}
199
200pub struct TaskExecutor {
202 registry: Arc<AgentRegistry>,
204 llm_client: Arc<dyn LlmClient>,
206 workspace: String,
208 mcp_manager: Option<Arc<McpManager>>,
210 parent_context: Option<crate::child_run::ChildRunContext>,
212 max_parallel_tasks: usize,
213 subagent_tracker: Option<Arc<crate::subagent_task_tracker::InMemorySubagentTaskTracker>>,
216}
217
218impl TaskExecutor {
219 pub fn new(
221 registry: Arc<AgentRegistry>,
222 llm_client: Arc<dyn LlmClient>,
223 workspace: String,
224 ) -> Self {
225 Self {
226 registry,
227 llm_client,
228 workspace,
229 mcp_manager: None,
230 parent_context: None,
231 max_parallel_tasks: crate::agent::DEFAULT_MAX_PARALLEL_TASKS,
232 subagent_tracker: None,
233 }
234 }
235
236 pub fn with_mcp(
238 registry: Arc<AgentRegistry>,
239 llm_client: Arc<dyn LlmClient>,
240 workspace: String,
241 mcp_manager: Arc<McpManager>,
242 ) -> Self {
243 Self {
244 registry,
245 llm_client,
246 workspace,
247 mcp_manager: Some(mcp_manager),
248 parent_context: None,
249 max_parallel_tasks: crate::agent::DEFAULT_MAX_PARALLEL_TASKS,
250 subagent_tracker: None,
251 }
252 }
253
254 pub fn with_parent_context(mut self, ctx: crate::child_run::ChildRunContext) -> Self {
256 if let Some(max_parallel_tasks) = ctx.max_parallel_tasks {
257 self.max_parallel_tasks = max_parallel_tasks.max(1);
258 }
259 self.parent_context = Some(ctx);
260 self
261 }
262
263 pub fn with_max_parallel_tasks(mut self, max_parallel_tasks: usize) -> Self {
264 self.max_parallel_tasks = max_parallel_tasks.max(1);
265 self
266 }
267
268 pub fn with_subagent_tracker(
272 mut self,
273 tracker: Arc<crate::subagent_task_tracker::InMemorySubagentTaskTracker>,
274 ) -> Self {
275 self.subagent_tracker = Some(tracker);
276 self
277 }
278
279 pub async fn execute(
284 &self,
285 params: TaskParams,
286 event_tx: Option<broadcast::Sender<AgentEvent>>,
287 parent_session_id: Option<&str>,
288 ) -> Result<TaskResult> {
289 let task_id = format!("task-{}", uuid::Uuid::new_v4());
290 self.execute_with_task_id(task_id, params, event_tx, parent_session_id, true)
291 .await
292 }
293
294 pub async fn execute_with_task_id(
299 &self,
300 task_id: String,
301 params: TaskParams,
302 event_tx: Option<broadcast::Sender<AgentEvent>>,
303 parent_session_id: Option<&str>,
304 emit_start: bool,
305 ) -> Result<TaskResult> {
306 let session_id = format!("task-run-{}", task_id);
307 let started_ms = epoch_ms();
308
309 let agent = self
310 .registry
311 .get(¶ms.agent)
312 .context(format!("Unknown agent type: '{}'", params.agent))?;
313
314 if emit_start {
315 if let Some(ref tx) = event_tx {
316 let _ = tx.send(AgentEvent::SubagentStart {
317 task_id: task_id.clone(),
318 session_id: session_id.clone(),
319 parent_session_id: parent_session_id.unwrap_or_default().to_string(),
320 agent: params.agent.clone(),
321 description: params.description.clone(),
322 started_ms,
323 });
324 }
325 }
326
327 let child_executor = if let Some(ref parent_ctx) = self.parent_context {
330 if let Some(ref services) = parent_ctx.workspace_services {
331 crate::tools::ToolExecutor::new_with_workspace_services_and_artifact_limits(
332 self.workspace.clone(),
333 Arc::clone(services),
334 crate::tools::ArtifactStoreLimits::default(),
335 )
336 } else {
337 crate::tools::ToolExecutor::new(self.workspace.clone())
338 }
339 } else {
340 crate::tools::ToolExecutor::new(self.workspace.clone())
341 };
342
343 if let Some(ref mcp) = self.mcp_manager {
345 let all_tools = mcp.get_all_tools().await;
346 let mut by_server: std::collections::HashMap<
347 String,
348 Vec<crate::mcp::protocol::McpTool>,
349 > = std::collections::HashMap::new();
350 for (server, tool) in all_tools {
351 by_server.entry(server).or_default().push(tool);
352 }
353 for (server_name, tools) in by_server {
354 let wrappers =
355 crate::mcp::tools::create_mcp_tools(&server_name, tools, Arc::clone(mcp));
356 for wrapper in wrappers {
357 child_executor.register_dynamic_tool(wrapper);
358 }
359 }
360 }
361
362 let child_executor = Arc::new(child_executor);
363
364 let mut child_config = AgentConfig {
365 tools: child_executor.definitions(),
366 ..AgentConfig::default()
367 };
368 agent.apply_to(&mut child_config);
369 if let Some(ref parent_ctx) = self.parent_context {
370 parent_ctx.apply_to(&mut child_config);
371 }
372 if let Some(max_steps) = params.max_steps {
373 child_config.max_tool_rounds = max_steps;
374 }
375
376 let mut tool_context =
377 ToolContext::new(PathBuf::from(&self.workspace)).with_session_id(session_id.clone());
378 if let Some(ref parent_ctx) = self.parent_context {
379 if let Some(ref services) = parent_ctx.workspace_services {
380 tool_context = tool_context.with_workspace_services(Arc::clone(services));
381 }
382 }
383
384 let agent_loop = AgentLoop::new(
385 Arc::clone(&self.llm_client),
386 child_executor,
387 tool_context,
388 child_config,
389 );
390
391 let child_event_tx = if let Some(ref broadcast_tx) = event_tx {
396 let (mpsc_tx, mut mpsc_rx) = tokio::sync::mpsc::channel(100);
397 let broadcast_tx_clone = broadcast_tx.clone();
398 let progress_task_id = task_id.clone();
399 let progress_session_id = session_id.clone();
400
401 tokio::spawn(async move {
402 while let Some(event) = mpsc_rx.recv().await {
403 if let Some(progress) = synthesize_subagent_progress(
404 &event,
405 &progress_task_id,
406 &progress_session_id,
407 ) {
408 let _ = broadcast_tx_clone.send(progress);
409 }
410 let _ = broadcast_tx_clone.send(event);
411 }
412 });
413
414 Some(mpsc_tx)
415 } else {
416 None
417 };
418
419 let cancel_token = tokio_util::sync::CancellationToken::new();
422 if let Some(ref tracker) = self.subagent_tracker {
423 tracker
424 .register_canceller(&task_id, cancel_token.clone())
425 .await;
426 }
427
428 let (output, success) = match agent_loop
429 .execute_with_session(
430 &[],
431 ¶ms.prompt,
432 Some(&session_id),
433 child_event_tx,
434 Some(&cancel_token),
435 )
436 .await
437 {
438 Ok(result) => (result.text, true),
439 Err(e) if cancel_token.is_cancelled() => {
440 (format!("Task cancelled by caller: {}", e), false)
441 }
442 Err(e) => (format!("Task failed: {}", e), false),
443 };
444
445 if let Some(ref tracker) = self.subagent_tracker {
446 tracker.clear_canceller(&task_id).await;
447 }
448
449 if let Some(ref tx) = event_tx {
450 let _ = tx.send(AgentEvent::SubagentEnd {
451 task_id: task_id.clone(),
452 session_id: session_id.clone(),
453 agent: params.agent.clone(),
454 output: output.clone(),
455 success,
456 finished_ms: epoch_ms(),
457 });
458 }
459
460 Ok(TaskResult {
461 output,
462 session_id,
463 agent: params.agent,
464 success,
465 task_id,
466 structured: None,
467 })
468 }
469
470 pub fn execute_background(
478 self: Arc<Self>,
479 params: TaskParams,
480 event_tx: Option<broadcast::Sender<AgentEvent>>,
481 parent_session_id: Option<String>,
482 ) -> String {
483 let task_id = format!("task-{}", uuid::Uuid::new_v4());
484 let session_id = format!("task-run-{}", task_id);
485
486 if let Some(ref tx) = event_tx {
487 let _ = tx.send(AgentEvent::SubagentStart {
488 task_id: task_id.clone(),
489 session_id,
490 parent_session_id: parent_session_id.clone().unwrap_or_default(),
491 agent: params.agent.clone(),
492 description: params.description.clone(),
493 started_ms: epoch_ms(),
494 });
495 }
496
497 let task_id_for_spawn = task_id.clone();
498 let task_id_for_log = task_id.clone();
499 tokio::spawn(async move {
500 if let Err(e) = self
501 .execute_with_task_id(
502 task_id_for_spawn,
503 params,
504 event_tx,
505 parent_session_id.as_deref(),
506 false,
507 )
508 .await
509 {
510 tracing::error!("Background task {} failed: {}", task_id_for_log, e);
511 }
512 });
513
514 task_id
515 }
516
517 pub async fn execute_parallel(
525 self: &Arc<Self>,
526 tasks: Vec<TaskParams>,
527 event_tx: Option<broadcast::Sender<AgentEvent>>,
528 parent_session_id: Option<&str>,
529 ) -> Vec<TaskResult> {
530 let parent = parent_session_id.map(|s| s.to_string());
531 let specs = tasks
532 .into_iter()
533 .map(|params| AgentStepSpec {
534 task_id: format!("task-{}", uuid::Uuid::new_v4()),
535 agent: params.agent,
536 description: params.description,
537 prompt: params.prompt,
538 max_steps: params.max_steps,
539 parent_session_id: parent.clone(),
540 output_schema: params.output_schema,
541 })
542 .collect();
543
544 let executor: Arc<dyn AgentExecutor> = Arc::<Self>::clone(self);
545 crate::orchestration::execute_steps_parallel(executor, specs, event_tx)
546 .await
547 .into_iter()
548 .map(TaskResult::from)
549 .collect()
550 }
551}
552
553impl From<TaskResult> for StepOutcome {
554 fn from(r: TaskResult) -> Self {
555 StepOutcome {
556 task_id: r.task_id,
557 session_id: r.session_id,
558 agent: r.agent,
559 output: r.output,
560 success: r.success,
561 structured: r.structured,
562 }
563 }
564}
565
566impl From<StepOutcome> for TaskResult {
567 fn from(o: StepOutcome) -> Self {
568 TaskResult {
569 output: o.output,
570 session_id: o.session_id,
571 agent: o.agent,
572 success: o.success,
573 task_id: o.task_id,
574 structured: o.structured,
575 }
576 }
577}
578
579#[async_trait]
583impl AgentExecutor for TaskExecutor {
584 async fn execute_step(
585 &self,
586 spec: AgentStepSpec,
587 event_tx: Option<broadcast::Sender<AgentEvent>>,
588 ) -> StepOutcome {
589 let agent = spec.agent.clone();
590 let task_id = spec.task_id.clone();
591 let output_schema = spec.output_schema.clone();
592 let params = TaskParams {
593 agent: spec.agent,
594 description: spec.description,
595 prompt: spec.prompt,
596 background: false,
597 max_steps: spec.max_steps,
598 output_schema: None,
599 };
600 let mut outcome: StepOutcome = match self
601 .execute_with_task_id(
602 task_id.clone(),
603 params,
604 event_tx,
605 spec.parent_session_id.as_deref(),
606 true,
607 )
608 .await
609 {
610 Ok(result) => result.into(),
611 Err(e) => return StepOutcome::failed(task_id, agent, format!("Task failed: {e}")),
612 };
613
614 if outcome.success {
619 if let Some(schema) = output_schema {
620 match self.coerce_to_schema(&outcome.output, schema).await {
621 Ok(object) => outcome.structured = Some(object),
622 Err(e) => {
623 outcome.success = false;
624 outcome.output =
625 format!("{}\n\n[structured output failed: {e}]", outcome.output);
626 }
627 }
628 }
629 }
630 outcome
631 }
632
633 fn concurrency_hint(&self) -> usize {
634 self.max_parallel_tasks
635 }
636}
637
638impl TaskExecutor {
639 async fn coerce_to_schema(
644 &self,
645 output: &str,
646 schema: serde_json::Value,
647 ) -> Result<serde_json::Value> {
648 let req = StructuredRequest {
649 prompt: format!(
650 "Convert the following task result into a single JSON object that conforms to \
651 the required schema. Use only information present in the result.\n\n\
652 --- TASK RESULT ---\n{output}"
653 ),
654 system: Some(
655 "You output exactly one JSON object matching the provided schema.".to_string(),
656 ),
657 schema,
658 schema_name: "step_output".to_string(),
659 schema_description: None,
660 mode: StructuredMode::Tool,
663 max_repair_attempts: 2,
664 };
665 let result = generate_blocking(&*self.llm_client, &req).await?;
666 Ok(result.object)
667 }
668}
669
670pub fn task_params_schema() -> serde_json::Value {
672 serde_json::json!({
673 "type": "object",
674 "additionalProperties": false,
675 "properties": {
676 "agent": {
677 "type": "string",
678 "description": "Required. Canonical agent type to use (for example: explore, general, plan, verification, review). Always provide this exact field name: 'agent'."
679 },
680 "description": {
681 "type": "string",
682 "description": "Required. Short task label for display and tracking. Always provide this exact field name: 'description'."
683 },
684 "prompt": {
685 "type": "string",
686 "description": "Required. Detailed instruction for the delegated child run. Always provide this exact field name: 'prompt'."
687 },
688 "background": {
689 "type": "boolean",
690 "description": "Optional. Run the task in the background. Default: false.",
691 "default": false
692 },
693 "max_steps": {
694 "type": "integer",
695 "description": "Optional. Maximum number of steps for this task."
696 },
697 "output_schema": {
698 "type": "object",
699 "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."
700 }
701 },
702 "required": ["agent", "description", "prompt"],
703 "examples": [
704 {
705 "agent": "explore",
706 "description": "Find Rust files",
707 "prompt": "Search the workspace for Rust files and summarize the layout."
708 },
709 {
710 "agent": "general",
711 "description": "Investigate test failure",
712 "prompt": "Inspect the failing tests and explain the root cause.",
713 "max_steps": 6
714 }
715 ]
716 })
717}
718
719pub struct TaskTool {
722 executor: Arc<TaskExecutor>,
723}
724
725impl TaskTool {
726 pub fn new(executor: Arc<TaskExecutor>) -> Self {
728 Self { executor }
729 }
730}
731
732#[async_trait]
733impl Tool for TaskTool {
734 fn name(&self) -> &str {
735 "task"
736 }
737
738 fn description(&self) -> &str {
739 "Delegate a bounded task to a specialized child run. Built-in agents: explore (read-only codebase and web evidence search), general/general-purpose (full access multi-step), plan (read-only planning), verification (adversarial validation), review (code review). Custom agents from agent_dirs and .a3s/agents are also available; .claude/agents is read for compatibility."
740 }
741
742 fn parameters(&self) -> serde_json::Value {
743 task_params_schema()
744 }
745
746 async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
747 let params: TaskParams =
748 serde_json::from_value(args.clone()).context("Invalid task parameters")?;
749
750 if params.background {
751 let task_id = Arc::clone(&self.executor).execute_background(
752 params,
753 ctx.agent_event_tx.clone(),
754 ctx.session_id.clone(),
755 );
756 return Ok(ToolOutput::success(format!(
757 "Task started in background. Task ID: {}",
758 task_id
759 )));
760 }
761
762 let result = self
763 .executor
764 .execute(
765 params,
766 ctx.agent_event_tx.clone(),
767 ctx.session_id.as_deref(),
768 )
769 .await?;
770 let (content, truncated) = format_task_result_for_context(&result);
771 let metadata = serde_json::json!({
772 "task_id": result.task_id,
773 "session_id": result.session_id,
774 "agent": result.agent,
775 "success": result.success,
776 "output_bytes": result.output.len(),
777 "truncated_for_context": truncated,
778 "artifact_id": task_artifact_id(&result),
779 "artifact_uri": task_artifact_uri(&result),
780 });
781
782 if result.success {
783 Ok(ToolOutput::success(content).with_metadata(metadata))
784 } else {
785 Ok(ToolOutput::error(content).with_metadata(metadata))
786 }
787 }
788}
789
790#[derive(Debug, Clone, Serialize, Deserialize)]
792#[serde(deny_unknown_fields)]
793pub struct ParallelTaskParams {
794 pub tasks: Vec<TaskParams>,
796 #[serde(default)]
799 pub allow_partial_failure: bool,
800}
801
802pub fn parallel_task_params_schema() -> serde_json::Value {
804 serde_json::json!({
805 "type": "object",
806 "additionalProperties": false,
807 "properties": {
808 "tasks": {
809 "type": "array",
810 "description": "List of tasks to execute in parallel. Each task runs as an independent delegated child run concurrently.",
811 "items": {
812 "type": "object",
813 "additionalProperties": false,
814 "properties": {
815 "agent": {
816 "type": "string",
817 "description": "Required. Canonical agent type for this task."
818 },
819 "description": {
820 "type": "string",
821 "description": "Required. Short task label for display and tracking."
822 },
823 "prompt": {
824 "type": "string",
825 "description": "Required. Detailed instruction for the delegated child run."
826 },
827 "background": {
828 "type": "boolean",
829 "description": "Optional. Run this delegated child task in the background. Default: false.",
830 "default": false
831 },
832 "max_steps": {
833 "type": "integer",
834 "description": "Optional. Maximum number of tool/model steps for this delegated child task."
835 },
836 "output_schema": {
837 "type": "object",
838 "description": "Optional. JSON Schema object the delegated child result must satisfy. When provided, the validated object is returned in each result's metadata."
839 }
840 },
841 "required": ["agent", "description", "prompt"]
842 },
843 "minItems": 1
844 },
845 "allow_partial_failure": {
846 "type": "boolean",
847 "description": "Optional. Defaults to false. When true, the parallel_task tool succeeds if at least one child task succeeds, while preserving failed child results in the output and metadata.",
848 "default": false
849 }
850 },
851 "required": ["tasks"],
852 "examples": [
853 {
854 "tasks": [
855 {
856 "agent": "explore",
857 "description": "Find Rust files",
858 "prompt": "List Rust files under src/."
859 },
860 {
861 "agent": "explore",
862 "description": "Find tests",
863 "prompt": "List test files and summarize their purpose."
864 }
865 ]
866 }
867 ]
868 })
869}
870
871pub struct ParallelTaskTool {
875 executor: Arc<TaskExecutor>,
876}
877
878impl ParallelTaskTool {
879 pub fn new(executor: Arc<TaskExecutor>) -> Self {
881 Self { executor }
882 }
883}
884
885#[async_trait]
886impl Tool for ParallelTaskTool {
887 fn name(&self) -> &str {
888 "parallel_task"
889 }
890
891 fn description(&self) -> &str {
892 "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. Use this only when the work genuinely splits into branches that can be investigated or implemented separately (e.g. inspect several unrelated modules at once, or run review and verification in parallel). Do NOT use it for trivial, conversational, or single-step requests, or for steps that depend on one another — handle those directly. Built-in agents: explore (read-only codebase and web evidence search), general/general-purpose (full access multi-step), plan (read-only planning), verification (adversarial validation), review (code review). Custom agents from agent_dirs and .a3s/agents are also available; .claude/agents is read for compatibility."
893 }
894
895 fn parameters(&self) -> serde_json::Value {
896 parallel_task_params_schema()
897 }
898
899 async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
900 let params: ParallelTaskParams =
901 serde_json::from_value(args.clone()).context("Invalid parallel task parameters")?;
902
903 if params.tasks.is_empty() {
904 return Ok(ToolOutput::error("No tasks provided".to_string()));
905 }
906
907 let task_count = params.tasks.len();
908
909 let results = self
910 .executor
911 .execute_parallel(
912 params.tasks,
913 ctx.agent_event_tx.clone(),
914 ctx.session_id.as_deref(),
915 )
916 .await;
917
918 let mut output = format!("Executed {} tasks in parallel:\n\n", task_count);
920 let mut metadata_results = Vec::new();
921 for (i, result) in results.iter().enumerate() {
922 let status = if result.success { "[OK]" } else { "[ERR]" };
923 let (formatted, truncated) = format_task_result_for_context(result);
924 metadata_results.push(serde_json::json!({
925 "task_id": result.task_id,
926 "session_id": result.session_id,
927 "agent": result.agent,
928 "success": result.success,
929 "output": formatted.clone(),
930 "structured": result.structured,
931 "output_bytes": result.output.len(),
932 "truncated_for_context": truncated,
933 "artifact_id": task_artifact_id(result),
934 "artifact_uri": task_artifact_uri(result),
935 }));
936 output.push_str(&format!(
937 "--- Task {} ({}) {} ---\n{}\n\n",
938 i + 1,
939 result.agent,
940 status,
941 formatted
942 ));
943 }
944
945 let success_count = results.iter().filter(|result| result.success).count();
946 let failed_count = results.len().saturating_sub(success_count);
947 let all_success = failed_count == 0;
948 let partial_failure = failed_count > 0 && success_count > 0;
949 if params.allow_partial_failure && partial_failure {
950 output.push_str(&format!(
951 "Partial failure tolerated: {success_count} succeeded, {failed_count} failed.\n"
952 ));
953 }
954
955 let tool_success = all_success || (params.allow_partial_failure && success_count > 0);
956 let output = if tool_success {
957 ToolOutput::success(output)
958 } else {
959 ToolOutput::error(output)
960 };
961
962 Ok(output.with_metadata(serde_json::json!({
963 "task_count": task_count,
964 "result_count": results.len(),
965 "success_count": success_count,
966 "failed_count": failed_count,
967 "all_success": all_success,
968 "partial_failure": partial_failure,
969 "allow_partial_failure": params.allow_partial_failure,
970 "results": metadata_results,
971 })))
972 }
973}
974
975#[cfg(test)]
976mod tests {
977 use super::*;
978
979 #[test]
980 fn test_task_params_deserialize() {
981 let json = r#"{
982 "agent": "explore",
983 "description": "Find auth code",
984 "prompt": "Search for authentication files"
985 }"#;
986
987 let params: TaskParams = serde_json::from_str(json).unwrap();
988 assert_eq!(params.agent, "explore");
989 assert_eq!(params.description, "Find auth code");
990 assert!(!params.background);
991 }
992
993 #[test]
994 fn test_task_params_with_background() {
995 let json = r#"{
996 "agent": "general",
997 "description": "Long task",
998 "prompt": "Do something complex",
999 "background": true
1000 }"#;
1001
1002 let params: TaskParams = serde_json::from_str(json).unwrap();
1003 assert!(params.background);
1004 }
1005
1006 #[test]
1007 fn test_task_params_with_max_steps() {
1008 let json = r#"{
1009 "agent": "plan",
1010 "description": "Planning task",
1011 "prompt": "Create a plan",
1012 "max_steps": 10
1013 }"#;
1014
1015 let params: TaskParams = serde_json::from_str(json).unwrap();
1016 assert_eq!(params.agent, "plan");
1017 assert_eq!(params.max_steps, Some(10));
1018 assert!(!params.background);
1019 }
1020
1021 #[test]
1022 fn test_task_params_all_fields() {
1023 let json = r#"{
1024 "agent": "general",
1025 "description": "Complex task",
1026 "prompt": "Do everything",
1027 "background": true,
1028 "max_steps": 20
1029 }"#;
1030
1031 let params: TaskParams = serde_json::from_str(json).unwrap();
1032 assert_eq!(params.agent, "general");
1033 assert_eq!(params.description, "Complex task");
1034 assert_eq!(params.prompt, "Do everything");
1035 assert!(params.background);
1036 assert_eq!(params.max_steps, Some(20));
1037 }
1038
1039 #[test]
1040 fn test_task_params_missing_required_field() {
1041 let json = r#"{
1042 "agent": "explore",
1043 "description": "Missing prompt"
1044 }"#;
1045
1046 let result: Result<TaskParams, _> = serde_json::from_str(json);
1047 assert!(result.is_err());
1048 }
1049
1050 #[test]
1051 fn test_task_params_serialize() {
1052 let params = TaskParams {
1053 agent: "explore".to_string(),
1054 description: "Test task".to_string(),
1055 prompt: "Test prompt".to_string(),
1056 background: false,
1057 max_steps: Some(5),
1058 output_schema: None,
1059 };
1060
1061 let json = serde_json::to_string(¶ms).unwrap();
1062 assert!(json.contains("explore"));
1063 assert!(json.contains("Test task"));
1064 assert!(json.contains("Test prompt"));
1065 }
1066
1067 #[test]
1068 fn test_task_params_clone() {
1069 let params = TaskParams {
1070 agent: "explore".to_string(),
1071 description: "Test".to_string(),
1072 prompt: "Prompt".to_string(),
1073 background: true,
1074 max_steps: None,
1075 output_schema: None,
1076 };
1077
1078 let cloned = params.clone();
1079 assert_eq!(params.agent, cloned.agent);
1080 assert_eq!(params.description, cloned.description);
1081 assert_eq!(params.background, cloned.background);
1082 }
1083
1084 #[test]
1085 fn test_task_result_serialize() {
1086 let result = TaskResult {
1087 output: "Found 5 files".to_string(),
1088 session_id: "session-123".to_string(),
1089 agent: "explore".to_string(),
1090 success: true,
1091 task_id: "task-456".to_string(),
1092 structured: None,
1093 };
1094
1095 let json = serde_json::to_string(&result).unwrap();
1096 assert!(json.contains("Found 5 files"));
1097 assert!(json.contains("explore"));
1098 }
1099
1100 #[test]
1101 fn test_task_result_deserialize() {
1102 let json = r#"{
1103 "output": "Task completed",
1104 "session_id": "sess-789",
1105 "agent": "general",
1106 "success": false,
1107 "task_id": "task-123"
1108 }"#;
1109
1110 let result: TaskResult = serde_json::from_str(json).unwrap();
1111 assert_eq!(result.output, "Task completed");
1112 assert_eq!(result.session_id, "sess-789");
1113 assert_eq!(result.agent, "general");
1114 assert!(!result.success);
1115 assert_eq!(result.task_id, "task-123");
1116 }
1117
1118 #[test]
1119 fn test_task_result_clone() {
1120 let result = TaskResult {
1121 output: "Output".to_string(),
1122 session_id: "session-1".to_string(),
1123 agent: "explore".to_string(),
1124 success: true,
1125 task_id: "task-1".to_string(),
1126 structured: None,
1127 };
1128
1129 let cloned = result.clone();
1130 assert_eq!(result.output, cloned.output);
1131 assert_eq!(result.success, cloned.success);
1132 }
1133
1134 #[test]
1135 fn test_compact_task_output_preserves_small_output() {
1136 let (output, truncated) = compact_task_output("short result");
1137 assert_eq!(output, "short result");
1138 assert!(!truncated);
1139 }
1140
1141 #[test]
1142 fn test_format_task_result_for_context_truncates_large_output() {
1143 let result = TaskResult {
1144 output: format!("{}TAIL", "x".repeat(TASK_OUTPUT_CONTEXT_LIMIT + 500)),
1145 session_id: "session-1".to_string(),
1146 agent: "explore".to_string(),
1147 success: true,
1148 task_id: "task-1".to_string(),
1149 structured: None,
1150 };
1151
1152 let (formatted, truncated) = format_task_result_for_context(&result);
1153 assert!(truncated);
1154 assert!(formatted.contains("Output excerpt"));
1155 assert!(formatted.contains("bytes omitted"));
1156 assert!(formatted.contains("Artifact ID: task-output:task-1"));
1157 assert!(formatted.contains("Artifact URI: a3s://tasks/session-1/runs/task-1/output"));
1158 assert!(formatted.contains("TAIL"));
1159 assert!(formatted.len() < result.output.len());
1160 }
1161
1162 #[test]
1163 fn test_task_artifact_reference_is_stable() {
1164 let result = TaskResult {
1165 output: "done".to_string(),
1166 session_id: "session-1".to_string(),
1167 agent: "explore".to_string(),
1168 success: true,
1169 task_id: "task-1".to_string(),
1170 structured: None,
1171 };
1172
1173 assert_eq!(task_artifact_id(&result), "task-output:task-1");
1174 assert_eq!(
1175 task_artifact_uri(&result),
1176 "a3s://tasks/session-1/runs/task-1/output"
1177 );
1178
1179 let (formatted, truncated) = format_task_result_for_context(&result);
1180 assert!(!truncated);
1181 assert!(formatted.contains("Artifact URI: a3s://tasks/session-1/runs/task-1/output"));
1182 }
1183
1184 #[test]
1185 fn test_task_params_schema() {
1186 let schema = task_params_schema();
1187 assert_eq!(schema["type"], "object");
1188 assert_eq!(schema["additionalProperties"], false);
1189 assert!(schema["properties"]["agent"].is_object());
1190 assert!(schema["properties"]["prompt"].is_object());
1191 }
1192
1193 #[test]
1194 fn test_task_params_schema_required_fields() {
1195 let schema = task_params_schema();
1196 let required = schema["required"].as_array().unwrap();
1197 assert!(required.contains(&serde_json::json!("agent")));
1198 assert!(required.contains(&serde_json::json!("description")));
1199 assert!(required.contains(&serde_json::json!("prompt")));
1200 }
1201
1202 #[test]
1203 fn test_task_params_schema_properties() {
1204 let schema = task_params_schema();
1205 let props = &schema["properties"];
1206
1207 assert_eq!(props["agent"]["type"], "string");
1208 assert_eq!(props["description"]["type"], "string");
1209 assert_eq!(props["prompt"]["type"], "string");
1210 assert_eq!(props["background"]["type"], "boolean");
1211 assert_eq!(props["background"]["default"], false);
1212 assert_eq!(props["max_steps"]["type"], "integer");
1213 assert_eq!(props["output_schema"]["type"], "object");
1214 }
1215
1216 #[test]
1217 fn test_task_params_schema_descriptions() {
1218 let schema = task_params_schema();
1219 let props = &schema["properties"];
1220
1221 assert!(props["agent"]["description"].is_string());
1222 assert!(props["description"]["description"].is_string());
1223 assert!(props["prompt"]["description"].is_string());
1224 assert!(props["background"]["description"].is_string());
1225 assert!(props["max_steps"]["description"].is_string());
1226 assert!(props["output_schema"]["description"].is_string());
1227 }
1228
1229 #[test]
1230 fn test_task_params_default_background() {
1231 let params = TaskParams {
1232 agent: "explore".to_string(),
1233 description: "Test".to_string(),
1234 prompt: "Test prompt".to_string(),
1235 background: false,
1236 max_steps: None,
1237 output_schema: None,
1238 };
1239 assert!(!params.background);
1240 }
1241
1242 #[test]
1243 fn test_task_params_serialize_skip_none() {
1244 let params = TaskParams {
1245 agent: "explore".to_string(),
1246 description: "Test".to_string(),
1247 prompt: "Test prompt".to_string(),
1248 background: false,
1249 max_steps: None,
1250 output_schema: None,
1251 };
1252 let json = serde_json::to_string(¶ms).unwrap();
1253 assert!(!json.contains("max_steps"));
1255 }
1256
1257 #[test]
1258 fn test_task_params_serialize_with_max_steps() {
1259 let params = TaskParams {
1260 agent: "explore".to_string(),
1261 description: "Test".to_string(),
1262 prompt: "Test prompt".to_string(),
1263 background: false,
1264 max_steps: Some(15),
1265 output_schema: None,
1266 };
1267 let json = serde_json::to_string(¶ms).unwrap();
1268 assert!(json.contains("max_steps"));
1269 assert!(json.contains("15"));
1270 }
1271
1272 #[test]
1273 fn test_task_result_success_true() {
1274 let result = TaskResult {
1275 output: "Success".to_string(),
1276 session_id: "sess-1".to_string(),
1277 agent: "explore".to_string(),
1278 success: true,
1279 task_id: "task-1".to_string(),
1280 structured: None,
1281 };
1282 assert!(result.success);
1283 }
1284
1285 #[test]
1286 fn test_task_result_success_false() {
1287 let result = TaskResult {
1288 output: "Failed".to_string(),
1289 session_id: "sess-1".to_string(),
1290 agent: "explore".to_string(),
1291 success: false,
1292 task_id: "task-1".to_string(),
1293 structured: None,
1294 };
1295 assert!(!result.success);
1296 }
1297
1298 #[test]
1299 fn test_task_params_empty_strings() {
1300 let params = TaskParams {
1301 agent: "".to_string(),
1302 description: "".to_string(),
1303 prompt: "".to_string(),
1304 background: false,
1305 max_steps: None,
1306 output_schema: None,
1307 };
1308 let json = serde_json::to_string(¶ms).unwrap();
1309 let deserialized: TaskParams = serde_json::from_str(&json).unwrap();
1310 assert_eq!(deserialized.agent, "");
1311 assert_eq!(deserialized.description, "");
1312 assert_eq!(deserialized.prompt, "");
1313 }
1314
1315 #[test]
1316 fn test_task_result_empty_output() {
1317 let result = TaskResult {
1318 output: "".to_string(),
1319 session_id: "sess-1".to_string(),
1320 agent: "explore".to_string(),
1321 success: true,
1322 task_id: "task-1".to_string(),
1323 structured: None,
1324 };
1325 assert_eq!(result.output, "");
1326 }
1327
1328 #[test]
1329 fn test_task_params_debug_format() {
1330 let params = TaskParams {
1331 agent: "explore".to_string(),
1332 description: "Test".to_string(),
1333 prompt: "Test prompt".to_string(),
1334 background: false,
1335 max_steps: None,
1336 output_schema: None,
1337 };
1338 let debug_str = format!("{:?}", params);
1339 assert!(debug_str.contains("explore"));
1340 assert!(debug_str.contains("Test"));
1341 }
1342
1343 #[test]
1344 fn test_task_result_debug_format() {
1345 let result = TaskResult {
1346 output: "Output".to_string(),
1347 session_id: "sess-1".to_string(),
1348 agent: "explore".to_string(),
1349 success: true,
1350 task_id: "task-1".to_string(),
1351 structured: None,
1352 };
1353 let debug_str = format!("{:?}", result);
1354 assert!(debug_str.contains("Output"));
1355 assert!(debug_str.contains("explore"));
1356 }
1357
1358 #[test]
1359 fn test_task_params_roundtrip() {
1360 let original = TaskParams {
1361 agent: "general".to_string(),
1362 description: "Roundtrip test".to_string(),
1363 prompt: "Test roundtrip serialization".to_string(),
1364 background: true,
1365 max_steps: Some(42),
1366 output_schema: None,
1367 };
1368 let json = serde_json::to_string(&original).unwrap();
1369 let deserialized: TaskParams = serde_json::from_str(&json).unwrap();
1370 assert_eq!(original.agent, deserialized.agent);
1371 assert_eq!(original.description, deserialized.description);
1372 assert_eq!(original.prompt, deserialized.prompt);
1373 assert_eq!(original.background, deserialized.background);
1374 assert_eq!(original.max_steps, deserialized.max_steps);
1375 }
1376
1377 #[test]
1378 fn test_task_result_roundtrip() {
1379 let original = TaskResult {
1380 output: "Roundtrip output".to_string(),
1381 session_id: "sess-roundtrip".to_string(),
1382 agent: "plan".to_string(),
1383 success: false,
1384 task_id: "task-roundtrip".to_string(),
1385 structured: None,
1386 };
1387 let json = serde_json::to_string(&original).unwrap();
1388 let deserialized: TaskResult = serde_json::from_str(&json).unwrap();
1389 assert_eq!(original.output, deserialized.output);
1390 assert_eq!(original.session_id, deserialized.session_id);
1391 assert_eq!(original.agent, deserialized.agent);
1392 assert_eq!(original.success, deserialized.success);
1393 assert_eq!(original.task_id, deserialized.task_id);
1394 }
1395
1396 #[test]
1397 fn test_parallel_task_params_deserialize() {
1398 let json = r#"{
1399 "tasks": [
1400 { "agent": "explore", "description": "Find auth", "prompt": "Search auth files" },
1401 { "agent": "general", "description": "Fix bug", "prompt": "Fix the login bug" }
1402 ]
1403 }"#;
1404
1405 let params: ParallelTaskParams = serde_json::from_str(json).unwrap();
1406 assert_eq!(params.tasks.len(), 2);
1407 assert_eq!(params.tasks[0].agent, "explore");
1408 assert_eq!(params.tasks[1].agent, "general");
1409 }
1410
1411 #[test]
1412 fn test_parallel_task_params_single_task() {
1413 let json = r#"{
1414 "tasks": [
1415 { "agent": "plan", "description": "Plan work", "prompt": "Create a plan" }
1416 ]
1417 }"#;
1418
1419 let params: ParallelTaskParams = serde_json::from_str(json).unwrap();
1420 assert_eq!(params.tasks.len(), 1);
1421 }
1422
1423 #[test]
1424 fn test_parallel_task_params_empty_tasks() {
1425 let json = r#"{ "tasks": [] }"#;
1426 let params: ParallelTaskParams = serde_json::from_str(json).unwrap();
1427 assert!(params.tasks.is_empty());
1428 assert!(!params.allow_partial_failure);
1429 }
1430
1431 #[test]
1432 fn test_parallel_task_params_missing_tasks() {
1433 let json = r#"{}"#;
1434 let result: Result<ParallelTaskParams, _> = serde_json::from_str(json);
1435 assert!(result.is_err());
1436 }
1437
1438 #[test]
1439 fn test_parallel_task_params_serialize() {
1440 let params = ParallelTaskParams {
1441 tasks: vec![
1442 TaskParams {
1443 agent: "explore".to_string(),
1444 description: "Task 1".to_string(),
1445 prompt: "Prompt 1".to_string(),
1446 background: false,
1447 max_steps: None,
1448 output_schema: None,
1449 },
1450 TaskParams {
1451 agent: "general".to_string(),
1452 description: "Task 2".to_string(),
1453 prompt: "Prompt 2".to_string(),
1454 background: false,
1455 max_steps: Some(10),
1456 output_schema: None,
1457 },
1458 ],
1459 allow_partial_failure: false,
1460 };
1461 let json = serde_json::to_string(¶ms).unwrap();
1462 assert!(json.contains("explore"));
1463 assert!(json.contains("general"));
1464 assert!(json.contains("Prompt 1"));
1465 assert!(json.contains("Prompt 2"));
1466 }
1467
1468 #[test]
1469 fn test_parallel_task_params_roundtrip() {
1470 let original = ParallelTaskParams {
1471 tasks: vec![
1472 TaskParams {
1473 agent: "explore".to_string(),
1474 description: "Explore".to_string(),
1475 prompt: "Find files".to_string(),
1476 background: false,
1477 max_steps: None,
1478 output_schema: None,
1479 },
1480 TaskParams {
1481 agent: "plan".to_string(),
1482 description: "Plan".to_string(),
1483 prompt: "Make plan".to_string(),
1484 background: false,
1485 max_steps: Some(5),
1486 output_schema: None,
1487 },
1488 ],
1489 allow_partial_failure: true,
1490 };
1491 let json = serde_json::to_string(&original).unwrap();
1492 let deserialized: ParallelTaskParams = serde_json::from_str(&json).unwrap();
1493 assert_eq!(original.tasks.len(), deserialized.tasks.len());
1494 assert_eq!(original.tasks[0].agent, deserialized.tasks[0].agent);
1495 assert_eq!(original.tasks[1].agent, deserialized.tasks[1].agent);
1496 assert_eq!(original.tasks[1].max_steps, deserialized.tasks[1].max_steps);
1497 assert_eq!(
1498 original.allow_partial_failure,
1499 deserialized.allow_partial_failure
1500 );
1501 }
1502
1503 #[test]
1504 fn test_parallel_task_params_clone() {
1505 let params = ParallelTaskParams {
1506 tasks: vec![TaskParams {
1507 agent: "explore".to_string(),
1508 description: "Test".to_string(),
1509 prompt: "Prompt".to_string(),
1510 background: false,
1511 max_steps: None,
1512 output_schema: None,
1513 }],
1514 allow_partial_failure: false,
1515 };
1516 let cloned = params.clone();
1517 assert_eq!(params.tasks.len(), cloned.tasks.len());
1518 assert_eq!(params.tasks[0].agent, cloned.tasks[0].agent);
1519 }
1520
1521 #[test]
1522 fn test_parallel_task_params_schema() {
1523 let schema = parallel_task_params_schema();
1524 assert_eq!(schema["type"], "object");
1525 assert_eq!(schema["additionalProperties"], false);
1526 assert!(schema["properties"]["tasks"].is_object());
1527 assert_eq!(schema["properties"]["tasks"]["type"], "array");
1528 assert_eq!(schema["properties"]["tasks"]["minItems"], 1);
1529 assert_eq!(
1530 schema["properties"]["allow_partial_failure"]["type"],
1531 "boolean"
1532 );
1533 assert_eq!(
1534 schema["properties"]["allow_partial_failure"]["default"],
1535 false
1536 );
1537 }
1538
1539 #[test]
1540 fn test_parallel_task_params_schema_required() {
1541 let schema = parallel_task_params_schema();
1542 let required = schema["required"].as_array().unwrap();
1543 assert!(required.contains(&serde_json::json!("tasks")));
1544 }
1545
1546 #[test]
1547 fn test_parallel_task_params_schema_items() {
1548 let schema = parallel_task_params_schema();
1549 let items = &schema["properties"]["tasks"]["items"];
1550 assert_eq!(items["type"], "object");
1551 assert_eq!(items["additionalProperties"], false);
1552 let item_required = items["required"].as_array().unwrap();
1553 assert!(item_required.contains(&serde_json::json!("agent")));
1554 assert!(item_required.contains(&serde_json::json!("description")));
1555 assert!(item_required.contains(&serde_json::json!("prompt")));
1556 assert_eq!(items["properties"]["background"]["type"], "boolean");
1557 assert_eq!(items["properties"]["background"]["default"], false);
1558 assert_eq!(items["properties"]["max_steps"]["type"], "integer");
1559 assert_eq!(items["properties"]["output_schema"]["type"], "object");
1560 }
1561
1562 #[test]
1563 fn test_task_schema_examples_use_delegation_core() {
1564 let task = task_params_schema();
1565 let task_examples = task["examples"].as_array().unwrap();
1566 assert_eq!(task_examples[0]["agent"], "explore");
1567 assert!(task_examples[0].get("task").is_none());
1568
1569 let parallel = parallel_task_params_schema();
1570 let parallel_examples = parallel["examples"].as_array().unwrap();
1571 assert!(!parallel_examples[0]["tasks"].as_array().unwrap().is_empty());
1572 }
1573
1574 #[test]
1575 fn test_parallel_task_params_debug() {
1576 let params = ParallelTaskParams {
1577 tasks: vec![TaskParams {
1578 agent: "explore".to_string(),
1579 description: "Debug test".to_string(),
1580 prompt: "Test".to_string(),
1581 background: false,
1582 max_steps: None,
1583 output_schema: None,
1584 }],
1585 allow_partial_failure: false,
1586 };
1587 let debug_str = format!("{:?}", params);
1588 assert!(debug_str.contains("explore"));
1589 assert!(debug_str.contains("Debug test"));
1590 }
1591
1592 #[test]
1593 fn test_parallel_task_params_large_count() {
1594 let tasks: Vec<TaskParams> = (0..150)
1596 .map(|i| TaskParams {
1597 agent: "explore".to_string(),
1598 description: format!("Task {}", i),
1599 prompt: format!("Prompt for task {}", i),
1600 background: false,
1601 max_steps: Some(10),
1602 output_schema: None,
1603 })
1604 .collect();
1605
1606 let params = ParallelTaskParams {
1607 tasks,
1608 allow_partial_failure: false,
1609 };
1610 let json = serde_json::to_string(¶ms).unwrap();
1611 let deserialized: ParallelTaskParams = serde_json::from_str(&json).unwrap();
1612 assert_eq!(deserialized.tasks.len(), 150);
1613 assert_eq!(deserialized.tasks[0].description, "Task 0");
1614 assert_eq!(deserialized.tasks[149].description, "Task 149");
1615 }
1616
1617 #[test]
1618 fn test_task_params_max_steps_zero() {
1619 let params = TaskParams {
1621 agent: "explore".to_string(),
1622 description: "Edge case".to_string(),
1623 prompt: "Zero steps".to_string(),
1624 background: false,
1625 max_steps: Some(0),
1626 output_schema: None,
1627 };
1628 let json = serde_json::to_string(¶ms).unwrap();
1629 let deserialized: TaskParams = serde_json::from_str(&json).unwrap();
1630 assert_eq!(deserialized.max_steps, Some(0));
1631 }
1632
1633 #[test]
1634 fn test_parallel_task_params_all_background() {
1635 let tasks: Vec<TaskParams> = (0..5)
1636 .map(|i| TaskParams {
1637 agent: "general".to_string(),
1638 description: format!("BG task {}", i),
1639 prompt: "Run in background".to_string(),
1640 background: true,
1641 max_steps: None,
1642 output_schema: None,
1643 })
1644 .collect();
1645 let params = ParallelTaskParams {
1646 tasks,
1647 allow_partial_failure: false,
1648 };
1649 for task in ¶ms.tasks {
1650 assert!(task.background);
1651 }
1652 }
1653
1654 #[test]
1655 fn test_task_params_rejects_permissive_field() {
1656 let json = r#"{
1657 "agent": "general",
1658 "description": "Legacy field rejection",
1659 "prompt": "Verify legacy fields are rejected",
1660 "permissive": true
1661 }"#;
1662
1663 let result: Result<TaskParams, _> = serde_json::from_str(json);
1664 assert!(result.is_err());
1665 }
1666
1667 #[test]
1668 fn test_task_params_schema_hides_permissive_field() {
1669 let schema = task_params_schema();
1670 let props = &schema["properties"];
1671
1672 assert!(props.get("permissive").is_none());
1673 }
1674
1675 use crate::agent::tests::MockLlmClient;
1680 use crate::llm::{ContentBlock, LlmResponse, Message, StreamEvent, TokenUsage, ToolDefinition};
1681 use crate::permissions::PermissionPolicy;
1682 use crate::subagent::AgentRegistry;
1683 use std::sync::atomic::{AtomicUsize, Ordering};
1684 use std::time::Duration;
1685 use tokio::sync::{mpsc, Barrier};
1686
1687 fn text_response(text: impl Into<String>) -> LlmResponse {
1688 LlmResponse {
1689 message: Message {
1690 role: "assistant".to_string(),
1691 content: vec![ContentBlock::Text { text: text.into() }],
1692 reasoning_content: None,
1693 },
1694 usage: TokenUsage {
1695 prompt_tokens: 10,
1696 completion_tokens: 5,
1697 total_tokens: 15,
1698 cache_read_tokens: None,
1699 cache_write_tokens: None,
1700 },
1701 stop_reason: Some("end_turn".to_string()),
1702 token_logprobs: Vec::new(),
1703 meta: None,
1704 }
1705 }
1706
1707 fn pre_analysis_response(messages: &[Message]) -> LlmResponse {
1708 let prompt = last_text(messages);
1709 let response = serde_json::json!({
1710 "intent": "GeneralPurpose",
1711 "requires_planning": false,
1712 "goal": {
1713 "description": prompt,
1714 "success_criteria": []
1715 },
1716 "execution_plan": {
1717 "complexity": "Simple",
1718 "steps": [{
1719 "id": "step-1",
1720 "description": prompt,
1721 "tool": null,
1722 "dependencies": [],
1723 "success_criteria": "Complete the request"
1724 }],
1725 "required_tools": []
1726 },
1727 "optimized_input": prompt
1728 });
1729 text_response(response.to_string())
1730 }
1731
1732 fn last_text(messages: &[Message]) -> String {
1733 messages
1734 .last()
1735 .and_then(|message| {
1736 message.content.iter().find_map(|block| {
1737 if let ContentBlock::Text { text } = block {
1738 Some(text.clone())
1739 } else {
1740 None
1741 }
1742 })
1743 })
1744 .unwrap_or_default()
1745 }
1746
1747 struct SchemaCoercionClient;
1752
1753 #[async_trait::async_trait]
1754 impl LlmClient for SchemaCoercionClient {
1755 async fn complete(
1756 &self,
1757 messages: &[Message],
1758 system: Option<&str>,
1759 tools: &[ToolDefinition],
1760 ) -> Result<LlmResponse> {
1761 if system == Some(crate::prompts::PRE_ANALYSIS_SYSTEM) {
1762 return Ok(pre_analysis_response(messages));
1763 }
1764 if tools.iter().any(|t| t.name == "emit_step_output") {
1767 return Ok(MockLlmClient::tool_call_response(
1768 "coerce-1",
1769 "emit_step_output",
1770 serde_json::json!({ "verdict": "ok" }),
1771 ));
1772 }
1773 Ok(text_response("The verdict is ok."))
1774 }
1775
1776 async fn complete_streaming(
1777 &self,
1778 _messages: &[Message],
1779 _system: Option<&str>,
1780 _tools: &[ToolDefinition],
1781 _cancel_token: tokio_util::sync::CancellationToken,
1782 ) -> Result<mpsc::Receiver<StreamEvent>> {
1783 anyhow::bail!("streaming is not used by schema coercion tests")
1784 }
1785 }
1786
1787 fn verdict_schema() -> serde_json::Value {
1788 serde_json::json!({
1789 "type": "object",
1790 "properties": { "verdict": { "type": "string" } },
1791 "required": ["verdict"]
1792 })
1793 }
1794
1795 #[tokio::test]
1796 async fn execute_step_with_schema_coerces_structured_output() {
1797 let workspace = tempfile::tempdir().unwrap();
1798 let executor = TaskExecutor::new(
1799 Arc::new(AgentRegistry::new()),
1800 Arc::new(SchemaCoercionClient),
1801 workspace.path().to_string_lossy().to_string(),
1802 );
1803 let spec = AgentStepSpec::new("step-1", "general", "assess", "Assess the thing.")
1804 .with_output_schema(verdict_schema());
1805
1806 let outcome = executor.execute_step(spec, None).await;
1807
1808 assert!(outcome.success, "step should succeed: {}", outcome.output);
1809 assert_eq!(
1810 outcome.structured,
1811 Some(serde_json::json!({ "verdict": "ok" })),
1812 "a schema'd step returns the validated object in `structured`"
1813 );
1814 }
1815
1816 #[tokio::test]
1817 async fn execute_step_without_schema_has_no_structured_output() {
1818 let workspace = tempfile::tempdir().unwrap();
1819 let executor = TaskExecutor::new(
1820 Arc::new(AgentRegistry::new()),
1821 Arc::new(SchemaCoercionClient),
1822 workspace.path().to_string_lossy().to_string(),
1823 );
1824 let spec = AgentStepSpec::new("step-2", "general", "assess", "Assess the thing.");
1825
1826 let outcome = executor.execute_step(spec, None).await;
1827
1828 assert!(outcome.success, "step should succeed: {}", outcome.output);
1829 assert_eq!(
1830 outcome.structured, None,
1831 "no schema requested → no structured output, no coercion call"
1832 );
1833 }
1834
1835 struct SchemaFailClient;
1839
1840 #[async_trait::async_trait]
1841 impl LlmClient for SchemaFailClient {
1842 async fn complete(
1843 &self,
1844 messages: &[Message],
1845 system: Option<&str>,
1846 tools: &[ToolDefinition],
1847 ) -> Result<LlmResponse> {
1848 if system == Some(crate::prompts::PRE_ANALYSIS_SYSTEM) {
1849 return Ok(pre_analysis_response(messages));
1850 }
1851 if tools.iter().any(|t| t.name == "emit_step_output") {
1852 return Ok(MockLlmClient::tool_call_response(
1854 "coerce-fail",
1855 "emit_step_output",
1856 serde_json::json!({}),
1857 ));
1858 }
1859 Ok(text_response("some answer"))
1860 }
1861
1862 async fn complete_streaming(
1863 &self,
1864 _messages: &[Message],
1865 _system: Option<&str>,
1866 _tools: &[ToolDefinition],
1867 _cancel_token: tokio_util::sync::CancellationToken,
1868 ) -> Result<mpsc::Receiver<StreamEvent>> {
1869 anyhow::bail!("streaming unused")
1870 }
1871 }
1872
1873 #[tokio::test]
1874 async fn execute_step_with_schema_demotes_step_on_coercion_failure() {
1875 let workspace = tempfile::tempdir().unwrap();
1876 let executor = TaskExecutor::new(
1877 Arc::new(AgentRegistry::new()),
1878 Arc::new(SchemaFailClient),
1879 workspace.path().to_string_lossy().to_string(),
1880 );
1881 let spec = AgentStepSpec::new("step-x", "general", "assess", "Assess the thing.")
1882 .with_output_schema(verdict_schema());
1883
1884 let outcome = executor.execute_step(spec, None).await;
1885
1886 assert!(
1887 !outcome.success,
1888 "a step whose output can't satisfy the schema is demoted to failure"
1889 );
1890 assert_eq!(outcome.structured, None, "no validated object on failure");
1891 assert!(
1892 outcome.output.contains("[structured output failed"),
1893 "the demotion marker is appended: {}",
1894 outcome.output
1895 );
1896 }
1897
1898 #[tokio::test]
1899 async fn parallel_isolates_schema_coercion_failure_from_sibling() {
1900 let workspace = tempfile::tempdir().unwrap();
1901 let executor: Arc<dyn AgentExecutor> = Arc::new(TaskExecutor::new(
1902 Arc::new(AgentRegistry::new()),
1903 Arc::new(SchemaFailClient),
1904 workspace.path().to_string_lossy().to_string(),
1905 ));
1906 let specs = vec![
1909 AgentStepSpec::new("plain", "general", "d", "p"),
1910 AgentStepSpec::new("schemad", "general", "d", "p").with_output_schema(verdict_schema()),
1911 ];
1912 let out = crate::orchestration::execute_steps_parallel(executor, specs, None).await;
1913
1914 assert_eq!(out.len(), 2);
1915 assert_eq!(out[0].task_id, "plain");
1916 assert!(out[0].success, "no-schema sibling unaffected");
1917 assert_eq!(out[0].structured, None);
1918 assert_eq!(out[1].task_id, "schemad");
1919 assert!(!out[1].success, "schema-failing step surfaces as failure");
1920 assert_eq!(out[1].structured, None);
1921 assert!(out[1].output.contains("[structured output failed"));
1922 }
1923
1924 #[tokio::test]
1925 async fn failed_step_with_schema_skips_coercion() {
1926 let workspace = tempfile::tempdir().unwrap();
1927 let executor = TaskExecutor::new(
1928 Arc::new(AgentRegistry::new()),
1929 Arc::new(SchemaCoercionClient),
1930 workspace.path().to_string_lossy().to_string(),
1931 );
1932 let spec = AgentStepSpec::new("step-y", "no-such-agent", "d", "p")
1935 .with_output_schema(verdict_schema());
1936
1937 let outcome = executor.execute_step(spec, None).await;
1938
1939 assert!(!outcome.success);
1940 assert_eq!(outcome.structured, None);
1941 assert!(
1942 !outcome.output.contains("[structured output failed"),
1943 "coercion never ran — failure is the run error, not a coercion failure: {}",
1944 outcome.output
1945 );
1946 }
1947
1948 struct StaticLlmClient {
1949 text: String,
1950 }
1951
1952 impl StaticLlmClient {
1953 fn new(text: impl Into<String>) -> Self {
1954 Self { text: text.into() }
1955 }
1956 }
1957
1958 #[async_trait::async_trait]
1959 impl LlmClient for StaticLlmClient {
1960 async fn complete(
1961 &self,
1962 messages: &[Message],
1963 system: Option<&str>,
1964 _tools: &[ToolDefinition],
1965 ) -> Result<LlmResponse> {
1966 if system == Some(crate::prompts::PRE_ANALYSIS_SYSTEM) {
1967 return Ok(pre_analysis_response(messages));
1968 }
1969 Ok(text_response(self.text.clone()))
1970 }
1971
1972 async fn complete_streaming(
1973 &self,
1974 _messages: &[Message],
1975 _system: Option<&str>,
1976 _tools: &[ToolDefinition],
1977 _cancel_token: tokio_util::sync::CancellationToken,
1978 ) -> Result<mpsc::Receiver<StreamEvent>> {
1979 anyhow::bail!("streaming is not used by task executor tests")
1980 }
1981 }
1982
1983 struct ConcurrentLlmClient {
1984 barrier: Arc<Barrier>,
1985 active: AtomicUsize,
1986 max_active: AtomicUsize,
1987 }
1988
1989 impl ConcurrentLlmClient {
1990 fn new(task_count: usize) -> Self {
1991 Self {
1992 barrier: Arc::new(Barrier::new(task_count)),
1993 active: AtomicUsize::new(0),
1994 max_active: AtomicUsize::new(0),
1995 }
1996 }
1997
1998 fn max_active(&self) -> usize {
1999 self.max_active.load(Ordering::SeqCst)
2000 }
2001
2002 fn record_active(&self) {
2003 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
2004 let mut observed = self.max_active.load(Ordering::SeqCst);
2005 while active > observed {
2006 match self.max_active.compare_exchange(
2007 observed,
2008 active,
2009 Ordering::SeqCst,
2010 Ordering::SeqCst,
2011 ) {
2012 Ok(_) => break,
2013 Err(next) => observed = next,
2014 }
2015 }
2016 }
2017 }
2018
2019 struct LimitedConcurrencyLlmClient {
2020 active: AtomicUsize,
2021 max_active: AtomicUsize,
2022 }
2023
2024 impl LimitedConcurrencyLlmClient {
2025 fn new() -> Self {
2026 Self {
2027 active: AtomicUsize::new(0),
2028 max_active: AtomicUsize::new(0),
2029 }
2030 }
2031
2032 fn max_active(&self) -> usize {
2033 self.max_active.load(Ordering::SeqCst)
2034 }
2035
2036 fn record_active(&self) {
2037 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
2038 self.max_active.fetch_max(active, Ordering::SeqCst);
2039 }
2040 }
2041
2042 #[async_trait::async_trait]
2043 impl LlmClient for LimitedConcurrencyLlmClient {
2044 async fn complete(
2045 &self,
2046 messages: &[Message],
2047 system: Option<&str>,
2048 _tools: &[ToolDefinition],
2049 ) -> Result<LlmResponse> {
2050 if system == Some(crate::prompts::PRE_ANALYSIS_SYSTEM) {
2051 return Ok(pre_analysis_response(messages));
2052 }
2053
2054 let prompt = last_text(messages);
2055 self.record_active();
2056 tokio::time::sleep(Duration::from_millis(40)).await;
2057 self.active.fetch_sub(1, Ordering::SeqCst);
2058 Ok(text_response(format!("completed: {prompt}")))
2059 }
2060
2061 async fn complete_streaming(
2062 &self,
2063 _messages: &[Message],
2064 _system: Option<&str>,
2065 _tools: &[ToolDefinition],
2066 _cancel_token: tokio_util::sync::CancellationToken,
2067 ) -> Result<mpsc::Receiver<StreamEvent>> {
2068 anyhow::bail!("streaming is not used by task executor tests")
2069 }
2070 }
2071
2072 #[async_trait::async_trait]
2073 impl LlmClient for ConcurrentLlmClient {
2074 async fn complete(
2075 &self,
2076 messages: &[Message],
2077 system: Option<&str>,
2078 _tools: &[ToolDefinition],
2079 ) -> Result<LlmResponse> {
2080 if system == Some(crate::prompts::PRE_ANALYSIS_SYSTEM) {
2081 return Ok(pre_analysis_response(messages));
2082 }
2083
2084 let prompt = last_text(messages);
2085 self.record_active();
2086 self.barrier.wait().await;
2087 if prompt.contains("slow") {
2088 tokio::time::sleep(Duration::from_millis(120)).await;
2089 } else {
2090 tokio::time::sleep(Duration::from_millis(10)).await;
2091 }
2092 self.active.fetch_sub(1, Ordering::SeqCst);
2093 Ok(text_response(format!("completed: {prompt}")))
2094 }
2095
2096 async fn complete_streaming(
2097 &self,
2098 _messages: &[Message],
2099 _system: Option<&str>,
2100 _tools: &[ToolDefinition],
2101 _cancel_token: tokio_util::sync::CancellationToken,
2102 ) -> Result<mpsc::Receiver<StreamEvent>> {
2103 anyhow::bail!("streaming is not used by task executor tests")
2104 }
2105 }
2106
2107 fn test_registry_with_writer() -> Arc<AgentRegistry> {
2108 let registry = AgentRegistry::new();
2109 let spec = crate::subagent::WorkerAgentSpec::custom("writer", "Write files")
2110 .with_permissions(PermissionPolicy::new().allow("write(*)").allow("read(*)"))
2111 .with_prompt("Write files when asked.")
2112 .with_max_steps(3);
2113 registry.register(spec.into_agent_definition());
2114 Arc::new(registry)
2115 }
2116
2117 fn test_registry_with_text_worker() -> Arc<AgentRegistry> {
2118 let registry = AgentRegistry::new();
2119 let spec = crate::subagent::WorkerAgentSpec::custom("worker", "Text worker")
2120 .with_prompt("Return a concise result.")
2121 .with_max_steps(1);
2122 registry.register(spec.into_agent_definition());
2123 Arc::new(registry)
2124 }
2125
2126 #[tokio::test]
2127 async fn task_child_run_permission_allow() {
2128 let workspace = tempfile::tempdir().unwrap();
2129 let mock = Arc::new(MockLlmClient::new(vec![
2130 MockLlmClient::tool_call_response(
2131 "t1",
2132 "write",
2133 serde_json::json!({
2134 "file_path": workspace.path().join("out.txt").to_string_lossy(),
2135 "content": "WRITTEN"
2136 }),
2137 ),
2138 MockLlmClient::text_response("Done."),
2139 ]));
2140
2141 let executor = TaskExecutor::new(
2142 test_registry_with_writer(),
2143 mock,
2144 workspace.path().to_string_lossy().to_string(),
2145 );
2146
2147 let result = executor
2148 .execute(
2149 TaskParams {
2150 agent: "writer".to_string(),
2151 description: "Write file".to_string(),
2152 prompt: "Write out.txt".to_string(),
2153 background: false,
2154 max_steps: Some(3),
2155 output_schema: None,
2156 },
2157 None,
2158 None,
2159 )
2160 .await
2161 .unwrap();
2162
2163 assert!(
2164 result.success,
2165 "child run should succeed: {}",
2166 result.output
2167 );
2168 assert!(
2169 !result.output.contains("Permission denied"),
2170 "no permission denial: {}",
2171 result.output
2172 );
2173 let content = std::fs::read_to_string(workspace.path().join("out.txt")).unwrap();
2174 assert_eq!(content, "WRITTEN");
2175 }
2176
2177 #[tokio::test]
2178 async fn task_child_run_permission_deny() {
2179 let workspace = tempfile::tempdir().unwrap();
2180 let registry = AgentRegistry::new();
2181 let spec = crate::subagent::WorkerAgentSpec::custom("restricted", "Restricted agent")
2182 .with_permissions(PermissionPolicy::new().allow("read(*)").deny("bash(*)"))
2183 .with_max_steps(3);
2184 registry.register(spec.into_agent_definition());
2185
2186 let mock = Arc::new(MockLlmClient::new(vec![
2187 MockLlmClient::tool_call_response(
2188 "t1",
2189 "bash",
2190 serde_json::json!({"command": "echo hello"}),
2191 ),
2192 MockLlmClient::text_response("Could not run bash."),
2193 ]));
2194
2195 let executor = TaskExecutor::new(
2196 Arc::new(registry),
2197 mock,
2198 workspace.path().to_string_lossy().to_string(),
2199 );
2200
2201 let result = executor
2202 .execute(
2203 TaskParams {
2204 agent: "restricted".to_string(),
2205 description: "Try bash".to_string(),
2206 prompt: "Run echo hello".to_string(),
2207 background: false,
2208 max_steps: Some(3),
2209 output_schema: None,
2210 },
2211 None,
2212 None,
2213 )
2214 .await
2215 .unwrap();
2216
2217 assert!(result.success, "agent should complete: {}", result.output);
2220 }
2221
2222 #[tokio::test]
2223 async fn task_child_run_confirmation_auto_approve() {
2224 let workspace = tempfile::tempdir().unwrap();
2225 let registry = AgentRegistry::new();
2226 let spec = crate::subagent::WorkerAgentSpec::custom("reader-writer", "Read and write")
2229 .with_permissions(PermissionPolicy::new().allow("read(*)"))
2230 .with_max_steps(3);
2231 registry.register(spec.into_agent_definition());
2232
2233 let mock = Arc::new(MockLlmClient::new(vec![
2234 MockLlmClient::tool_call_response(
2235 "t1",
2236 "write",
2237 serde_json::json!({
2238 "file_path": workspace.path().join("auto.txt").to_string_lossy(),
2239 "content": "AUTO_APPROVED"
2240 }),
2241 ),
2242 MockLlmClient::text_response("Written."),
2243 ]));
2244
2245 let executor = TaskExecutor::new(
2246 Arc::new(registry),
2247 mock,
2248 workspace.path().to_string_lossy().to_string(),
2249 );
2250
2251 let result = executor
2252 .execute(
2253 TaskParams {
2254 agent: "reader-writer".to_string(),
2255 description: "Write via auto-approve".to_string(),
2256 prompt: "Write auto.txt".to_string(),
2257 background: false,
2258 max_steps: Some(3),
2259 output_schema: None,
2260 },
2261 None,
2262 None,
2263 )
2264 .await
2265 .unwrap();
2266
2267 assert!(
2268 result.success,
2269 "Ask should be auto-approved: {}",
2270 result.output
2271 );
2272 assert!(
2273 !result.output.contains("MissingConfirmationManager"),
2274 "no MissingConfirmationManager: {}",
2275 result.output
2276 );
2277 }
2278
2279 #[tokio::test]
2280 async fn task_child_run_step_budget_enforced() {
2281 let workspace = tempfile::tempdir().unwrap();
2282 let mock = Arc::new(MockLlmClient::new(vec![
2283 MockLlmClient::tool_call_response(
2284 "t1",
2285 "read",
2286 serde_json::json!({"file_path": "/tmp/a.txt"}),
2287 ),
2288 MockLlmClient::tool_call_response(
2289 "t2",
2290 "read",
2291 serde_json::json!({"file_path": "/tmp/b.txt"}),
2292 ),
2293 MockLlmClient::tool_call_response(
2294 "t3",
2295 "read",
2296 serde_json::json!({"file_path": "/tmp/c.txt"}),
2297 ),
2298 MockLlmClient::text_response("Should not reach here."),
2299 ]));
2300
2301 let executor = TaskExecutor::new(
2302 test_registry_with_writer(),
2303 mock,
2304 workspace.path().to_string_lossy().to_string(),
2305 );
2306
2307 let result = executor
2308 .execute(
2309 TaskParams {
2310 agent: "writer".to_string(),
2311 description: "Exceed budget".to_string(),
2312 prompt: "Read many files".to_string(),
2313 background: false,
2314 max_steps: Some(2),
2315 output_schema: None,
2316 },
2317 None,
2318 None,
2319 )
2320 .await
2321 .unwrap();
2322
2323 assert!(
2325 !result.success,
2326 "should fail when exceeding step budget: {}",
2327 result.output
2328 );
2329 assert!(
2330 result.output.contains("Max tool rounds") || result.output.contains("max tool rounds"),
2331 "error should mention tool rounds: {}",
2332 result.output
2333 );
2334 }
2335
2336 #[tokio::test]
2337 async fn parallel_task_executor_runs_children_concurrently_and_preserves_input_order() {
2338 let workspace = tempfile::tempdir().unwrap();
2339 let client = Arc::new(ConcurrentLlmClient::new(2));
2340 let executor = Arc::new(TaskExecutor::new(
2341 test_registry_with_text_worker(),
2342 client.clone(),
2343 workspace.path().to_string_lossy().to_string(),
2344 ));
2345
2346 let tasks = vec![
2347 TaskParams {
2348 agent: "worker".to_string(),
2349 description: "Slow task".to_string(),
2350 prompt: "slow branch".to_string(),
2351 background: false,
2352 max_steps: Some(1),
2353 output_schema: None,
2354 },
2355 TaskParams {
2356 agent: "worker".to_string(),
2357 description: "Fast task".to_string(),
2358 prompt: "fast branch".to_string(),
2359 background: false,
2360 max_steps: Some(1),
2361 output_schema: None,
2362 },
2363 ];
2364
2365 let results = tokio::time::timeout(
2366 Duration::from_secs(2),
2367 executor.execute_parallel(tasks, None, None),
2368 )
2369 .await
2370 .expect("parallel children should reach the barrier and complete");
2371
2372 assert_eq!(results.len(), 2);
2373 assert!(
2374 client.max_active() >= 2,
2375 "expected concurrent child execution, max_active={}",
2376 client.max_active()
2377 );
2378 assert!(results[0].success);
2379 assert!(results[0].output.contains("slow branch"));
2380 assert!(results[1].success);
2381 assert!(results[1].output.contains("fast branch"));
2382 }
2383
2384 #[tokio::test]
2385 async fn parallel_task_executor_respects_configured_concurrency_limit() {
2386 let workspace = tempfile::tempdir().unwrap();
2387 let client = Arc::new(LimitedConcurrencyLlmClient::new());
2388 let executor = Arc::new(
2389 TaskExecutor::new(
2390 test_registry_with_text_worker(),
2391 client.clone(),
2392 workspace.path().to_string_lossy().to_string(),
2393 )
2394 .with_max_parallel_tasks(2),
2395 );
2396
2397 let tasks = (0..5)
2398 .map(|idx| TaskParams {
2399 agent: "worker".to_string(),
2400 description: format!("Task {idx}"),
2401 prompt: format!("branch {idx}"),
2402 background: false,
2403 max_steps: Some(1),
2404 output_schema: None,
2405 })
2406 .collect::<Vec<_>>();
2407
2408 let results = executor.execute_parallel(tasks, None, None).await;
2409
2410 assert_eq!(results.len(), 5);
2411 assert!(results.iter().all(|result| result.success));
2412 assert_eq!(client.max_active(), 2);
2413 }
2414
2415 #[tokio::test]
2416 async fn parallel_task_executor_isolates_unknown_agent_failure() {
2417 let workspace = tempfile::tempdir().unwrap();
2418 let executor = Arc::new(TaskExecutor::new(
2419 test_registry_with_text_worker(),
2420 Arc::new(StaticLlmClient::new("valid branch done")),
2421 workspace.path().to_string_lossy().to_string(),
2422 ));
2423
2424 let tasks = vec![
2425 TaskParams {
2426 agent: "missing-agent".to_string(),
2427 description: "Missing".to_string(),
2428 prompt: "should fail".to_string(),
2429 background: false,
2430 max_steps: Some(1),
2431 output_schema: None,
2432 },
2433 TaskParams {
2434 agent: "worker".to_string(),
2435 description: "Valid".to_string(),
2436 prompt: "should succeed".to_string(),
2437 background: false,
2438 max_steps: Some(1),
2439 output_schema: None,
2440 },
2441 ];
2442
2443 let results = executor.execute_parallel(tasks, None, None).await;
2444
2445 assert_eq!(results.len(), 2);
2446 assert!(!results[0].success);
2447 assert_eq!(results[0].agent, "missing-agent");
2448 assert!(results[0].output.contains("Unknown agent type"));
2449 assert!(results[1].success);
2450 assert_eq!(results[1].agent, "worker");
2451 assert!(results[1].output.contains("valid branch done"));
2452 }
2453
2454 #[tokio::test]
2455 async fn parallel_task_executor_emits_subagent_events_for_each_child() {
2456 let workspace = tempfile::tempdir().unwrap();
2457 let executor = Arc::new(TaskExecutor::new(
2458 test_registry_with_text_worker(),
2459 Arc::new(StaticLlmClient::new("done")),
2460 workspace.path().to_string_lossy().to_string(),
2461 ));
2462 let (tx, mut rx) = broadcast::channel(64);
2463
2464 let tasks = vec![
2465 TaskParams {
2466 agent: "worker".to_string(),
2467 description: "One".to_string(),
2468 prompt: "first".to_string(),
2469 background: false,
2470 max_steps: Some(1),
2471 output_schema: None,
2472 },
2473 TaskParams {
2474 agent: "worker".to_string(),
2475 description: "Two".to_string(),
2476 prompt: "second".to_string(),
2477 background: false,
2478 max_steps: Some(1),
2479 output_schema: None,
2480 },
2481 ];
2482
2483 let results = executor.execute_parallel(tasks, Some(tx), None).await;
2484 assert_eq!(results.len(), 2);
2485 tokio::time::sleep(Duration::from_millis(20)).await;
2486
2487 let mut starts = Vec::new();
2488 let mut ends = Vec::new();
2489 let mut progress_statuses: Vec<String> = Vec::new();
2490 while let Ok(event) = rx.try_recv() {
2491 match event {
2492 AgentEvent::SubagentStart { description, .. } => starts.push(description),
2493 AgentEvent::SubagentEnd { agent, success, .. } => ends.push((agent, success)),
2494 AgentEvent::SubagentProgress { status, .. } => progress_statuses.push(status),
2495 _ => {}
2496 }
2497 }
2498
2499 starts.sort();
2500 assert_eq!(starts, vec!["One".to_string(), "Two".to_string()]);
2501 assert_eq!(ends.len(), 2);
2502 assert!(ends
2503 .iter()
2504 .all(|(agent, success)| agent == "worker" && *success));
2505 assert!(
2508 progress_statuses
2509 .iter()
2510 .filter(|s| s == &"turn_completed")
2511 .count()
2512 >= 2,
2513 "expected at least two turn_completed progress events, got {:?}",
2514 progress_statuses
2515 );
2516 }
2517
2518 #[tokio::test]
2519 async fn parallel_task_tool_reports_error_when_any_child_fails() {
2520 let workspace = tempfile::tempdir().unwrap();
2521 let executor = Arc::new(TaskExecutor::new(
2522 test_registry_with_text_worker(),
2523 Arc::new(StaticLlmClient::new("valid branch done")),
2524 workspace.path().to_string_lossy().to_string(),
2525 ));
2526 let tool = ParallelTaskTool::new(executor);
2527 let ctx = ToolContext::new(workspace.path().to_path_buf());
2528
2529 let output = tool
2530 .execute(
2531 &serde_json::json!({
2532 "tasks": [
2533 {
2534 "agent": "missing-agent",
2535 "description": "Missing",
2536 "prompt": "should fail"
2537 },
2538 {
2539 "agent": "worker",
2540 "description": "Valid",
2541 "prompt": "should succeed"
2542 }
2543 ]
2544 }),
2545 &ctx,
2546 )
2547 .await
2548 .unwrap();
2549
2550 assert!(
2551 !output.success,
2552 "parallel_task should fail when any child result fails"
2553 );
2554 assert!(output.content.contains("[ERR]"));
2555 assert!(output.content.contains("[OK]"));
2556 let metadata = output.metadata.expect("metadata");
2557 assert_eq!(metadata["task_count"], 2);
2558 assert_eq!(metadata["success_count"], 1);
2559 assert_eq!(metadata["failed_count"], 1);
2560 assert_eq!(metadata["all_success"], false);
2561 assert_eq!(metadata["partial_failure"], true);
2562 assert_eq!(metadata["allow_partial_failure"], false);
2563 assert_eq!(metadata["results"][0]["success"], false);
2564 assert_eq!(metadata["results"][1]["success"], true);
2565 }
2566
2567 #[tokio::test]
2568 async fn parallel_task_tool_allows_partial_failure_when_requested() {
2569 let workspace = tempfile::tempdir().unwrap();
2570 let executor = Arc::new(TaskExecutor::new(
2571 test_registry_with_text_worker(),
2572 Arc::new(StaticLlmClient::new("valid branch done")),
2573 workspace.path().to_string_lossy().to_string(),
2574 ));
2575 let tool = ParallelTaskTool::new(executor);
2576 let ctx = ToolContext::new(workspace.path().to_path_buf());
2577
2578 let output = tool
2579 .execute(
2580 &serde_json::json!({
2581 "allow_partial_failure": true,
2582 "tasks": [
2583 {
2584 "agent": "missing-agent",
2585 "description": "Missing",
2586 "prompt": "should fail"
2587 },
2588 {
2589 "agent": "worker",
2590 "description": "Valid",
2591 "prompt": "should succeed"
2592 }
2593 ]
2594 }),
2595 &ctx,
2596 )
2597 .await
2598 .unwrap();
2599
2600 assert!(
2601 output.success,
2602 "parallel_task should continue when partial failures are allowed"
2603 );
2604 assert!(output.content.contains("[ERR]"));
2605 assert!(output.content.contains("[OK]"));
2606 assert!(output.content.contains("Partial failure tolerated"));
2607 let metadata = output.metadata.expect("metadata");
2608 assert_eq!(metadata["task_count"], 2);
2609 assert_eq!(metadata["success_count"], 1);
2610 assert_eq!(metadata["failed_count"], 1);
2611 assert_eq!(metadata["all_success"], false);
2612 assert_eq!(metadata["partial_failure"], true);
2613 assert_eq!(metadata["allow_partial_failure"], true);
2614 assert_eq!(metadata["results"][0]["success"], false);
2615 assert_eq!(metadata["results"][1]["success"], true);
2616 }
2617
2618 #[tokio::test]
2619 async fn parallel_task_tool_returns_structured_child_output_when_schema_requested() {
2620 let workspace = tempfile::tempdir().unwrap();
2621 let executor = Arc::new(TaskExecutor::new(
2622 test_registry_with_text_worker(),
2623 Arc::new(SchemaCoercionClient),
2624 workspace.path().to_string_lossy().to_string(),
2625 ));
2626 let tool = ParallelTaskTool::new(executor);
2627 let ctx = ToolContext::new(workspace.path().to_path_buf());
2628
2629 let output = tool
2630 .execute(
2631 &serde_json::json!({
2632 "tasks": [{
2633 "agent": "worker",
2634 "description": "Structured verdict",
2635 "prompt": "Return a verdict.",
2636 "output_schema": verdict_schema()
2637 }]
2638 }),
2639 &ctx,
2640 )
2641 .await
2642 .unwrap();
2643
2644 assert!(
2645 output.success,
2646 "parallel_task should succeed: {}",
2647 output.content
2648 );
2649 let metadata = output.metadata.expect("metadata");
2650 assert_eq!(metadata["results"][0]["success"], true);
2651 assert_eq!(metadata["results"][0]["structured"]["verdict"], "ok");
2652 }
2653
2654 #[tokio::test]
2655 async fn parallel_task_tool_still_fails_when_all_children_fail() {
2656 let workspace = tempfile::tempdir().unwrap();
2657 let executor = Arc::new(TaskExecutor::new(
2658 test_registry_with_text_worker(),
2659 Arc::new(StaticLlmClient::new("unused")),
2660 workspace.path().to_string_lossy().to_string(),
2661 ));
2662 let tool = ParallelTaskTool::new(executor);
2663 let ctx = ToolContext::new(workspace.path().to_path_buf());
2664
2665 let output = tool
2666 .execute(
2667 &serde_json::json!({
2668 "allow_partial_failure": true,
2669 "tasks": [
2670 {
2671 "agent": "missing-one",
2672 "description": "Missing one",
2673 "prompt": "should fail"
2674 },
2675 {
2676 "agent": "missing-two",
2677 "description": "Missing two",
2678 "prompt": "should also fail"
2679 }
2680 ]
2681 }),
2682 &ctx,
2683 )
2684 .await
2685 .unwrap();
2686
2687 assert!(
2688 !output.success,
2689 "parallel_task should fail if every child task fails"
2690 );
2691 let metadata = output.metadata.expect("metadata");
2692 assert_eq!(metadata["success_count"], 0);
2693 assert_eq!(metadata["failed_count"], 2);
2694 assert_eq!(metadata["all_success"], false);
2695 assert_eq!(metadata["partial_failure"], false);
2696 assert_eq!(metadata["allow_partial_failure"], true);
2697 }
2698
2699 #[tokio::test]
2700 async fn parallel_task_both_inherit_permissions() {
2701 let workspace = tempfile::tempdir().unwrap();
2702 let mock = Arc::new(MockLlmClient::new(vec![
2703 MockLlmClient::tool_call_response(
2705 "t1",
2706 "write",
2707 serde_json::json!({
2708 "file_path": workspace.path().join("p1.txt").to_string_lossy(),
2709 "content": "P1"
2710 }),
2711 ),
2712 MockLlmClient::text_response("Done 1."),
2713 MockLlmClient::tool_call_response(
2715 "t2",
2716 "write",
2717 serde_json::json!({
2718 "file_path": workspace.path().join("p2.txt").to_string_lossy(),
2719 "content": "P2"
2720 }),
2721 ),
2722 MockLlmClient::text_response("Done 2."),
2723 ]));
2724
2725 let executor = Arc::new(TaskExecutor::new(
2726 test_registry_with_writer(),
2727 mock,
2728 workspace.path().to_string_lossy().to_string(),
2729 ));
2730
2731 let tasks = vec![
2732 TaskParams {
2733 agent: "writer".to_string(),
2734 description: "Write p1".to_string(),
2735 prompt: "Write p1.txt".to_string(),
2736 background: false,
2737 max_steps: Some(3),
2738 output_schema: None,
2739 },
2740 TaskParams {
2741 agent: "writer".to_string(),
2742 description: "Write p2".to_string(),
2743 prompt: "Write p2.txt".to_string(),
2744 background: false,
2745 max_steps: Some(3),
2746 output_schema: None,
2747 },
2748 ];
2749
2750 let results = executor.execute_parallel(tasks, None, None).await;
2751 assert_eq!(results.len(), 2);
2752
2753 for result in &results {
2754 assert!(
2755 result.success,
2756 "parallel child should succeed: {}",
2757 result.output
2758 );
2759 }
2760 }
2761
2762 #[test]
2763 fn synthesize_progress_emits_tool_completed_for_tool_end() {
2764 let event = AgentEvent::ToolEnd {
2765 id: "call-1".to_string(),
2766 name: "bash".to_string(),
2767 output: "hello".to_string(),
2768 exit_code: 0,
2769 metadata: None,
2770 error_kind: None,
2771 };
2772 let progress =
2773 synthesize_subagent_progress(&event, "task-1", "task-run-task-1").expect("some");
2774 match progress {
2775 AgentEvent::SubagentProgress {
2776 task_id,
2777 session_id,
2778 status,
2779 metadata,
2780 } => {
2781 assert_eq!(task_id, "task-1");
2782 assert_eq!(session_id, "task-run-task-1");
2783 assert_eq!(status, "tool_completed");
2784 assert_eq!(metadata["tool"], "bash");
2785 assert_eq!(metadata["exit_code"], 0);
2786 assert_eq!(metadata["output_bytes"], 5);
2787 assert!(metadata.get("error_kind").is_none());
2788 }
2789 other => panic!("expected SubagentProgress, got {:?}", other),
2790 }
2791 }
2792
2793 #[test]
2794 fn synthesize_progress_includes_error_kind_when_present() {
2795 let event = AgentEvent::ToolEnd {
2796 id: "call-2".to_string(),
2797 name: "edit".to_string(),
2798 output: "boom".to_string(),
2799 exit_code: 1,
2800 metadata: None,
2801 error_kind: Some(crate::tools::ToolErrorKind::NotFound {
2802 path: "missing.txt".to_string(),
2803 }),
2804 };
2805 let progress =
2806 synthesize_subagent_progress(&event, "task-x", "task-run-task-x").expect("some");
2807 if let AgentEvent::SubagentProgress { metadata, .. } = progress {
2808 assert!(
2809 metadata.get("error_kind").is_some(),
2810 "error_kind should propagate into metadata"
2811 );
2812 } else {
2813 panic!("expected SubagentProgress");
2814 }
2815 }
2816
2817 #[test]
2818 fn synthesize_progress_emits_turn_completed_for_turn_end() {
2819 let event = AgentEvent::TurnEnd {
2820 turn: 3,
2821 usage: crate::llm::TokenUsage {
2822 prompt_tokens: 100,
2823 completion_tokens: 25,
2824 total_tokens: 125,
2825 cache_read_tokens: None,
2826 cache_write_tokens: None,
2827 },
2828 };
2829 let progress =
2830 synthesize_subagent_progress(&event, "task-1", "task-run-task-1").expect("some");
2831 if let AgentEvent::SubagentProgress {
2832 status, metadata, ..
2833 } = progress
2834 {
2835 assert_eq!(status, "turn_completed");
2836 assert_eq!(metadata["turn"], 3);
2837 assert_eq!(metadata["total_tokens"], 125);
2838 assert_eq!(metadata["prompt_tokens"], 100);
2839 assert_eq!(metadata["completion_tokens"], 25);
2840 } else {
2841 panic!("expected SubagentProgress");
2842 }
2843 }
2844
2845 #[test]
2846 fn synthesize_progress_ignores_unrelated_events() {
2847 let ignored = [
2848 AgentEvent::TextDelta {
2849 text: "hi".to_string(),
2850 },
2851 AgentEvent::ToolStart {
2852 id: "x".to_string(),
2853 name: "bash".to_string(),
2854 },
2855 AgentEvent::TurnStart { turn: 1 },
2856 AgentEvent::SubagentStart {
2857 task_id: "nested".to_string(),
2858 session_id: "nested-run".to_string(),
2859 parent_session_id: "parent".to_string(),
2860 agent: "explore".to_string(),
2861 description: "nested".to_string(),
2862 started_ms: 0,
2863 },
2864 ];
2865 for event in &ignored {
2866 assert!(
2867 synthesize_subagent_progress(event, "task", "session").is_none(),
2868 "{:?} should not emit progress",
2869 event
2870 );
2871 }
2872 }
2873}