1use crate::audit::{AuditEventType, AuditLog};
7use crate::config::Config;
8use crate::error::Result;
9use crate::healing::SelfHealingEngine;
10use crate::llm::{
11 ChatMessage, Choice, LLMProviderTrait, MultiModelManager, ProviderFallbackChain, RetryConfig,
12 TokenBudget,
13};
14use crate::mcp::McpClient;
15use crate::policy::{Decision, PolicyEngine};
16use crate::ravenfabric::RavenFabricClient;
17use crate::sandbox::Sandbox;
18use crate::tools::{ToolCall, ToolRegistry, ToolResult};
19use serde::{Deserialize, Serialize};
20use std::path::PathBuf;
21use std::sync::Arc;
22use tokio::sync::RwLock;
23use tracing::{debug, info, instrument, warn};
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct ConversationMemory {
31 max_messages: usize,
33 messages: Vec<ChatMessage>,
35}
36
37impl ConversationMemory {
38 pub fn new(system_prompt: &str, max_messages: usize) -> Self {
41 Self {
42 max_messages,
43 messages: vec![ChatMessage::new("system", system_prompt.to_string())],
44 }
45 }
46
47 pub fn add_user_message(&mut self, content: &str) -> &[ChatMessage] {
49 self.messages
50 .push(ChatMessage::new("user", content.to_string()));
51 self.trim_to_max();
52 &self.messages
53 }
54
55 pub fn add_user_message_with_images(
57 &mut self,
58 text: &str,
59 image_data_uris: Vec<String>,
60 ) -> &[ChatMessage] {
61 self.messages.push(ChatMessage::with_images(
62 "user",
63 text.to_string(),
64 image_data_uris,
65 ));
66 self.trim_to_max();
67 &self.messages
68 }
69
70 pub fn add_assistant_message(&mut self, content: &str) {
72 self.messages
73 .push(ChatMessage::new("assistant", content.to_string()));
74 self.trim_to_max();
75 }
76
77 pub fn history(&self) -> &[ChatMessage] {
79 &self.messages
80 }
81
82 pub fn from_history(messages: Vec<ChatMessage>, max_messages: usize) -> Self {
85 Self {
86 max_messages,
87 messages,
88 }
89 }
90
91 #[allow(dead_code)]
93 pub fn len(&self) -> usize {
94 self.messages.len()
95 }
96
97 #[allow(dead_code)]
99 pub fn is_empty(&self) -> bool {
100 self.messages.is_empty()
101 }
102
103 fn trim_to_max(&mut self) {
105 if self.max_messages == 0 {
106 return;
107 }
108 while self.messages.len() > self.max_messages {
109 if self.messages.len() > 1 {
111 self.messages.remove(1);
112 } else {
113 break;
114 }
115 }
116 }
117}
118
119#[derive(Debug, Clone, Serialize, Deserialize)]
125pub struct CheckpointState {
126 pub session_id: String,
128 pub iteration: usize,
130 pub max_iterations: usize,
132 pub messages: Vec<ChatMessage>,
134 pub initial_prompt: String,
136 pub system_prompt: String,
138 pub provider: String,
140 pub model: String,
142 pub enable_tools: bool,
144 pub last_checkpoint: String,
146}
147
148impl CheckpointState {
149 #[allow(clippy::too_many_arguments)]
151 pub fn new(
152 session_id: String,
153 iteration: usize,
154 max_iterations: usize,
155 messages: Vec<ChatMessage>,
156 initial_prompt: &str,
157 system_prompt: &str,
158 provider: &str,
159 model: &str,
160 enable_tools: bool,
161 ) -> Self {
162 Self {
163 session_id,
164 iteration,
165 max_iterations,
166 messages,
167 initial_prompt: initial_prompt.to_string(),
168 system_prompt: system_prompt.to_string(),
169 provider: provider.to_string(),
170 model: model.to_string(),
171 enable_tools,
172 last_checkpoint: chrono::Utc::now().to_rfc3339(),
173 }
174 }
175}
176
177pub fn save_checkpoint(
182 checkpoint_dir: &std::path::Path,
183 state: &CheckpointState,
184) -> std::result::Result<std::path::PathBuf, String> {
185 let path = checkpoint_dir.join(format!("{}.json", state.session_id));
186
187 std::fs::create_dir_all(checkpoint_dir)
189 .map_err(|e| format!("Failed to create checkpoint directory: {}", e))?;
190
191 let content = serde_json::to_string_pretty(state)
192 .map_err(|e| format!("Failed to serialize checkpoint: {}", e))?;
193
194 let tmp_path = path.with_extension("json.tmp");
196 std::fs::write(&tmp_path, &content)
197 .map_err(|e| format!("Failed to write checkpoint: {}", e))?;
198 std::fs::rename(&tmp_path, &path)
199 .map_err(|e| format!("Failed to finalize checkpoint: {}", e))?;
200
201 Ok(path)
202}
203
204pub fn load_checkpoint(
209 checkpoint_dir: &std::path::Path,
210 session_id: &str,
211) -> Option<CheckpointState> {
212 let path = checkpoint_dir.join(format!("{}.json", session_id));
213
214 match std::fs::read_to_string(&path) {
215 Ok(content) => match serde_json::from_str::<CheckpointState>(&content) {
216 Ok(state) => {
217 info!(
218 session_id = %session_id,
219 iteration = state.iteration,
220 max_iterations = state.max_iterations,
221 "Loaded checkpoint"
222 );
223 Some(state)
224 }
225 Err(e) => {
226 warn!(
227 session_id = %session_id,
228 error = %e,
229 "Failed to deserialize checkpoint"
230 );
231 None
232 }
233 },
234 Err(e) => {
235 if e.kind() != std::io::ErrorKind::NotFound {
236 warn!(
237 session_id = %session_id,
238 error = %e,
239 "Failed to read checkpoint"
240 );
241 }
242 None
243 }
244 }
245}
246
247pub fn delete_checkpoint(checkpoint_dir: &std::path::Path, session_id: &str) {
251 let path = checkpoint_dir.join(format!("{}.json", session_id));
252 if path.exists() {
253 if let Err(e) = std::fs::remove_file(&path) {
254 warn!(
255 session_id = %session_id,
256 error = %e,
257 "Failed to delete checkpoint"
258 );
259 } else {
260 debug!(
261 session_id = %session_id,
262 "Deleted checkpoint"
263 );
264 }
265 }
266}
267
268pub struct AgentLoopConfig {
273 pub max_iterations: usize,
275 pub enable_tools: bool,
277 pub require_approval: bool,
279 pub prompt_injection_protection: bool,
281 pub token_lifetime_secs: u64,
285 pub no_final_required: bool,
287 pub fallback_chain: Option<Arc<std::sync::Mutex<ProviderFallbackChain>>>,
289 pub token_budget: Option<Arc<std::sync::Mutex<TokenBudget>>>,
291 pub ravenfabric: Option<RavenFabricClient>,
293 pub checkpoint_dir: Option<PathBuf>,
297 pub session_id: Option<String>,
300 pub metrics_callback: Option<Box<dyn Fn(u64, u64) + Send + Sync>>,
304
305 pub load_manager: Option<Arc<crate::load::LoadManager>>,
308
309 pub retry_config: Option<RetryConfig>,
314
315 pub healing_engine: Option<Arc<std::sync::Mutex<SelfHealingEngine>>>,
320}
321
322impl Default for AgentLoopConfig {
323 fn default() -> Self {
324 Self {
325 max_iterations: 10,
326 enable_tools: false,
327 require_approval: false,
328 prompt_injection_protection: true,
329 token_lifetime_secs: 0,
330 no_final_required: true,
331 fallback_chain: None,
332 token_budget: None,
333 ravenfabric: None,
334 checkpoint_dir: None,
335 session_id: None,
336 metrics_callback: None,
337 load_manager: None,
338 retry_config: None,
339 healing_engine: None,
340 }
341 }
342}
343
344impl std::fmt::Debug for AgentLoopConfig {
346 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
347 f.debug_struct("AgentLoopConfig")
348 .field("max_iterations", &self.max_iterations)
349 .field("enable_tools", &self.enable_tools)
350 .field("require_approval", &self.require_approval)
351 .field(
352 "prompt_injection_protection",
353 &self.prompt_injection_protection,
354 )
355 .field("token_lifetime_secs", &self.token_lifetime_secs)
356 .field("no_final_required", &self.no_final_required)
357 .field("fallback_chain", &self.fallback_chain)
358 .field("token_budget", &self.token_budget)
359 .field("ravenfabric", &self.ravenfabric)
360 .field("checkpoint_dir", &self.checkpoint_dir)
361 .field("session_id", &self.session_id)
362 .field(
363 "metrics_callback",
364 &self.metrics_callback.as_ref().map(|_| "Box<Fn>"),
365 )
366 .field(
367 "load_manager",
368 &self.load_manager.as_ref().map(|_| "Arc<LoadManager>"),
369 )
370 .field(
371 "healing_engine",
372 &self
373 .healing_engine
374 .as_ref()
375 .map(|_| "Arc<Mutex<SelfHealingEngine>>"),
376 )
377 .finish()
378 }
379}
380
381impl Clone for AgentLoopConfig {
384 fn clone(&self) -> Self {
385 Self {
386 max_iterations: self.max_iterations,
387 enable_tools: self.enable_tools,
388 require_approval: self.require_approval,
389 prompt_injection_protection: self.prompt_injection_protection,
390 token_lifetime_secs: self.token_lifetime_secs,
391 no_final_required: self.no_final_required,
392 fallback_chain: self.fallback_chain.clone(),
393 token_budget: self.token_budget.clone(),
394 ravenfabric: self.ravenfabric.clone(),
395 checkpoint_dir: self.checkpoint_dir.clone(),
396 session_id: self.session_id.clone(),
397 metrics_callback: None,
398 load_manager: self.load_manager.clone(),
399 retry_config: self.retry_config.clone(),
400 healing_engine: self.healing_engine.clone(),
401 }
402 }
403}
404
405#[instrument(skip_all, fields(provider = %llm.provider_name(), model = %llm.model()))]
412pub async fn run_agent_loop(
413 llm: Arc<dyn LLMProviderTrait>,
414 initial_prompt: &str,
415 system_prompt: &str,
416 config: AgentLoopConfig,
417) -> Result<String> {
418 run_agent_loop_with_registry(llm, initial_prompt, system_prompt, config, None).await
419}
420
421#[instrument(skip_all, fields(provider = %llm.provider_name(), model = %llm.model()))]
426pub async fn run_agent_loop_with_registry(
427 llm: Arc<dyn LLMProviderTrait>,
428 initial_prompt: &str,
429 system_prompt: &str,
430 config: AgentLoopConfig,
431 tool_registry: Option<ToolRegistry>,
432) -> Result<String> {
433 let registry = tool_registry.unwrap_or_else(ToolRegistry::with_default_tools);
434 run_agent_loop_inner(
435 llm,
436 initial_prompt,
437 system_prompt,
438 config,
439 registry,
440 "security integration",
441 false,
442 Vec::new(),
443 )
444 .await
445}
446
447#[allow(dead_code)]
452#[instrument(skip_all, fields(provider = %llm.provider_name(), model = %llm.model(), image_count = image_data_uris.len()))]
453pub async fn run_agent_loop_with_images(
454 llm: Arc<dyn LLMProviderTrait>,
455 initial_prompt: &str,
456 system_prompt: &str,
457 config: AgentLoopConfig,
458 tool_registry: Option<ToolRegistry>,
459 image_data_uris: Vec<String>,
460) -> Result<String> {
461 let registry = tool_registry.unwrap_or_else(ToolRegistry::with_default_tools);
462 run_agent_loop_inner(
463 llm,
464 initial_prompt,
465 system_prompt,
466 config,
467 registry,
468 "security integration",
469 false,
470 image_data_uris,
471 )
472 .await
473}
474
475async fn call_llm_with_retry(
491 llm: &Arc<dyn LLMProviderTrait>,
492 messages: Vec<ChatMessage>,
493 retry_config: Option<&RetryConfig>,
494 audit_log: &AuditLog,
495 session_id: &str,
496 checkpoint_dir: &Option<PathBuf>,
497 iteration: usize,
498) -> std::result::Result<crate::llm::ChatResponse, crate::llm::LLMError> {
499 let max_attempts = match retry_config {
500 Some(cfg) => cfg.max_retries + 1, None => 1,
502 };
503
504 let mut last_error = None;
505
506 for attempt in 0..max_attempts {
507 if attempt > 0 {
508 let delay = retry_config.unwrap().delay_for_attempt(attempt - 1);
509 info!(
510 attempt = attempt + 1,
511 max_attempts = max_attempts,
512 delay_ms = delay.as_millis(),
513 iteration = iteration,
514 "Retrying LLM call after transient error"
515 );
516 tokio::time::sleep(delay).await;
517 }
518
519 match llm.chat(messages.clone()).await {
520 Ok(response) => {
521 if attempt > 0 {
522 info!(
523 attempt = attempt + 1,
524 iteration = iteration,
525 "LLM call succeeded on retry"
526 );
527 let _ = audit_log.append(
528 AuditEventType::Custom("Info".to_string()),
529 "llm_retry",
530 &format!("LLM call succeeded on retry attempt {}", attempt + 1),
531 None,
532 );
533 }
534 return Ok(response);
535 }
536 Err(e) => {
537 let is_transient = matches!(
538 &e,
539 crate::llm::LLMError::RequestFailed(_)
540 | crate::llm::LLMError::RateLimited
541 | crate::llm::LLMError::CircuitBreakerOpen(_)
542 );
543
544 if is_transient && attempt + 1 < max_attempts {
545 warn!(
546 error = %e,
547 attempt = attempt + 1,
548 max_attempts = max_attempts,
549 iteration = iteration,
550 "Transient LLM error, will retry"
551 );
552 let _ = audit_log.append(
553 AuditEventType::Error,
554 "llm_retry",
555 &format!("Transient LLM error on attempt {}: {}", attempt + 1, e),
556 None,
557 );
558 last_error = Some(e);
559 continue;
561 }
562
563 if attempt + 1 >= max_attempts && is_transient {
565 warn!(
566 error = %e,
567 attempts = max_attempts,
568 iteration = iteration,
569 "All retry attempts exhausted"
570 );
571 let _ = audit_log.append(
572 AuditEventType::Error,
573 "llm_retry",
574 &format!("All {} retry attempts exhausted: {}", max_attempts, e),
575 None,
576 );
577 }
578
579 if let Some(ref cp_dir) = checkpoint_dir {
581 delete_checkpoint(cp_dir, session_id);
582 }
583 return Err(e);
584 }
585 }
586 }
587
588 Err(last_error.unwrap_or(crate::llm::LLMError::RequestFailed(
590 "All retry attempts exhausted".to_string(),
591 )))
592}
593
594#[allow(clippy::too_many_arguments)]
606#[instrument(skip_all, fields(provider = %llm.provider_name(), model = %llm.model()))]
607async fn run_agent_loop_inner(
608 llm: Arc<dyn LLMProviderTrait>,
609 initial_prompt: &str,
610 system_prompt: &str,
611 config: AgentLoopConfig,
612 registry: ToolRegistry,
613 loop_label: &str,
614 mcp_enabled: bool,
615 image_data_uris: Vec<String>,
616) -> Result<String> {
617 let policy_engine = PolicyEngine::default_secure();
619 let mut sandbox = Sandbox::default();
620 sandbox.init().await.map_err(|e| {
621 crate::error::RavenClawsError::CommandExecution(format!("Sandbox init failed: {}", e))
622 })?;
623 let audit_log = AuditLog::new(format!("agent-{}", std::process::id()));
624
625 let injection_detector = if config.prompt_injection_protection {
627 Some(crate::policy::InjectionDetector::new())
628 } else {
629 None
630 };
631
632 let session_start = std::time::Instant::now();
634
635 info!(
636 provider = llm.provider_name(),
637 model = llm.model(),
638 max_iterations = config.max_iterations,
639 enable_tools = config.enable_tools,
640 tool_count = registry.len(),
641 require_approval = config.require_approval,
642 prompt_injection_protection = config.prompt_injection_protection,
643 token_lifetime_secs = config.token_lifetime_secs,
644 "Agent loop starting with {}",
645 loop_label
646 );
647
648 let _ = audit_log.append(
650 AuditEventType::AgentStart,
651 "agent",
652 &format!(
653 "Agent loop started with {} (model: {})",
654 llm.provider_name(),
655 llm.model()
656 ),
657 Some(serde_json::json!({
658 "provider": llm.provider_name(),
659 "model": llm.model(),
660 "max_iterations": config.max_iterations,
661 "enable_tools": config.enable_tools,
662 "mcp_enabled": mcp_enabled,
663 "tool_count": registry.len(),
664 "require_approval": config.require_approval,
665 "prompt_injection_protection": config.prompt_injection_protection,
666 "token_lifetime_secs": config.token_lifetime_secs,
667 })),
668 );
669
670 let (mut memory, start_iteration) = if let Some(ref checkpoint_dir) = config.checkpoint_dir {
674 if let Some(ref session_id) = config.session_id {
675 if let Some(checkpoint) = load_checkpoint(checkpoint_dir, session_id) {
676 info!(
677 session_id = %session_id,
678 iteration = checkpoint.iteration,
679 max_iterations = checkpoint.max_iterations,
680 "Resuming agent loop from checkpoint"
681 );
682 (
683 ConversationMemory::from_history(checkpoint.messages, 0),
684 checkpoint.iteration + 1, )
686 } else {
687 info!(
688 session_id = %session_id,
689 "No checkpoint found, starting fresh"
690 );
691 let mut m = ConversationMemory::new(system_prompt, 0);
692 if image_data_uris.is_empty() {
693 m.add_user_message(initial_prompt);
694 } else {
695 m.add_user_message_with_images(initial_prompt, image_data_uris.clone());
696 }
697 (m, 0)
698 }
699 } else {
700 let mut m = ConversationMemory::new(system_prompt, 0);
701 if image_data_uris.is_empty() {
702 m.add_user_message(initial_prompt);
703 } else {
704 m.add_user_message_with_images(initial_prompt, image_data_uris.clone());
705 }
706 (m, 0)
707 }
708 } else {
709 let mut m = ConversationMemory::new(system_prompt, 0);
710 if image_data_uris.is_empty() {
711 m.add_user_message(initial_prompt);
712 } else {
713 m.add_user_message_with_images(initial_prompt, image_data_uris.clone());
714 }
715 (m, 0)
716 };
717
718 let session_id = config
720 .session_id
721 .clone()
722 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
723
724 for iteration in start_iteration..config.max_iterations {
725 if config.token_lifetime_secs > 0 {
727 let elapsed = session_start.elapsed().as_secs();
728 if elapsed >= config.token_lifetime_secs {
729 warn!(
730 iteration = iteration,
731 elapsed_secs = elapsed,
732 token_lifetime_secs = config.token_lifetime_secs,
733 "Agent loop reached token lifetime limit"
734 );
735 let _ = audit_log.append(
736 AuditEventType::SecurityViolation,
737 "token_lifetime",
738 &format!(
739 "Session expired after {} seconds (limit: {}s)",
740 elapsed, config.token_lifetime_secs
741 ),
742 Some(serde_json::json!({
743 "elapsed_secs": elapsed,
744 "token_lifetime_secs": config.token_lifetime_secs,
745 "iteration": iteration,
746 })),
747 );
748 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
750 delete_checkpoint(checkpoint_dir, &session_id);
751 }
752 return Err(crate::error::RavenClawsError::SecurityViolation(format!(
753 "Session token expired after {} seconds (limit: {}s)",
754 elapsed, config.token_lifetime_secs
755 )));
756 }
757 }
758 let messages = memory.history().to_vec();
759
760 if let Some(ref budget) = config.token_budget {
762 let budget = budget.lock().unwrap();
763 if budget.remaining() < 100 {
764 warn!(
765 iteration = iteration,
766 remaining = budget.remaining(),
767 "Token budget exhausted"
768 );
769 let _ = audit_log.append(
770 AuditEventType::SecurityViolation,
771 "token_budget",
772 &format!("Token budget exhausted (remaining: {})", budget.remaining()),
773 Some(serde_json::json!({
774 "remaining": budget.remaining(),
775 "used": budget.used_tokens,
776 "iteration": iteration,
777 })),
778 );
779 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
781 delete_checkpoint(checkpoint_dir, &session_id);
782 }
783 return Err(crate::error::RavenClawsError::SecurityViolation(
784 "Token budget exhausted".to_string(),
785 ));
786 }
787 }
788
789 if let Some(ref load_manager) = config.load_manager {
791 let admission = load_manager.check_admission();
792 if !admission.is_allowed() {
793 warn!(
794 ?admission,
795 iteration = iteration,
796 "Admission denied before LLM call"
797 );
798 let _ = audit_log.append(
799 AuditEventType::Error,
800 "load_manager",
801 &format!("Admission denied: {:?}", admission),
802 None,
803 );
804 load_manager.record_outcome(crate::load::RequestOutcome::Failure);
805 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
807 delete_checkpoint(checkpoint_dir, &session_id);
808 }
809 return Err(crate::error::RavenClawsError::SecurityViolation(format!(
810 "Admission denied: {:?}",
811 admission
812 )));
813 }
814 }
815
816 if let Some(ref healing) = config.healing_engine {
818 let agent_id = llm.provider_name().to_string();
819 let healthy = {
820 let mut engine = healing.lock().unwrap();
821 engine.is_healthy(&agent_id)
822 };
823 if !healthy {
824 warn!(
825 agent_id = %agent_id,
826 iteration = iteration,
827 "Agent blocked by self-healing circuit breaker"
828 );
829 let _ = audit_log.append(
830 AuditEventType::Error,
831 "healing",
832 &format!("Agent '{}' blocked by circuit breaker", agent_id),
833 None,
834 );
835 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
837 delete_checkpoint(checkpoint_dir, &session_id);
838 }
839 return Err(crate::error::RavenClawsError::HealingError(format!(
840 "Agent '{}' blocked by circuit breaker",
841 agent_id
842 )));
843 }
844 }
845
846 let response = match call_llm_with_retry(
850 &llm,
851 messages.clone(),
852 config.retry_config.as_ref(),
853 &audit_log,
854 &session_id,
855 &config.checkpoint_dir,
856 iteration,
857 )
858 .await
859 {
860 Ok(r) => {
861 if let Some(ref load_manager) = config.load_manager {
863 load_manager.record_outcome(crate::load::RequestOutcome::Success);
864 }
865 if let Some(ref healing) = config.healing_engine {
867 let agent_id = llm.provider_name().to_string();
868 let mut engine = healing.lock().unwrap();
869 engine.record_success(&agent_id);
870 }
871 r
872 }
873 Err(e) => {
874 if let Some(ref load_manager) = config.load_manager {
876 load_manager.record_outcome(crate::load::RequestOutcome::Failure);
877 }
878 if let Some(ref healing) = config.healing_engine {
880 let agent_id = llm.provider_name().to_string();
881 let mut engine = healing.lock().unwrap();
882 engine.record_failure(&agent_id, &e.to_string());
883 }
884 if let Some(ref chain) = config.fallback_chain {
886 warn!(error = %e, "Primary LLM failed after retries, trying fallback chain");
887 let _ = audit_log.append(
888 AuditEventType::Error,
889 "llm",
890 &format!("Primary LLM failed after retries, trying fallback: {}", e),
891 None,
892 );
893 let configs = {
895 let c = chain.lock().unwrap();
896 c.configs.clone()
897 };
898 let mut temp_chain = ProviderFallbackChain::new(configs);
899 match temp_chain.chat_with_fallback(messages).await {
900 Ok(r) => {
901 info!("Fallback chain succeeded");
902 if let Some(ref healing) = config.healing_engine {
904 let agent_id = llm.provider_name().to_string();
905 let mut engine = healing.lock().unwrap();
906 engine.record_success(&agent_id);
907 }
908 if let Some(ref budget) = config.token_budget {
910 if let Some(usage) = &r.usage {
911 let mut b = budget.lock().unwrap();
912 b.record_usage(usage.total_tokens);
913 }
914 }
915 r
916 }
917 Err(fallback_e) => {
918 warn!(error = %fallback_e, "Fallback chain also failed");
919 let _ = audit_log.append(
920 AuditEventType::Error,
921 "llm",
922 &format!("All providers failed: {}", fallback_e),
923 None,
924 );
925 return Err(crate::error::RavenClawsError::Llm(fallback_e));
927 }
928 }
929 } else {
930 warn!(error = %e, "LLM request failed after retries");
931 let _ = audit_log.append(
932 AuditEventType::Error,
933 "llm",
934 &format!("LLM request failed after retries: {}", e),
935 None,
936 );
937 return Err(crate::error::RavenClawsError::Llm(e));
939 }
940 }
941 };
942
943 let mut iteration_tokens: u64 = 0;
945 if let Some(ref budget) = config.token_budget {
946 if let Some(usage) = &response.usage {
947 let mut b = budget.lock().unwrap();
948 b.record_usage(usage.total_tokens);
949 iteration_tokens = usage.total_tokens as u64;
950 debug!(
951 iteration = iteration,
952 tokens_used = usage.total_tokens,
953 total_used = b.used_tokens,
954 remaining = b.remaining(),
955 "Token usage recorded"
956 );
957 }
958 } else if let Some(usage) = &response.usage {
959 iteration_tokens = usage.total_tokens as u64;
960 }
961
962 if let Some(ref cb) = config.metrics_callback {
964 cb(iteration_tokens, 0);
965 }
966
967 if let Some(ref rf) = config.ravenfabric {
969 if rf.is_enabled() {
970 let _ = rf.health().await;
971 info!(
972 iteration = iteration,
973 ravenfabric = true,
974 "RavenFabric health check completed"
975 );
976 }
977 }
978
979 let first_choice = response.choices.first();
980 let content = first_choice
981 .map(|c| c.message.content.clone())
982 .unwrap_or_default();
983
984 debug!(
985 iteration = iteration,
986 response_length = content.len(),
987 response_preview = %content[..content.len().min(500)],
988 "LLM response received"
989 );
990
991 if let Some(ref detector) = injection_detector {
993 match detector.check(&content) {
994 crate::policy::InjectionVerdict::Suspicious(reason) => {
995 warn!(
996 iteration = iteration,
997 reason = %reason,
998 "Prompt-injection detected in LLM response"
999 );
1000 let _ = audit_log.append(
1001 AuditEventType::SecurityViolation,
1002 "injection_detector",
1003 &format!("Prompt-injection detected: {}", reason),
1004 Some(serde_json::json!({
1005 "reason": reason,
1006 "iteration": iteration,
1007 "content_preview": &content[..content.len().min(200)],
1008 })),
1009 );
1010 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
1012 delete_checkpoint(checkpoint_dir, &session_id);
1013 }
1014 return Err(crate::error::RavenClawsError::SecurityViolation(format!(
1015 "LLM response blocked: potential prompt injection ({})",
1016 reason
1017 )));
1018 }
1019 crate::policy::InjectionVerdict::Clean => {}
1020 }
1021 }
1022
1023 if config.enable_tools {
1025 if let Some((tool_name, args)) = first_choice.and_then(parse_structured_tool_call) {
1026 info!(tool = %tool_name, "Structured tool call detected");
1027
1028 if let Some(tool_result) = execute_parsed_tool_call(
1030 tool_name,
1031 args,
1032 ®istry,
1033 &policy_engine,
1034 &sandbox,
1035 &audit_log,
1036 config.require_approval,
1037 )
1038 .await
1039 {
1040 let observation = if tool_result.success {
1041 format!("OBSERVATION: {}", tool_result.output)
1042 } else {
1043 format!(
1044 "OBSERVATION: Tool failed with error: {}",
1045 tool_result.error.as_deref().unwrap_or("unknown error")
1046 )
1047 };
1048
1049 memory.add_user_message(&observation);
1050
1051 if let Some(ref cb) = config.metrics_callback {
1053 cb(0, 1);
1054 }
1055
1056 info!(
1057 iteration = iteration,
1058 tool = %tool_result.tool_name,
1059 success = tool_result.success,
1060 "Structured tool executed"
1061 );
1062 continue;
1063 }
1064 }
1065 }
1066
1067 if content.contains("FINAL:") {
1069 let final_response = content
1070 .split("FINAL:")
1071 .nth(1)
1072 .unwrap_or("")
1073 .trim()
1074 .to_string();
1075
1076 memory.add_assistant_message(&content);
1077
1078 let _ = audit_log.append(
1080 AuditEventType::AgentFinish,
1081 "agent",
1082 "Agent loop completed successfully",
1083 Some(serde_json::json!({
1084 "iterations": iteration + 1,
1085 "final_response_length": final_response.len(),
1086 })),
1087 );
1088
1089 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
1091 delete_checkpoint(checkpoint_dir, &session_id);
1092 }
1093
1094 return Ok(final_response);
1095 }
1096
1097 if config.enable_tools {
1099 if let Some(tool_result) = execute_tool_call_with_security(
1100 &content,
1101 ®istry,
1102 &policy_engine,
1103 &sandbox,
1104 &audit_log,
1105 )
1106 .await
1107 {
1108 let observation = if tool_result.success {
1109 format!("OBSERVATION: {}", tool_result.output)
1110 } else {
1111 format!(
1112 "OBSERVATION: Tool failed with error: {}",
1113 tool_result.error.as_deref().unwrap_or("unknown error")
1114 )
1115 };
1116
1117 memory.add_assistant_message(&content);
1118 memory.add_user_message(&observation);
1119
1120 if let Some(ref cb) = config.metrics_callback {
1122 cb(0, 1);
1123 }
1124
1125 info!(
1126 iteration = iteration,
1127 tool = %tool_result.tool_name,
1128 success = tool_result.success,
1129 "Tool executed"
1130 );
1131 continue;
1132 }
1133 }
1134
1135 memory.add_assistant_message(&content);
1137
1138 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
1140 let checkpoint = CheckpointState::new(
1141 session_id.clone(),
1142 iteration,
1143 config.max_iterations,
1144 memory.history().to_vec(),
1145 initial_prompt,
1146 system_prompt,
1147 llm.provider_name(),
1148 llm.model(),
1149 config.enable_tools,
1150 );
1151 if let Err(e) = save_checkpoint(checkpoint_dir, &checkpoint) {
1152 warn!(
1153 session_id = %session_id,
1154 iteration = iteration,
1155 error = %e,
1156 "Failed to save checkpoint"
1157 );
1158 } else {
1159 debug!(
1160 session_id = %session_id,
1161 iteration = iteration,
1162 "Checkpoint saved"
1163 );
1164 }
1165 }
1166
1167 if config.no_final_required {
1169 info!(
1170 iteration = iteration,
1171 response_length = content.len(),
1172 "no_final_required: treating response as completion"
1173 );
1174 let _ = audit_log.append(
1175 AuditEventType::AgentFinish,
1176 "agent",
1177 "Agent loop completed (no_final_required)",
1178 Some(serde_json::json!({
1179 "iterations": iteration + 1,
1180 "final_response_length": content.len(),
1181 })),
1182 );
1183 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
1185 delete_checkpoint(checkpoint_dir, &session_id);
1186 }
1187 return Ok(content);
1188 }
1189
1190 info!(
1191 iteration = iteration,
1192 thought = %content.lines().find(|l| l.starts_with("THOUGHT:")).unwrap_or("<no thought>"),
1193 "Agent loop progress"
1194 );
1195 }
1196
1197 warn!(
1199 max_iterations = config.max_iterations,
1200 "Agent loop reached max iterations"
1201 );
1202
1203 let _ = audit_log.append(
1204 AuditEventType::Error,
1205 "agent",
1206 "Agent loop reached max iterations without completing",
1207 Some(serde_json::json!({
1208 "max_iterations": config.max_iterations,
1209 })),
1210 );
1211
1212 if let Some(ref checkpoint_dir) = config.checkpoint_dir {
1214 delete_checkpoint(checkpoint_dir, &session_id);
1215 }
1216
1217 let history = memory.history();
1218 if history.len() > 1 {
1219 if let Some(last) = history.last() {
1220 return Ok(last.content.clone());
1221 }
1222 }
1223
1224 Err(crate::error::RavenClawsError::CommandExecution(
1225 "Agent loop reached max iterations without completing the task".to_string(),
1226 ))
1227}
1228
1229#[allow(dead_code)]
1235#[instrument(skip_all, fields(provider = %llm.provider_name(), model = %llm.model()))]
1236pub async fn run_agent_loop_with_mcp(
1237 llm: Arc<dyn LLMProviderTrait>,
1238 initial_prompt: &str,
1239 system_prompt: &str,
1240 config: AgentLoopConfig,
1241 mcp_client: Option<Arc<RwLock<McpClient>>>,
1242) -> Result<String> {
1243 run_agent_loop_with_mcp_and_registry(
1244 llm,
1245 initial_prompt,
1246 system_prompt,
1247 config,
1248 mcp_client,
1249 None,
1250 )
1251 .await
1252}
1253
1254#[instrument(skip_all, fields(provider = %llm.provider_name(), model = %llm.model()))]
1256pub async fn run_agent_loop_with_mcp_and_registry(
1257 llm: Arc<dyn LLMProviderTrait>,
1258 initial_prompt: &str,
1259 system_prompt: &str,
1260 config: AgentLoopConfig,
1261 mcp_client: Option<Arc<RwLock<McpClient>>>,
1262 tool_registry: Option<ToolRegistry>,
1263) -> Result<String> {
1264 let mut registry = tool_registry.unwrap_or_else(ToolRegistry::with_default_tools);
1266
1267 if let Some(client) = &mcp_client {
1269 match crate::mcp::register_mcp_tools(&mut registry, client.clone()).await {
1270 Ok(count) => {
1271 info!(count, "MCP tools registered");
1272 }
1273 Err(e) => {
1274 warn!(error = %e, "Failed to register MCP tools");
1275 }
1276 }
1277 }
1278
1279 let mcp_enabled = mcp_client.is_some();
1280 run_agent_loop_inner(
1281 llm,
1282 initial_prompt,
1283 system_prompt,
1284 config,
1285 registry,
1286 "MCP integration",
1287 mcp_enabled,
1288 Vec::new(),
1289 )
1290 .await
1291}
1292
1293#[instrument(skip_all, fields(provider = %llm.provider_name(), model = %llm.model(), image_count = image_data_uris.len()))]
1295pub async fn run_agent_loop_with_mcp_and_images(
1296 llm: Arc<dyn LLMProviderTrait>,
1297 initial_prompt: &str,
1298 system_prompt: &str,
1299 config: AgentLoopConfig,
1300 mcp_client: Option<Arc<RwLock<McpClient>>>,
1301 tool_registry: Option<ToolRegistry>,
1302 image_data_uris: Vec<String>,
1303) -> Result<String> {
1304 let mut registry = tool_registry.unwrap_or_else(ToolRegistry::with_default_tools);
1305
1306 if let Some(client) = &mcp_client {
1307 match crate::mcp::register_mcp_tools(&mut registry, client.clone()).await {
1308 Ok(count) => {
1309 info!(count, "MCP tools registered");
1310 }
1311 Err(e) => {
1312 warn!(error = %e, "Failed to register MCP tools");
1313 }
1314 }
1315 }
1316
1317 let mcp_enabled = mcp_client.is_some();
1318 run_agent_loop_inner(
1319 llm,
1320 initial_prompt,
1321 system_prompt,
1322 config,
1323 registry,
1324 "MCP integration",
1325 mcp_enabled,
1326 image_data_uris,
1327 )
1328 .await
1329}
1330
1331async fn prompt_for_approval(tool_name: &str, args: &serde_json::Value) -> bool {
1336 use std::io::{IsTerminal, Write};
1337
1338 let args_str = serde_json::to_string_pretty(args).unwrap_or_default();
1339
1340 if !std::io::stdin().is_terminal() {
1342 warn!(
1343 tool = %tool_name,
1344 "stdin is not a TTY — auto-approving tool call (use --require-approval only in interactive mode)"
1345 );
1346 return true;
1347 }
1348
1349 eprintln!("\n⚠️ Tool requires approval:");
1351 eprintln!(" Tool: {}", tool_name);
1352 for line in args_str.lines() {
1353 eprintln!(" {}", line);
1354 }
1355 eprint!(" Approve? [y/N] ");
1356 std::io::stderr().flush().ok();
1357
1358 let mut input = String::new();
1359 match std::io::stdin().read_line(&mut input) {
1360 Ok(_) => {
1361 let trimmed = input.trim().to_lowercase();
1362 trimmed == "y" || trimmed == "yes"
1363 }
1364 Err(e) => {
1365 warn!(error = %e, "Failed to read approval input — denying by default");
1366 false
1367 }
1368 }
1369}
1370
1371#[cfg(test)]
1374async fn prompt_for_approval_with_input(
1375 tool_name: &str,
1376 args: &serde_json::Value,
1377 input: &str,
1378) -> bool {
1379 use std::io::Write;
1380
1381 let args_str = serde_json::to_string_pretty(args).unwrap_or_default();
1382
1383 eprintln!("\n⚠️ Tool requires approval:");
1384 eprintln!(" Tool: {}", tool_name);
1385 for line in args_str.lines() {
1386 eprintln!(" {}", line);
1387 }
1388 eprint!(" Approve? [y/N] ");
1389 std::io::stderr().flush().ok();
1390
1391 let trimmed = input.trim().to_lowercase();
1392 trimmed == "y" || trimmed == "yes"
1393}
1394
1395async fn execute_parsed_tool_call(
1404 tool_name: String,
1405 args: serde_json::Value,
1406 registry: &ToolRegistry,
1407 policy_engine: &PolicyEngine,
1408 _sandbox: &Sandbox,
1409 audit_log: &AuditLog,
1410 require_approval: bool,
1411) -> Option<ToolResult> {
1412 info!(tool = %tool_name, "Executing parsed tool call");
1413
1414 let _ = audit_log.tool_call(&tool_name, &args);
1416
1417 if require_approval && policy_engine.requires_approval(&tool_name) {
1419 let _ = audit_log.append(
1420 AuditEventType::ApprovalRequested,
1421 "approval",
1422 &format!("Approval required for tool: {}", tool_name),
1423 Some(serde_json::json!({"tool": tool_name, "args": args})),
1424 );
1425
1426 let granted = prompt_for_approval(&tool_name, &args).await;
1428
1429 if !granted {
1430 let _ = audit_log.approval(&tool_name, false, Some("Denied by user"));
1431 warn!(tool = %tool_name, "Tool call denied by user");
1432 return Some(ToolResult {
1433 tool_name: tool_name.clone(),
1434 success: false,
1435 output: String::new(),
1436 error: Some(format!("Approval denied by user for tool: {}", tool_name)),
1437 exit_code: Some(-1),
1438 duration_ms: None,
1439 });
1440 }
1441
1442 let _ = audit_log.approval(&tool_name, true, Some("Approved by user"));
1443 info!(tool = %tool_name, "Tool call approved by user");
1444 }
1445
1446 let policy_decision = policy_engine.check_tool_call(&tool_name, &args);
1448
1449 match &policy_decision {
1451 Decision::Allow => {
1452 let _ = audit_log.policy_decision(&tool_name, true, None);
1453 }
1454 Decision::Deny(reason) => {
1455 let _ = audit_log.policy_decision(&tool_name, false, Some(reason));
1456 warn!(tool = %tool_name, reason = %reason, "Tool call denied by policy");
1457 return Some(ToolResult {
1458 tool_name: tool_name.clone(),
1459 success: false,
1460 output: String::new(),
1461 error: Some(format!("Policy denied: {}", reason)),
1462 exit_code: Some(-1),
1463 duration_ms: None,
1464 });
1465 }
1466 }
1467
1468 let tool_name_clone = tool_name.clone();
1470 let call = ToolCall {
1471 name: tool_name.clone(),
1472 arguments: args,
1473 id: None,
1474 };
1475
1476 let result = match registry.execute(call).await {
1477 Ok(result) => {
1478 let _ = audit_log.append(
1480 AuditEventType::ToolResult,
1481 &tool_name_clone,
1482 &format!(
1483 "Tool executed: {} (success: {})",
1484 tool_name_clone, result.success
1485 ),
1486 Some(serde_json::json!({
1487 "success": result.success,
1488 "exit_code": result.exit_code,
1489 "duration_ms": result.duration_ms,
1490 })),
1491 );
1492 result
1493 }
1494 Err(e) => {
1495 let _ = audit_log.append(
1497 AuditEventType::Error,
1498 &tool_name_clone,
1499 &format!("Tool execution failed: {}", e),
1500 None,
1501 );
1502 ToolResult {
1503 tool_name: tool_name_clone,
1504 success: false,
1505 output: String::new(),
1506 error: Some(e.to_string()),
1507 exit_code: Some(-1),
1508 duration_ms: None,
1509 }
1510 }
1511 };
1512
1513 Some(result)
1514}
1515
1516async fn execute_tool_call_with_security(
1525 content: &str,
1526 registry: &ToolRegistry,
1527 policy_engine: &PolicyEngine,
1528 _sandbox: &Sandbox,
1529 audit_log: &AuditLog,
1530) -> Option<ToolResult> {
1531 let (tool_name, args) = parse_tool_call(content)?;
1533
1534 execute_parsed_tool_call(
1536 tool_name,
1537 args,
1538 registry,
1539 policy_engine,
1540 _sandbox,
1541 audit_log,
1542 false, )
1544 .await
1545}
1546
1547fn parse_structured_tool_call(choice: &Choice) -> Option<(String, serde_json::Value)> {
1551 let tool_calls = choice.tool_calls.as_ref()?;
1552 let first_call = tool_calls.first()?;
1553
1554 let tool_name = first_call.function.name.clone();
1555 let args: serde_json::Value = serde_json::from_str(&first_call.function.arguments).ok()?;
1556
1557 Some((tool_name, args))
1558}
1559
1560fn parse_tool_call(content: &str) -> Option<(String, serde_json::Value)> {
1562 let mut lines = content.lines();
1563 let tool_call_line = lines.find(|l| l.trim().starts_with("TOOL_CALL:"))?;
1564
1565 let tool_name = tool_call_line
1566 .trim()
1567 .strip_prefix("TOOL_CALL:")
1568 .map(|s| s.trim())
1569 .filter(|s| !s.is_empty())?
1570 .to_string();
1571
1572 let args_line = lines.find(|l| l.trim().starts_with("ARGS:"))?;
1574 let args_str = args_line.trim().strip_prefix("ARGS:").map(|s| s.trim())?;
1575
1576 let args: serde_json::Value = serde_json::from_str(args_str).ok()?;
1577
1578 Some((tool_name, args))
1579}
1580
1581pub async fn run_single(
1583 llm: Arc<dyn LLMProviderTrait>,
1584 config: Config,
1585 ravenfabric: Option<RavenFabricClient>,
1586) -> Result<()> {
1587 info!(
1588 "Starting single agent mode with provider: {}",
1589 llm.provider_name()
1590 );
1591
1592 if let Some(ref rf) = ravenfabric {
1594 if rf.is_enabled() {
1595 info!("RavenFabric remote execution available");
1596 match rf.health().await {
1597 Ok(true) => info!("RavenFabric mesh is healthy"),
1598 Ok(false) => warn!("RavenFabric mesh returned unhealthy status"),
1599 Err(e) => warn!(error = %e, "RavenFabric health check failed"),
1600 }
1601 }
1602 }
1603
1604 let system_prompt = &config.llm.system_prompt;
1605
1606 let messages = vec![
1607 ChatMessage::new("system", system_prompt.to_string()),
1608 ChatMessage::new("user", "Ready. Awaiting instructions."),
1609 ];
1610
1611 match llm.chat(messages).await {
1612 Ok(response) => {
1613 if let Some(choice) = response.choices.first() {
1614 info!(provider = llm.provider_name(), model = llm.model(), response = %choice.message.content, "Agent response received");
1615
1616 if let Some(ref rf) = ravenfabric {
1618 if rf.is_enabled() {
1619 let preview = choice.message.content.chars().take(500).collect::<String>();
1620 let _ = rf.broadcast(&preview, 30).await;
1621 info!("Agent result broadcast to RavenFabric mesh");
1622 }
1623 }
1624 }
1625 }
1626 Err(e) => {
1627 warn!(error = %e, provider = llm.provider_name(), "LLM request failed");
1628 }
1629 }
1630
1631 Ok(())
1632}
1633
1634pub async fn run_swarm(
1639 llm: Arc<dyn LLMProviderTrait>,
1640 config: Config,
1641 ravenfabric: Option<RavenFabricClient>,
1642) -> Result<()> {
1643 info!("Starting swarm mode (single-provider) — 3 parallel agents");
1644
1645 if let Some(ref rf) = ravenfabric {
1647 if rf.is_enabled() {
1648 info!("RavenFabric remote execution available for swarm coordination");
1649 match rf.health().await {
1650 Ok(true) => info!("RavenFabric mesh is healthy"),
1651 Ok(false) => warn!("RavenFabric mesh returned unhealthy status"),
1652 Err(e) => warn!(error = %e, "RavenFabric health check failed"),
1653 }
1654 }
1655 }
1656
1657 let _system_prompt = &config.llm.system_prompt;
1658 let num_agents = 3;
1659 let mut handles = Vec::new();
1660
1661 let personas = [
1663 "You are an analytical agent. Focus on logic, structure, and precision.",
1664 "You are a creative agent. Focus on innovation, alternatives, and possibilities.",
1665 "You are a pragmatic agent. Focus on simplicity, efficiency, and practicality.",
1666 ];
1667
1668 for (i, persona) in personas.iter().enumerate().take(num_agents) {
1669 let llm_clone = llm.clone();
1670 let persona = persona.to_string();
1671 let task = "Analyze the given task and provide your solution.".to_string();
1672
1673 let handle = tokio::spawn(async move {
1674 let mut memory = ConversationMemory::new(&persona, 10);
1675 memory.add_user_message(&task);
1676
1677 let messages = memory.history().to_vec();
1678 match llm_clone.chat(messages).await {
1679 Ok(response) => {
1680 let content = response
1681 .choices
1682 .first()
1683 .map(|c| c.message.content.clone())
1684 .unwrap_or_default();
1685 Ok((i, content))
1686 }
1687 Err(e) => Err(format!("Agent {} failed: {}", i, e)),
1688 }
1689 });
1690
1691 handles.push(handle);
1692 }
1693
1694 let mut results: Vec<(usize, String)> = Vec::new();
1696 for handle in handles {
1697 match handle.await {
1698 Ok(Ok((idx, result))) => {
1699 info!("Agent {} completed: {} chars", idx, result.len());
1700 results.push((idx, result));
1701 }
1702 Ok(Err(e)) => warn!("Agent failed: {}", e),
1703 Err(e) => warn!("Agent join failed: {}", e),
1704 }
1705 }
1706
1707 println!("\n🐦⬛ Swarm Results ({} agents):", results.len());
1709 for (idx, result) in &results {
1710 println!(
1711 "\n── Agent {} ({}) ──",
1712 idx + 1,
1713 personas[*idx].split('.').next().unwrap_or("Unknown")
1714 );
1715 println!("{}", result);
1716 }
1717
1718 if let Some(ref rf) = ravenfabric {
1720 if rf.is_enabled() {
1721 let summary = format!(
1722 "Swarm completed: {} agents, results: {}",
1723 results.len(),
1724 results
1725 .iter()
1726 .map(|(i, r)| format!("Agent {}: {} chars", i, r.len()))
1727 .collect::<Vec<_>>()
1728 .join(", ")
1729 );
1730 let _ = rf.broadcast(&summary, 30).await;
1731 info!("Swarm results broadcast to RavenFabric mesh");
1732 }
1733 }
1734
1735 Ok(())
1736}
1737
1738pub async fn run_supervisor(
1743 llm: Arc<dyn LLMProviderTrait>,
1744 config: Config,
1745 ravenfabric: Option<RavenFabricClient>,
1746) -> Result<()> {
1747 info!("Starting supervisor mode (single-provider)");
1748
1749 if let Some(ref rf) = ravenfabric {
1751 if rf.is_enabled() {
1752 info!("RavenFabric remote execution available for supervisor coordination");
1753 match rf.health().await {
1754 Ok(true) => info!("RavenFabric mesh is healthy"),
1755 Ok(false) => warn!("RavenFabric mesh returned unhealthy status"),
1756 Err(e) => warn!(error = %e, "RavenFabric health check failed"),
1757 }
1758 }
1759 }
1760
1761 let system_prompt = &config.llm.system_prompt;
1762 let policy_engine = PolicyEngine::default_secure();
1763 let mut sandbox = Sandbox::default();
1764 sandbox.init().await.map_err(|e| {
1765 crate::error::RavenClawsError::CommandExecution(format!("Sandbox init failed: {}", e))
1766 })?;
1767 let audit_log = AuditLog::new(format!("supervisor-{}", std::process::id()));
1768 let registry = ToolRegistry::with_default_tools();
1769
1770 let supervisor_prompt = format!(
1772 "You are a supervisor agent. Your task is to decompose complex tasks into subtasks \
1773 and coordinate sub-agents to complete them. \
1774 \n\nFor each subtask, respond with:\n\
1775 SUBTASK: <description>\n\
1776 AGENT: <agent_number>\n\
1777 \nWhen all subtasks are complete, respond with:\n\
1778 FINAL: <aggregated result>\n\
1779 \nTask: {}",
1780 "Coordinate the completion of the assigned task."
1781 );
1782
1783 let mut memory = ConversationMemory::new(system_prompt, 20);
1784 memory.add_user_message(&supervisor_prompt);
1785
1786 let mut subtask_results: Vec<String> = Vec::new();
1787 let mut iteration = 0;
1788 let max_iterations = 15;
1789
1790 loop {
1791 iteration += 1;
1792 if iteration > max_iterations {
1793 warn!("Supervisor reached max iterations");
1794 break;
1795 }
1796
1797 let messages = memory.history().to_vec();
1798 let response = match llm.chat(messages).await {
1799 Ok(r) => r,
1800 Err(e) => {
1801 warn!(error = %e, "Supervisor LLM request failed");
1802 continue;
1803 }
1804 };
1805
1806 let content = response
1807 .choices
1808 .first()
1809 .map(|c| c.message.content.clone())
1810 .unwrap_or_default();
1811
1812 if content.contains("FINAL:") {
1814 let final_response = content
1815 .split("FINAL:")
1816 .nth(1)
1817 .unwrap_or("")
1818 .trim()
1819 .to_string();
1820 info!("Supervisor completed task: {} chars", final_response.len());
1821
1822 let _ = audit_log.append(
1823 AuditEventType::AgentFinish,
1824 "supervisor",
1825 "Supervisor completed task coordination",
1826 Some(serde_json::json!({
1827 "iterations": iteration,
1828 "subtasks_completed": subtask_results.len(),
1829 })),
1830 );
1831
1832 println!("\n🐦⬛ Supervisor Result:\n{}", final_response);
1833 return Ok(());
1834 }
1835
1836 if content.contains("SUBTASK:") {
1838 let subtask_block = content.split("SUBTASK:").nth(1).unwrap_or("");
1839 let subtask_lines: Vec<&str> = subtask_block.lines().take(3).collect();
1840
1841 let subtask_desc = subtask_lines.first().unwrap_or(&"").trim();
1842 let agent_num = subtask_lines
1843 .iter()
1844 .find(|l| l.starts_with("AGENT:"))
1845 .and_then(|l| l.split(':').nth(1))
1846 .unwrap_or("1")
1847 .trim();
1848
1849 if !subtask_desc.is_empty() {
1850 info!("Subtask {}: {}", agent_num, subtask_desc);
1851
1852 let subtask_result = run_subtask_agent(
1854 llm.clone(),
1855 subtask_desc,
1856 system_prompt,
1857 &policy_engine,
1858 &sandbox,
1859 &audit_log,
1860 ®istry,
1861 )
1862 .await;
1863
1864 match subtask_result {
1865 Ok(result) => {
1866 info!("Subtask {} completed: {} chars", agent_num, result.len());
1867 subtask_results.push(format!("Agent {} result: {}", agent_num, result));
1868
1869 memory.add_assistant_message(&format!(
1870 "Decomposed subtask {}: {}",
1871 agent_num, subtask_desc
1872 ));
1873 memory
1874 .add_user_message(&format!("Subtask {} result: {}", agent_num, result));
1875 }
1876 Err(e) => {
1877 warn!("Subtask {} failed: {}", agent_num, e);
1878 memory
1879 .add_assistant_message(&format!("Subtask {} failed: {}", agent_num, e));
1880 }
1881 }
1882 }
1883 } else {
1884 memory.add_assistant_message(&content);
1885 }
1886 }
1887
1888 if !subtask_results.is_empty() {
1890 let aggregated = subtask_results.join("\n\n");
1891 info!(
1892 "Supervisor aggregated {} subtask results",
1893 subtask_results.len()
1894 );
1895
1896 if let Some(ref rf) = ravenfabric {
1898 if rf.is_enabled() {
1899 let summary = format!(
1900 "Supervisor completed: {} subtasks, result: {} chars",
1901 subtask_results.len(),
1902 aggregated.len()
1903 );
1904 let _ = rf.broadcast(&summary, 30).await;
1905 info!("Supervisor result broadcast to RavenFabric mesh");
1906 }
1907 }
1908
1909 println!("\n🐦⬛ Supervisor Aggregated Result:\n{}", aggregated);
1910 return Ok(());
1911 }
1912
1913 Err(crate::error::RavenClawsError::CommandExecution(
1914 "Supervisor mode completed without results".to_string(),
1915 ))
1916}
1917
1918async fn run_subtask_agent(
1920 llm: Arc<dyn LLMProviderTrait>,
1921 subtask: &str,
1922 system_prompt: &str,
1923 policy_engine: &PolicyEngine,
1924 sandbox: &Sandbox,
1925 audit_log: &AuditLog,
1926 registry: &ToolRegistry,
1927) -> Result<String> {
1928 let mut memory = ConversationMemory::new(system_prompt, 10);
1929 memory.add_user_message(&format!("Execute this subtask: {}", subtask));
1930
1931 for i in 0..5 {
1932 let messages = memory.history().to_vec();
1933 let response = match llm.chat(messages).await {
1934 Ok(r) => r,
1935 Err(e) => {
1936 warn!(error = %e, iteration = i, "Subtask agent LLM failed");
1937 continue;
1938 }
1939 };
1940
1941 let content = response
1942 .choices
1943 .first()
1944 .map(|c| c.message.content.clone())
1945 .unwrap_or_default();
1946
1947 if content.contains("FINAL:") || content.contains("DONE:") {
1948 return Ok(content
1949 .replace("FINAL:", "")
1950 .replace("DONE:", "")
1951 .trim()
1952 .to_string());
1953 }
1954
1955 if let Some(tool_result) =
1957 execute_tool_call_with_security(&content, registry, policy_engine, sandbox, audit_log)
1958 .await
1959 {
1960 memory.add_assistant_message(&content);
1961 memory.add_user_message(&format!("Tool result: {}", tool_result.output));
1962 } else {
1963 memory.add_assistant_message(&content);
1964 memory.add_user_message("Continue with next step.");
1965 }
1966 }
1967
1968 Ok("Subtask completed".to_string())
1969}
1970
1971pub async fn run_single_multi(
1973 multi_llm: MultiModelManager,
1974 config: Config,
1975 ravenfabric: Option<RavenFabricClient>,
1976) -> Result<()> {
1977 info!(
1978 "Starting single agent mode (multi-model) with {} providers",
1979 multi_llm.client_count()
1980 );
1981
1982 if let Some(ref rf) = ravenfabric {
1984 if rf.is_enabled() {
1985 info!("RavenFabric remote execution available");
1986 match rf.health().await {
1987 Ok(true) => info!("RavenFabric mesh is healthy"),
1988 Ok(false) => warn!("RavenFabric mesh returned unhealthy status"),
1989 Err(e) => warn!(error = %e, "RavenFabric health check failed"),
1990 }
1991 }
1992 }
1993
1994 let system_prompt = &config.llm.system_prompt;
1995
1996 let messages = vec![
1997 ChatMessage::new("system", system_prompt.to_string()),
1998 ChatMessage::new("user", "Ready. Awaiting instructions."),
1999 ];
2000
2001 let mut last_index = 0;
2003 for i in 0..multi_llm.client_count() {
2004 let client = if i == 0 {
2005 multi_llm.get_client(0)
2006 } else {
2007 multi_llm.next_client(last_index)
2008 };
2009
2010 if let Some(client) = client {
2011 match client.chat(messages.clone()).await {
2012 Ok(response) => {
2013 if let Some(choice) = response.choices.first() {
2014 info!(provider = client.provider_name(), model = client.model(), response = %choice.message.content, "Provider response received");
2015 }
2016 }
2017 Err(e) => {
2018 warn!(error = %e, provider = client.provider_name(), model = client.model(), "Provider request failed");
2019 }
2020 }
2021 last_index = i;
2022 }
2023 }
2024
2025 if let Some(ref rf) = ravenfabric {
2027 if rf.is_enabled() {
2028 let _ = rf
2029 .broadcast("Single agent (multi-model) completed", 30)
2030 .await;
2031 info!("Multi-model result broadcast to RavenFabric mesh");
2032 }
2033 }
2034
2035 Ok(())
2036}
2037
2038pub async fn run_swarm_multi(
2043 multi_llm: MultiModelManager,
2044 config: Config,
2045 ravenfabric: Option<RavenFabricClient>,
2046) -> Result<()> {
2047 info!(
2048 "Starting swarm mode (multi-model) — {} parallel agents",
2049 multi_llm.client_count()
2050 );
2051
2052 if let Some(ref rf) = ravenfabric {
2054 if rf.is_enabled() {
2055 info!("RavenFabric remote execution available for swarm coordination");
2056 match rf.health().await {
2057 Ok(true) => info!("RavenFabric mesh is healthy"),
2058 Ok(false) => warn!("RavenFabric mesh returned unhealthy status"),
2059 Err(e) => warn!(error = %e, "RavenFabric health check failed"),
2060 }
2061 }
2062 }
2063
2064 let _system_prompt = &config.llm.system_prompt;
2065 let num_agents = multi_llm.client_count().min(3); let mut handles = Vec::new();
2067
2068 let personas = [
2070 "You are an analytical agent. Focus on logic, structure, and precision.",
2071 "You are a creative agent. Focus on innovation, alternatives, and possibilities.",
2072 "You are a pragmatic agent. Focus on simplicity, efficiency, and practicality.",
2073 ];
2074
2075 for i in 0..num_agents {
2076 let client = multi_llm.get_client(i).unwrap().clone();
2077 let persona = personas.get(i).unwrap_or(&personas[0]).to_string();
2078 let task = "Analyze the given task and provide your solution.".to_string();
2079
2080 let handle = tokio::spawn(async move {
2081 let mut memory = ConversationMemory::new(&persona, 10);
2082 memory.add_user_message(&task);
2083
2084 let messages = memory.history().to_vec();
2085 match client.chat(messages).await {
2086 Ok(response) => {
2087 let content = response
2088 .choices
2089 .first()
2090 .map(|c| c.message.content.clone())
2091 .unwrap_or_default();
2092 Ok((
2093 i,
2094 client.provider_name().to_string(),
2095 client.model().to_string(),
2096 content,
2097 ))
2098 }
2099 Err(e) => Err(format!("Agent {} failed: {}", i, e)),
2100 }
2101 });
2102
2103 handles.push(handle);
2104 }
2105
2106 let mut results: Vec<(usize, String, String, String)> = Vec::new();
2108 for handle in handles {
2109 match handle.await {
2110 Ok(Ok((idx, provider, model, result))) => {
2111 info!(
2112 "Agent {} ({}:{}) completed: {} chars",
2113 idx,
2114 provider,
2115 model,
2116 result.len()
2117 );
2118 results.push((idx, provider, model, result));
2119 }
2120 Ok(Err(e)) => warn!("Agent failed: {}", e),
2121 Err(e) => warn!("Agent join failed: {}", e),
2122 }
2123 }
2124
2125 println!(
2127 "\n🐦⬛ Swarm Results ({} agents, multi-model):",
2128 results.len()
2129 );
2130 for (idx, provider, model, result) in &results {
2131 println!("\n── Agent {} ({}:{}) ──", idx + 1, provider, model);
2132 println!("{}", result);
2133 }
2134
2135 if let Some(ref rf) = ravenfabric {
2137 if rf.is_enabled() {
2138 let summary = format!("Multi-model swarm completed: {} agents", results.len());
2139 let _ = rf.broadcast(&summary, 30).await;
2140 info!("Multi-model swarm results broadcast to RavenFabric mesh");
2141 }
2142 }
2143
2144 Ok(())
2145}
2146
2147pub async fn run_supervisor_multi(
2152 multi_llm: MultiModelManager,
2153 config: Config,
2154 ravenfabric: Option<RavenFabricClient>,
2155) -> Result<()> {
2156 info!(
2157 "Starting supervisor mode (multi-model) with {} providers",
2158 multi_llm.client_count()
2159 );
2160
2161 if let Some(ref rf) = ravenfabric {
2163 if rf.is_enabled() {
2164 info!("RavenFabric remote execution available for supervisor coordination");
2165 match rf.health().await {
2166 Ok(true) => info!("RavenFabric mesh is healthy"),
2167 Ok(false) => warn!("RavenFabric mesh returned unhealthy status"),
2168 Err(e) => warn!(error = %e, "RavenFabric health check failed"),
2169 }
2170 }
2171 }
2172
2173 let system_prompt = &config.llm.system_prompt;
2174 let policy_engine = PolicyEngine::default_secure();
2175 let mut sandbox = Sandbox::default();
2176 sandbox.init().await.map_err(|e| {
2177 crate::error::RavenClawsError::CommandExecution(format!("Sandbox init failed: {}", e))
2178 })?;
2179 let audit_log = AuditLog::new(format!("supervisor-multi-{}", std::process::id()));
2180 let registry = ToolRegistry::with_default_tools();
2181
2182 let supervisor_prompt = format!(
2184 "You are a supervisor agent coordinating multiple LLM providers. \
2185 Decompose tasks and assign them to appropriate providers based on their strengths. \
2186 \n\nFor each subtask, respond with:\n\
2187 SUBTASK: <description>\n\
2188 PROVIDER: <provider_index 0-{}>\n\
2189 \nWhen complete, respond with:\n\
2190 FINAL: <aggregated result>\n\
2191 \nTask: {}",
2192 multi_llm.client_count() - 1,
2193 "Coordinate the completion of the assigned task using available providers."
2194 );
2195
2196 let mut memory = ConversationMemory::new(system_prompt, 20);
2197 memory.add_user_message(&supervisor_prompt);
2198
2199 let mut subtask_results: Vec<String> = Vec::new();
2200 let mut iteration = 0;
2201 let max_iterations = 15;
2202
2203 loop {
2204 iteration += 1;
2205 if iteration > max_iterations {
2206 warn!("Supervisor reached max iterations");
2207 break;
2208 }
2209
2210 let supervisor_client = multi_llm
2212 .get_client(iteration % multi_llm.client_count())
2213 .or_else(|| multi_llm.get_client(0))
2214 .cloned();
2215
2216 let messages = memory.history().to_vec();
2217 let response =
2218 match supervisor_client.map(|c| tokio::spawn(async move { c.chat(messages).await })) {
2219 Some(handle) => match handle.await {
2220 Ok(Ok(r)) => r,
2221 Ok(Err(e)) => {
2222 warn!(error = %e, "Supervisor LLM request failed");
2223 continue;
2224 }
2225 Err(e) => {
2226 warn!(error = %e, "Supervisor task join failed");
2227 continue;
2228 }
2229 },
2230 None => {
2231 warn!("No LLM clients available");
2232 break;
2233 }
2234 };
2235
2236 let content = response
2237 .choices
2238 .first()
2239 .map(|c| c.message.content.clone())
2240 .unwrap_or_default();
2241
2242 if content.contains("FINAL:") {
2244 let final_response = content
2245 .split("FINAL:")
2246 .nth(1)
2247 .unwrap_or("")
2248 .trim()
2249 .to_string();
2250 info!("Supervisor completed task: {} chars", final_response.len());
2251
2252 let _ = audit_log.append(
2253 AuditEventType::AgentFinish,
2254 "supervisor",
2255 "Supervisor completed task coordination",
2256 Some(serde_json::json!({
2257 "iterations": iteration,
2258 "subtasks_completed": subtask_results.len(),
2259 "providers_used": multi_llm.client_count(),
2260 })),
2261 );
2262
2263 println!("\n🐦⬛ Supervisor Result (multi-model):\n{}", final_response);
2264 return Ok(());
2265 }
2266
2267 if content.contains("SUBTASK:") && content.contains("PROVIDER:") {
2269 let subtask_block = content.split("SUBTASK:").nth(1).unwrap_or("");
2270 let subtask_lines: Vec<&str> = subtask_block.lines().take(4).collect();
2271
2272 let subtask_desc = subtask_lines.first().unwrap_or(&"").trim();
2273 let provider_idx = subtask_lines
2274 .iter()
2275 .find(|l| l.starts_with("PROVIDER:"))
2276 .and_then(|l| l.split(':').nth(1))
2277 .and_then(|s| s.trim().parse::<usize>().ok())
2278 .unwrap_or(0);
2279
2280 if !subtask_desc.is_empty() {
2281 info!("Subtask for provider {}: {}", provider_idx, subtask_desc);
2282
2283 let client = multi_llm
2284 .get_client(provider_idx)
2285 .or_else(|| multi_llm.get_client(0));
2286
2287 if let Some(client) = client {
2288 let subtask_result = run_subtask_agent(
2289 client.clone(),
2290 subtask_desc,
2291 system_prompt,
2292 &policy_engine,
2293 &sandbox,
2294 &audit_log,
2295 ®istry,
2296 )
2297 .await;
2298
2299 match subtask_result {
2300 Ok(result) => {
2301 info!("Subtask {} completed: {} chars", provider_idx, result.len());
2302 subtask_results.push(format!(
2303 "Provider {} ({}): {}",
2304 provider_idx,
2305 client.provider_name(),
2306 result
2307 ));
2308
2309 memory.add_assistant_message(&format!(
2310 "Assigned subtask to provider {}: {}",
2311 provider_idx, subtask_desc
2312 ));
2313 memory.add_user_message(&format!(
2314 "Provider {} result: {}",
2315 provider_idx, result
2316 ));
2317 }
2318 Err(e) => {
2319 warn!("Subtask {} failed: {}", provider_idx, e);
2320 memory.add_assistant_message(&format!(
2321 "Provider {} subtask failed: {}",
2322 provider_idx, e
2323 ));
2324 }
2325 }
2326 }
2327 }
2328 } else {
2329 memory.add_assistant_message(&content);
2330 }
2331 }
2332
2333 if !subtask_results.is_empty() {
2335 let aggregated = subtask_results.join("\n\n");
2336 info!(
2337 "Supervisor aggregated {} subtask results",
2338 subtask_results.len()
2339 );
2340
2341 if let Some(ref rf) = ravenfabric {
2343 if rf.is_enabled() {
2344 let summary = format!(
2345 "Multi-model supervisor completed: {} subtasks, result: {} chars",
2346 subtask_results.len(),
2347 aggregated.len()
2348 );
2349 let _ = rf.broadcast(&summary, 30).await;
2350 info!("Multi-model supervisor result broadcast to RavenFabric mesh");
2351 }
2352 }
2353
2354 println!(
2355 "\n🐦⬛ Supervisor Aggregated Result (multi-model):\n{}",
2356 aggregated
2357 );
2358 return Ok(());
2359 }
2360
2361 Err(crate::error::RavenClawsError::CommandExecution(
2362 "Supervisor mode completed without results".to_string(),
2363 ))
2364}
2365
2366pub async fn run_repl(llm: Arc<dyn LLMProviderTrait>, config: Config) -> Result<()> {
2368 use tokio::io::{AsyncBufReadExt, BufReader};
2369
2370 info!("Starting interactive REPL mode");
2371
2372 let system_prompt = &config.llm.system_prompt;
2373 let mut memory = ConversationMemory::new(system_prompt, 0);
2374
2375 let stdin = BufReader::new(tokio::io::stdin());
2376 let mut lines = stdin.lines();
2377
2378 println!("RavenClaws REPL — type /exit to quit, /reset to clear history");
2379
2380 loop {
2381 print!("\n> ");
2382 use tokio::io::AsyncWriteExt;
2383 tokio::io::stdout().flush().await?;
2384
2385 let line = match lines.next_line().await {
2386 Ok(Some(l)) => l,
2387 Ok(None) => break, Err(e) => {
2389 warn!(error = %e, "REPL read error");
2390 break;
2391 }
2392 };
2393
2394 let input = line.trim();
2395
2396 if input.is_empty() {
2397 continue;
2398 }
2399
2400 match input {
2401 "/exit" | "/quit" => {
2402 println!("Exiting REPL.");
2403 break;
2404 }
2405 "/reset" => {
2406 memory = ConversationMemory::new(system_prompt, 0);
2407 println!("Conversation history reset.");
2408 continue;
2409 }
2410 _ => {}
2411 }
2412
2413 memory.add_user_message(input);
2414 let messages = memory.history().to_vec();
2415
2416 match llm.chat(messages).await {
2417 Ok(response) => {
2418 if let Some(choice) = response.choices.first() {
2419 let content = &choice.message.content;
2420 println!("{}", content);
2421 memory.add_assistant_message(content);
2422 }
2423 }
2424 Err(e) => {
2425 warn!(error = %e, "LLM request failed");
2426 println!("Error: {}", e);
2427 }
2428 }
2429 }
2430
2431 Ok(())
2432}
2433
2434#[cfg(test)]
2435mod tests {
2436 use super::*;
2437
2438 #[test]
2439 fn test_swarm_function_exists() {
2440 let _fn_ptr: fn(Arc<dyn LLMProviderTrait>, Config, Option<RavenFabricClient>) -> _ =
2442 run_swarm;
2443 }
2444
2445 #[test]
2446 fn test_supervisor_function_exists() {
2447 let _fn_ptr: fn(Arc<dyn LLMProviderTrait>, Config, Option<RavenFabricClient>) -> _ =
2449 run_supervisor;
2450 }
2451
2452 #[test]
2453 fn test_conversation_memory_new() {
2454 let mem = ConversationMemory::new("system prompt", 10);
2455 assert_eq!(mem.messages.len(), 1);
2456 assert_eq!(mem.messages[0].role, "system");
2457 assert_eq!(mem.messages[0].content, "system prompt");
2458 }
2459
2460 #[test]
2461 fn test_conversation_memory_add_user() {
2462 let mut mem = ConversationMemory::new("system", 10);
2463 mem.add_user_message("hello");
2464 assert_eq!(mem.messages.len(), 2);
2465 assert_eq!(mem.messages[1].role, "user");
2466 assert_eq!(mem.messages[1].content, "hello");
2467 }
2468
2469 #[test]
2470 fn test_conversation_memory_trim() {
2471 let mut mem = ConversationMemory::new("system", 3);
2472 mem.add_user_message("msg1");
2473 mem.add_assistant_message("resp1");
2474 mem.add_user_message("msg2");
2475 mem.add_assistant_message("resp2");
2476 assert!(mem.messages.len() <= 3);
2478 }
2479
2480 #[test]
2481 fn test_parse_tool_call_valid() {
2482 let content = "THOUGHT: I need to run a command\nTOOL_CALL: shell_exec\nARGS: {\"command\": \"echo hello\"}";
2483 let (name, args) = parse_tool_call(content).unwrap();
2484 assert_eq!(name, "shell_exec");
2485 assert_eq!(args["command"], "echo hello");
2486 }
2487
2488 #[test]
2489 fn test_parse_tool_call_missing_tool() {
2490 let content = "THOUGHT: no tool here";
2491 assert!(parse_tool_call(content).is_none());
2492 }
2493
2494 #[test]
2495 fn test_parse_tool_call_missing_args() {
2496 let content = "TOOL_CALL: shell_exec\nNo args line";
2497 assert!(parse_tool_call(content).is_none());
2498 }
2499
2500 #[test]
2501 fn test_parse_tool_call_invalid_json() {
2502 let content = "TOOL_CALL: shell_exec\nARGS: not valid json";
2503 assert!(parse_tool_call(content).is_none());
2504 }
2505
2506 #[test]
2507 fn test_agent_loop_config_default() {
2508 let config = AgentLoopConfig::default();
2509 assert_eq!(config.max_iterations, 10);
2510 assert!(!config.enable_tools);
2511 assert!(!config.require_approval);
2512 }
2513
2514 #[test]
2515 fn test_agent_loop_config_require_approval() {
2516 let config = AgentLoopConfig {
2517 max_iterations: 5,
2518 enable_tools: true,
2519 require_approval: true,
2520 prompt_injection_protection: true,
2521 token_lifetime_secs: 0,
2522 no_final_required: false,
2523 fallback_chain: None,
2524 token_budget: None,
2525 ravenfabric: None,
2526 checkpoint_dir: None,
2527 session_id: None,
2528 metrics_callback: None,
2529 load_manager: None,
2530 retry_config: None,
2531 healing_engine: None,
2532 };
2533 assert_eq!(config.max_iterations, 5);
2534 assert!(config.enable_tools);
2535 assert!(config.require_approval);
2536 assert!(config.prompt_injection_protection);
2537 assert_eq!(config.token_lifetime_secs, 0);
2538 }
2539
2540 #[test]
2541 fn test_prompt_for_approval_yes() {
2542 let args = serde_json::json!({"command": "echo hello"});
2543 let result = tokio_test::block_on(prompt_for_approval_with_input("shell_exec", &args, "y"));
2544 assert!(result, "Should approve for 'y'");
2545 }
2546
2547 #[test]
2548 fn test_prompt_for_approval_yes_full() {
2549 let args = serde_json::json!({"command": "echo hello"});
2550 let result =
2551 tokio_test::block_on(prompt_for_approval_with_input("shell_exec", &args, "yes"));
2552 assert!(result, "Should approve for 'yes'");
2553 }
2554
2555 #[test]
2556 fn test_prompt_for_approval_no() {
2557 let args = serde_json::json!({"command": "echo hello"});
2558 let result = tokio_test::block_on(prompt_for_approval_with_input("shell_exec", &args, "n"));
2559 assert!(!result, "Should deny for 'n'");
2560 }
2561
2562 #[test]
2563 fn test_prompt_for_approval_no_full() {
2564 let args = serde_json::json!({"command": "echo hello"});
2565 let result =
2566 tokio_test::block_on(prompt_for_approval_with_input("shell_exec", &args, "no"));
2567 assert!(!result, "Should deny for 'no'");
2568 }
2569
2570 #[test]
2571 fn test_prompt_for_approval_empty() {
2572 let args = serde_json::json!({"command": "echo hello"});
2573 let result = tokio_test::block_on(prompt_for_approval_with_input("shell_exec", &args, ""));
2574 assert!(!result, "Should deny for empty input (default N)");
2575 }
2576
2577 #[test]
2578 fn test_prompt_for_approval_uppercase() {
2579 let args = serde_json::json!({"command": "echo hello"});
2580 let result = tokio_test::block_on(prompt_for_approval_with_input("shell_exec", &args, "Y"));
2581 assert!(result, "Should approve for uppercase 'Y'");
2582 }
2583
2584 #[test]
2585 fn test_prompt_for_approval_auto_approves_non_tty() {
2586 #[allow(clippy::let_underscore_future)]
2592 let _ = prompt_for_approval_with_input("test", &serde_json::json!({}), "y");
2593 }
2594
2595 #[test]
2596 fn test_execute_parsed_tool_call_skips_approval_when_not_required() {
2597 let registry = ToolRegistry::with_default_tools();
2598 let policy_engine = PolicyEngine::default_secure();
2599 let sandbox = Sandbox::default();
2600 let audit_log = AuditLog::new("test-session".to_string());
2601
2602 let args = serde_json::json!({"command": "echo hello"});
2603 let result = tokio_test::block_on(execute_parsed_tool_call(
2604 "shell_exec".to_string(),
2605 args,
2606 ®istry,
2607 &policy_engine,
2608 &sandbox,
2609 &audit_log,
2610 false, ));
2612
2613 assert!(result.is_some());
2614 let tool_result = result.unwrap();
2615 assert_eq!(tool_result.tool_name, "shell_exec");
2616 }
2617
2618 #[test]
2619 fn test_execute_parsed_tool_call_approval_not_needed_for_read_only_tools() {
2620 let registry = ToolRegistry::with_default_tools();
2623 let policy_engine = PolicyEngine::default_secure();
2624 let sandbox = Sandbox::default();
2625 let audit_log = AuditLog::new("test-session".to_string());
2626
2627 let args = serde_json::json!({"path": "/tmp/test.txt"});
2628 let result = tokio_test::block_on(execute_parsed_tool_call(
2629 "read_file".to_string(),
2630 args,
2631 ®istry,
2632 &policy_engine,
2633 &sandbox,
2634 &audit_log,
2635 true, ));
2637
2638 assert!(result.is_some());
2640 let tool_result = result.unwrap();
2641 assert_eq!(tool_result.tool_name, "read_file");
2642 }
2643
2644 #[test]
2645 fn test_agent_loop_config_token_lifetime_zero_disabled() {
2646 let config = AgentLoopConfig {
2647 max_iterations: 10,
2648 enable_tools: false,
2649 require_approval: false,
2650 prompt_injection_protection: false,
2651 token_lifetime_secs: 0,
2652 no_final_required: false,
2653 fallback_chain: None,
2654 token_budget: None,
2655 ravenfabric: None,
2656 checkpoint_dir: None,
2657 session_id: None,
2658 metrics_callback: None,
2659 load_manager: None,
2660 retry_config: None,
2661 healing_engine: None,
2662 };
2663 assert_eq!(config.token_lifetime_secs, 0);
2664 }
2666
2667 #[test]
2668 fn test_agent_loop_config_token_lifetime_nonzero() {
2669 let config = AgentLoopConfig {
2670 max_iterations: 10,
2671 enable_tools: false,
2672 require_approval: false,
2673 prompt_injection_protection: false,
2674 token_lifetime_secs: 3600,
2675 no_final_required: false,
2676 fallback_chain: None,
2677 token_budget: None,
2678 ravenfabric: None,
2679 checkpoint_dir: None,
2680 session_id: None,
2681 metrics_callback: None,
2682 load_manager: None,
2683 retry_config: None,
2684 healing_engine: None,
2685 };
2686 assert_eq!(config.token_lifetime_secs, 3600);
2687 }
2688
2689 #[test]
2690 fn test_agent_loop_config_default_includes_token_lifetime() {
2691 let config = AgentLoopConfig::default();
2692 assert_eq!(config.token_lifetime_secs, 0);
2693 }
2694
2695 #[test]
2696 fn test_agent_loop_config_retry_config_default_none() {
2697 let config = AgentLoopConfig::default();
2698 assert!(config.retry_config.is_none());
2699 }
2700
2701 #[test]
2702 fn test_agent_loop_config_retry_config_custom() {
2703 let config = AgentLoopConfig {
2704 max_iterations: 10,
2705 enable_tools: false,
2706 require_approval: false,
2707 prompt_injection_protection: false,
2708 token_lifetime_secs: 0,
2709 no_final_required: false,
2710 fallback_chain: None,
2711 token_budget: None,
2712 ravenfabric: None,
2713 checkpoint_dir: None,
2714 session_id: None,
2715 metrics_callback: None,
2716 load_manager: None,
2717 retry_config: Some(RetryConfig {
2718 max_retries: 5,
2719 base_delay_ms: 50,
2720 max_delay_ms: 5000,
2721 jitter: 0.1,
2722 }),
2723 healing_engine: None,
2724 };
2725 assert!(config.retry_config.is_some());
2726 assert_eq!(config.retry_config.as_ref().unwrap().max_retries, 5);
2727 assert_eq!(config.retry_config.as_ref().unwrap().base_delay_ms, 50);
2728 }
2729
2730 #[test]
2731 fn test_retry_config_delay_calculation() {
2732 let config = RetryConfig {
2733 max_retries: 3,
2734 base_delay_ms: 100,
2735 max_delay_ms: 10000,
2736 jitter: 0.0, };
2738
2739 let d0 = config.delay_for_attempt(0);
2741 assert_eq!(d0.as_millis(), 100);
2742
2743 let d1 = config.delay_for_attempt(1);
2745 assert_eq!(d1.as_millis(), 200);
2746
2747 let d2 = config.delay_for_attempt(2);
2749 assert_eq!(d2.as_millis(), 400);
2750 }
2751
2752 #[test]
2753 fn test_retry_config_delay_capped() {
2754 let config = RetryConfig {
2755 max_retries: 10,
2756 base_delay_ms: 1000,
2757 max_delay_ms: 5000,
2758 jitter: 0.0, };
2760
2761 let d5 = config.delay_for_attempt(5);
2763 assert_eq!(d5.as_millis(), 5000);
2764 }
2765
2766 #[test]
2767 fn test_retry_config_delay_with_jitter() {
2768 let config = RetryConfig {
2769 max_retries: 3,
2770 base_delay_ms: 100,
2771 max_delay_ms: 10000,
2772 jitter: 0.5,
2773 };
2774
2775 let d = config.delay_for_attempt(0);
2777 assert!(d.as_millis() >= 100);
2778 assert!(d.as_millis() <= 150);
2780 }
2781}