1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde_json::{json, Value};
5
6use crate::engine::{PlanGenerator, PlanStore, StepExecutor};
7use crate::tool::{Tool, ToolContext, ToolControlFlow, ToolOutput};
8use crate::types::{AgentError, AgentEvent, AgentResult, PlanStatus, StepStatus};
9
10#[derive(Clone)]
14pub struct PlanOrchestrator {
15 plan_generator: Arc<dyn PlanGenerator>,
16 step_executor: Arc<dyn StepExecutor>,
17 plan_store: Arc<dyn PlanStore>,
18}
19
20impl PlanOrchestrator {
21 pub fn new(
22 plan_generator: Arc<dyn PlanGenerator>,
23 step_executor: Arc<dyn StepExecutor>,
24 plan_store: Arc<dyn PlanStore>,
25 ) -> Self {
26 Self {
27 plan_generator,
28 step_executor,
29 plan_store,
30 }
31 }
32
33 pub fn with_step_executor(&mut self, step_executor: Arc<dyn StepExecutor>) {
34 self.step_executor = step_executor;
35 }
36}
37
38#[async_trait]
39impl Tool for PlanOrchestrator {
40 fn name(&self) -> &'static str {
41 "create_plan"
42 }
43
44 fn definition(&self) -> Value {
45 json!({
46 "type": "function",
47 "function": {
48 "name": "create_plan",
49 "description": "Analyze a task and generate an execution plan (without executing commands). Used for complex tasks that require multiple steps. The system will analyze the objective and generate a plan; after user review and confirmation, use execute_plan to execute it.",
50 "parameters": {
51 "type": "object",
52 "properties": {
53 "objective": {
54 "type": "string",
55 "description": "The overall goal of the task, e.g. 'check disk space', 'troubleshoot network issues'"
56 },
57 "context": {
58 "type": "string",
59 "description": "Additional context information, such as target host, environment variables, etc."
60 }
61 },
62 "required": ["objective"]
63 }
64 }
65 })
66 }
67
68 async fn call(&self, args: &Value, ctx: &ToolContext) -> AgentResult<ToolOutput> {
69 let objective = args
70 .get("objective")
71 .and_then(Value::as_str)
72 .unwrap_or("unnamed task")
73 .to_string();
74 let context = args
75 .get("context")
76 .and_then(Value::as_str)
77 .unwrap_or("")
78 .to_string();
79
80 let plan_id = {
81 let timestamp = std::time::SystemTime::now()
82 .duration_since(std::time::UNIX_EPOCH)
83 .unwrap_or_default()
84 .as_millis();
85 static COUNTER: std::sync::atomic::AtomicU64 =
86 std::sync::atomic::AtomicU64::new(0);
87 let count = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
88 format!("plan-{timestamp}-{count}")
89 };
90
91 let event_bus_g = ctx.event_bus.clone();
92 let session_id_g = ctx.session_id.clone();
93 let plan_id_g = plan_id.clone();
94 let on_generating = Box::new(move || {
95 let _ = event_bus_g.send(AgentEvent::PlanGenerating {
96 session_id: session_id_g.clone(),
97 plan_id: plan_id_g.clone(),
98 });
99 });
100
101 let event_bus_s = ctx.event_bus.clone();
102 let session_id_s = ctx.session_id.clone();
103 let plan_id_s = plan_id.clone();
104 let on_step_parsed = Box::new(move |index: usize, step_id: String, description: String| {
105 let _ = event_bus_s.send(AgentEvent::PlanStepParsed {
106 session_id: session_id_s.clone(),
107 plan_id: plan_id_s.clone(),
108 step_index: index,
109 step_id,
110 step_description: description,
111 });
112 });
113
114 let event_bus_t = ctx.event_bus.clone();
115 let session_id_t = ctx.session_id.clone();
116 let on_raw_chunk = Box::new(move |text: String| {
117 let _ = event_bus_t.send(AgentEvent::ThoughtDelta {
118 session_id: session_id_t.clone(),
119 text,
120 });
121 });
122
123 let tools = vec![];
124
125 match self
126 .plan_generator
127 .generate_plan_streaming(
128 &objective,
129 &context,
130 &tools,
131 on_generating,
132 on_step_parsed,
133 on_raw_chunk,
134 )
135 .await
136 {
137 Ok(mut plan) => {
138 plan.id = plan_id.clone();
139 plan.objective = objective.clone();
140
141 self.plan_store
142 .save_plan(&plan, json!({"session_id": ctx.session_id.to_string()}))
143 .await?;
144
145 let _ = ctx.event_bus.send(AgentEvent::PlanGenerated {
146 session_id: ctx.session_id.clone(),
147 plan: plan.clone(),
148 });
149
150 let step_details: Vec<Value> = plan
151 .steps
152 .iter()
153 .map(|s| {
154 json!({
155 "id": s.id,
156 "description": s.description,
157 })
158 })
159 .collect();
160
161 let summary = if ctx.language == crate::types::Language::Zh {
162 format!(
163 "计划已生成，包含 {} 个步骤，等待用户确认。计划ID: {}",
164 plan.steps.len(),
165 plan_id
166 )
167 } else {
168 format!(
169 "Plan generated with {} steps, awaiting user confirmation. plan_id: {}",
170 plan.steps.len(),
171 plan_id
172 )
173 };
174
175 Ok(ToolOutput {
176 summary,
177 raw: Some(json!({
178 "objective": objective,
179 "plan_id": plan_id,
180 "steps_count": plan.steps.len(),
181 "steps": step_details,
182 "success": true,
183 "status": "awaiting_confirmation",
184 })),
185 control_flow: ToolControlFlow::Continue,
186 truncation: None,
187 })
188 }
189 Err(e) => {
190 let _ = ctx.event_bus.send(AgentEvent::PlanFailed {
191 session_id: ctx.session_id.clone(),
192 plan_id: plan_id.clone(),
193 error: e.to_string(),
194 });
195
196 let summary = if ctx.language == crate::types::Language::Zh {
197 format!("计划生成失败: {e}")
198 } else {
199 format!("Plan generation failed: {e}")
200 };
201
202 Ok(ToolOutput {
203 summary,
204 raw: Some(json!({
205 "objective": objective,
206 "plan_id": plan_id,
207 "success": false,
208 "error": e.to_string(),
209 })),
210 control_flow: ToolControlFlow::Continue,
211 truncation: None,
212 })
213 }
214 }
215 }
216}
217
218#[derive(Clone)]
220pub struct PlanExecTool {
221 step_executor: Arc<dyn StepExecutor>,
222 plan_store: Arc<dyn PlanStore>,
223 recovery: Arc<dyn crate::engine::RecoveryStrategy>,
224}
225
226impl PlanExecTool {
227 pub fn new(
228 step_executor: Arc<dyn StepExecutor>,
229 plan_store: Arc<dyn PlanStore>,
230 recovery: Arc<dyn crate::engine::RecoveryStrategy>,
231 ) -> Self {
232 Self {
233 step_executor,
234 plan_store,
235 recovery,
236 }
237 }
238}
239
240#[async_trait]
241impl Tool for PlanExecTool {
242 fn name(&self) -> &'static str {
243 "execute_plan"
244 }
245
246 fn definition(&self) -> Value {
247 json!({
248 "type": "function",
249 "function": {
250 "name": "execute_plan",
251 "description": "Execute a previously generated plan. First use create_plan to generate a plan, then after user review and confirmation, use this tool to execute it.",
252 "parameters": {
253 "type": "object",
254 "properties": {
255 "plan_id": {
256 "type": "string",
257 "description": "The plan ID to execute (obtained from the create_plan result)"
258 }
259 },
260 "required": ["plan_id"]
261 }
262 }
263 })
264 }
265
266 async fn call(&self, args: &Value, ctx: &ToolContext) -> AgentResult<ToolOutput> {
267 let plan_id = args
268 .get("plan_id")
269 .and_then(Value::as_str)
270 .unwrap_or("")
271 .to_string();
272
273 let is_zh = ctx.language == crate::types::Language::Zh;
274
275 if plan_id.is_empty() {
276 let summary = if is_zh {
277 "缺少 plan_id 参数".to_string()
278 } else {
279 "Missing plan_id parameter".to_string()
280 };
281 return Ok(ToolOutput {
282 summary,
283 raw: Some(json!({
284 "success": false,
285 "error": "plan_id is required",
286 })),
287 control_flow: ToolControlFlow::Continue,
288 truncation: None,
289 });
290 }
291
292 let plan_data = self
293 .plan_store
294 .load_plan(&plan_id)
295 .await?
296 .ok_or_else(|| AgentError::plan_storage(
297 if is_zh {
298 format!("计划 {plan_id} 不存在，可能已过期或从未创建")
299 } else {
300 format!("Plan {plan_id} does not exist, it may have expired or never been created")
301 }
302 ))?;
303
304 let mut plan = plan_data.plan;
305 let objective = plan.objective.clone();
306 let mut execution_summary = if is_zh {
307 format!("计划 '{}' 的执行结果:\n", objective)
308 } else {
309 format!("Execution results for plan '{}':\n", objective)
310 };
311 let mut step_results = Vec::new();
312 let mut overall_success = true;
313 let mut failed_step_name: Option<String> = None;
314 let mut _completed_count = 0usize;
315
316 for (index, step) in plan.steps.iter_mut().enumerate() {
317 if step.status != StepStatus::Pending {
318 continue;
319 }
320
321 step.status = StepStatus::Running;
322
323 let _ = ctx.event_bus.send(AgentEvent::PlanStepStarted {
324 session_id: ctx.session_id.clone(),
325 step_id: step.id.clone(),
326 step_description: step.description.clone(),
327 });
328
329 execution_summary.push_str(&if is_zh {
330 format!("步骤 {}: {}\n", index + 1, step.description)
331 } else {
332 format!("Step {}: {}\n", index + 1, step.description)
333 });
334
335 match self
336 .step_executor
337 .execute_step(step, &plan_data.metadata)
338 .await
339 {
340 Ok(result) => {
341 let step_success = result.success;
342 step.status = if step_success {
343 StepStatus::Completed
344 } else {
345 StepStatus::Failed
346 };
347 step.result = Some(result.clone());
348
349 execution_summary.push_str(&if is_zh {
350 format!(
351 " 结果: {}\n",
352 if step_success { "成功" } else { "失败" }
353 )
354 } else {
355 format!(
356 " Result: {}\n",
357 if step_success { "OK" } else { "FAILED" }
358 )
359 });
360
361 let _ = ctx.event_bus.send(AgentEvent::PlanStepCompleted {
362 session_id: ctx.session_id.clone(),
363 step_id: step.id.clone(),
364 success: step_success,
365 result: result.output.clone(),
366 });
367
368 if step_success {
369 _completed_count += 1;
370 } else {
371 if let Some(ref output) = result.output {
372 execution_summary.push_str(&if is_zh {
373 format!(" 错误: {output}\n")
374 } else {
375 format!(" Error: {output}\n")
376 });
377 }
378
379 match self
380 .recovery
381 .handle_step_failure(
382 step,
383 result.output.as_deref().unwrap_or(""),
384 0,
385 )
386 .await
387 {
388 Ok(action) => match action {
389 crate::types::RecoveryAction::Retry => {
390 execution_summary.push_str(
391 if is_zh {
392 " [重试] 系统建议重试该步骤（计划标记为未完全成功）\n"
393 } else {
394 " [Retry] System suggests retrying this step (plan marked as not fully successful)\n"
395 },
396 );
397 overall_success = false;
398 }
399 crate::types::RecoveryAction::Skip => {
400 execution_summary.push_str(
401 if is_zh {
402 " [跳过] 系统建议跳过该步骤（计划标记为未完全成功）\n"
403 } else {
404 " [Skip] System suggests skipping this step (plan marked as not fully successful)\n"
405 },
406 );
407 step.status = StepStatus::Skipped;
408 overall_success = false;
409 }
410 crate::types::RecoveryAction::Abort => {
411 execution_summary.push_str(
412 if is_zh {
413 " [中止] 系统建议中止计划\n"
414 } else {
415 " [Abort] System suggests aborting the plan\n"
416 },
417 );
418 overall_success = false;
419 failed_step_name = Some(step.description.clone());
420 break;
421 }
422 },
423 Err(_e) => {
424 overall_success = false;
425 failed_step_name = Some(step.description.clone());
426 break;
427 }
428 }
429 }
430
431 step_results.push(json!({
432 "step": step.description,
433 "success": step_success,
434 "output": result.output,
435 }));
436 }
437 Err(e) => {
438 step.status = StepStatus::Failed;
439 execution_summary.push_str(&if is_zh {
440 format!(" 执行错误: {e}\n")
441 } else {
442 format!(" Execution error: {e}\n")
443 });
444
445 let _ = ctx.event_bus.send(AgentEvent::PlanStepCompleted {
446 session_id: ctx.session_id.clone(),
447 step_id: step.id.clone(),
448 success: false,
449 result: Some(e.to_string()),
450 });
451
452 overall_success = false;
453 failed_step_name = Some(step.description.clone());
454 break;
455 }
456 }
457 }
458
459 plan.status = if overall_success && plan.is_completed() {
460 PlanStatus::Completed
461 } else if plan.has_failed() {
462 PlanStatus::Failed
463 } else {
464 PlanStatus::Executing
465 };
466
467 self.plan_store
468 .save_plan(&plan, plan_data.metadata)
469 .await?;
470
471 if overall_success {
472 execution_summary.push_str(
473 if is_zh { "所有步骤执行完毕。" } else { "All steps completed." }
474 );
475 }
476
477 let _ = ctx.event_bus.send(AgentEvent::PlanCompleted {
478 session_id: ctx.session_id.clone(),
479 plan_id: plan_id.clone(),
480 success: overall_success,
481 });
482
483 Ok(ToolOutput {
484 summary: execution_summary,
485 raw: Some(json!({
486 "objective": objective,
487 "plan_id": plan_id,
488 "steps": step_results,
489 "success": overall_success,
490 "failed_step": failed_step_name,
491 })),
492 control_flow: ToolControlFlow::Continue,
493 truncation: None,
494 })
495 }
496}