diff --git a/crates/chat-cli/src/cli/chat/cli.rs b/crates/chat-cli/src/cli/chat/cli.rs index d4d10a450b..d435a4c12e 100644 --- a/crates/chat-cli/src/cli/chat/cli.rs +++ b/crates/chat-cli/src/cli/chat/cli.rs @@ -17,9 +17,9 @@ pub struct Chat { /// prompt requests permissions to use a tool, unless --trust-all-tools is also used. #[arg(long)] pub no_interactive: bool, - /// Start a new conversation and overwrites any previous conversation from this directory. - #[arg(long)] - pub new: bool, + /// Resumes the previous conversation from this directory. + #[arg(short, long)] + pub resume: bool, /// The first question to ask pub input: Option, /// Context profile to use diff --git a/crates/chat-cli/src/cli/chat/command.rs b/crates/chat-cli/src/cli/chat/command.rs index 006eb9b151..f2a4262b60 100644 --- a/crates/chat-cli/src/cli/chat/command.rs +++ b/crates/chat-cli/src/cli/chat/command.rs @@ -51,10 +51,10 @@ pub enum Command { subcommand: Option, }, Usage, - Import { + Load { path: String, }, - Export { + Save { path: String, force: bool, }, @@ -818,15 +818,15 @@ impl Command { } }, "usage" => Self::Usage, - "import" => { + "load" => { let Some(path) = parts.get(1) else { return Err("path is required".to_string()); }; - Self::Import { + Self::Load { path: (*path).to_string(), } }, - "export" => { + "save" => { let force = parts.contains(&"-f") || parts.contains(&"--force"); let Some(path) = parts.get(1) else { return Err("path is required".to_string()); @@ -835,7 +835,7 @@ impl Command { if !path.ends_with(".json") { path.push_str(".json"); } - Self::Export { path, force } + Self::Save { path, force } }, unknown_command => { let looks_like_path = { diff --git a/crates/chat-cli/src/cli/chat/conversation_state.rs b/crates/chat-cli/src/cli/chat/conversation_state.rs index 23d5fdb434..2feab5bbd0 100644 --- a/crates/chat-cli/src/cli/chat/conversation_state.rs +++ b/crates/chat-cli/src/cli/chat/conversation_state.rs @@ -1,5 +1,6 @@ use std::collections::{ HashMap, + HashSet, VecDeque, }; use std::sync::Arc; @@ -32,7 +33,6 @@ use super::hooks::{ }; use super::message::{ AssistantMessage, - AssistantToolUse, ToolUseResult, ToolUseResultBlock, UserMessage, @@ -60,7 +60,6 @@ use crate::api_client::model::{ ToolInputSchema, ToolResult, ToolResultContentBlock, - ToolResultStatus, ToolSpecification, ToolUse, UserInputMessage, @@ -347,7 +346,7 @@ impl ConversationState { } } - self.enforce_tool_use_history_invariants(true); + self.enforce_tool_use_history_invariants(); } /// Here we also need to make sure that the tool result corresponds to one of the tools @@ -362,105 +361,52 @@ impl ConversationState { /// intervention here is to substitute the ambiguous, partial name with a dummy. /// 3. The model had decided to call a tool that does not exist. The intervention here is to /// substitute the non-existent tool name with a dummy. - pub fn enforce_tool_use_history_invariants(&mut self, last_only: bool) { - let tool_name_list = self.tool_manager.tn_map.keys().map(String::as_str).collect::>(); - // We need to first determine what the range of interest is. There are two places where we - // would call this function: - // 1. When there are changes to the list of available tools, in which case we comb through the - // entire conversation - // 2. When we send a message, in which case we only examine the most recent entry - let (tool_use_results, mut tool_uses) = if last_only { - if let (Some((_, AssistantMessage::ToolUse { ref mut tool_uses, .. })), Some(user_msg)) = ( - self.history - .range_mut(self.valid_history_range.0..self.valid_history_range.1) - .last(), - &mut self.next_message, - ) { - let tool_use_results = user_msg - .tool_use_results() - .map_or(Vec::new(), |results| results.iter().collect::>()); - let tool_uses = tool_uses.iter_mut().collect::>(); - (tool_use_results, tool_uses) - } else { - (Vec::new(), Vec::new()) - } - } else { - let tool_use_results = self.next_message.as_ref().map_or(Vec::new(), |user_msg| { - user_msg - .tool_use_results() - .map_or(Vec::new(), |results| results.iter().collect::>()) - }); - self.history - .iter_mut() - .filter_map(|(user_msg, asst_msg)| { - if let (Some(tool_use_results), AssistantMessage::ToolUse { ref mut tool_uses, .. }) = - (user_msg.tool_use_results(), asst_msg) - { - Some((tool_use_results, tool_uses)) - } else { - None - } + pub fn enforce_tool_use_history_invariants(&mut self) { + let tool_names: HashSet<_> = self + .tools + .values() + .flat_map(|tools| { + tools.iter().map(|tool| match tool { + Tool::ToolSpecification(tool_specification) => tool_specification.name.as_str(), }) - .fold( - (tool_use_results, Vec::<&mut AssistantToolUse>::new()), - |(mut tool_use_results, mut tool_uses), (results, uses)| { - let mut results = results.iter().collect::>(); - let mut uses = uses.iter_mut().collect::>(); - tool_use_results.append(&mut results); - tool_uses.append(&mut uses); - (tool_use_results, tool_uses) - }, - ) - }; + }) + .filter(|name| *name != DUMMY_TOOL_NAME) + .collect(); + + for (_, assistant) in &mut self.history { + if let AssistantMessage::ToolUse { ref mut tool_uses, .. } = assistant { + for tool_use in tool_uses { + if tool_names.contains(tool_use.name.as_str()) { + continue; + } - // Replace tool uses associated with tools that does not exist / no longer exists with - // dummy (i.e. put them to sleep / dormant) - for result in tool_use_results { - let tool_use_id = result.tool_use_id.as_str(); - let corresponding_tool_use = tool_uses.iter_mut().find(|tool_use| tool_use_id == tool_use.id); - if let Some(tool_use) = corresponding_tool_use { - if tool_name_list.contains(&tool_use.name.as_str()) { - // If this tool matches of the tools in our list, this is not our - // concern, error or not. - continue; - } - if let ToolResultStatus::Error = result.status { - // case 2 and 3 - tool_use.name = DUMMY_TOOL_NAME.to_string(); - tool_use.args = serde_json::json!({}); - } else { - // case 1 - let full_name = tool_name_list.iter().find(|name| name.ends_with(&tool_use.name)); - // We should be able to find a match but if not we'll just treat it as - // a dummy and move on - if let Some(full_name) = full_name { - tool_use.name = (*full_name).to_string(); - } else { - tool_use.name = DUMMY_TOOL_NAME.to_string(); - tool_use.args = serde_json::json!({}); + if tool_names.contains(tool_use.orig_name.as_str()) { + tool_use.name = tool_use.orig_name.clone(); + tool_use.args = tool_use.orig_args.clone(); + continue; } - } - } - } - // Revive tools that were previously dormant if they now corresponds to one of the tools in - // our list of available tools. Note that this check only works because tn_map does NOT - // contain names of native tools. - for tool_use in tool_uses { - if tool_use.name == DUMMY_TOOL_NAME - && tool_use - .orig_name - .as_ref() - .is_some_and(|name| tool_name_list.contains(&(*name).as_str())) - { - tool_use.name = tool_use - .orig_name - .as_ref() - .map_or(DUMMY_TOOL_NAME.to_string(), |name| name.clone()); - tool_use.args = tool_use - .orig_args - .as_ref() - .map_or(serde_json::json!({}), |args| args.clone()); + let names: Vec<&str> = tool_names + .iter() + .filter_map(|name| { + if name.ends_with(&tool_use.name) { + Some(*name) + } else { + None + } + }) + .collect(); + + // There's only one tool use matching, so we can just replace it with the + // found name. + if names.len() == 1 { + tool_use.name = (*names.first().unwrap()).to_string(); + continue; + } + + // Otherwise, we have to replace it with a dummy. + tool_use.name = DUMMY_TOOL_NAME.to_string(); + } } } } @@ -514,8 +460,8 @@ impl ConversationState { .expect("unable to construct conversation state") } - pub async fn update_state(&mut self) { - let needs_update = self.tool_manager.has_new_stuff.load(Ordering::Acquire); + pub async fn update_state(&mut self, force_update: bool) { + let needs_update = self.tool_manager.has_new_stuff.load(Ordering::Acquire) || force_update; if !needs_update { return; } @@ -540,12 +486,13 @@ impl ConversationState { // We call this in [Self::enforce_conversation_invariants] as well. But we need to call it // here as well because when it's being called in [Self::enforce_conversation_invariants] // it is only checking the last entry. - self.enforce_tool_use_history_invariants(false); + self.enforce_tool_use_history_invariants(); } /// Returns a conversation state representation which reflects the exact conversation to send /// back to the model. pub async fn backend_conversation_state(&mut self, run_hooks: bool, quiet: bool) -> BackendConversationState<'_> { + self.update_state(false).await; self.enforce_conversation_invariants(); // Run hooks and add to conversation start and next user message. diff --git a/crates/chat-cli/src/cli/chat/message.rs b/crates/chat-cli/src/cli/chat/message.rs index 733a18f581..46bad55d18 100644 --- a/crates/chat-cli/src/cli/chat/message.rs +++ b/crates/chat-cli/src/cli/chat/message.rs @@ -349,11 +349,11 @@ pub struct AssistantToolUse { /// The name for the tool as exposed to the model pub name: String, /// Original name of the tool - pub orig_name: Option, + pub orig_name: String, /// The input to pass to the tool as exposed to the model pub args: serde_json::Value, /// Original input passed to the tool - pub orig_args: Option, + pub orig_args: serde_json::Value, } impl From for ToolUse { diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs index 2fd1e5dbe7..0ca68b33ca 100644 --- a/crates/chat-cli/src/cli/chat/mod.rs +++ b/crates/chat-cli/src/cli/chat/mod.rs @@ -203,7 +203,8 @@ const WELCOME_TEXT: &str = color_print::cstr! {" const SMALL_SCREEN_WELCOME_TEXT: &str = color_print::cstr! {"Welcome to Amazon Q!"}; const RESUME_TEXT: &str = color_print::cstr! {"Picking up where we left off..."}; -const ROTATING_TIPS: [&str; 11] = [ +const ROTATING_TIPS: [&str; 12] = [ + color_print::cstr! {"You can resume the last conversation from your current directory by launching with q chat --resume"}, color_print::cstr! {"Get notified whenever Q CLI finishes responding. Just run q settings chat.enableNotifications true"}, color_print::cstr! {"You can use /editor to edit your prompt with a vim-like experience"}, color_print::cstr! {"/usage shows you a visual breakdown of your current context window usage"}, @@ -263,8 +264,8 @@ const HELP_TEXT: &str = color_print::cstr! {" clear Clear all files from current context [--global] hooks View and manage context hooks /usage Show current session's context window usage -/import Import conversation state from a JSON file -/export Export conversation state to a JSON file +/load Load conversation state from a JSON file +/save Save conversation state to a JSON file MCP: You can now configure the Amazon Q CLI to use MCP servers. \nLearn how: https://docs.aws.amazon.com/en_us/amazonq/latest/qdeveloper-ug/command-line-mcp.html @@ -299,7 +300,7 @@ pub async fn launch_chat(database: &mut Database, telemetry: &TelemetryThread, a telemetry, args.input, args.no_interactive, - args.new, + args.resume, args.accept_all, args.profile, args.trust_all_tools, @@ -314,7 +315,7 @@ pub async fn chat( telemetry: &TelemetryThread, input: Option, no_interactive: bool, - new_conversation: bool, + resume_conversation: bool, accept_all: bool, profile: Option, trust_all_tools: bool, @@ -443,7 +444,7 @@ pub async fn chat( input, InputSource::new(database, prompt_request_sender, prompt_response_receiver)?, interactive, - new_conversation, + resume_conversation, client, || terminal::window_size().map(|s| s.columns.into()).ok(), tool_manager, @@ -530,7 +531,7 @@ impl ChatContext { mut input: Option, input_source: InputSource, interactive: bool, - new_conversation: bool, + resume_conversation: bool, client: StreamingClient, terminal_width_provider: fn() -> Option, tool_manager: ToolManager, @@ -542,22 +543,7 @@ impl ChatContext { let output_clone = output.clone(); let mut existing_conversation = false; - let conversation_state = if new_conversation { - let new_state = ConversationState::new( - ctx_clone, - conversation_id, - tool_config, - profile, - Some(output_clone), - tool_manager, - ) - .await; - - std::env::current_dir() - .ok() - .and_then(|cwd| database.set_conversation_by_path(cwd, &new_state).ok()); - new_state - } else { + let conversation_state = if resume_conversation { let prior = std::env::current_dir() .ok() .and_then(|cwd| database.get_conversation_by_path(cwd).ok()) @@ -572,7 +558,8 @@ impl ChatContext { cs.reload_serialized_state(Arc::clone(&ctx), Some(output.clone())).await; input = Some(input.unwrap_or("In a few words, summarize our conversation so far.".to_owned())); cs.tool_manager = tool_manager; - cs.enforce_tool_use_history_invariants(false); + cs.update_state(true).await; + cs.enforce_tool_use_history_invariants(); cs } else { ConversationState::new( @@ -585,6 +572,16 @@ impl ChatContext { ) .await } + } else { + ConversationState::new( + ctx_clone, + conversation_id, + tool_config, + profile, + Some(output_clone), + tool_manager, + ) + .await }; Ok(Self { @@ -816,7 +813,7 @@ impl ChatContext { debug!(?chat_state, "changing to state"); // Update conversation state with new tool information - self.conversation_state.update_state().await; + self.conversation_state.update_state(false).await; let result = match chat_state { ChatState::PromptUser { @@ -2812,7 +2809,7 @@ impl ChatContext { skip_printing_tools: true, } }, - Command::Import { path } => { + Command::Load { path } => { macro_rules! tri { ($v:expr) => { match $v { @@ -2854,7 +2851,7 @@ impl ChatContext { skip_printing_tools: true, } }, - Command::Export { path, force } => { + Command::Save { path, force } => { macro_rules! tri { ($v:expr) => { match $v { diff --git a/crates/chat-cli/src/cli/chat/parser.rs b/crates/chat-cli/src/cli/chat/parser.rs index ab3bcd208b..8211271697 100644 --- a/crates/chat-cli/src/cli/chat/parser.rs +++ b/crates/chat-cli/src/cli/chat/parser.rs @@ -222,9 +222,9 @@ impl ResponseParser { self.tool_uses.push(AssistantToolUse { id: id.clone(), name: name.clone(), - orig_name: Some(name.clone()), + orig_name: name.clone(), args: args.clone(), - orig_args: Some(args.clone()), + orig_args: args.clone(), }); let message = Box::new(AssistantMessage::new_tool_use( Some(self.message_id.clone()), @@ -244,11 +244,14 @@ impl ResponseParser { // if the tool just does not need any input _ => serde_json::json!({}), }; + let orig_name = name.clone(); + let orig_args = args.clone(); Ok(AssistantToolUse { id, name, + orig_name, args, - ..Default::default() + orig_args, }) } diff --git a/crates/chat-cli/src/cli/chat/prompt.rs b/crates/chat-cli/src/cli/chat/prompt.rs index 782731b2ce..1b36d041f5 100644 --- a/crates/chat-cli/src/cli/chat/prompt.rs +++ b/crates/chat-cli/src/cli/chat/prompt.rs @@ -76,8 +76,8 @@ pub const COMMANDS: &[&str] = &[ "/compact", "/compact help", "/usage", - "/import", - "/export", + "/save", + "/load", ]; pub fn generate_prompt(current_profile: Option<&str>, warning: bool) -> String { diff --git a/crates/chat-cli/src/cli/mod.rs b/crates/chat-cli/src/cli/mod.rs index c03ae41d43..ea32659d0b 100644 --- a/crates/chat-cli/src/cli/mod.rs +++ b/crates/chat-cli/src/cli/mod.rs @@ -369,7 +369,7 @@ mod test { subcommand: Some(CliRootCommands::Chat(Chat { accept_all: false, no_interactive: false, - new: false, + resume: false, input: None, profile: None, trust_all_tools: false, @@ -408,7 +408,7 @@ mod test { CliRootCommands::Chat(Chat { accept_all: false, no_interactive: false, - new: false, + resume: false, input: None, profile: Some("my-profile".to_string()), trust_all_tools: false, @@ -424,7 +424,7 @@ mod test { CliRootCommands::Chat(Chat { accept_all: false, no_interactive: false, - new: false, + resume: false, input: Some("Hello".to_string()), profile: Some("my-profile".to_string()), trust_all_tools: false, @@ -440,7 +440,7 @@ mod test { CliRootCommands::Chat(Chat { accept_all: true, no_interactive: false, - new: false, + resume: false, input: None, profile: Some("my-profile".to_string()), trust_all_tools: false, @@ -450,13 +450,25 @@ mod test { } #[test] - fn test_chat_with_no_interactive_new() { + fn test_chat_with_no_interactive_and_resume() { assert_parse!( - ["chat", "--no-interactive", "--new"], + ["chat", "--no-interactive", "--resume"], CliRootCommands::Chat(Chat { accept_all: false, no_interactive: true, - new: true, + resume: true, + input: None, + profile: None, + trust_all_tools: false, + trust_tools: None, + }) + ); + assert_parse!( + ["chat", "--no-interactive", "-r"], + CliRootCommands::Chat(Chat { + accept_all: false, + no_interactive: true, + resume: true, input: None, profile: None, trust_all_tools: false, @@ -472,7 +484,7 @@ mod test { CliRootCommands::Chat(Chat { accept_all: false, no_interactive: false, - new: false, + resume: false, input: None, profile: None, trust_all_tools: true, @@ -488,7 +500,7 @@ mod test { CliRootCommands::Chat(Chat { accept_all: false, no_interactive: false, - new: false, + resume: false, input: None, profile: None, trust_all_tools: false, @@ -504,7 +516,7 @@ mod test { CliRootCommands::Chat(Chat { accept_all: false, no_interactive: false, - new: false, + resume: false, input: None, profile: None, trust_all_tools: false, diff --git a/crates/q_cli/src/cli/mod.rs b/crates/q_cli/src/cli/mod.rs index 4e6e50a65c..cb9c43bca2 100644 --- a/crates/q_cli/src/cli/mod.rs +++ b/crates/q_cli/src/cli/mod.rs @@ -344,16 +344,18 @@ impl Cli { CliRootCommands::Telemetry(subcommand) => subcommand.execute().await, CliRootCommands::Version { changelog } => Self::print_version(changelog), CliRootCommands::Dashboard => launch_dashboard(false).await, - CliRootCommands::Chat { args } => Self::execute_chat(Some(args)).await, + CliRootCommands::Chat { args } => Self::execute_chat(Some(args), true).await, CliRootCommands::Inline(subcommand) => subcommand.execute(&cli_context).await, }, // Root command - None => Self::execute_chat(None).await, + None => Self::execute_chat(None, true).await, } } - async fn execute_chat(args: Option>) -> Result { - assert_logged_in().await?; + pub async fn execute_chat(args: Option>, enforce_login: bool) -> Result { + if enforce_login { + assert_logged_in().await?; + } let secret_store = SecretStore::new().await.ok(); if let Some(secret_store) = secret_store { diff --git a/crates/q_cli/src/main.rs b/crates/q_cli/src/main.rs index d8eaaa7420..eabc5119c6 100644 --- a/crates/q_cli/src/main.rs +++ b/crates/q_cli/src/main.rs @@ -26,11 +26,26 @@ fn main() -> Result { fig_telemetry::set_dispatch_mode(fig_telemetry::DispatchMode::On); fig_telemetry::init_global_telemetry_emitter(); + let mut args = std::env::args(); + let subcommand = args.nth(1); let multithread = matches!( - std::env::args().nth(1).as_deref(), + subcommand.as_deref(), Some("init" | "_" | "internal" | "completion" | "hook" | "chat") ); + let runtime = if multithread { + tokio::runtime::Builder::new_multi_thread() + } else { + tokio::runtime::Builder::new_current_thread() + } + .enable_all() + .build()?; + + // Hack as clap doesn't expose a custom command help. + if subcommand.as_deref() == Some("chat") && args.any(|arg| ["--help", "-h"].contains(&arg.as_str())) { + runtime.block_on(cli::Cli::execute_chat(Some(vec!["--help".to_owned()]), true))?; + } + let parsed = match cli::Cli::try_parse() { Ok(cli) => cli, Err(err) => { @@ -58,14 +73,6 @@ fn main() -> Result { let verbose = parsed.verbose > 0; - let runtime = if multithread { - tokio::runtime::Builder::new_multi_thread() - } else { - tokio::runtime::Builder::new_current_thread() - } - .enable_all() - .build()?; - let result = runtime.block_on(async { let result = parsed.execute().await; fig_telemetry::finish_telemetry().await;