1use super::super::{Tool, ToolContext};
10use crate::runtime::subagent::SubagentStatus;
11use crate::{Result, RuntimeError};
12use serde_json::{json, Value};
13
14pub struct SubagentCollectTool;
15
16#[async_trait::async_trait]
17impl Tool for SubagentCollectTool {
18 fn origin(&self) -> crate::tools::ToolOrigin {
19 crate::tools::ToolOrigin::Builtin
20 }
21
22 fn name(&self) -> &str {
23 "subagent_collect"
24 }
25
26 fn description(&self) -> &str {
27 "Retrieve a reactive subagent's result. Non-blocking — returns immediately \
28 (status + partial output if still running). Standard flow: dispatch with \
29 subagent_start, carry on (or end your turn) — the subagent's completion event \
30 wakes you with a terminal status (completed/failed/timed_out/cancelled); then \
31 call this ONCE with reconciled=true — that single call both collects the result and attests \
32 reconciliation, clearing the completion gate. Calling without reconciled=true \
33 is an inspection-only read: it leaves the completion gate blocked and forces \
34 a second call. reconciled=true is idempotent and safe on repeats."
35 }
36
37 fn parameters(&self) -> Value {
38 json!({
39 "type": "object",
40 "properties": {
41 "handle_id": {
42 "type": "string",
43 "description": "Handle ID returned by subagent_start (e.g. \"sa_3\")."
44 },
45 "reconciled": {
46 "type": "boolean",
47 "default": false,
48 "description": "Set true on your terminal collect (after subagent_status reports done) to attest the result is reconciled with foreground work and unblock completion — one call collects AND reconciles. Leave false only for a deliberate inspection-only read; a later collect with reconciled=true still reconciles (idempotent)."
49 }
50 },
51 "required": ["handle_id"]
52 })
53 }
54
55 async fn execute(&self, params: Value, ctx: ToolContext) -> Result<String> {
56 let handle_id = params["handle_id"]
57 .as_str()
58 .ok_or_else(|| RuntimeError::Tool("Missing 'handle_id' parameter".to_string()))?
59 .to_string();
60
61 let registry = ctx.capabilities.subagent_registry.as_ref().ok_or_else(|| {
62 RuntimeError::Tool("SubagentRegistry not available on this ToolContext".to_string())
63 })?;
64
65 let mut reg = registry.lock().unwrap();
66 let archive_path = reg.archive_path(&handle_id);
67 let Some(handle) = reg.get_mut(&handle_id) else {
68 drop(reg);
69 let archived = archive_path
71 .as_ref()
72 .and_then(|path| super::archive::read_archive(path).ok());
73 if let Some(orchestration) = &ctx.capabilities.orchestration {
74 if orchestration.is_unreconciled(&handle_id) {
75 orchestration
76 .terminal_and_collect(
77 &handle_id,
78 agent_core::orchestration::WorkerTerminal::Failed,
79 )
80 .map_err(RuntimeError::Tool)?;
81 if params["reconciled"].as_bool().unwrap_or(false) {
82 orchestration
83 .reconcile(&handle_id)
84 .map_err(RuntimeError::Tool)?;
85 }
86 if archived.is_none() {
87 return Ok(json!({
88 "handle_id": handle_id,
89 "status": "expired",
90 "note": "Subagent output expired; orchestration lifecycle remains recoverable.",
91 "collected": false
92 })
93 .to_string());
94 }
95 }
96 }
97 if let (Some(archive), Some(path)) = (archived, archive_path) {
98 let mut body = json!({
99 "handle_id": handle_id,
100 "status": archive.status,
101 "error": archive.error,
102 "output": archive.output,
103 "model": archive.meta.model,
104 "collected": true,
105 "archive_path": path,
106 "note": "Handle was reaped; result read from its resume archive. \
107 subagent_resume can still continue this worker."
108 });
109 if let Some(context) =
110 crate::runtime::subagent::archived_context_summary(&archive.history)
111 {
112 body["context"] = context;
113 }
114 return Ok(body.to_string());
115 }
116 return Err(RuntimeError::Tool(format!(
117 "No subagent found with handle_id '{}'. Finished handles are retained for {} minutes after completion and then garbage-collected.",
118 handle_id,
119 crate::runtime::subagent::FINISHED_HANDLE_TTL.as_secs() / 60
120 )));
121 };
122
123 super::archive::retry_failed_archive(&handle.state, &handle_id);
124 let status = handle.status();
125 let output: String = handle.partial_output();
126 let elapsed = handle.elapsed_secs();
127 let model = handle.model.clone();
128 let terminal = handle.terminal_diagnostic();
129 let authorization = handle.authorization.clone();
130 let archive_error = handle.archive_error();
131 let context = handle
132 .state
133 .read()
134 .unwrap_or_else(|p| p.into_inner())
135 .context_summary();
136
137 let already_collected = handle.is_collected();
139 if status != SubagentStatus::Running {
140 handle.mark_collected();
141 }
142 drop(reg);
143
144 if status == SubagentStatus::Running {
145 let char_count = output.chars().count();
147 let output_so_far: String = if char_count > 500 {
148 output.chars().skip(char_count - 500).collect()
149 } else {
150 output
151 };
152 return Ok(json!({
153 "handle_id": handle_id,
154 "status": "running",
155 "elapsed_secs": (elapsed * 10.0).round() / 10.0,
156 "output_so_far": output_so_far
157 })
158 .to_string());
159 }
160
161 if let Some(orchestration) = &ctx.capabilities.orchestration {
170 if !already_collected {
171 let terminal = match status {
172 SubagentStatus::Completed => {
173 agent_core::orchestration::WorkerTerminal::Completed
174 }
175 SubagentStatus::TimedOut => agent_core::orchestration::WorkerTerminal::TimedOut,
176 _ => agent_core::orchestration::WorkerTerminal::Failed,
177 };
178 orchestration
179 .terminal_and_collect(&handle_id, terminal)
180 .map_err(RuntimeError::Tool)?;
181 }
182 if params["reconciled"].as_bool().unwrap_or(false) {
183 orchestration
184 .reconcile(&handle_id)
185 .map_err(RuntimeError::Tool)?;
186 }
187 }
188
189 let reconciled_now = params["reconciled"].as_bool().unwrap_or(false);
193 let mut body = json!({
194 "archiveError": archive_error,
195 "handle_id": handle_id,
196 "status": status.as_str(),
197 "output": output,
198 "model": model,
199 "terminal_cause": terminal,
200 "authorization": authorization,
201 "collected": already_collected,
202 });
203 if let Some(context) = context {
205 body["context"] = context;
206 }
207 if !reconciled_now && ctx.capabilities.orchestration.is_some() {
208 body["note"] = json!(
213 "UNRECONCILED read — the completion gate is still blocked. \
214 Call subagent_collect again with reconciled=true after \
215 integrating this result."
216 );
217 }
218 if let Some(reason) = status.failure_reason() {
219 body["error"] = json!(reason);
220 }
221 if let Some(path) = archive_path {
222 body["archive_path"] = json!(path);
224 }
225 Ok(body.to_string())
226 }
227}
228
229#[cfg(test)]
230mod tests {
231 use super::*;
232 use crate::runtime::subagent::{
233 SubagentHandle, SubagentRegistry, SubagentState, SubagentStatus,
234 };
235 use crate::tools::test_helpers::create_tool_context;
236 use crate::tools::Tool;
237 use serde_json::json;
238 use std::sync::{Arc, Mutex, RwLock};
239 use tokio::sync::{mpsc, oneshot};
240
241 fn make_finished_handle(id: &str, output: &str) -> SubagentHandle {
242 let state = Arc::new(RwLock::new(SubagentState::new()));
243 {
244 let mut s = state.write().unwrap();
245 s.status = SubagentStatus::Completed;
246 s.partial_text = output.to_string();
247 s.finished_at = Some(std::time::Instant::now());
248 }
249 let (steer_tx, _steer_rx) = mpsc::unbounded_channel();
250 let (shutdown_tx, _shutdown_rx) = oneshot::channel();
251 let (_result_tx, result_rx) = oneshot::channel();
252 SubagentHandle::new(
253 id.to_string(),
254 0,
255 "test-agent".to_string(),
256 "test task".to_string(),
257 "claude-sonnet-4-6".to_string(),
258 "system prompt".to_string(),
259 300,
260 state,
261 Some(steer_tx),
262 Some(shutdown_tx),
263 Some(result_rx),
264 )
265 }
266
267 #[tokio::test]
268 async fn managed_terminal_waits_for_finalization_before_collection() {
269 let registry = Arc::new(Mutex::new(SubagentRegistry::new()));
270 let h = make_finished_handle("sa_pending", "still draining");
271 let state = Arc::clone(&h.state);
272 {
273 let mut s = state.write().unwrap();
274 s.finished_at = None;
275 s.archive_meta = Some(crate::runtime::subagent::SubagentArchiveMeta {
276 agent_name: "worker".into(),
277 model: "model".into(),
278 system_prompt: "system".into(),
279 timeout_secs: 1,
280 cwd: None,
281 thinking: None,
282 });
283 }
284 registry.lock().unwrap().register(h);
285 let response = SubagentCollectTool
286 .execute(
287 json!({"handle_id":"sa_pending", "reconciled":true}),
288 make_ctx_with_registry(Arc::clone(®istry)),
289 )
290 .await
291 .unwrap();
292 assert_eq!(
293 serde_json::from_str::<Value>(&response).unwrap()["status"],
294 "running"
295 );
296 let mut reg = registry.lock().unwrap();
297 assert!(!reg.get("sa_pending").unwrap().is_collected());
298 assert_eq!(reg.display_rows()[0].status, SubagentStatus::Running);
299 reg.release_finished_resources("sa_pending");
300 reg.cleanup_finished_with_ttl(std::time::Duration::ZERO);
301 assert!(reg.get("sa_pending").is_some());
302 assert!(state.read().unwrap().finished_at.is_none());
303 }
304
305 fn make_ctx_with_registry(registry: Arc<Mutex<SubagentRegistry>>) -> crate::tools::ToolContext {
306 let mut ctx = create_tool_context();
307 ctx.capabilities.subagent_registry = Some(registry);
308 ctx.capabilities.orchestration = None;
314 ctx
315 }
316
317 #[tokio::test]
318 async fn collect_after_reap_reads_the_archive() {
319 let dir = tempfile::tempdir().unwrap();
320 let meta = crate::runtime::subagent::SubagentArchiveMeta {
321 agent_name: "test-agent".into(),
322 model: "anthropic/claude-opus-5-5".into(),
323 system_prompt: "s".into(),
324 timeout_secs: 60,
325 cwd: None,
326 thinking: None,
327 };
328 let path = super::super::archive::write_archive(
329 dir.path(),
330 "sa_gone",
331 &meta,
332 "timed_out",
333 Some("worker timed out"),
334 "[TIMED OUT after 60s] latest progress",
335 &[Arc::new(json!({"role": "user", "content": "task"}))],
336 )
337 .unwrap();
338 let registry = Arc::new(Mutex::new(SubagentRegistry::new()));
339 {
340 let mut reg = registry.lock().unwrap();
341 let state = Arc::new(RwLock::new(SubagentState::new()));
342 {
343 let mut st = state.write().unwrap();
344 st.status = SubagentStatus::TimedOut;
345 st.finished_at = Some(std::time::Instant::now());
346 st.archive_path = Some(path.clone());
347 }
348 let (steer_tx, _steer_rx) = mpsc::unbounded_channel();
349 let (shutdown_tx, _shutdown_rx) = oneshot::channel();
350 let (_result_tx, result_rx) = oneshot::channel();
351 let mut h = SubagentHandle::new(
352 "sa_gone".into(),
353 0,
354 "test-agent".into(),
355 "task".into(),
356 "anthropic/claude-opus-5-5".into(),
357 "s".into(),
358 60,
359 state,
360 Some(steer_tx),
361 Some(shutdown_tx),
362 Some(result_rx),
363 );
364 h.mark_collected();
365 reg.register(h);
366 reg.cleanup_finished_with_ttl(std::time::Duration::ZERO);
367 assert!(reg.get("sa_gone").is_none());
368 }
369 let body: serde_json::Value = serde_json::from_str(
370 &SubagentCollectTool
371 .execute(
372 json!({"handle_id": "sa_gone"}),
373 make_ctx_with_registry(registry),
374 )
375 .await
376 .unwrap(),
377 )
378 .unwrap();
379 assert_eq!(body["status"], "timed_out");
380 assert_eq!(body["error"], "worker timed out");
381 assert_eq!(body["output"], "[TIMED OUT after 60s] latest progress");
382 assert_eq!(body["archive_path"], json!(path));
383 }
384
385 #[tokio::test]
387 async fn collect_marks_collected_and_is_idempotent() {
388 let tool = SubagentCollectTool;
389
390 let registry = Arc::new(Mutex::new(SubagentRegistry::new()));
391 let handle = make_finished_handle("sa_42", "result text");
392 registry.lock().unwrap().register(handle);
393
394 let result1 = tool
396 .execute(
397 json!({"handle_id": "sa_42"}),
398 make_ctx_with_registry(registry.clone()),
399 )
400 .await
401 .unwrap();
402 let body1: serde_json::Value = serde_json::from_str(&result1).unwrap();
403 assert_eq!(body1["status"], "completed");
404 assert_eq!(body1["output"], "result text");
405 assert_eq!(
406 body1["collected"], false,
407 "first collect should report collected=false"
408 );
409
410 let result2 = tool
412 .execute(
413 json!({"handle_id": "sa_42"}),
414 make_ctx_with_registry(registry.clone()),
415 )
416 .await
417 .unwrap();
418 let body2: serde_json::Value = serde_json::from_str(&result2).unwrap();
419 assert_eq!(body2["status"], "completed");
420 assert_eq!(
421 body2["collected"], true,
422 "second collect should report collected=true"
423 );
424
425 let reg = registry.lock().unwrap();
427 let h = reg.get("sa_42").unwrap();
428 assert!(
429 h.is_collected(),
430 "handle must be marked collected in registry"
431 );
432 }
433
434 fn make_orch_ctx(
437 handle_id: &str,
438 output: &str,
439 ) -> (
440 Arc<Mutex<SubagentRegistry>>,
441 Arc<crate::orchestration::OrchestrationRuntime>,
442 crate::tools::ToolContext,
443 ) {
444 use crate::orchestration::OrchestrationRuntime;
445 use agent_core::orchestration::CompletionGate;
446
447 let foreground = agent_core::prompt::QualifiedModelId::parse("anthropic/claude-sonnet-4-6")
448 .expect("test foreground is qualified");
449 let orch =
450 Arc::new(OrchestrationRuntime::baseline(foreground, 8, 64).expect("baseline runtime"));
451 orch.authorize(handle_id, "anthropic/claude-sonnet-4-6")
452 .expect("authorize worker");
453 orch.terminal_and_collect(
456 handle_id,
457 agent_core::orchestration::WorkerTerminal::Completed,
458 )
459 .expect("mark terminal + collected");
460 assert!(
461 matches!(orch.completion_gate(), CompletionGate::Blocked { .. }),
462 "terminal worker must block completion until reconciled"
463 );
464
465 let registry = Arc::new(Mutex::new(SubagentRegistry::new()));
466 registry
467 .lock()
468 .unwrap()
469 .register(make_finished_handle(handle_id, output));
470
471 let mut ctx = create_tool_context();
472 ctx.capabilities.subagent_registry = Some(registry.clone());
473 ctx.capabilities.orchestration = Some(orch.clone());
474 (registry, orch, ctx)
475 }
476
477 #[tokio::test]
480 async fn deferred_reconcile_on_repeat_collect_unblocks_completion() {
481 use agent_core::orchestration::CompletionGate;
482
483 let tool = SubagentCollectTool;
484 let (registry, orch, ctx) = make_orch_ctx("sa_deferred", "result text");
485
486 let result1 = tool
487 .execute(
488 json!({"handle_id": "sa_deferred", "reconciled": false}),
489 ctx,
490 )
491 .await
492 .unwrap();
493 let body1: serde_json::Value = serde_json::from_str(&result1).unwrap();
494 assert_eq!(body1["collected"], false);
495 match orch.completion_gate() {
496 CompletionGate::Blocked { workers } => {
497 assert_eq!(workers, vec!["sa_deferred".to_string()]);
498 assert!(
499 workers.iter().all(|id| id.starts_with("sa_")),
500 "blocked IDs must be subagent_collect handles, got {workers:?}"
501 );
502 assert!(
503 workers.iter().all(|id| !id.starts_with("worker-")),
504 "blocked IDs must never leak policy WorkerHandle, got {workers:?}"
505 );
506 }
507 other => panic!(
508 "first collect without reconcile must leave completion blocked, got {other:?}"
509 ),
510 }
511
512 let mut ctx2 = create_tool_context();
515 ctx2.capabilities.subagent_registry = Some(registry.clone());
516 ctx2.capabilities.orchestration = Some(orch.clone());
517
518 let result2 = tool
519 .execute(
520 json!({"handle_id": "sa_deferred", "reconciled": true}),
521 ctx2,
522 )
523 .await
524 .unwrap();
525 let body2: serde_json::Value = serde_json::from_str(&result2).unwrap();
526 assert_eq!(body2["collected"], true);
527 assert_eq!(
528 orch.completion_gate(),
529 CompletionGate::Allowed,
530 "repeat collect with reconciled=true must unblock completion"
531 );
532 }
533
534 #[tokio::test]
536 async fn first_collect_with_reconcile_allows_completion() {
537 use agent_core::orchestration::CompletionGate;
538
539 let tool = SubagentCollectTool;
540 let (_registry, orch, ctx) = make_orch_ctx("sa_first_true", "result text");
541
542 let result = tool
543 .execute(
544 json!({"handle_id": "sa_first_true", "reconciled": true}),
545 ctx,
546 )
547 .await
548 .unwrap();
549 let body: serde_json::Value = serde_json::from_str(&result).unwrap();
550 assert_eq!(body["collected"], false);
551 assert_eq!(
552 orch.completion_gate(),
553 CompletionGate::Allowed,
554 "first collect with reconciled=true must allow completion"
555 );
556 }
557
558 #[tokio::test]
563 async fn unreconciled_terminal_collect_carries_reconcile_nudge() {
564 let tool = SubagentCollectTool;
565 let (_registry, _orch, ctx) = make_orch_ctx("sa_nudge", "result text");
566
567 let result = tool
568 .execute(json!({"handle_id": "sa_nudge", "reconciled": false}), ctx)
569 .await
570 .unwrap();
571 let body: serde_json::Value = serde_json::from_str(&result).unwrap();
572 let note = body["note"].as_str().expect("unreconciled read must nudge");
573 assert!(note.contains("reconciled=true"), "{note}");
574 }
575
576 #[tokio::test]
577 async fn reconciled_collect_carries_no_nudge() {
578 let tool = SubagentCollectTool;
579 let (_registry, _orch, ctx) = make_orch_ctx("sa_no_nudge", "result text");
580
581 let result = tool
582 .execute(json!({"handle_id": "sa_no_nudge", "reconciled": true}), ctx)
583 .await
584 .unwrap();
585 let body: serde_json::Value = serde_json::from_str(&result).unwrap();
586 assert!(
587 body.get("note").is_none(),
588 "single-call collect+reconcile must not carry a nudge: {body}"
589 );
590 }
591
592 #[tokio::test]
593 async fn forced_expiry_remains_collectible_and_clears_gate() {
594 use agent_core::orchestration::CompletionGate;
595
596 let tool = SubagentCollectTool;
597 let (registry, orch, ctx) = make_orch_ctx("sa_expired", &"result".repeat(10_000));
598 crate::runtime::subagent::reap_finished_with_ttl(
599 ®istry,
600 Some(orch.as_ref()),
601 std::time::Duration::ZERO,
602 );
603 assert!(registry
604 .lock()
605 .unwrap()
606 .get("sa_expired")
607 .unwrap()
608 .is_tombstone());
609
610 let result = tool
611 .execute(json!({"handle_id": "sa_expired", "reconciled": true}), ctx)
612 .await
613 .unwrap();
614 assert_ne!(
615 serde_json::from_str::<Value>(&result).unwrap()["status"],
616 "expired"
617 );
618 assert_eq!(orch.completion_gate(), CompletionGate::Allowed);
619 }
620
621 #[tokio::test]
622 async fn collected_handle_reaped_then_degraded_reconcile_clears_gate() {
623 use agent_core::orchestration::CompletionGate;
624
625 let tool = SubagentCollectTool;
626 let (registry, orch, ctx) = make_orch_ctx("sa_collected_expired", "result text");
627 tool.execute(
628 json!({"handle_id": "sa_collected_expired", "reconciled": false}),
629 ctx,
630 )
631 .await
632 .unwrap();
633 assert!(matches!(
634 orch.completion_gate(),
635 CompletionGate::Blocked { .. }
636 ));
637
638 crate::runtime::subagent::reap_finished_with_ttl(
639 ®istry,
640 Some(orch.as_ref()),
641 std::time::Duration::ZERO,
642 );
643 assert!(registry
644 .lock()
645 .unwrap()
646 .get("sa_collected_expired")
647 .is_none());
648
649 let mut ctx2 = create_tool_context();
650 ctx2.capabilities.subagent_registry = Some(registry);
651 ctx2.capabilities.orchestration = Some(orch.clone());
652 let response = tool
653 .execute(
654 json!({"handle_id": "sa_collected_expired", "reconciled": true}),
655 ctx2,
656 )
657 .await
658 .unwrap();
659 assert_eq!(
660 serde_json::from_str::<Value>(&response).unwrap()["status"],
661 "expired"
662 );
663 assert_eq!(orch.completion_gate(), CompletionGate::Allowed);
664 }
665
666 #[tokio::test]
667 async fn missing_mapped_handle_has_degraded_collect_not_404() {
668 use agent_core::orchestration::CompletionGate;
669
670 let tool = SubagentCollectTool;
671 let (registry, orch, ctx) = make_orch_ctx("sa_missing", "lost");
672 registry.lock().unwrap().remove("sa_missing");
673 let result = tool
674 .execute(json!({"handle_id": "sa_missing", "reconciled": true}), ctx)
675 .await
676 .unwrap();
677 let body: Value = serde_json::from_str(&result).unwrap();
678 assert_eq!(body["status"], "expired");
679 assert_eq!(orch.completion_gate(), CompletionGate::Allowed);
680 }
681
682 #[tokio::test]
684 async fn repeated_reconcile_true_is_idempotent() {
685 use agent_core::orchestration::CompletionGate;
686
687 let tool = SubagentCollectTool;
688 let (registry, orch, ctx) = make_orch_ctx("sa_repeat_true", "result text");
689
690 let result1 = tool
691 .execute(
692 json!({"handle_id": "sa_repeat_true", "reconciled": true}),
693 ctx,
694 )
695 .await
696 .unwrap();
697 let body1: serde_json::Value = serde_json::from_str(&result1).unwrap();
698 assert_eq!(body1["collected"], false);
699 assert_eq!(orch.completion_gate(), CompletionGate::Allowed);
700
701 let mut ctx2 = create_tool_context();
702 ctx2.capabilities.subagent_registry = Some(registry.clone());
703 ctx2.capabilities.orchestration = Some(orch.clone());
704
705 let result2 = tool
706 .execute(
707 json!({"handle_id": "sa_repeat_true", "reconciled": true}),
708 ctx2,
709 )
710 .await
711 .unwrap();
712 let body2: serde_json::Value = serde_json::from_str(&result2).unwrap();
713 assert_eq!(body2["collected"], true);
714 assert_eq!(
715 orch.completion_gate(),
716 CompletionGate::Allowed,
717 "repeat reconciled collect must stay Allowed (idempotent)"
718 );
719 let reg = registry.lock().unwrap();
721 assert!(reg.get("sa_repeat_true").unwrap().is_collected());
722 }
723}