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