1use std::sync::atomic::{AtomicUsize, Ordering};
18use std::sync::Arc;
19
20use async_trait::async_trait;
21use serde_json::Value;
22use tokio_util::sync::CancellationToken;
23use uuid::Uuid;
24
25use crate::error::{CoreError, Result as CoreResult};
26use crate::event::EventSink;
27use crate::message::{AgentMessage, ContentBlock, Role};
28use crate::model::{Model, ModelProvider};
29use crate::runner::{run_agent, RunConfig, RunOutcome};
30use crate::tool::{InvokeContext, Tool, ToolDefinition, ToolResult};
31
32pub const DEFAULT_MAX_DEPTH: usize = 5;
36
37pub const DEFAULT_MAX_DELEGATIONS: usize = 64;
44
45#[derive(Clone)]
52pub struct SubagentProfile {
53 pub name: String,
55 pub description: String,
57 pub instructions: String,
59 pub model: Option<Model>,
61 pub config: Option<RunConfig>,
63 pub tools: Vec<Arc<dyn Tool>>,
65 pub subagents: Vec<SubagentProfile>,
68}
69
70impl SubagentProfile {
71 #[must_use]
76 pub fn new(name: impl Into<String>, instructions: impl Into<String>) -> Self {
77 Self {
78 name: name.into(),
79 description: String::new(),
82 instructions: instructions.into(),
83 model: None,
84 config: None,
85 tools: Vec::new(),
86 subagents: Vec::new(),
87 }
88 }
89
90 #[must_use]
92 pub fn with_description(mut self, description: impl Into<String>) -> Self {
93 self.description = description.into();
94 self
95 }
96
97 #[must_use]
99 pub fn with_model(mut self, model: Model) -> Self {
100 self.model = Some(model);
101 self
102 }
103
104 #[must_use]
106 pub fn with_config(mut self, config: RunConfig) -> Self {
107 self.config = Some(config);
108 self
109 }
110
111 #[must_use]
113 pub fn with_tool(mut self, tool: Arc<dyn Tool>) -> Self {
114 self.tools.push(tool);
115 self
116 }
117
118 #[must_use]
120 pub fn with_subagent(mut self, subagent: SubagentProfile) -> Self {
121 self.subagents.push(subagent);
122 self
123 }
124}
125
126#[derive(Clone, Copy, Debug)]
128pub struct SubagentOptions {
129 pub max_depth: usize,
134 pub max_delegations: usize,
140}
141
142impl Default for SubagentOptions {
143 fn default() -> Self {
144 Self {
145 max_depth: DEFAULT_MAX_DEPTH,
146 max_delegations: DEFAULT_MAX_DELEGATIONS,
147 }
148 }
149}
150
151pub struct TaskTool {
163 provider: Arc<dyn ModelProvider>,
165 parent_model: Model,
167 parent_config: RunConfig,
169 subagents: Vec<SubagentProfile>,
171 max_depth: usize,
173 depth: usize,
175 cancel: CancellationToken,
177 event_sink: Option<Arc<dyn EventSink>>,
180 remaining_delegations: Arc<AtomicUsize>,
183}
184
185impl TaskTool {
186 #[must_use]
191 pub fn new(
192 provider: Arc<dyn ModelProvider>,
193 parent_model: Model,
194 parent_config: RunConfig,
195 subagents: Vec<SubagentProfile>,
196 options: SubagentOptions,
197 cancel: CancellationToken,
198 event_sink: Option<Arc<dyn EventSink>>,
199 ) -> Self {
200 Self {
201 provider,
202 parent_model,
203 parent_config,
204 subagents,
205 max_depth: options.max_depth,
206 depth: 0,
207 cancel,
208 event_sink,
209 remaining_delegations: Arc::new(AtomicUsize::new(options.max_delegations)),
210 }
211 }
212
213 fn child(
215 &self,
216 subagents: Vec<SubagentProfile>,
217 parent_model: Model,
218 parent_config: RunConfig,
219 ) -> Self {
220 Self {
221 provider: Arc::clone(&self.provider),
222 parent_model,
223 parent_config,
224 subagents,
225 max_depth: self.max_depth,
226 depth: self.depth + 1,
227 cancel: self.cancel.clone(),
228 event_sink: self.event_sink.as_ref().map(Arc::clone),
229 remaining_delegations: Arc::clone(&self.remaining_delegations),
231 }
232 }
233
234 fn resolve(&self, name: &str) -> Option<&SubagentProfile> {
236 self.subagents.iter().find(|s| s.name == name)
237 }
238
239 fn max_delegations_hint(&self) -> usize {
244 self.remaining_delegations.load(Ordering::Relaxed) + 1
247 }
248
249 async fn delegate(&self, profile: &SubagentProfile, prompt: String) -> CoreResult<RunOutcome> {
251 let child_model = profile
253 .model
254 .clone()
255 .unwrap_or_else(|| self.parent_model.clone());
256 let child_config = profile
257 .config
258 .clone()
259 .unwrap_or_else(|| self.parent_config.clone());
260
261 let mut child_tools: Vec<Arc<dyn Tool>> = Vec::new();
268 if !profile.subagents.is_empty() {
269 let child_task = self.child(
270 profile.subagents.clone(),
271 child_model.clone(),
274 child_config.clone(),
275 );
276 child_tools.push(Arc::new(child_task));
277 }
278 child_tools.extend(profile.tools.clone());
279
280 let child_session = Uuid::new_v4();
282 let mut child_messages = vec![
283 AgentMessage {
284 role: Role::System,
285 content: vec![ContentBlock::Text {
286 text: profile.instructions.clone(),
287 }],
288 },
289 AgentMessage {
290 role: Role::User,
291 content: vec![ContentBlock::Text { text: prompt }],
292 },
293 ];
294
295 let child_hooks = crate::event::RunHooks {
298 session_id: Some(child_session),
299 turn_sink: None,
300 event_sink: self.event_sink.as_deref(),
301 policy: None,
302 };
303
304 run_agent(
308 self.provider.as_ref(),
309 &child_tools,
310 &mut child_messages,
311 &child_model,
312 &child_config,
313 &self.cancel,
314 &child_hooks,
315 )
316 .await
317 }
318}
319
320#[async_trait]
321impl Tool for TaskTool {
322 fn definition(&self) -> ToolDefinition {
323 let mut desc = String::from(
324 "Delegate a focused subtask to a named subagent. The subagent runs \
325 in a fresh context and its answer is returned to you. Call this \
326 only when a declared subagent is well-suited to the work. \
327 Available subagents:",
328 );
329 if self.subagents.is_empty() {
330 desc.push_str(" (none declared)");
331 } else {
332 for s in &self.subagents {
333 let guidance = if s.description.trim().is_empty() {
334 "(no description provided)"
335 } else {
336 s.description.trim()
337 };
338 desc.push_str(&format!("\n - \"{}\": {}", s.name, guidance));
339 }
340 }
341
342 let mut fields = serde_json::Map::new();
344 fields.insert("type".into(), Value::String("object".into()));
345 fields.insert(
346 "properties".into(),
347 serde_json::json!({
348 "agent": {
349 "type": "string",
350 "description": "The name of the declared subagent to delegate to."
351 },
352 "prompt": {
353 "type": "string",
354 "description": "The task to give the subagent (it sees this, not your conversation history)."
355 }
356 }),
357 );
358 fields.insert(
359 "required".into(),
360 Value::Array(vec![
361 Value::String("agent".into()),
362 Value::String("prompt".into()),
363 ]),
364 );
365
366 ToolDefinition {
367 name: "task".into(),
368 label: "Task".into(),
369 description: desc,
370 parameters: crate::tool::ParameterSchema {
371 fields: fields.into_iter().collect(),
372 },
373 }
374 }
375
376 async fn execute(&self, ctx: InvokeContext, input: Value) -> CoreResult<ToolResult> {
377 let obj = input.as_object().ok_or_else(|| {
379 CoreError::ToolInputValidation("task tool expects an object input".into())
380 })?;
381 let agent = obj.get("agent").and_then(Value::as_str).ok_or_else(|| {
382 CoreError::ToolInputValidation("task tool requires a string `agent`".into())
383 })?;
384 let prompt = obj.get("prompt").and_then(Value::as_str).ok_or_else(|| {
385 CoreError::ToolInputValidation("task tool requires a string `prompt`".into())
386 })?;
387
388 let profile = match self.resolve(agent) {
390 Some(p) => p,
391 None => {
392 let known: Vec<&str> = self.subagents.iter().map(|s| s.name.as_str()).collect();
393 return Err(CoreError::ToolInputValidation(format!(
394 "subagent not declared: \"{agent}\" (known: {})",
395 known.join(", ")
396 )));
397 }
398 };
399
400 let prev = self.remaining_delegations.fetch_sub(1, Ordering::Relaxed);
405 if prev == 0 {
406 self.remaining_delegations.fetch_add(1, Ordering::Relaxed);
408 return Err(CoreError::ToolInputValidation(format!(
409 "delegation budget exhausted (max {} total delegations across the tree)",
410 self.max_delegations_hint()
411 )));
412 }
413
414 if self.depth >= self.max_depth {
416 self.remaining_delegations.fetch_add(1, Ordering::Relaxed);
419 return Err(CoreError::ToolInputValidation(format!(
420 "delegation depth exceeded (depth {} >= max_depth {})",
421 self.depth, self.max_depth
422 )));
423 }
424
425 if ctx.cancel.is_cancelled() {
428 return Err(CoreError::Cancelled("task delegation cancelled".into()));
429 }
430
431 let outcome = self.delegate(profile, prompt.to_string()).await?;
435 Ok(ToolResult {
436 content: vec![serde_json::json!({
437 "type": "text",
438 "text": if outcome.final_text.trim().is_empty() {
439 "(subagent returned no text)".to_string()
440 } else {
441 outcome.final_text
442 },
443 })],
444 details: None,
445 })
446 }
447}
448
449#[cfg(test)]
450mod tests {
451 use super::*;
452 use crate::model::{ModelRequest, ModelResponse};
453
454 fn dummy_profile(name: &str) -> SubagentProfile {
455 SubagentProfile::new(name, "you are a helper")
456 }
457
458 fn top_level_tool(profiles: Vec<SubagentProfile>, max_depth: usize) -> TaskTool {
459 TaskTool::new(
460 test_provider(),
465 Model {
466 id: "test/model".into(),
467 },
468 RunConfig::default(),
469 profiles,
470 SubagentOptions {
471 max_depth,
472 max_delegations: DEFAULT_MAX_DELEGATIONS,
473 },
474 CancellationToken::new(),
475 None,
476 )
477 }
478
479 fn test_provider() -> Arc<dyn ModelProvider> {
481 use async_trait::async_trait;
482 struct TestProvider;
483 #[async_trait]
484 impl ModelProvider for TestProvider {
485 async fn invoke(
486 &self,
487 _request: crate::model::ModelRequest,
488 ) -> CoreResult<crate::model::ModelResponse> {
489 Ok(crate::model::ModelResponse {
491 messages: vec![crate::message::AgentMessage {
492 role: crate::message::Role::Assistant,
493 content: vec![crate::message::ContentBlock::Text {
494 text: "child done".into(),
495 }],
496 }],
497 })
498 }
499 }
500 Arc::new(TestProvider)
501 }
502
503 #[test]
504 fn definition_lists_declared_subagents() {
505 let profiles = vec![
506 dummy_profile("reviewer").with_description("Review changes."),
507 dummy_profile("classifier").with_description("Classify issues."),
508 ];
509 let tool = top_level_tool(profiles, DEFAULT_MAX_DEPTH);
510 let def = tool.definition();
511 assert_eq!(def.name, "task");
512 assert_eq!(def.label, "Task");
513 assert!(def.description.contains("\"reviewer\""), "missing reviewer");
514 assert!(def.description.contains("Review changes."));
515 assert!(def.description.contains("\"classifier\""));
516 assert!(def.description.contains("Classify issues."));
517 }
518
519 #[test]
520 fn definition_handles_no_subagents() {
521 let tool = top_level_tool(vec![], DEFAULT_MAX_DEPTH);
522 let def = tool.definition();
523 assert!(def.description.contains("(none declared)"));
524 }
525
526 #[test]
527 fn definition_schema_requires_agent_and_prompt() {
528 let tool = top_level_tool(vec![dummy_profile("a")], DEFAULT_MAX_DEPTH);
529 let def = tool.definition();
530 let required = def
531 .parameters
532 .fields
533 .get("required")
534 .and_then(|v| v.as_array())
535 .expect("required array");
536 let names: Vec<&str> = required.iter().filter_map(Value::as_str).collect();
537 assert!(names.contains(&"agent"));
538 assert!(names.contains(&"prompt"));
539 }
540
541 #[tokio::test]
542 async fn unknown_agent_returns_error() {
543 let tool = top_level_tool(vec![dummy_profile("reviewer")], DEFAULT_MAX_DEPTH);
544 let ctx = InvokeContext {
545 tool_call_id: "c1".into(),
546 cancel: CancellationToken::new(),
547 };
548 let err = tool
549 .execute(ctx, serde_json::json!({ "agent": "ghost", "prompt": "hi" }))
550 .await
551 .expect_err("unknown agent should error");
552 let msg = err.to_string();
553 assert!(msg.contains("not declared"), "msg: {msg}");
554 assert!(msg.contains("ghost"));
555 assert!(msg.contains("reviewer"));
557 }
558
559 #[tokio::test]
560 async fn depth_exceeded_at_max_zero() {
561 let tool = top_level_tool(vec![dummy_profile("a")], 0);
563 let ctx = InvokeContext {
564 tool_call_id: "c2".into(),
565 cancel: CancellationToken::new(),
566 };
567 let err = tool
568 .execute(ctx, serde_json::json!({ "agent": "a", "prompt": "hi" }))
569 .await
570 .expect_err("depth should exceed");
571 let msg = err.to_string();
572 assert!(msg.contains("depth exceeded"), "msg: {msg}");
573 assert!(msg.contains("max_depth 0"));
574 }
575
576 #[tokio::test]
577 async fn budget_exhaustion_blocks_delegation() {
578 let tool = TaskTool::new(
581 test_provider(),
582 Model {
583 id: "test/model".into(),
584 },
585 RunConfig::default(),
586 vec![dummy_profile("worker")],
587 SubagentOptions {
588 max_depth: DEFAULT_MAX_DEPTH,
589 max_delegations: 1,
590 },
591 CancellationToken::new(),
592 None,
593 );
594 let ctx1 = InvokeContext {
595 tool_call_id: "b1".into(),
596 cancel: CancellationToken::new(),
597 };
598 let r1 = tool
600 .execute(
601 ctx1,
602 serde_json::json!({ "agent": "worker", "prompt": "go" }),
603 )
604 .await
605 .expect("first delegation succeeds");
606 assert_eq!(r1.content.len(), 1);
607
608 let ctx2 = InvokeContext {
610 tool_call_id: "b2".into(),
611 cancel: CancellationToken::new(),
612 };
613 let err = tool
614 .execute(
615 ctx2,
616 serde_json::json!({ "agent": "worker", "prompt": "again" }),
617 )
618 .await
619 .expect_err("budget should be exhausted");
620 let msg = err.to_string();
621 assert!(msg.contains("budget exhausted"), "msg: {msg}");
622 }
623
624 #[tokio::test]
625 async fn delegate_runs_child_and_returns_text() {
626 let tool = top_level_tool(vec![dummy_profile("worker")], DEFAULT_MAX_DEPTH);
627 let ctx = InvokeContext {
628 tool_call_id: "c3".into(),
629 cancel: CancellationToken::new(),
630 };
631 let result = tool
632 .execute(
633 ctx,
634 serde_json::json!({ "agent": "worker", "prompt": "do it" }),
635 )
636 .await
637 .expect("delegation should succeed");
638 assert_eq!(result.content.len(), 1);
639 let text = result.content[0]
640 .get("text")
641 .and_then(Value::as_str)
642 .expect("text");
643 assert_eq!(text, "child done");
644 }
645
646 #[tokio::test]
647 async fn cancellation_aborts_before_child_spawn() {
648 let tool = top_level_tool(vec![dummy_profile("a")], DEFAULT_MAX_DEPTH);
649 let cancel = CancellationToken::new();
650 let ctx = InvokeContext {
651 tool_call_id: "c4".into(),
652 cancel: cancel.clone(),
653 };
654 cancel.cancel();
655 let err = tool
656 .execute(ctx, serde_json::json!({ "agent": "a", "prompt": "hi" }))
657 .await
658 .expect_err("should be cancelled");
659 assert!(matches!(err, CoreError::Cancelled(_)), "err: {err}");
660 }
661
662 struct ScriptedProvider {
666 responses: std::sync::Mutex<std::collections::VecDeque<ModelResponse>>,
667 }
668
669 impl ScriptedProvider {
670 fn new(responses: Vec<Vec<AgentMessage>>) -> Self {
671 let responses = responses
672 .into_iter()
673 .map(|msgs| ModelResponse { messages: msgs })
674 .collect();
675 Self {
676 responses: std::sync::Mutex::new(responses),
677 }
678 }
679 }
680
681 #[async_trait]
682 impl ModelProvider for ScriptedProvider {
683 async fn invoke(&self, _request: ModelRequest) -> CoreResult<ModelResponse> {
684 let next = self
685 .responses
686 .lock()
687 .unwrap()
688 .pop_front()
689 .unwrap_or(ModelResponse { messages: vec![] });
690 Ok(next)
691 }
692 }
693
694 fn assistant_text(t: &str) -> AgentMessage {
695 AgentMessage {
696 role: Role::Assistant,
697 content: vec![ContentBlock::Text { text: t.into() }],
698 }
699 }
700
701 fn parent_task_call(agent: &str, prompt: &str) -> AgentMessage {
703 AgentMessage {
704 role: Role::Assistant,
705 content: vec![ContentBlock::ToolUse {
706 id: "call_1".into(),
707 call: crate::tool::ToolCall {
708 name: "task".into(),
709 input: serde_json::json!({ "agent": agent, "prompt": prompt }),
710 },
711 }],
712 }
713 }
714
715 #[tokio::test]
716 async fn integration_parent_delegates_and_child_answers() {
717 let provider: Arc<dyn ModelProvider> = Arc::new(ScriptedProvider::new(vec![
720 vec![parent_task_call("worker", "do the work")],
722 vec![assistant_text("child done")],
724 vec![assistant_text("got: child done")],
727 ]));
728 let cancel = CancellationToken::new();
729 let task = Arc::new(TaskTool::new(
730 Arc::clone(&provider),
731 Model {
732 id: "test/m".into(),
733 },
734 RunConfig::default(),
735 vec![dummy_profile("worker")],
736 SubagentOptions::default(),
737 cancel.clone(),
738 None,
739 ));
740 let tools: Vec<Arc<dyn Tool>> = vec![task];
741 let mut messages = vec![
742 AgentMessage {
743 role: Role::System,
744 content: vec![ContentBlock::Text {
745 text: "be brief".into(),
746 }],
747 },
748 AgentMessage {
749 role: Role::User,
750 content: vec![ContentBlock::Text {
751 text: "delegate the work".into(),
752 }],
753 },
754 ];
755 let outcome = run_agent(
756 provider.as_ref(),
757 &tools,
758 &mut messages,
759 &Model {
760 id: "test/m".into(),
761 },
762 &RunConfig::default(),
763 &cancel,
764 &crate::event::RunHooks::default(),
765 )
766 .await
767 .expect("parent run");
768 assert_eq!(outcome.turns, 2);
769 assert_eq!(outcome.final_text, "got: child done");
770 }
771
772 #[tokio::test]
773 async fn integration_nested_delegation_stops_at_max_depth() {
774 let grandchild = dummy_profile("grandchild");
777 let child = SubagentProfile::new("child", "you delegate").with_subagent(grandchild);
778
779 let provider: Arc<dyn ModelProvider> = Arc::new(ScriptedProvider::new(vec![
782 vec![parent_task_call("child", "sub-delegate")],
784 vec![AgentMessage {
786 role: Role::Assistant,
787 content: vec![ContentBlock::ToolUse {
788 id: "cchild".into(),
789 call: crate::tool::ToolCall {
790 name: "task".into(),
791 input: serde_json::json!({
792 "agent": "grandchild",
793 "prompt": "too deep"
794 }),
795 },
796 }],
797 }],
798 vec![assistant_text("grandchild was unreachable")],
801 vec![assistant_text("done")],
803 ]));
804 let cancel = CancellationToken::new();
805 let task = Arc::new(TaskTool::new(
806 Arc::clone(&provider),
807 Model {
808 id: "test/m".into(),
809 },
810 RunConfig::default(),
811 vec![child],
812 SubagentOptions {
814 max_depth: 1,
815 max_delegations: DEFAULT_MAX_DELEGATIONS,
816 },
817 cancel.clone(),
818 None,
819 ));
820 let tools: Vec<Arc<dyn Tool>> = vec![task];
821 let mut messages = vec![AgentMessage {
822 role: Role::User,
823 content: vec![ContentBlock::Text { text: "go".into() }],
824 }];
825 let outcome = run_agent(
826 provider.as_ref(),
827 &tools,
828 &mut messages,
829 &Model {
830 id: "test/m".into(),
831 },
832 &RunConfig::default(),
833 &cancel,
834 &crate::event::RunHooks::default(),
835 )
836 .await
837 .expect("parent run");
838 assert_eq!(outcome.turns, 2);
842 }
843}