diff --git a/crates/chat-cli/src/cli/chat/conversation.rs b/crates/chat-cli/src/cli/chat/conversation.rs index 5a27fc803..93bca65a6 100644 --- a/crates/chat-cli/src/cli/chat/conversation.rs +++ b/crates/chat-cli/src/cli/chat/conversation.rs @@ -6,11 +6,8 @@ use std::collections::{ use std::io::Write; use std::sync::atomic::Ordering; +use crossterm::style; use crossterm::style::Color; -use crossterm::{ - execute, - style, -}; use serde::{ Deserialize, Serialize, @@ -64,6 +61,7 @@ use crate::cli::agent::hook::{ HookTrigger, }; use crate::cli::chat::ChatError; +use crate::execute_conditional; use crate::mcp_client::Prompt; use crate::os::Os; @@ -332,6 +330,7 @@ impl ConversationState { &mut self, os: &Os, stderr: &mut impl Write, + enable_syling: bool, run_perprompt_hooks: bool, ) -> Result { debug_assert!(self.next_message.is_some()); @@ -341,7 +340,8 @@ impl ConversationState { let context = self.backend_conversation_state(os, run_perprompt_hooks, stderr).await?; if !context.dropped_context_files.is_empty() { - execute!( + execute_conditional!( + enable_syling, stderr, style::SetForegroundColor(Color::DarkYellow), style::Print("\nSome context files are dropped due to size limit, please run "), @@ -1069,7 +1069,7 @@ mod tests { conversation.set_next_user_message("start".to_string()).await; for i in 0..=(MAX_CONVERSATION_STATE_HISTORY_LEN + 100) { let s = conversation - .as_sendable_conversation_state(&os, &mut vec![], true) + .as_sendable_conversation_state(&os, &mut vec![], true, true) .await .unwrap(); assert_conversation_state_invariants(s, i); @@ -1097,7 +1097,7 @@ mod tests { conversation.set_next_user_message("start".to_string()).await; for i in 0..=(MAX_CONVERSATION_STATE_HISTORY_LEN + 100) { let s = conversation - .as_sendable_conversation_state(&os, &mut vec![], true) + .as_sendable_conversation_state(&os, &mut vec![], true, true) .await .unwrap(); assert_conversation_state_invariants(s, i); @@ -1125,7 +1125,7 @@ mod tests { conversation.set_next_user_message("start".to_string()).await; for i in 0..=(MAX_CONVERSATION_STATE_HISTORY_LEN + 100) { let s = conversation - .as_sendable_conversation_state(&os, &mut vec![], true) + .as_sendable_conversation_state(&os, &mut vec![], true, true) .await .unwrap(); assert_conversation_state_invariants(s, i); @@ -1181,7 +1181,7 @@ mod tests { conversation.set_next_user_message("start".to_string()).await; for i in 0..=(MAX_CONVERSATION_STATE_HISTORY_LEN + 100) { let s = conversation - .as_sendable_conversation_state(&os, &mut vec![], true) + .as_sendable_conversation_state(&os, &mut vec![], true, true) .await .unwrap(); diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index b4f701ce7..8958cb38b 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -156,6 +156,10 @@ use crate::telemetry::{ get_error_reason, }; use crate::util::MCP_SERVER_TOOL_DELIMITER; +use crate::{ + execute_conditional, + queue_conditional, +}; const LIMIT_REACHED_TEXT: &str = color_print::cstr! { "You've used all your free requests for this month. You have two options: 1. Upgrade to a paid subscription for increased limits. See our Pricing page for what's included> https://aws.amazon.com/q/developer/pricing/ @@ -204,6 +208,7 @@ pub struct ChatArgs { impl ChatArgs { pub async fn execute(mut self, os: &mut Os) -> Result { let mut input = self.input; + let enable_styling = !self.no_interactive; if self.no_interactive && input.is_none() { if !std::io::stdin().is_terminal() { @@ -233,7 +238,8 @@ impl ChatArgs { .iter() .any(|arg| arg == "--profile" || arg.starts_with("--profile=")) { - execute!( + execute_conditional!( + enable_styling, stderr, style::SetForegroundColor(Color::Yellow), style::Print("WARNING: "), @@ -271,7 +277,8 @@ impl ChatArgs { .is_some_and(|a| !a.mcp_servers.mcp_servers.is_empty()) { if !self.no_interactive && !os.database.settings.get_bool(Setting::McpLoadedBefore).unwrap_or(false) { - execute!( + execute_conditional!( + enable_styling, stderr, style::Print( "To learn more about MCP safety, see https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/command-line-mcp-security.html\n\n" @@ -557,6 +564,8 @@ pub struct ChatSession { /// Pending prompts to be sent pending_prompts: VecDeque, interactive: bool, + /// Whether text styling should be enabled for this session + enable_styling: bool, inner: Option, ctrlc_rx: broadcast::Receiver<()>, } @@ -621,7 +630,8 @@ impl ChatSession { cs.tool_manager = tool_manager; if let Some(profile) = cs.current_profile() { if agents.switch(profile).is_err() { - execute!( + execute_conditional!( + !interactive, stderr, style::SetForegroundColor(Color::Red), style::Print("Error"), @@ -678,6 +688,7 @@ impl ChatSession { failed_request_ids: Vec::new(), pending_prompts: VecDeque::new(), interactive, + enable_styling: interactive, inner: Some(ChatState::default()), ctrlc_rx, }) @@ -686,6 +697,7 @@ impl ChatSession { pub async fn next(&mut self, os: &mut Os) -> Result<(), ChatError> { // Update conversation state with new tool information self.conversation.update_state(false).await; + let enable_styling = self.enable_styling(); let mut ctrl_c_stream = self.ctrlc_rx.resubscribe(); let result = match self.inner.take().expect("state must always be Some") { @@ -773,7 +785,8 @@ impl ChatSession { if self.spinner.is_some() { drop(self.spinner.take()); - queue!( + queue_conditional!( + enable_styling, self.stderr, terminal::Clear(terminal::ClearType::CurrentLine), cursor::MoveToColumn(0), @@ -782,7 +795,7 @@ impl ChatSession { let (context, report, display_err_message) = match err { ChatError::Interrupted { tool_uses: ref inter } => { - execute!(self.stderr, style::Print("\n\n"))?; + execute_conditional!(enable_styling, self.stderr, style::Print("\n\n"))?; // If there was an interrupt during tool execution, then we add fake // messages to "reset" the chat state. @@ -792,7 +805,7 @@ impl ChatSession { .abandon_tool_use(tool_uses, "The user interrupted the tool execution.".to_string()); let _ = self .conversation - .as_sendable_conversation_state(os, &mut self.stderr, false) + .as_sendable_conversation_state(os, &mut self.stderr, self.enable_styling, false) .await?; self.conversation.push_assistant_message( os, @@ -811,7 +824,8 @@ impl ChatSession { ChatError::CompactHistoryFailure => { // This error is not retryable - the user must take manual intervention to manage // their context. - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Red), style::Print("Your conversation is too large to continue.\n"), @@ -838,7 +852,8 @@ impl ChatSession { .get_bool(Setting::ChatDisableAutoCompaction) .unwrap_or(false) { - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Red), style::Print("The conversation history has overflowed.\n"), @@ -862,8 +877,8 @@ impl ChatSession { ..Default::default() }, }); - - execute!( + execute_conditional!( + enable_styling, self.stdout, style::SetForegroundColor(Color::Yellow), style::Print("The context window has overflowed, summarizing the history..."), @@ -880,7 +895,8 @@ impl ChatSession { } => { let err = "Request quota exceeded. Please wait a moment and try again.".to_string(); self.conversation.append_transcript(err.clone()); - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetAttribute(Attribute::Bold), style::SetForegroundColor(Color::Red), @@ -925,7 +941,8 @@ impl ChatSession { } ); self.conversation.append_transcript(err.clone()); - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetAttribute(Attribute::Bold), style::SetForegroundColor(Color::Red), @@ -939,7 +956,8 @@ impl ChatSession { ApiClientError::MonthlyLimitReached { .. } => { let subscription_status = get_subscription_status(os).await; if subscription_status.is_err() { - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Red), style::Print(format!( @@ -950,7 +968,8 @@ impl ChatSession { )?; } - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Yellow), style::Print("Monthly request limit reached"), @@ -965,7 +984,8 @@ impl ChatSession { if subscription_status.is_err() || subscription_status.is_ok_and(|s| s == ActualSubscriptionStatus::None) { - execute!( + execute_conditional!( + enable_styling, self.stderr, style::Print(format!("\n\n{LIMIT_REACHED_TEXT} {limits_text}")), style::SetForegroundColor(Color::DarkGrey), @@ -977,7 +997,8 @@ impl ChatSession { style::SetForegroundColor(Color::Reset), )?; } else { - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Yellow), style::Print(format!(" - {limits_text}\n\n")), @@ -1008,7 +1029,8 @@ impl ChatSession { // Remove non-ASCII and ANSI characters. let re = Regex::new(r"((\x9B|\x1B\[)[0-?]*[ -\/]*[@-~])|([^\x00-\x7F]+)").unwrap(); - queue!( + queue_conditional!( + enable_styling, self.stderr, style::SetAttribute(Attribute::Bold), style::SetForegroundColor(Color::Red), @@ -1016,10 +1038,11 @@ impl ChatSession { let text = re.replace_all(&format!("{}: {:?}\n", context, report), "").into_owned(); - queue!(self.stderr, style::Print(&text),)?; + queue_conditional!(enable_styling, self.stderr, style::Print(&text),)?; self.conversation.append_transcript(text); - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetAttribute(Attribute::Reset), style::SetForegroundColor(Color::Reset), @@ -1046,7 +1069,9 @@ impl Drop for ChatSession { spinner.stop(); } - execute!( + let enable_styling = self.enable_styling(); + execute_conditional!( + enable_styling, self.stderr, cursor::MoveToColumn(0), style::SetAttribute(Attribute::Reset), @@ -1132,11 +1157,13 @@ impl ChatSession { async fn spawn(&mut self, os: &mut Os) -> Result<()> { let is_small_screen = self.terminal_width() < GREETING_BREAK_POINT; + let enable_styling = self.enable_styling(); if os .database .settings .get_bool(Setting::ChatGreetingEnabled) .unwrap_or(true) + && self.interactive { let welcome_text = match self.existing_conversation { true => RESUME_TEXT, @@ -1185,7 +1212,8 @@ impl ChatSession { } if self.all_tools_trusted() { - queue!( + queue_conditional!( + enable_styling, self.stderr, style::Print(format!( "{}{TRUST_ALL_TEXT}\n\n", @@ -1198,7 +1226,8 @@ impl ChatSession { if let Some(ref id) = self.conversation.model { let model_options = get_model_options(os).await?; if let Some(model_option) = model_options.iter().find(|option| option.model_id == *id) { - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Cyan), style::Print(format!("🤖 You are chatting with {}\n", model_option.name)), @@ -1271,10 +1300,12 @@ impl ChatSession { request_metadata_lock: Arc>>, ) -> Result { let hist = self.conversation.history(); + let enable_styling = self.enable_styling(); debug!(?strategy, ?hist, "compacting history"); if self.conversation.history().is_empty() { - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Yellow), style::Print("\nConversation too short to compact.\n\n"), @@ -1288,7 +1319,8 @@ impl ChatSession { if strategy.truncate_large_messages { info!("truncating large messages"); - execute!( + execute_conditional!( + enable_styling, self.stderr, terminal::Clear(terminal::ClearType::CurrentLine), cursor::MoveToColumn(0), @@ -1445,7 +1477,8 @@ impl ChatSession { // Print output to the user. { - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Green), style::Print("✔ Conversation history has been compacted successfully!\n\n"), @@ -1454,7 +1487,8 @@ impl ChatSession { let mut output = Vec::new(); if let Some(custom_prompt) = &custom_prompt { - execute!( + execute_conditional!( + enable_styling, output, style::Print(format!("• Custom prompt applied: {}\n", custom_prompt)) )?; @@ -1466,7 +1500,8 @@ impl ChatSession { // Add a border around the summary for better visual separation let terminal_width = self.terminal_width(); let border = "═".repeat(terminal_width.min(80)); - execute!( + execute_conditional!( + enable_styling, self.stderr, style::Print("\n"), style::SetForegroundColor(Color::Cyan), @@ -1480,7 +1515,8 @@ impl ChatSession { style::Print("\n\n"), )?; - execute!( + execute_conditional!( + enable_styling, output, style::Print(&summary), style::Print("\n\n"), @@ -1490,7 +1526,8 @@ impl ChatSession { )?; animate_output(&mut self.stderr, &output)?; - execute!( + execute_conditional!( + enable_styling, self.stderr, style::Print(&border), style::Print("\n\n"), @@ -1502,7 +1539,7 @@ impl ChatSession { if should_retry { Ok(ChatState::HandleResponseStream( self.conversation - .as_sendable_conversation_state(os, &mut self.stderr, false) + .as_sendable_conversation_state(os, &mut self.stderr, self.enable_styling, false) .await?, )) } else { @@ -1588,7 +1625,8 @@ impl ChatSession { } async fn handle_input(&mut self, os: &mut Os, mut user_input: String) -> Result { - queue!(self.stderr, style::Print('\n'))?; + let enable_styling = self.enable_styling(); + queue_conditional!(enable_styling, self.stderr, style::Print('\n'))?; user_input = sanitize_unicode_tags(&user_input); let input = user_input.trim(); @@ -1632,7 +1670,8 @@ impl ChatSession { } }, Err(err) => { - queue!( + queue_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Red), style::Print(format!("\nFailed to execute command: {}\n", err)), @@ -1717,7 +1756,8 @@ impl ChatSession { match result { Ok(status) => { if !status.success() { - queue!( + queue_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Yellow), style::Print(format!("Self exited with status: {}\n", status)), @@ -1726,7 +1766,8 @@ impl ChatSession { } }, Err(e) => { - queue!( + queue_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Red), style::Print(format!("\nFailed to execute command: {}\n", e)), @@ -1795,13 +1836,13 @@ impl ChatSession { let conv_state = self .conversation - .as_sendable_conversation_state(os, &mut self.stderr, true) + .as_sendable_conversation_state(os, &mut self.stderr, self.enable_styling, true) .await?; self.send_tool_use_telemetry(os).await; - queue!(self.stderr, style::SetForegroundColor(Color::Magenta))?; - queue!(self.stderr, style::SetForegroundColor(Color::Reset))?; - queue!(self.stderr, cursor::Hide)?; + queue_conditional!(enable_styling, self.stderr, style::SetForegroundColor(Color::Magenta))?; + queue_conditional!(enable_styling, self.stderr, style::SetForegroundColor(Color::Reset))?; + queue_conditional!(enable_styling, self.stderr, cursor::Hide)?; if self.interactive { self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_owned())); @@ -1878,6 +1919,8 @@ impl ChatSession { let mut tool_results = vec![]; let mut image_blocks: Vec = Vec::new(); + let enable_styling = self.enable_styling(); + for tool in &self.tool_uses { let tool_start = std::time::Instant::now(); let mut tool_telemetry = self.tool_use_telemetry_events.entry(tool.id.clone()); @@ -1900,14 +1943,15 @@ impl ChatSession { let invoke_result = tool.tool.invoke(os, &mut self.stdout).await; if self.spinner.is_some() { - queue!( + queue_conditional!( + enable_styling, self.stderr, terminal::Clear(terminal::ClearType::CurrentLine), cursor::MoveToColumn(0), cursor::Show )?; } - execute!(self.stdout, style::Print("\n"))?; + execute_conditional!(enable_styling, self.stdout, style::Print("\n"))?; let tool_end_time = Instant::now(); let tool_time = tool_end_time.duration_since(tool_start); @@ -1943,7 +1987,8 @@ impl ChatSession { } debug!("tool result output: {:#?}", result); - execute!( + execute_conditional!( + enable_styling, self.stdout, style::Print(CONTINUATION_LINE), style::Print("\n"), @@ -1967,7 +2012,8 @@ impl ChatSession { }, Err(err) => { error!(?err, "An error occurred processing the tool"); - execute!( + execute_conditional!( + enable_styling, self.stderr, style::Print(CONTINUATION_LINE), style::Print("\n"), @@ -2007,7 +2053,8 @@ impl ChatSession { if !image_blocks.is_empty() { let images = image_blocks.into_iter().map(|(block, _)| block).collect(); self.conversation.add_tool_results_with_images(tool_results, images); - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetAttribute(Attribute::Reset), style::SetForegroundColor(Color::Reset), @@ -2017,8 +2064,13 @@ impl ChatSession { self.conversation.add_tool_results(tool_results); } - execute!(self.stderr, cursor::Hide)?; - execute!(self.stderr, style::Print("\n"), style::SetAttribute(Attribute::Reset))?; + execute_conditional!(enable_styling, self.stderr, cursor::Hide)?; + execute_conditional!( + enable_styling, + self.stderr, + style::Print("\n"), + style::SetAttribute(Attribute::Reset) + )?; if self.interactive { self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_string())); } @@ -2028,7 +2080,7 @@ impl ChatSession { self.send_tool_use_telemetry(os).await; return Ok(ChatState::HandleResponseStream( self.conversation - .as_sendable_conversation_state(os, &mut self.stderr, false) + .as_sendable_conversation_state(os, &mut self.stderr, self.enable_styling, false) .await?, )); } @@ -2049,6 +2101,7 @@ impl ChatSession { let mut rx = self.send_message(os, state, request_metadata_lock, None).await?; let request_id = rx.request_id().map(String::from); + let enable_styling = self.enable_styling(); let mut buf = String::new(); let mut offset = 0; @@ -2064,7 +2117,8 @@ impl ChatSession { if self.spinner.is_some() { drop(self.spinner.take()); - queue!( + queue_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Reset), cursor::MoveToColumn(0), @@ -2086,8 +2140,9 @@ impl ChatSession { }, parser::ResponseEvent::AssistantText(text) => { // Add Q response prefix before the first assistant text. - if !response_prefix_printed && !text.trim().is_empty() { - queue!( + if !response_prefix_printed && !text.trim().is_empty() && enable_styling { + queue_conditional!( + enable_styling, self.stdout, style::SetForegroundColor(Color::Green), style::Print("> "), @@ -2100,7 +2155,8 @@ impl ChatSession { parser::ResponseEvent::ToolUse(tool_use) => { if self.spinner.is_some() { drop(self.spinner.take()); - queue!( + queue_conditional!( + enable_styling, self.stderr, terminal::Clear(terminal::ClearType::CurrentLine), cursor::MoveToColumn(0), @@ -2154,8 +2210,11 @@ impl ChatSession { duration.as_secs() ); - execute!(self.stderr, cursor::Hide)?; - self.spinner = Some(Spinner::new(Spinners::Dots, "Dividing up the work...".to_string())); + execute_conditional!(enable_styling, self.stderr, cursor::Hide)?; + if self.interactive { + self.spinner = + Some(Spinner::new(Spinners::Dots, "Dividing up the work...".to_string())); + } // For stream timeouts, we'll tell the model to try and split its response into // smaller chunks. @@ -2173,7 +2232,7 @@ impl ChatSession { self.send_tool_use_telemetry(os).await; return Ok(ChatState::HandleResponseStream( self.conversation - .as_sendable_conversation_state(os, &mut self.stderr, false) + .as_sendable_conversation_state(os, &mut self.stderr, self.enable_styling, false) .await?, )); }, @@ -2210,7 +2269,7 @@ impl ChatSession { self.send_tool_use_telemetry(os).await; return Ok(ChatState::HandleResponseStream( self.conversation - .as_sendable_conversation_state(os, &mut self.stderr, false) + .as_sendable_conversation_state(os, &mut self.stderr, self.enable_styling, false) .await?, )); }, @@ -2244,7 +2303,8 @@ impl ChatSession { if tool_name_being_recvd.is_none() && !buf.is_empty() && self.spinner.is_some() { drop(self.spinner.take()); - queue!( + queue_conditional!( + enable_styling, self.stderr, terminal::Clear(terminal::ClearType::CurrentLine), cursor::MoveToColumn(0), @@ -2253,29 +2313,39 @@ impl ChatSession { } // Print the response for normal cases - loop { - let input = Partial::new(&buf[offset..]); - match interpret_markdown(input, &mut self.stdout, &mut state) { - Ok(parsed) => { - offset += parsed.offset_from(&input); - self.stdout.flush()?; - state.newline = state.set_newline; - state.set_newline = false; - }, - Err(err) => match err.into_inner() { - Some(err) => return Err(ChatError::Custom(err.to_string().into())), - None => break, // Data was incomplete - }, - } + if enable_styling { + loop { + let input = Partial::new(&buf[offset..]); + match interpret_markdown(input, &mut self.stdout, &mut state) { + Ok(parsed) => { + offset += parsed.offset_from(&input); + self.stdout.flush()?; + state.newline = state.set_newline; + state.set_newline = false; + }, + Err(err) => match err.into_inner() { + Some(err) => return Err(ChatError::Custom(err.to_string().into())), + None => break, // Data was incomplete + }, + } - // TODO: We should buffer output based on how much we have to parse, not as a constant - // Do not remove unless you are nabochay :) - tokio::time::sleep(Duration::from_millis(8)).await; + // TODO: We should buffer output based on how much we have to parse, not as a constant + // Do not remove unless you are nabochay :) + tokio::time::sleep(Duration::from_millis(8)).await; + } + } else { + // When styling is disabled, output raw text without markdown processing + let remaining = &buf[offset..]; + if !remaining.is_empty() { + write!(self.stdout, "{}", remaining)?; + self.stdout.flush()?; + offset = buf.len(); + } } // Set spinner after showing all of the assistant text content so far. if tool_name_being_recvd.is_some() { - queue!(self.stderr, cursor::Hide)?; + queue_conditional!(enable_styling, self.stderr, cursor::Hide)?; if self.interactive { self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_string())); } @@ -2292,11 +2362,17 @@ impl ChatSession { play_notification_bell(tool_uses.is_empty()); } - queue!(self.stderr, style::ResetColor, style::SetAttribute(Attribute::Reset))?; - execute!(self.stdout, style::Print("\n"))?; + queue_conditional!( + enable_styling, + self.stderr, + style::ResetColor, + style::SetAttribute(Attribute::Reset) + )?; + execute_conditional!(enable_styling, self.stdout, style::Print("\n"))?; for (i, citation) in &state.citations { - queue!( + queue_conditional!( + enable_styling, self.stdout, style::Print("\n"), style::SetForegroundColor(Color::Blue), @@ -2329,6 +2405,7 @@ impl ChatSession { async fn validate_tools(&mut self, os: &Os, tool_uses: Vec) -> Result { let conv_id = self.conversation.conversation_id().to_owned(); + let enable_styling = self.enable_styling(); debug!(?tool_uses, "Validating tool uses"); let mut queued_tools: Vec = Vec::new(); let mut tool_results: Vec = Vec::new(); @@ -2379,7 +2456,8 @@ impl ChatSession { // If we have any validation errors, then return them immediately to the model. if !tool_results.is_empty() { debug!(?tool_results, "Error found in the model tools"); - queue!( + queue_conditional!( + enable_styling, self.stderr, style::SetAttribute(Attribute::Bold), style::Print("Tool validation failed: "), @@ -2395,7 +2473,8 @@ impl ChatSession { .ok(), }; if let Some(content) = content { - queue!( + queue_conditional!( + enable_styling, self.stderr, style::Print("\n"), style::SetForegroundColor(Color::Red), @@ -2420,7 +2499,7 @@ impl ChatSession { return Ok(ChatState::HandleResponseStream( self.conversation - .as_sendable_conversation_state(os, &mut self.stderr, false) + .as_sendable_conversation_state(os, &mut self.stderr, self.enable_styling, false) .await?, )); } @@ -2453,7 +2532,7 @@ impl ChatSession { Ok(ChatState::HandleResponseStream( self.conversation - .as_sendable_conversation_state(os, &mut self.stderr, true) + .as_sendable_conversation_state(os, &mut self.stderr, self.enable_styling, true) .await?, )) } @@ -2484,8 +2563,10 @@ impl ChatSession { async fn print_tool_description(&mut self, os: &Os, tool_index: usize, trusted: bool) -> Result<(), ChatError> { let tool_use = &self.tool_uses[tool_index]; + let enable_styling = self.enable_styling(); - queue!( + queue_conditional!( + enable_styling, self.stdout, style::SetForegroundColor(Color::Magenta), style::Print(format!( @@ -2496,7 +2577,8 @@ impl ChatSession { style::SetForegroundColor(Color::Reset) )?; if let Tool::Custom(ref tool) = tool_use.tool { - queue!( + queue_conditional!( + enable_styling, self.stdout, style::SetForegroundColor(Color::Reset), style::Print(" from mcp server "), @@ -2506,7 +2588,8 @@ impl ChatSession { )?; } - execute!( + execute_conditional!( + enable_styling, self.stdout, style::Print("\n"), style::Print(CONTINUATION_LINE), @@ -2526,6 +2609,7 @@ impl ChatSession { /// Helper function to read user input with a prompt and Ctrl+C handling fn read_user_input(&mut self, prompt: &str, exit_on_single_ctrl_c: bool) -> Option { let mut ctrl_c = false; + let enable_styling = self.enable_styling(); loop { match (self.input_source.read_line(Some(prompt)), ctrl_c) { (Ok(Some(line)), _) => { @@ -2538,7 +2622,8 @@ impl ChatSession { if exit_on_single_ctrl_c { return None; } - execute!( + execute_conditional!( + enable_styling, self.stderr, style::Print(format!( "\n(To exit the CLI, press Ctrl+C or Ctrl+D again or type {})\n\n", @@ -2581,14 +2666,21 @@ impl ChatSession { self.conversation.agents.trust_all_tools } + /// Returns true if text styling should be disabled for this session + fn enable_styling(&self) -> bool { + self.enable_styling + } + /// Display character limit warnings based on current conversation size async fn display_char_warnings(&mut self, os: &Os) -> Result<(), ChatError> { let warning_level = self.conversation.get_token_warning_level(os).await?; + let enable_styling = self.enable_styling(); match warning_level { TokenWarningLevel::Critical => { // Memory constraint warning with gentler wording - execute!( + execute_conditional!( + enable_styling, self.stderr, style::SetForegroundColor(Color::Yellow), style::SetAttribute(Attribute::Bold), diff --git a/crates/chat-cli/src/cli/chat/util/mod.rs b/crates/chat-cli/src/cli/chat/util/mod.rs index cc4b739a9..5f80c1f5c 100644 --- a/crates/chat-cli/src/cli/chat/util/mod.rs +++ b/crates/chat-cli/src/cli/chat/util/mod.rs @@ -1,5 +1,6 @@ pub mod images; pub mod issue; +pub mod styling; #[cfg(test)] pub mod test; pub mod ui; diff --git a/crates/chat-cli/src/cli/chat/util/styling.rs b/crates/chat-cli/src/cli/chat/util/styling.rs new file mode 100644 index 000000000..35ea023a9 --- /dev/null +++ b/crates/chat-cli/src/cli/chat/util/styling.rs @@ -0,0 +1,214 @@ +use std::fmt::Display; + +use crossterm::{ + cursor, + style, + terminal, +}; + +/// Trait to identify commands that perform terminal styling operations. +/// Used by conditional macros to skip styling commands when styling is disabled. +pub trait IsStyling { + /// Returns true if this command performs styling (colors, attributes, cursor movement, etc.) + fn is_styling(&self) -> bool; +} + +/// Macro to mark crossterm command types as styling operations. +/// This allows the conditional execution macros to skip these commands +/// when styling is disabled (e.g., in non-interactive mode). +#[macro_export] +macro_rules! declare_as_styling { + ($type_name:path) => { + impl $crate::cli::chat::util::styling::IsStyling for $type_name { + fn is_styling(&self) -> bool { + true + } + } + }; +} + +/// Conditionally executes crossterm commands based on styling preference. +/// Styling commands (colors, cursor movement, etc.) are skipped when styling is disabled, +/// while non-styling commands (like Print) are always executed. +/// Equivalent to crossterm::execute! but with conditional styling support. +#[macro_export] +macro_rules! execute_conditional { + ($enable_styling:expr, $writer:expr $(, $command:expr)* $(,)? ) => {{ + // Queue each command, then flush + $crate::queue_conditional!($enable_styling, $writer $(, $command)*) + .and_then(|()| { + ::std::io::Write::flush($writer.by_ref()) + }) + }} +} + +/// Conditionally queues crossterm commands based on styling preference. +/// When styling is enabled, all commands are queued normally. +/// When styling is disabled, only non-styling commands (like Print) are queued, +/// while styling commands are skipped to avoid terminal formatting in non-interactive mode. +/// Equivalent to crossterm::queue! but with conditional styling support. +#[macro_export] +macro_rules! queue_conditional { + ($enable_styling:expr, $writer:expr $(, $command:expr)* $(,)?) => {{ + // This allows the macro to take both mut impl Write and &mut impl Write. + Ok($writer.by_ref()) + $(.and_then(|writer| { + if ($enable_styling || !$crate::cli::chat::util::styling::IsStyling::is_styling(&$command)) { + crossterm::QueueableCommand::queue(writer, $command) + } else { + std::io::Result::Ok(writer) + } + }))* + .map(|_| ()) + }} +} + +// Styling commands - these are skipped when styling is disabled +declare_as_styling!(style::SetForegroundColor); +declare_as_styling!(style::ResetColor); +declare_as_styling!(style::SetAttribute); +declare_as_styling!(terminal::Clear); +declare_as_styling!(cursor::MoveToColumn); +declare_as_styling!(cursor::Hide); +declare_as_styling!(cursor::Show); + +// Non-styling commands - these are always executed +impl IsStyling for style::Print { + fn is_styling(&self) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use std::io::Write; + + use crossterm::style::Color; + use crossterm::{ + Command, + style, + }; + + #[test] + fn test_queue_conditional_macros() { + let mut output: Vec = Vec::new(); + + // With styling enabled - should include ANSI codes + queue_conditional!( + true, + &mut output, + style::SetForegroundColor(Color::Red), + style::Print("test"), + style::ResetColor + ) + .unwrap(); + + let styled_output = String::from_utf8(output.clone()).unwrap(); + let mut expected_red = String::new(); + style::SetForegroundColor(Color::Red) + .write_ansi(&mut expected_red) + .unwrap(); + let mut expected_reset = String::new(); + style::ResetColor.write_ansi(&mut expected_reset).unwrap(); + + assert!(styled_output.contains(&expected_red)); + assert!(styled_output.contains(&expected_reset)); + assert!(styled_output.contains("test")); + + // With styling disabled - should only contain text + output.clear(); + queue_conditional!( + false, + &mut output, + style::SetForegroundColor(Color::Red), + style::Print("test"), + style::ResetColor + ) + .unwrap(); + + let unstyled_output = String::from_utf8(output).unwrap(); + assert_eq!(unstyled_output, "test"); + assert!(!unstyled_output.contains(&expected_red)); + assert!(!unstyled_output.contains(&expected_reset)); + } + + #[test] + fn test_execute_conditional_macro() { + use std::sync::Arc; + use std::sync::atomic::{ + AtomicUsize, + Ordering, + }; + + // Mock writer that tracks flush calls + struct MockWriter { + buffer: Vec, + flush_count: Arc, + } + + impl std::io::Write for MockWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.buffer.extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.flush_count.fetch_add(1, Ordering::Relaxed); + Ok(()) + } + } + + let flush_count = Arc::new(AtomicUsize::new(0)); + let mut mock_writer = MockWriter { + buffer: Vec::new(), + flush_count: flush_count.clone(), + }; + + let mut expected_red = String::new(); + style::SetForegroundColor(Color::Red) + .write_ansi(&mut expected_red) + .unwrap(); + let mut expected_reset = String::new(); + style::ResetColor.write_ansi(&mut expected_reset).unwrap(); + + // Test execute_conditional with styling enabled + execute_conditional!( + true, + &mut mock_writer, + style::SetForegroundColor(Color::Red), + style::Print("test"), + style::ResetColor + ) + .unwrap(); + + let styled_output = String::from_utf8(mock_writer.buffer.clone()).unwrap(); + assert!(styled_output.contains(&expected_red)); + assert!(styled_output.contains(&expected_reset)); + assert!(styled_output.contains("test")); + assert_eq!(flush_count.load(Ordering::Relaxed), 1, "flush should be called once"); + + // Reset for next test + mock_writer.buffer.clear(); + flush_count.store(0, Ordering::Relaxed); + + // Test execute_conditional with styling disabled + execute_conditional!( + false, + &mut mock_writer, + style::SetForegroundColor(Color::Red), + style::Print("test"), + style::ResetColor + ) + .unwrap(); + + let unstyled_output = String::from_utf8(mock_writer.buffer).unwrap(); + assert_eq!(unstyled_output, "test"); + assert!(!unstyled_output.contains(&expected_red)); + assert!(!unstyled_output.contains(&expected_reset)); + assert_eq!( + flush_count.load(Ordering::Relaxed), + 1, + "flush should still be called once even when styling is disabled" + ); + } +}