1pub mod commands;
34pub mod compact;
35pub mod history;
36pub mod ide;
37pub mod persistence;
38pub mod prompts;
39pub mod session;
40pub mod tools;
41pub mod ui;
42use colored::Colorize;
43use commands::TokenUsage;
44use history::{ConversationHistory, ToolCallRecord};
45use ide::IdeClient;
46use rig::{
47 client::{CompletionClient, ProviderClient},
48 completion::Prompt,
49 providers::{anthropic, openai},
50};
51use session::{ChatSession, PlanMode};
52use std::path::Path;
53use std::sync::Arc;
54use tokio::sync::Mutex as TokioMutex;
55use ui::{ResponseFormatter, ToolDisplayHook};
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
59pub enum ProviderType {
60 #[default]
61 OpenAI,
62 Anthropic,
63 Bedrock,
64}
65
66impl std::fmt::Display for ProviderType {
67 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68 match self {
69 ProviderType::OpenAI => write!(f, "openai"),
70 ProviderType::Anthropic => write!(f, "anthropic"),
71 ProviderType::Bedrock => write!(f, "bedrock"),
72 }
73 }
74}
75
76impl std::str::FromStr for ProviderType {
77 type Err = String;
78
79 fn from_str(s: &str) -> Result<Self, Self::Err> {
80 match s.to_lowercase().as_str() {
81 "openai" => Ok(ProviderType::OpenAI),
82 "anthropic" => Ok(ProviderType::Anthropic),
83 "bedrock" | "aws" | "aws-bedrock" => Ok(ProviderType::Bedrock),
84 _ => Err(format!(
85 "Unknown provider: {}. Use: openai, anthropic, or bedrock",
86 s
87 )),
88 }
89 }
90}
91
92#[derive(Debug, thiserror::Error)]
94pub enum AgentError {
95 #[error("Missing API key. Set {0} environment variable.")]
96 MissingApiKey(String),
97
98 #[error("Provider error: {0}")]
99 ProviderError(String),
100
101 #[error("Tool error: {0}")]
102 ToolError(String),
103}
104
105pub type AgentResult<T> = Result<T, AgentError>;
106
107fn get_system_prompt(project_path: &Path, query: Option<&str>, plan_mode: PlanMode) -> String {
109 if plan_mode.is_planning() {
111 return prompts::get_planning_prompt(project_path);
112 }
113
114 if let Some(q) = query {
115 if prompts::is_code_development_query(q) {
117 return prompts::get_code_development_prompt(project_path);
118 }
119 if prompts::is_generation_query(q) {
121 return prompts::get_devops_prompt(project_path, Some(q));
122 }
123 }
124 prompts::get_analysis_prompt(project_path)
126}
127
128pub async fn run_interactive(
130 project_path: &Path,
131 provider: ProviderType,
132 model: Option<String>,
133) -> AgentResult<()> {
134 use tools::*;
135
136 let mut session = ChatSession::new(project_path, provider, model);
137
138 let mut conversation_history = ConversationHistory::new();
144
145 let ide_client: Option<Arc<TokioMutex<IdeClient>>> = {
147 let mut client = IdeClient::new().await;
148 if client.is_ide_available() {
149 match client.connect().await {
150 Ok(()) => {
151 println!(
152 "{} Connected to {} IDE companion",
153 "โ".green(),
154 client.ide_name().unwrap_or("VS Code")
155 );
156 Some(Arc::new(TokioMutex::new(client)))
157 }
158 Err(e) => {
159 println!("{} IDE companion not connected: {}", "!".yellow(), e);
161 None
162 }
163 }
164 } else {
165 println!(
166 "{} No IDE detected (TERM_PROGRAM={})",
167 "ยท".dimmed(),
168 std::env::var("TERM_PROGRAM").unwrap_or_default()
169 );
170 None
171 }
172 };
173
174 ChatSession::load_api_key_to_env(session.provider);
176
177 if !ChatSession::has_api_key(session.provider) {
179 ChatSession::prompt_api_key(session.provider)?;
180 }
181
182 session.print_banner();
183
184 let mut raw_chat_history: Vec<rig::completion::Message> = Vec::new();
200
201 let mut pending_input: Option<String> = None;
203 let mut auto_accept_writes = false;
205
206 let mut session_recorder = persistence::SessionRecorder::new(project_path);
208
209 loop {
210 if !conversation_history.is_empty() {
212 println!(
213 "{}",
214 format!(" ๐ฌ Context: {}", conversation_history.status()).dimmed()
215 );
216 }
217
218 let input = if let Some(pending) = pending_input.take() {
220 println!("{} {}", "โ".cyan(), pending.dimmed());
222 pending
223 } else {
224 auto_accept_writes = false;
226
227 let input_result = match session.read_input() {
229 Ok(result) => result,
230 Err(_) => break,
231 };
232
233 match input_result {
235 ui::InputResult::Submit(text) => ChatSession::process_submitted_text(&text),
236 ui::InputResult::Cancel | ui::InputResult::Exit => break,
237 ui::InputResult::TogglePlanMode => {
238 let new_mode = session.toggle_plan_mode();
240 if new_mode.is_planning() {
241 println!("{}", "โ
plan mode".yellow());
242 } else {
243 println!("{}", "โถ standard mode".green());
244 }
245 continue;
246 }
247 }
248 };
249
250 if input.is_empty() {
251 continue;
252 }
253
254 if ChatSession::is_command(&input) {
256 if input.trim().to_lowercase() == "/clear" || input.trim().to_lowercase() == "/c" {
258 conversation_history.clear();
259 raw_chat_history.clear();
260 }
261 match session.process_command(&input) {
262 Ok(true) => continue,
263 Ok(false) => break, Err(e) => {
265 eprintln!("{}", format!("Error: {}", e).red());
266 continue;
267 }
268 }
269 }
270
271 if !ChatSession::has_api_key(session.provider) {
273 eprintln!(
274 "{}",
275 "No API key configured. Use /provider to set one.".yellow()
276 );
277 continue;
278 }
279
280 if conversation_history.needs_compaction() {
282 println!("{}", " ๐ฆ Compacting conversation history...".dimmed());
283 if let Some(summary) = conversation_history.compact() {
284 println!(
285 "{}",
286 format!(" โ Compressed {} turns", summary.matches("Turn").count()).dimmed()
287 );
288 }
289 }
290
291 let estimated_input_tokens = estimate_raw_history_tokens(&raw_chat_history)
295 + input.len() / 4 + 5000; if estimated_input_tokens > 150_000 {
299 println!(
300 "{}",
301 " โ Large context detected. Pre-truncating...".yellow()
302 );
303
304 let old_count = raw_chat_history.len();
305 if raw_chat_history.len() > 20 {
307 let drain_count = raw_chat_history.len() - 20;
308 raw_chat_history.drain(0..drain_count);
309 conversation_history.clear(); println!(
311 "{}",
312 format!(
313 " โ Truncated {} โ {} messages",
314 old_count,
315 raw_chat_history.len()
316 )
317 .dimmed()
318 );
319 }
320 }
321
322 const MAX_RETRIES: u32 = 3;
328 const MAX_CONTINUATIONS: u32 = 10;
329 const _TOOL_CALL_CHECKPOINT: usize = 50;
330 const MAX_TOOL_CALLS: usize = 300;
331 let mut retry_attempt = 0;
332 let mut continuation_count = 0;
333 let mut total_tool_calls: usize = 0;
334 let mut auto_continue_tools = false; let mut current_input = input.clone();
336 let mut succeeded = false;
337
338 while retry_attempt < MAX_RETRIES && continuation_count < MAX_CONTINUATIONS && !succeeded {
339 if continuation_count > 0 {
341 eprintln!("{}", " ๐ก Sending continuation request...".dimmed());
342 }
343
344 let hook = ToolDisplayHook::new();
346
347 let progress = ui::GenerationIndicator::new();
349 hook.set_progress_state(progress.state()).await;
352
353 let project_path_buf = session.project_path.clone();
354 let preamble = get_system_prompt(
356 &session.project_path,
357 Some(¤t_input),
358 session.plan_mode,
359 );
360 let is_generation = prompts::is_generation_query(¤t_input);
361 let is_planning = session.plan_mode.is_planning();
362
363 let progress_state = progress.state();
368
369 let mut user_interrupted = false;
372
373 let response = tokio::select! {
375 biased; _ = tokio::signal::ctrl_c() => {
378 user_interrupted = true;
379 Err::<String, String>("User cancelled".to_string())
380 }
381
382 result = async {
383 match session.provider {
384 ProviderType::OpenAI => {
385 let client = openai::Client::from_env();
386 let reasoning_params =
389 if session.model.starts_with("gpt-5") || session.model.starts_with("o1") {
390 Some(serde_json::json!({
391 "reasoning": {
392 "effort": "medium",
393 "summary": "detailed"
394 }
395 }))
396 } else {
397 None
398 };
399
400 let mut builder = client
401 .agent(&session.model)
402 .preamble(&preamble)
403 .max_tokens(4096)
404 .tool(AnalyzeTool::new(project_path_buf.clone()))
405 .tool(SecurityScanTool::new(project_path_buf.clone()))
406 .tool(VulnerabilitiesTool::new(project_path_buf.clone()))
407 .tool(HadolintTool::new(project_path_buf.clone()))
408 .tool(DclintTool::new(project_path_buf.clone()))
409 .tool(KubelintTool::new(project_path_buf.clone()))
410 .tool(HelmlintTool::new(project_path_buf.clone()))
411 .tool(TerraformFmtTool::new(project_path_buf.clone()))
412 .tool(TerraformValidateTool::new(project_path_buf.clone()))
413 .tool(TerraformInstallTool::new())
414 .tool(ReadFileTool::new(project_path_buf.clone()))
415 .tool(ListDirectoryTool::new(project_path_buf.clone()))
416 .tool(WebFetchTool::new());
417
418 if is_planning {
420 builder = builder
422 .tool(ShellTool::new(project_path_buf.clone()).with_read_only(true))
423 .tool(PlanCreateTool::new(project_path_buf.clone()))
424 .tool(PlanListTool::new(project_path_buf.clone()));
425 } else if is_generation {
426 let (mut write_file_tool, mut write_files_tool) =
428 if let Some(ref client) = ide_client {
429 (
430 WriteFileTool::new(project_path_buf.clone())
431 .with_ide_client(client.clone()),
432 WriteFilesTool::new(project_path_buf.clone())
433 .with_ide_client(client.clone()),
434 )
435 } else {
436 (
437 WriteFileTool::new(project_path_buf.clone()),
438 WriteFilesTool::new(project_path_buf.clone()),
439 )
440 };
441 if auto_accept_writes {
443 write_file_tool = write_file_tool.without_confirmation();
444 write_files_tool = write_files_tool.without_confirmation();
445 }
446 builder = builder
447 .tool(write_file_tool)
448 .tool(write_files_tool)
449 .tool(ShellTool::new(project_path_buf.clone()))
450 .tool(PlanListTool::new(project_path_buf.clone()))
451 .tool(PlanNextTool::new(project_path_buf.clone()))
452 .tool(PlanUpdateTool::new(project_path_buf.clone()));
453 }
454
455 if let Some(params) = reasoning_params {
456 builder = builder.additional_params(params);
457 }
458
459 let agent = builder.build();
460 agent
464 .prompt(¤t_input)
465 .with_history(&mut raw_chat_history)
466 .with_hook(hook.clone())
467 .multi_turn(50)
468 .await
469 }
470 ProviderType::Anthropic => {
471 let client = anthropic::Client::from_env();
472
473 let mut builder = client
480 .agent(&session.model)
481 .preamble(&preamble)
482 .max_tokens(4096)
483 .tool(AnalyzeTool::new(project_path_buf.clone()))
484 .tool(SecurityScanTool::new(project_path_buf.clone()))
485 .tool(VulnerabilitiesTool::new(project_path_buf.clone()))
486 .tool(HadolintTool::new(project_path_buf.clone()))
487 .tool(DclintTool::new(project_path_buf.clone()))
488 .tool(KubelintTool::new(project_path_buf.clone()))
489 .tool(HelmlintTool::new(project_path_buf.clone()))
490 .tool(TerraformFmtTool::new(project_path_buf.clone()))
491 .tool(TerraformValidateTool::new(project_path_buf.clone()))
492 .tool(TerraformInstallTool::new())
493 .tool(ReadFileTool::new(project_path_buf.clone()))
494 .tool(ListDirectoryTool::new(project_path_buf.clone()))
495 .tool(WebFetchTool::new());
496
497 if is_planning {
499 builder = builder
501 .tool(ShellTool::new(project_path_buf.clone()).with_read_only(true))
502 .tool(PlanCreateTool::new(project_path_buf.clone()))
503 .tool(PlanListTool::new(project_path_buf.clone()));
504 } else if is_generation {
505 let (mut write_file_tool, mut write_files_tool) =
507 if let Some(ref client) = ide_client {
508 (
509 WriteFileTool::new(project_path_buf.clone())
510 .with_ide_client(client.clone()),
511 WriteFilesTool::new(project_path_buf.clone())
512 .with_ide_client(client.clone()),
513 )
514 } else {
515 (
516 WriteFileTool::new(project_path_buf.clone()),
517 WriteFilesTool::new(project_path_buf.clone()),
518 )
519 };
520 if auto_accept_writes {
522 write_file_tool = write_file_tool.without_confirmation();
523 write_files_tool = write_files_tool.without_confirmation();
524 }
525 builder = builder
526 .tool(write_file_tool)
527 .tool(write_files_tool)
528 .tool(ShellTool::new(project_path_buf.clone()))
529 .tool(PlanListTool::new(project_path_buf.clone()))
530 .tool(PlanNextTool::new(project_path_buf.clone()))
531 .tool(PlanUpdateTool::new(project_path_buf.clone()));
532 }
533
534 let agent = builder.build();
535
536 agent
540 .prompt(¤t_input)
541 .with_history(&mut raw_chat_history)
542 .with_hook(hook.clone())
543 .multi_turn(50)
544 .await
545 }
546 ProviderType::Bedrock => {
547 let client = crate::bedrock::client::Client::from_env();
549
550 let thinking_params = serde_json::json!({
556 "thinking": {
557 "type": "enabled",
558 "budget_tokens": 8000
559 }
560 });
561
562 let mut builder = client
563 .agent(&session.model)
564 .preamble(&preamble)
565 .max_tokens(64000) .tool(AnalyzeTool::new(project_path_buf.clone()))
567 .tool(SecurityScanTool::new(project_path_buf.clone()))
568 .tool(VulnerabilitiesTool::new(project_path_buf.clone()))
569 .tool(HadolintTool::new(project_path_buf.clone()))
570 .tool(DclintTool::new(project_path_buf.clone()))
571 .tool(KubelintTool::new(project_path_buf.clone()))
572 .tool(HelmlintTool::new(project_path_buf.clone()))
573 .tool(TerraformFmtTool::new(project_path_buf.clone()))
574 .tool(TerraformValidateTool::new(project_path_buf.clone()))
575 .tool(TerraformInstallTool::new())
576 .tool(ReadFileTool::new(project_path_buf.clone()))
577 .tool(ListDirectoryTool::new(project_path_buf.clone()))
578 .tool(WebFetchTool::new());
579
580 if is_planning {
582 builder = builder
584 .tool(ShellTool::new(project_path_buf.clone()).with_read_only(true))
585 .tool(PlanCreateTool::new(project_path_buf.clone()))
586 .tool(PlanListTool::new(project_path_buf.clone()));
587 } else if is_generation {
588 let (mut write_file_tool, mut write_files_tool) =
590 if let Some(ref client) = ide_client {
591 (
592 WriteFileTool::new(project_path_buf.clone())
593 .with_ide_client(client.clone()),
594 WriteFilesTool::new(project_path_buf.clone())
595 .with_ide_client(client.clone()),
596 )
597 } else {
598 (
599 WriteFileTool::new(project_path_buf.clone()),
600 WriteFilesTool::new(project_path_buf.clone()),
601 )
602 };
603 if auto_accept_writes {
605 write_file_tool = write_file_tool.without_confirmation();
606 write_files_tool = write_files_tool.without_confirmation();
607 }
608 builder = builder
609 .tool(write_file_tool)
610 .tool(write_files_tool)
611 .tool(ShellTool::new(project_path_buf.clone()))
612 .tool(PlanListTool::new(project_path_buf.clone()))
613 .tool(PlanNextTool::new(project_path_buf.clone()))
614 .tool(PlanUpdateTool::new(project_path_buf.clone()));
615 }
616
617 builder = builder.additional_params(thinking_params);
619
620 let agent = builder.build();
621
622 agent
624 .prompt(¤t_input)
625 .with_history(&mut raw_chat_history)
626 .with_hook(hook.clone())
627 .multi_turn(50)
628 .await
629 }
630 }.map_err(|e| e.to_string())
631 } => result
632 };
633
634 progress.stop().await;
636
637 let _ = (&progress_state, user_interrupted);
639
640 match response {
641 Ok(text) => {
642 println!();
644 ResponseFormatter::print_response(&text);
645
646 let hook_usage = hook.get_usage().await;
648 if hook_usage.has_data() {
649 session
651 .token_usage
652 .add_actual(hook_usage.input_tokens, hook_usage.output_tokens);
653 } else {
654 let prompt_tokens = TokenUsage::estimate_tokens(&input);
656 let completion_tokens = TokenUsage::estimate_tokens(&text);
657 session
658 .token_usage
659 .add_estimated(prompt_tokens, completion_tokens);
660 }
661 hook.reset_usage().await;
663
664 let model_short = session
666 .model
667 .split('/')
668 .next_back()
669 .unwrap_or(&session.model)
670 .split(':')
671 .next()
672 .unwrap_or(&session.model);
673 println!();
674 println!(
675 " {}[{}/{}]{}",
676 ui::colors::ansi::DIM,
677 model_short,
678 session.token_usage.format_compact(),
679 ui::colors::ansi::RESET
680 );
681
682 let tool_calls = extract_tool_calls_from_hook(&hook).await;
684 let batch_tool_count = tool_calls.len();
685 total_tool_calls += batch_tool_count;
686
687 if batch_tool_count > 10 {
689 println!(
690 "{}",
691 format!(
692 " โ Completed with {} tool calls ({} total this session)",
693 batch_tool_count, total_tool_calls
694 )
695 .dimmed()
696 );
697 }
698
699 conversation_history.add_turn(input.clone(), text.clone(), tool_calls.clone());
701
702 if conversation_history.needs_compaction() {
705 println!("{}", " ๐ฆ Compacting conversation history...".dimmed());
706 if let Some(summary) = conversation_history.compact() {
707 println!(
708 "{}",
709 format!(" โ Compressed {} turns", summary.matches("Turn").count())
710 .dimmed()
711 );
712 }
713 }
714
715 session.history.push(("user".to_string(), input.clone()));
717 session
718 .history
719 .push(("assistant".to_string(), text.clone()));
720
721 session_recorder.record_user_message(&input);
723 session_recorder.record_assistant_message(&text, Some(&tool_calls));
724 if let Err(e) = session_recorder.save() {
725 eprintln!(
726 "{}",
727 format!(" Warning: Failed to save session: {}", e).dimmed()
728 );
729 }
730
731 if let Some(plan_info) = find_plan_create_call(&tool_calls) {
733 println!(); match ui::show_plan_action_menu(&plan_info.0, plan_info.1) {
737 ui::PlanActionResult::ExecuteAutoAccept => {
738 if session.plan_mode.is_planning() {
740 session.plan_mode = session.plan_mode.toggle();
741 }
742 auto_accept_writes = true;
743 pending_input = Some(format!(
744 "Execute the plan at '{}'. Use plan_next to get tasks and execute them in order. Auto-accept all file writes.",
745 plan_info.0
746 ));
747 succeeded = true;
748 }
749 ui::PlanActionResult::ExecuteWithReview => {
750 if session.plan_mode.is_planning() {
752 session.plan_mode = session.plan_mode.toggle();
753 }
754 pending_input = Some(format!(
755 "Execute the plan at '{}'. Use plan_next to get tasks and execute them in order.",
756 plan_info.0
757 ));
758 succeeded = true;
759 }
760 ui::PlanActionResult::ChangePlan(feedback) => {
761 pending_input = Some(format!(
763 "Please modify the plan at '{}'. User feedback: {}",
764 plan_info.0, feedback
765 ));
766 succeeded = true;
767 }
768 ui::PlanActionResult::Cancel => {
769 succeeded = true;
771 }
772 }
773 } else {
774 succeeded = true;
775 }
776 }
777 Err(e) => {
778 let err_str = e.to_string();
779
780 println!();
781
782 if err_str.contains("cancelled") || err_str.contains("Cancelled") {
784 let completed_tools = extract_tool_calls_from_hook(&hook).await;
786 let tool_count = completed_tools.len();
787
788 eprintln!("{}", "โ Generation interrupted.".yellow());
789 if tool_count > 0 {
790 eprintln!(
791 "{}",
792 format!(" {} tool calls completed before interrupt.", tool_count)
793 .dimmed()
794 );
795 conversation_history.add_turn(
797 current_input.clone(),
798 format!("[Interrupted after {} tool calls]", tool_count),
799 completed_tools,
800 );
801 }
802 eprintln!("{}", " Type your next message to continue.".dimmed());
803
804 break;
806 }
807
808 if err_str.contains("MaxDepth")
810 || err_str.contains("max_depth")
811 || err_str.contains("reached limit")
812 {
813 let completed_tools = extract_tool_calls_from_hook(&hook).await;
815 let agent_thinking = extract_agent_messages_from_hook(&hook).await;
816 let batch_tool_count = completed_tools.len();
817 total_tool_calls += batch_tool_count;
818
819 eprintln!("{}", format!(
820 "โ Reached {} tool calls this batch ({} total). Maximum allowed: {}",
821 batch_tool_count, total_tool_calls, MAX_TOOL_CALLS
822 ).yellow());
823
824 if total_tool_calls >= MAX_TOOL_CALLS {
826 eprintln!(
827 "{}",
828 format!("Maximum tool call limit ({}) reached.", MAX_TOOL_CALLS)
829 .red()
830 );
831 eprintln!(
832 "{}",
833 "The task is too complex. Try breaking it into smaller parts."
834 .dimmed()
835 );
836 break;
837 }
838
839 let should_continue = if auto_continue_tools {
841 eprintln!(
842 "{}",
843 " Auto-continuing (you selected 'always')...".dimmed()
844 );
845 true
846 } else {
847 eprintln!(
848 "{}",
849 "Excessive tool calls used. Want to continue?".yellow()
850 );
851 eprintln!(
852 "{}",
853 " [y] Yes, continue [n] No, stop [a] Always continue".dimmed()
854 );
855 print!(" > ");
856 let _ = std::io::Write::flush(&mut std::io::stdout());
857
858 let mut response = String::new();
860 match std::io::stdin().read_line(&mut response) {
861 Ok(_) => {
862 let resp = response.trim().to_lowercase();
863 if resp == "a" || resp == "always" {
864 auto_continue_tools = true;
865 true
866 } else {
867 resp == "y" || resp == "yes" || resp.is_empty()
868 }
869 }
870 Err(_) => false,
871 }
872 };
873
874 if !should_continue {
875 eprintln!(
876 "{}",
877 "Stopped by user. Type 'continue' to resume later.".dimmed()
878 );
879 if !completed_tools.is_empty() {
881 conversation_history.add_turn(
882 current_input.clone(),
883 format!(
884 "[Stopped at checkpoint - {} tools completed]",
885 batch_tool_count
886 ),
887 vec![],
888 );
889 }
890 break;
891 }
892
893 eprintln!(
895 "{}",
896 format!(
897 " โ Continuing... {} remaining tool calls available",
898 MAX_TOOL_CALLS - total_tool_calls
899 )
900 .dimmed()
901 );
902
903 conversation_history.add_turn(
905 current_input.clone(),
906 format!(
907 "[Checkpoint - {} tools completed, continuing...]",
908 batch_tool_count
909 ),
910 vec![],
911 );
912
913 current_input =
915 build_continuation_prompt(&input, &completed_tools, &agent_thinking);
916
917 tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
919 continue; } else if err_str.contains("rate")
921 || err_str.contains("Rate")
922 || err_str.contains("429")
923 || err_str.contains("Too many tokens")
924 || err_str.contains("please wait")
925 || err_str.contains("throttl")
926 || err_str.contains("Throttl")
927 {
928 eprintln!("{}", "โ Rate limited by API provider.".yellow());
929 retry_attempt += 1;
931 let wait_secs = if err_str.contains("Too many tokens") {
932 30
933 } else {
934 5
935 };
936 eprintln!(
937 "{}",
938 format!(
939 " Waiting {} seconds before retry ({}/{})...",
940 wait_secs, retry_attempt, MAX_RETRIES
941 )
942 .dimmed()
943 );
944 tokio::time::sleep(tokio::time::Duration::from_secs(wait_secs)).await;
945 } else if is_input_too_long_error(&err_str) {
946 eprintln!(
950 "{}",
951 "โ Context too large for model. Truncating history...".yellow()
952 );
953
954 let old_token_count = estimate_raw_history_tokens(&raw_chat_history);
955 let old_msg_count = raw_chat_history.len();
956
957 let keep_count = match retry_attempt {
960 0 => 10,
961 1 => 6,
962 _ => 4,
963 };
964
965 if raw_chat_history.len() > keep_count {
966 let drain_count = raw_chat_history.len() - keep_count;
968 raw_chat_history.drain(0..drain_count);
969 }
970
971 let new_token_count = estimate_raw_history_tokens(&raw_chat_history);
972 eprintln!("{}", format!(
973 " โ Truncated: {} messages (~{} tokens) โ {} messages (~{} tokens)",
974 old_msg_count, old_token_count, raw_chat_history.len(), new_token_count
975 ).green());
976
977 conversation_history.clear();
979
980 retry_attempt += 1;
982 if retry_attempt < MAX_RETRIES {
983 eprintln!(
984 "{}",
985 format!(
986 " โ Retrying with truncated context ({}/{})...",
987 retry_attempt, MAX_RETRIES
988 )
989 .dimmed()
990 );
991 tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
992 } else {
993 eprintln!(
994 "{}",
995 "Context still too large after truncation. Try /clear to reset."
996 .red()
997 );
998 break;
999 }
1000 } else if is_truncation_error(&err_str) {
1001 let completed_tools = extract_tool_calls_from_hook(&hook).await;
1003 let agent_thinking = extract_agent_messages_from_hook(&hook).await;
1004
1005 let completed_count = completed_tools
1007 .iter()
1008 .filter(|t| !t.result_summary.contains("IN PROGRESS"))
1009 .count();
1010 let in_progress_count = completed_tools.len() - completed_count;
1011
1012 if !completed_tools.is_empty() && continuation_count < MAX_CONTINUATIONS {
1013 continuation_count += 1;
1015 let status_msg = if in_progress_count > 0 {
1016 format!(
1017 "โ Response truncated. {} completed, {} in-progress. Auto-continuing ({}/{})...",
1018 completed_count,
1019 in_progress_count,
1020 continuation_count,
1021 MAX_CONTINUATIONS
1022 )
1023 } else {
1024 format!(
1025 "โ Response truncated. {} tool calls completed. Auto-continuing ({}/{})...",
1026 completed_count, continuation_count, MAX_CONTINUATIONS
1027 )
1028 };
1029 eprintln!("{}", status_msg.yellow());
1030
1031 conversation_history.add_turn(
1036 current_input.clone(),
1037 format!("[Partial response - {} tools completed, {} in-progress before truncation. See continuation prompt for details.]",
1038 completed_count, in_progress_count),
1039 vec![] );
1041
1042 if conversation_history.needs_compaction() {
1045 eprintln!(
1046 "{}",
1047 " ๐ฆ Compacting history before continuation...".dimmed()
1048 );
1049 if let Some(summary) = conversation_history.compact() {
1050 eprintln!(
1051 "{}",
1052 format!(
1053 " โ Compressed {} turns",
1054 summary.matches("Turn").count()
1055 )
1056 .dimmed()
1057 );
1058 }
1059 }
1060
1061 current_input = build_continuation_prompt(
1063 &input,
1064 &completed_tools,
1065 &agent_thinking,
1066 );
1067
1068 eprintln!("{}", format!(
1070 " โ Continuing with {} files read, {} written, {} other actions tracked",
1071 completed_tools.iter().filter(|t| t.tool_name == "read_file").count(),
1072 completed_tools.iter().filter(|t| t.tool_name == "write_file" || t.tool_name == "write_files").count(),
1073 completed_tools.iter().filter(|t| t.tool_name != "read_file" && t.tool_name != "write_file" && t.tool_name != "write_files" && t.tool_name != "list_directory").count()
1074 ).dimmed());
1075
1076 tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
1078 } else if retry_attempt < MAX_RETRIES {
1080 retry_attempt += 1;
1082 eprintln!(
1083 "{}",
1084 format!(
1085 "โ Response error (attempt {}/{}). Retrying...",
1086 retry_attempt, MAX_RETRIES
1087 )
1088 .yellow()
1089 );
1090 tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
1091 } else {
1092 eprintln!("{}", format!("Error: {}", e).red());
1094 if continuation_count >= MAX_CONTINUATIONS {
1095 eprintln!("{}", format!("Max continuations ({}) reached. The task is too complex for one request.", MAX_CONTINUATIONS).dimmed());
1096 } else {
1097 eprintln!(
1098 "{}",
1099 "Max retries reached. The response may be too complex."
1100 .dimmed()
1101 );
1102 }
1103 eprintln!(
1104 "{}",
1105 "Try breaking your request into smaller parts.".dimmed()
1106 );
1107 break;
1108 }
1109 } else if err_str.contains("timeout") || err_str.contains("Timeout") {
1110 retry_attempt += 1;
1112 if retry_attempt < MAX_RETRIES {
1113 eprintln!(
1114 "{}",
1115 format!(
1116 "โ Request timed out (attempt {}/{}). Retrying...",
1117 retry_attempt, MAX_RETRIES
1118 )
1119 .yellow()
1120 );
1121 tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
1122 } else {
1123 eprintln!("{}", "Request timed out. Please try again.".red());
1124 break;
1125 }
1126 } else {
1127 eprintln!("{}", format!("Error: {}", e).red());
1129 if continuation_count > 0 {
1130 eprintln!(
1131 "{}",
1132 format!(
1133 " (occurred during continuation attempt {})",
1134 continuation_count
1135 )
1136 .dimmed()
1137 );
1138 }
1139 eprintln!("{}", "Error details for debugging:".dimmed());
1140 eprintln!(
1141 "{}",
1142 format!(" - retry_attempt: {}/{}", retry_attempt, MAX_RETRIES)
1143 .dimmed()
1144 );
1145 eprintln!(
1146 "{}",
1147 format!(
1148 " - continuation_count: {}/{}",
1149 continuation_count, MAX_CONTINUATIONS
1150 )
1151 .dimmed()
1152 );
1153 break;
1154 }
1155 }
1156 }
1157 }
1158 println!();
1159 }
1160
1161 Ok(())
1170}
1171
1172async fn extract_tool_calls_from_hook(hook: &ToolDisplayHook) -> Vec<ToolCallRecord> {
1178 let state = hook.state();
1179 let guard = state.lock().await;
1180
1181 guard
1182 .tool_calls
1183 .iter()
1184 .enumerate()
1185 .map(|(i, tc)| {
1186 let result = if tc.is_running {
1187 "[IN PROGRESS - may need to be re-run]".to_string()
1189 } else if let Some(output) = &tc.output {
1190 truncate_string(output, 200)
1191 } else {
1192 "completed".to_string()
1193 };
1194
1195 ToolCallRecord {
1196 tool_name: tc.name.clone(),
1197 args_summary: truncate_string(&tc.args, 100),
1198 result_summary: result,
1199 tool_id: Some(format!("tool_{}_{}", tc.name, i)),
1201 droppable: matches!(
1203 tc.name.as_str(),
1204 "read_file" | "list_directory" | "analyze_project"
1205 ),
1206 }
1207 })
1208 .collect()
1209}
1210
1211async fn extract_agent_messages_from_hook(hook: &ToolDisplayHook) -> Vec<String> {
1213 let state = hook.state();
1214 let guard = state.lock().await;
1215 guard.agent_messages.clone()
1216}
1217
1218fn truncate_string(s: &str, max_len: usize) -> String {
1220 if s.len() <= max_len {
1221 s.to_string()
1222 } else {
1223 format!("{}...", &s[..max_len.saturating_sub(3)])
1224 }
1225}
1226
1227fn estimate_raw_history_tokens(messages: &[rig::completion::Message]) -> usize {
1231 use rig::completion::message::{AssistantContent, UserContent};
1232
1233 messages
1234 .iter()
1235 .map(|msg| -> usize {
1236 match msg {
1237 rig::completion::Message::User { content } => {
1238 content
1239 .iter()
1240 .map(|c| -> usize {
1241 match c {
1242 UserContent::Text(t) => t.text.len() / 4,
1243 _ => 100, }
1245 })
1246 .sum::<usize>()
1247 }
1248 rig::completion::Message::Assistant { content, .. } => {
1249 content
1250 .iter()
1251 .map(|c| -> usize {
1252 match c {
1253 AssistantContent::Text(t) => t.text.len() / 4,
1254 AssistantContent::ToolCall(tc) => {
1255 let args_len = tc.function.arguments.to_string().len();
1257 (tc.function.name.len() + args_len) / 4
1258 }
1259 _ => 100,
1260 }
1261 })
1262 .sum::<usize>()
1263 }
1264 }
1265 })
1266 .sum()
1267}
1268
1269fn find_plan_create_call(tool_calls: &[ToolCallRecord]) -> Option<(String, usize)> {
1272 for tc in tool_calls {
1273 if tc.tool_name == "plan_create" {
1274 let plan_path =
1277 if let Ok(result) = serde_json::from_str::<serde_json::Value>(&tc.result_summary) {
1278 result
1279 .get("plan_path")
1280 .and_then(|v| v.as_str())
1281 .map(|s| s.to_string())
1282 } else {
1283 None
1284 };
1285
1286 let plan_path = plan_path.unwrap_or_else(|| {
1289 find_most_recent_plan_file().unwrap_or_else(|| "plans/plan.md".to_string())
1290 });
1291
1292 let task_count = count_tasks_in_plan_file(&plan_path).unwrap_or(0);
1294
1295 return Some((plan_path, task_count));
1296 }
1297 }
1298 None
1299}
1300
1301fn find_most_recent_plan_file() -> Option<String> {
1303 let plans_dir = std::env::current_dir().ok()?.join("plans");
1304 if !plans_dir.exists() {
1305 return None;
1306 }
1307
1308 let mut newest: Option<(std::path::PathBuf, std::time::SystemTime)> = None;
1309
1310 for entry in std::fs::read_dir(&plans_dir).ok()?.flatten() {
1311 let path = entry.path();
1312 if path.extension().is_some_and(|e| e == "md")
1313 && let Ok(metadata) = entry.metadata()
1314 && let Ok(modified) = metadata.modified()
1315 && newest.as_ref().map(|(_, t)| modified > *t).unwrap_or(true)
1316 {
1317 newest = Some((path, modified));
1318 }
1319 }
1320
1321 newest.map(|(path, _)| {
1322 path.strip_prefix(std::env::current_dir().unwrap_or_default())
1324 .map(|p| p.display().to_string())
1325 .unwrap_or_else(|_| path.display().to_string())
1326 })
1327}
1328
1329fn count_tasks_in_plan_file(plan_path: &str) -> Option<usize> {
1331 use regex::Regex;
1332
1333 let path = std::path::Path::new(plan_path);
1335 let content = if path.exists() {
1336 std::fs::read_to_string(path).ok()?
1337 } else {
1338 std::fs::read_to_string(std::env::current_dir().ok()?.join(plan_path)).ok()?
1340 };
1341
1342 let task_regex = Regex::new(r"^\s*-\s*\[[ x~!]\]").ok()?;
1344 let count = content
1345 .lines()
1346 .filter(|line| task_regex.is_match(line))
1347 .count();
1348
1349 Some(count)
1350}
1351
1352fn is_truncation_error(err_str: &str) -> bool {
1354 err_str.contains("JsonError")
1355 || err_str.contains("EOF while parsing")
1356 || err_str.contains("JSON")
1357 || err_str.contains("unexpected end")
1358}
1359
1360fn is_input_too_long_error(err_str: &str) -> bool {
1364 err_str.contains("too long")
1365 || err_str.contains("Too long")
1366 || err_str.contains("context length")
1367 || err_str.contains("maximum context")
1368 || err_str.contains("exceeds the model")
1369 || err_str.contains("Input is too long")
1370}
1371
1372fn build_continuation_prompt(
1375 original_task: &str,
1376 completed_tools: &[ToolCallRecord],
1377 agent_thinking: &[String],
1378) -> String {
1379 use std::collections::HashSet;
1380
1381 let mut files_read: HashSet<String> = HashSet::new();
1383 let mut files_written: HashSet<String> = HashSet::new();
1384 let mut dirs_listed: HashSet<String> = HashSet::new();
1385 let mut other_tools: Vec<String> = Vec::new();
1386 let mut in_progress: Vec<String> = Vec::new();
1387
1388 for tool in completed_tools {
1389 let is_in_progress = tool.result_summary.contains("IN PROGRESS");
1390
1391 if is_in_progress {
1392 in_progress.push(format!("{}({})", tool.tool_name, tool.args_summary));
1393 continue;
1394 }
1395
1396 match tool.tool_name.as_str() {
1397 "read_file" => {
1398 files_read.insert(tool.args_summary.clone());
1400 }
1401 "write_file" | "write_files" => {
1402 files_written.insert(tool.args_summary.clone());
1403 }
1404 "list_directory" => {
1405 dirs_listed.insert(tool.args_summary.clone());
1406 }
1407 _ => {
1408 other_tools.push(format!(
1409 "{}({})",
1410 tool.tool_name,
1411 truncate_string(&tool.args_summary, 40)
1412 ));
1413 }
1414 }
1415 }
1416
1417 let mut prompt = format!(
1418 "[CONTINUE] Your previous response was interrupted. DO NOT repeat completed work.\n\n\
1419 Original task: {}\n",
1420 truncate_string(original_task, 500)
1421 );
1422
1423 if !files_read.is_empty() {
1425 prompt.push_str("\n== FILES ALREADY READ (do NOT read again) ==\n");
1426 for file in &files_read {
1427 prompt.push_str(&format!(" - {}\n", file));
1428 }
1429 }
1430
1431 if !dirs_listed.is_empty() {
1432 prompt.push_str("\n== DIRECTORIES ALREADY LISTED ==\n");
1433 for dir in &dirs_listed {
1434 prompt.push_str(&format!(" - {}\n", dir));
1435 }
1436 }
1437
1438 if !files_written.is_empty() {
1439 prompt.push_str("\n== FILES ALREADY WRITTEN ==\n");
1440 for file in &files_written {
1441 prompt.push_str(&format!(" - {}\n", file));
1442 }
1443 }
1444
1445 if !other_tools.is_empty() {
1446 prompt.push_str("\n== OTHER COMPLETED ACTIONS ==\n");
1447 for tool in other_tools.iter().take(20) {
1448 prompt.push_str(&format!(" - {}\n", tool));
1449 }
1450 if other_tools.len() > 20 {
1451 prompt.push_str(&format!(" ... and {} more\n", other_tools.len() - 20));
1452 }
1453 }
1454
1455 if !in_progress.is_empty() {
1456 prompt.push_str("\n== INTERRUPTED (may need re-run) ==\n");
1457 for tool in &in_progress {
1458 prompt.push_str(&format!(" โ {}\n", tool));
1459 }
1460 }
1461
1462 if let Some(last_thought) = agent_thinking.last() {
1464 prompt.push_str(&format!(
1465 "\n== YOUR LAST THOUGHTS ==\n\"{}\"\n",
1466 truncate_string(last_thought, 300)
1467 ));
1468 }
1469
1470 prompt.push_str("\n== INSTRUCTIONS ==\n");
1471 prompt.push_str("IMPORTANT: Your previous response was too long and got cut off.\n");
1472 prompt.push_str("1. Do NOT re-read files listed above - they are already in context.\n");
1473 prompt.push_str("2. If writing a document, write it in SECTIONS - complete one section now, then continue.\n");
1474 prompt.push_str("3. Keep your response SHORT and focused. Better to complete small chunks than fail on large ones.\n");
1475 prompt.push_str("4. If the task involves writing a file, START WRITING NOW - don't explain what you'll do.\n");
1476
1477 prompt
1478}
1479
1480pub async fn run_query(
1482 project_path: &Path,
1483 query: &str,
1484 provider: ProviderType,
1485 model: Option<String>,
1486) -> AgentResult<String> {
1487 use tools::*;
1488
1489 let project_path_buf = project_path.to_path_buf();
1490 let preamble = get_system_prompt(project_path, Some(query), PlanMode::default());
1493 let is_generation = prompts::is_generation_query(query);
1494
1495 match provider {
1496 ProviderType::OpenAI => {
1497 let client = openai::Client::from_env();
1498 let model_name = model.as_deref().unwrap_or("gpt-5.2");
1499
1500 let reasoning_params =
1502 if model_name.starts_with("gpt-5") || model_name.starts_with("o1") {
1503 Some(serde_json::json!({
1504 "reasoning": {
1505 "effort": "medium",
1506 "summary": "detailed"
1507 }
1508 }))
1509 } else {
1510 None
1511 };
1512
1513 let mut builder = client
1514 .agent(model_name)
1515 .preamble(&preamble)
1516 .max_tokens(4096)
1517 .tool(AnalyzeTool::new(project_path_buf.clone()))
1518 .tool(SecurityScanTool::new(project_path_buf.clone()))
1519 .tool(VulnerabilitiesTool::new(project_path_buf.clone()))
1520 .tool(HadolintTool::new(project_path_buf.clone()))
1521 .tool(DclintTool::new(project_path_buf.clone()))
1522 .tool(KubelintTool::new(project_path_buf.clone()))
1523 .tool(HelmlintTool::new(project_path_buf.clone()))
1524 .tool(TerraformFmtTool::new(project_path_buf.clone()))
1525 .tool(TerraformValidateTool::new(project_path_buf.clone()))
1526 .tool(TerraformInstallTool::new())
1527 .tool(ReadFileTool::new(project_path_buf.clone()))
1528 .tool(ListDirectoryTool::new(project_path_buf.clone()))
1529 .tool(WebFetchTool::new());
1530
1531 if is_generation {
1533 builder = builder
1534 .tool(WriteFileTool::new(project_path_buf.clone()))
1535 .tool(WriteFilesTool::new(project_path_buf.clone()))
1536 .tool(ShellTool::new(project_path_buf.clone()));
1537 }
1538
1539 if let Some(params) = reasoning_params {
1540 builder = builder.additional_params(params);
1541 }
1542
1543 let agent = builder.build();
1544
1545 agent
1546 .prompt(query)
1547 .multi_turn(50)
1548 .await
1549 .map_err(|e| AgentError::ProviderError(e.to_string()))
1550 }
1551 ProviderType::Anthropic => {
1552 let client = anthropic::Client::from_env();
1553 let model_name = model.as_deref().unwrap_or("claude-sonnet-4-5-20250929");
1554
1555 let mut builder = client
1560 .agent(model_name)
1561 .preamble(&preamble)
1562 .max_tokens(4096)
1563 .tool(AnalyzeTool::new(project_path_buf.clone()))
1564 .tool(SecurityScanTool::new(project_path_buf.clone()))
1565 .tool(VulnerabilitiesTool::new(project_path_buf.clone()))
1566 .tool(HadolintTool::new(project_path_buf.clone()))
1567 .tool(DclintTool::new(project_path_buf.clone()))
1568 .tool(KubelintTool::new(project_path_buf.clone()))
1569 .tool(HelmlintTool::new(project_path_buf.clone()))
1570 .tool(TerraformFmtTool::new(project_path_buf.clone()))
1571 .tool(TerraformValidateTool::new(project_path_buf.clone()))
1572 .tool(TerraformInstallTool::new())
1573 .tool(ReadFileTool::new(project_path_buf.clone()))
1574 .tool(ListDirectoryTool::new(project_path_buf.clone()))
1575 .tool(WebFetchTool::new());
1576
1577 if is_generation {
1579 builder = builder
1580 .tool(WriteFileTool::new(project_path_buf.clone()))
1581 .tool(WriteFilesTool::new(project_path_buf.clone()))
1582 .tool(ShellTool::new(project_path_buf.clone()));
1583 }
1584
1585 let agent = builder.build();
1586
1587 agent
1588 .prompt(query)
1589 .multi_turn(50)
1590 .await
1591 .map_err(|e| AgentError::ProviderError(e.to_string()))
1592 }
1593 ProviderType::Bedrock => {
1594 let client = crate::bedrock::client::Client::from_env();
1596 let model_name = model
1597 .as_deref()
1598 .unwrap_or("global.anthropic.claude-sonnet-4-5-20250929-v1:0");
1599
1600 let thinking_params = serde_json::json!({
1602 "thinking": {
1603 "type": "enabled",
1604 "budget_tokens": 16000
1605 }
1606 });
1607
1608 let mut builder = client
1609 .agent(model_name)
1610 .preamble(&preamble)
1611 .max_tokens(64000) .tool(AnalyzeTool::new(project_path_buf.clone()))
1613 .tool(SecurityScanTool::new(project_path_buf.clone()))
1614 .tool(VulnerabilitiesTool::new(project_path_buf.clone()))
1615 .tool(HadolintTool::new(project_path_buf.clone()))
1616 .tool(DclintTool::new(project_path_buf.clone()))
1617 .tool(KubelintTool::new(project_path_buf.clone()))
1618 .tool(HelmlintTool::new(project_path_buf.clone()))
1619 .tool(TerraformFmtTool::new(project_path_buf.clone()))
1620 .tool(TerraformValidateTool::new(project_path_buf.clone()))
1621 .tool(TerraformInstallTool::new())
1622 .tool(ReadFileTool::new(project_path_buf.clone()))
1623 .tool(ListDirectoryTool::new(project_path_buf.clone()))
1624 .tool(WebFetchTool::new());
1625
1626 if is_generation {
1628 builder = builder
1629 .tool(WriteFileTool::new(project_path_buf.clone()))
1630 .tool(WriteFilesTool::new(project_path_buf.clone()))
1631 .tool(ShellTool::new(project_path_buf.clone()));
1632 }
1633
1634 let agent = builder.additional_params(thinking_params).build();
1635
1636 agent
1637 .prompt(query)
1638 .multi_turn(50)
1639 .await
1640 .map_err(|e| AgentError::ProviderError(e.to_string()))
1641 }
1642 }
1643}