agent_engine/tools/subagent/
resume.rs1use serde_json::{json, Value};
8use std::sync::atomic::Ordering;
9use std::sync::{Arc, RwLock};
10use std::time::Duration;
11use tokio::sync::{mpsc, oneshot};
12
13use crate::{Result, RuntimeError, LlmEvent, SessionEvent, AgentEvent};
14use super::super::{Tool, ToolContext, NEXT_SUBAGENT_ID};
15use crate::runtime::subagent::{SubagentHandle, SubagentResult, SubagentStatus, SubagentState};
16
17pub struct SubagentResumeTool;
18
19#[async_trait::async_trait]
20impl Tool for SubagentResumeTool {
21 fn name(&self) -> &str { "subagent_resume" }
22
23 fn description(&self) -> &str {
24 "Resume a finished or timed-out reactive subagent with new instructions. \
25 The previous subagent's conversation state is prepended as context so the \
26 new run has full history. Returns a new handle_id — the original handle \
27 remains readable for comparison. Only works on subagents in \
28 finished/timed_out/failed state."
29 }
30
31 fn parameters(&self) -> Value {
32 json!({
33 "type": "object",
34 "properties": {
35 "handle_id": {
36 "type": "string",
37 "description": "Handle ID of the completed subagent to resume (e.g. \"sa_3\")."
38 },
39 "instructions": {
40 "type": "string",
41 "description": "New task or context to prepend to the resumed subagent. \
42 Injected before the prior conversation history."
43 }
44 },
45 "required": ["handle_id", "instructions"]
46 })
47 }
48
49 async fn execute(&self, params: Value, ctx: ToolContext) -> Result<String> {
50 let prior_handle_id = params["handle_id"].as_str()
51 .ok_or_else(|| RuntimeError::Tool("Missing 'handle_id' parameter".to_string()))?
52 .to_string();
53
54 let instructions = params["instructions"].as_str()
55 .ok_or_else(|| RuntimeError::Tool("Missing 'instructions' parameter".to_string()))?
56 .to_string();
57
58 let registry = ctx.capabilities.subagent_registry.as_ref()
59 .ok_or_else(|| RuntimeError::Tool(
60 "SubagentRegistry not available on this ToolContext".to_string()
61 ))?;
62
63 let (agent_name, model, prior_context, prior_system_prompt, prior_timeout) = {
65 let reg = registry.lock().unwrap();
66 let handle = reg.get(&prior_handle_id)
67 .ok_or_else(|| RuntimeError::Tool(
68 format!("No subagent found with handle_id '{}'", prior_handle_id)
69 ))?;
70
71 if handle.status() == SubagentStatus::Running {
72 return Err(RuntimeError::Tool(format!(
73 "Subagent '{}' is still running. Call subagent_collect first, \
74 or wait until it finishes.",
75 prior_handle_id
76 )));
77 }
78
79 let prior = {
80 let state = handle.conversation_state();
81 if state.is_empty() {
82 handle.partial_output()
83 } else {
84 serde_json::to_string(&state).unwrap_or_else(|_| handle.partial_output())
85 }
86 };
87
88 (handle.agent_name.clone(), handle.model.clone(), prior, handle.system_prompt.clone(), handle.timeout_secs)
89 };
90
91 let resumed_task = format!(
93 "{instructions}\n\n\
94 ---\n\
95 [Prior conversation context from handle {prior_handle_id}]\n\
96 {prior_context}"
97 );
98
99 let system_prompt = prior_system_prompt;
101 let timeout_secs = prior_timeout;
102 let label = agent_name.clone();
103 let task_preview: String = resumed_task.chars().take(80).collect();
104 let task_full = resumed_task.clone();
105 let subagent_id = NEXT_SUBAGENT_ID.fetch_add(1, Ordering::Relaxed);
106 let handle_id = format!("sa_{}", subagent_id);
107
108 tracing::info!(
109 "subagent_resume: dispatching '{}' (id={}, resumed_from={}) model={}",
110 label, handle_id, prior_handle_id, model
111 );
112
113 let state = Arc::new(RwLock::new(SubagentState::new()));
114
115 let (steer_tx, steer_rx) = mpsc::unbounded_channel::<String>();
116 let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
117 let (result_tx, result_rx) = oneshot::channel::<SubagentResult>();
118
119 if let Some(ref tx) = ctx.channels.tx_events {
120 let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentStart {
121 subagent_id,
122 agent_name: label.clone(),
123 task_preview: task_preview.clone(),
124 }));
125 }
126
127 let state_t = Arc::clone(&state);
128 let task_full_a = task_full.clone();
129 let label_inner = label.clone();
130 let model_inner = model.clone();
131 let tx_events_inner = ctx.channels.tx_events.clone();
132 let start_time = std::time::Instant::now();
133
134 let system_prompt_for_handle = system_prompt.clone();
135 let thread_handle = std::thread::spawn(move || {
136 let panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
137 let rt = match tokio::runtime::Builder::new_current_thread()
138 .enable_all()
139 .build()
140 {
141 Ok(rt) => rt,
142 Err(e) => {
143 state_t.write().unwrap().status =
144 SubagentStatus::Failed(format!("tokio runtime: {}", e));
145 return;
146 }
147 };
148
149 let state_a = Arc::clone(&state_t);
150 let label_a = label_inner.clone();
151 let model_a = model_inner.clone();
152 let tx_events_a = tx_events_inner.clone();
153 let task_for_timeout = task_full_a.clone();
154 let task_for_complete = task_full_a.clone();
155 let task_for_stream = task_full_a;
156
157 let outcome: std::result::Result<SubagentResult, String> = rt.block_on(async move {
158 use futures::StreamExt;
159
160 let mut runtime = match crate::Runtime::new().await {
161 Ok(r) => r,
162 Err(e) => return Err(format!("Failed to create subagent runtime: {}", e)),
163 };
164
165 super::apply_subagent_runtime_policy(&mut runtime, &crate::config::load_config());
170 runtime.set_system_prompt(system_prompt);
171 runtime.set_model(model_a.clone());
172 runtime.set_tools(crate::ToolRegistry::without_subagent());
173
174 let cancel = crate::CancellationToken::new();
175 let cancel_inner = cancel.clone();
176 tokio::spawn(async move {
177 let _ = shutdown_rx.await;
178 cancel_inner.cancel();
179 });
180
181 let mut stream = runtime.run_stream_with_messages(
182 vec![serde_json::json!({"role": "user", "content": task_for_stream})],
183 cancel,
184 Some(steer_rx),
185 None,
186 false,
187 ).await;
188
189 let mut tool_count = 0u32;
190 let mut total_input_tokens = 0u64;
191 let mut total_output_tokens = 0u64;
192 let mut total_cache_read = 0u64;
193 let mut total_cache_creation = 0u64;
194 let mut total_cache_5m: Option<u64> = None;
196 let mut total_cache_1h: Option<u64> = None;
197
198 let timeout_fut = tokio::time::sleep(Duration::from_secs(timeout_secs));
199 tokio::pin!(timeout_fut);
200
201 loop {
202 tokio::select! {
203 event = stream.next() => {
204 let Some(event) = event else { break };
205 match event {
206 crate::StreamEvent::Llm(LlmEvent::Thinking(_)) => {
207 if let Some(ref tx) = tx_events_a {
208 let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentUpdate {
209 subagent_id,
210 agent_name: label_a.clone(),
211 status: "💭 thinking...".to_string(),
212 }));
213 }
214 }
215 crate::StreamEvent::Llm(LlmEvent::Text(text)) => {
216 state_a.write().unwrap().partial_text.push_str(&text);
217 }
218 crate::StreamEvent::Llm(LlmEvent::ToolUseStart { tool_name: name, .. }) => {
219 tool_count += 1;
220 if let Some(ref tx) = tx_events_a {
221 let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentUpdate {
222 subagent_id,
223 agent_name: label_a.clone(),
224 status: format!("⚙ {} (tool #{})", name, tool_count),
225 }));
226 }
227 }
228 crate::StreamEvent::Llm(LlmEvent::ToolUse { tool_name, input, .. }) => {
229 let input_str = input.to_string();
230 let input_preview: String = input_str.chars().take(200).collect();
231 state_a.write().unwrap().tool_log
232 .push(format!("[tool_use]: {} — {}", tool_name, input_preview));
233 let detail = match tool_name.as_str() {
234 "bash" => {
235 let cmd = input["command"].as_str().unwrap_or("");
236 let preview: String = cmd.chars().take(60).collect();
237 format!("$ {}", preview)
238 }
239 "read" => format!("reading {}", input["path"].as_str().unwrap_or("?").rsplit('/').next().unwrap_or("?")),
240 "write" => format!("writing {}", input["path"].as_str().unwrap_or("?").rsplit('/').next().unwrap_or("?")),
241 "edit" => format!("editing {}", input["path"].as_str().unwrap_or("?").rsplit('/').next().unwrap_or("?")),
242 "grep" => format!("grep /{}/", input["pattern"].as_str().unwrap_or("?").chars().take(30).collect::<String>()),
243 "find" => format!("find {}", input["pattern"].as_str().unwrap_or("?")),
244 "ls" => format!("ls {}", input["path"].as_str().unwrap_or(".").rsplit('/').next().unwrap_or(".")),
245 other => {
246 if other.starts_with("ext__") {
247 other.splitn(3, "__").last().unwrap_or(other).to_string()
248 } else {
249 other.to_string()
250 }
251 }
252 };
253 if let Some(ref tx) = tx_events_a {
254 let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentUpdate {
255 subagent_id,
256 agent_name: label_a.clone(),
257 status: detail,
258 }));
259 }
260 }
261 crate::StreamEvent::Llm(LlmEvent::ToolResult { result, .. }) => {
262 let preview: String = result.chars().take(300).collect();
263 state_a.write().unwrap().tool_log
264 .push(format!("[tool_result]: {}", preview));
265 }
266 crate::StreamEvent::Session(SessionEvent::Usage {
267 input_tokens, output_tokens,
268 cache_read_input_tokens, cache_creation_input_tokens,
269 cache_creation_5m, cache_creation_1h,
270 model: _,
271 }) => {
272 total_input_tokens += input_tokens;
273 total_output_tokens += output_tokens;
274 total_cache_read += cache_read_input_tokens;
275 total_cache_creation += cache_creation_input_tokens;
276 crate::core::rpc_dispatch::merge_split(&mut total_cache_5m, cache_creation_5m);
277 crate::core::rpc_dispatch::merge_split(&mut total_cache_1h, cache_creation_1h);
278 }
279 crate::StreamEvent::Session(SessionEvent::Error(e)) => return Err(e),
280 crate::StreamEvent::Session(SessionEvent::Done) => break,
281 _ => {}
282 }
283 }
284 _ = &mut timeout_fut => {
285 let (partial, log) = {
286 let mut s = state_a.write().unwrap();
287 s.status = SubagentStatus::TimedOut;
288 s.conversation_state = vec![
289 serde_json::json!({"role": "user", "content": task_for_timeout.clone()}),
290 serde_json::json!({"role": "assistant", "content": &s.partial_text}),
291 ];
292 (s.partial_text.clone(), s.tool_log.clone())
293 };
294 let mut text = format!("[TIMED OUT after {}s — partial results below]\n\n", timeout_secs);
295 if !log.is_empty() {
296 text.push_str(&log.join("\n"));
297 text.push('\n');
298 }
299 if !partial.is_empty() {
300 text.push_str("\n[partial response]:\n");
301 text.push_str(&partial);
302 }
303 return Ok(SubagentResult {
304 text,
305 model: model_a.clone(),
306 input_tokens: total_input_tokens,
307 output_tokens: total_output_tokens,
308 cache_read: total_cache_read,
309 cache_creation: total_cache_creation,
310 cache_creation_5m: total_cache_5m,
311 cache_creation_1h: total_cache_1h,
312 tool_count,
313 });
314 }
315 }
316 }
317
318 Ok(SubagentResult {
319 text: state_a.write().unwrap().partial_text.clone(),
320 model: model_a.clone(),
321 input_tokens: total_input_tokens,
322 output_tokens: total_output_tokens,
323 cache_read: total_cache_read,
324 cache_creation: total_cache_creation,
325 cache_creation_5m: total_cache_5m,
326 cache_creation_1h: total_cache_1h,
327 tool_count,
328 })
329 });
330
331 match outcome {
332 Ok(sa_result) => {
333 {
334 let mut s = state_t.write().unwrap();
335 if matches!(s.status, SubagentStatus::Running) {
336 s.status = SubagentStatus::Completed;
337 s.conversation_state = vec![
338 serde_json::json!({"role": "user", "content": task_for_complete.clone()}),
339 serde_json::json!({"role": "assistant", "content": sa_result.text.clone()}),
340 ];
341 }
342 }
343 let elapsed = start_time.elapsed().as_secs_f64();
344 let preview: String = sa_result.text.chars().take(120).collect();
345 if let Some(ref tx) = tx_events_inner {
346 let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentDone {
347 subagent_id,
348 agent_name: label_inner.clone(),
349 result_preview: preview,
350 duration_secs: elapsed,
351 }));
352 }
353 let _ = result_tx.send(sa_result);
354 }
355 Err(e) => {
356 state_t.write().unwrap().status = SubagentStatus::Failed(e.clone());
357 let elapsed = start_time.elapsed().as_secs_f64();
358 if let Some(ref tx) = tx_events_inner {
359 let _ = tx.send(crate::StreamEvent::Agent(AgentEvent::SubagentDone {
360 subagent_id,
361 agent_name: label_inner.clone(),
362 result_preview: format!("ERROR: {}", e),
363 duration_secs: elapsed,
364 }));
365 }
366 }
367 }
368 }));
369
370 if let Err(panic_info) = panic_result {
371 let msg = if let Some(s) = panic_info.downcast_ref::<&str>() {
372 s.to_string()
373 } else if let Some(s) = panic_info.downcast_ref::<String>() {
374 s.clone()
375 } else {
376 "unknown panic".to_string()
377 };
378 tracing::error!("Resumed subagent thread panicked: {}", msg);
379 state_t.write().unwrap().status = SubagentStatus::Failed(format!("panic: {}", msg));
380 }
381 });
382
383 let handle = SubagentHandle::new(
384 handle_id.clone(),
385 label.clone(),
386 task_preview,
387 model,
388 system_prompt_for_handle,
389 timeout_secs,
390 state,
391 Some(steer_tx),
392 Some(shutdown_tx),
393 Some(result_rx),
394 );
395
396 {
397 let mut reg = registry.lock().unwrap();
398 reg.register(handle);
399 if let Some(h) = reg.get_mut(&handle_id) {
400 h.set_thread_handle(thread_handle);
401 }
402 }
403
404 Ok(json!({
405 "handle_id": handle_id,
406 "resumed_from": prior_handle_id,
407 "agent_name": label,
408 "status": "running"
409 }).to_string())
410 }
411}