diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index fd301913e7..d0018610a8 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -31,6 +31,8 @@ use crate::bottom_pane::paste_burst::FlushResult; use crate::bottom_pane::prompt_args::expand_custom_prompt; use crate::bottom_pane::prompt_args::expand_if_numeric_with_positional_args; use crate::bottom_pane::prompt_args::parse_slash_name; +use crate::bottom_pane::prompt_args::prompt_argument_names; +use crate::bottom_pane::prompt_args::prompt_command_with_arg_placeholders; use crate::bottom_pane::prompt_args::prompt_has_numeric_placeholders; use crate::slash_command::SlashCommand; use crate::style::user_message_style; @@ -44,6 +46,7 @@ use crate::bottom_pane::textarea::TextArea; use crate::bottom_pane::textarea::TextAreaState; use crate::clipboard_paste::normalize_pasted_path; use crate::clipboard_paste::pasted_image_format; +use crate::history_cell; use crate::ui_consts::LIVE_PREFIX_COLS; use codex_file_search::FileMatch; use std::cell::RefCell; @@ -71,6 +74,16 @@ struct AttachedImage { path: PathBuf, } +enum PromptSelectionMode { + Completion, + Submit, +} + +enum PromptSelectionAction { + Insert { text: String, cursor: Option }, + Submit { text: String }, +} + pub(crate) struct ChatComposer { textarea: TextArea, textarea_state: RefCell, @@ -436,17 +449,17 @@ impl ChatComposer { } CommandItem::UserPrompt(idx) => { if let Some(prompt) = popup.prompt(idx) { - let name = prompt.name.clone(); - let starts_with_cmd = first_line - .trim_start() - .starts_with(format!("/{PROMPTS_CMD_PREFIX}:{name}").as_str()); - if !starts_with_cmd { - self.textarea.set_text( - format!("/{PROMPTS_CMD_PREFIX}:{name} ").as_str(), - ); - } - if !self.textarea.text().is_empty() { - cursor_target = Some(self.textarea.text().len()); + match prompt_selection_action( + prompt, + first_line, + PromptSelectionMode::Completion, + ) { + PromptSelectionAction::Insert { text, cursor } => { + let target = cursor.unwrap_or(text.len()); + self.textarea.set_text(&text); + cursor_target = Some(target); + } + PromptSelectionAction::Submit { .. } => {} } } } @@ -484,28 +497,21 @@ impl ChatComposer { } CommandItem::UserPrompt(idx) => { if let Some(prompt) = popup.prompt(idx) { - let has_numeric = prompt_has_numeric_placeholders(&prompt.content); - - if !has_numeric { - // No placeholders at all: auto-submit the literal content - self.textarea.set_text(""); - return (InputResult::Submitted(prompt.content.clone()), true); - } - // Numeric placeholders present. - // If the user already typed positional args on the first line, - // expand immediately and submit; otherwise insert "/name " so - // they can type args. - let first_line = self.textarea.text().lines().next().unwrap_or(""); - if let Some(expanded) = - expand_if_numeric_with_positional_args(prompt, first_line) - { - self.textarea.set_text(""); - return (InputResult::Submitted(expanded), true); - } else { - let name = prompt.name.clone(); - let text = format!("/{PROMPTS_CMD_PREFIX}:{name} "); - self.textarea.set_text(&text); - self.textarea.set_cursor(self.textarea.text().len()); + match prompt_selection_action( + prompt, + first_line, + PromptSelectionMode::Submit, + ) { + PromptSelectionAction::Submit { text } => { + self.textarea.set_text(""); + return (InputResult::Submitted(text), true); + } + PromptSelectionAction::Insert { text, cursor } => { + let target = cursor.unwrap_or(text.len()); + self.textarea.set_text(&text); + self.textarea.set_cursor(target); + return (InputResult::None, true); + } } } return (InputResult::None, true); @@ -919,6 +925,7 @@ impl ChatComposer { return (InputResult::None, true); } let mut text = self.textarea.text().to_string(); + let original_input = text.clone(); self.textarea.set_text(""); // Replace all pending pastes in the text @@ -932,13 +939,20 @@ impl ChatComposer { // If there is neither text nor attachments, suppress submission entirely. let has_attachments = !self.attached_images.is_empty(); text = text.trim().to_string(); - - if let Some(expanded) = - expand_custom_prompt(&text, &self.custom_prompts).unwrap_or_default() - { + let expanded_prompt = match expand_custom_prompt(&text, &self.custom_prompts) { + Ok(expanded) => expanded, + Err(err) => { + self.app_event_tx.send(AppEvent::InsertHistoryCell(Box::new( + history_cell::new_error_event(err.user_message()), + ))); + self.textarea.set_text(&original_input); + self.textarea.set_cursor(original_input.len()); + return (InputResult::None, true); + } + }; + if let Some(expanded) = expanded_prompt { text = expanded; } - if text.is_empty() && !has_attachments { return (InputResult::None, true); } @@ -1472,6 +1486,54 @@ impl WidgetRef for ChatComposer { } } +fn prompt_selection_action( + prompt: &CustomPrompt, + first_line: &str, + mode: PromptSelectionMode, +) -> PromptSelectionAction { + let named_args = prompt_argument_names(&prompt.content); + let has_numeric = prompt_has_numeric_placeholders(&prompt.content); + + match mode { + PromptSelectionMode::Completion => { + if !named_args.is_empty() { + let (text, cursor) = + prompt_command_with_arg_placeholders(&prompt.name, &named_args); + return PromptSelectionAction::Insert { + text, + cursor: Some(cursor), + }; + } + if has_numeric { + let text = format!("/{PROMPTS_CMD_PREFIX}:{} ", prompt.name); + return PromptSelectionAction::Insert { text, cursor: None }; + } + let text = format!("/{PROMPTS_CMD_PREFIX}:{}", prompt.name); + PromptSelectionAction::Insert { text, cursor: None } + } + PromptSelectionMode::Submit => { + if !named_args.is_empty() { + let (text, cursor) = + prompt_command_with_arg_placeholders(&prompt.name, &named_args); + return PromptSelectionAction::Insert { + text, + cursor: Some(cursor), + }; + } + if has_numeric { + if let Some(expanded) = expand_if_numeric_with_positional_args(prompt, first_line) { + return PromptSelectionAction::Submit { text: expanded }; + } + let text = format!("/{PROMPTS_CMD_PREFIX}:{} ", prompt.name); + return PromptSelectionAction::Insert { text, cursor: None }; + } + PromptSelectionAction::Submit { + text: prompt.content.clone(), + } + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -1487,7 +1549,6 @@ mod tests { use crate::bottom_pane::InputResult; use crate::bottom_pane::chat_composer::AttachedImage; use crate::bottom_pane::chat_composer::LARGE_PASTE_CHAR_THRESHOLD; - use crate::bottom_pane::footer::footer_height; use crate::bottom_pane::prompt_args::extract_positional_args_for_prompt_line; use crate::bottom_pane::textarea::TextArea; use tokio::sync::mpsc::unbounded_channel; @@ -2625,6 +2686,174 @@ mod tests { assert!(composer.textarea.is_empty()); } + #[test] + fn custom_prompt_submission_expands_arguments() { + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + true, + sender, + false, + "Ask Codex to do anything".to_string(), + false, + ); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $USER changes on $BRANCH".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text("/prompts:my-prompt USER=Alice BRANCH=main"); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!( + InputResult::Submitted("Review Alice changes on main".to_string()), + result + ); + assert!(composer.textarea.is_empty()); + } + + #[test] + fn custom_prompt_submission_accepts_quoted_values() { + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + true, + sender, + false, + "Ask Codex to do anything".to_string(), + false, + ); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Pair $USER with $BRANCH".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text("/prompts:my-prompt USER=\"Alice Smith\" BRANCH=dev-main"); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!( + InputResult::Submitted("Pair Alice Smith with dev-main".to_string()), + result + ); + assert!(composer.textarea.is_empty()); + } + + #[test] + fn custom_prompt_invalid_args_reports_error() { + let (tx, mut rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + true, + sender, + false, + "Ask Codex to do anything".to_string(), + false, + ); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $USER changes".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text("/prompts:my-prompt USER=Alice stray"); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!(InputResult::None, result); + assert_eq!( + "/prompts:my-prompt USER=Alice stray", + composer.textarea.text() + ); + + let mut found_error = false; + while let Ok(event) = rx.try_recv() { + if let AppEvent::InsertHistoryCell(cell) = event { + let message = cell + .display_lines(80) + .into_iter() + .map(|line| line.to_string()) + .collect::>() + .join("\n"); + assert!(message.contains("expected key=value")); + found_error = true; + break; + } + } + assert!(found_error, "expected error history cell to be sent"); + } + + #[test] + fn custom_prompt_missing_required_args_reports_error() { + let (tx, mut rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + true, + sender, + false, + "Ask Codex to do anything".to_string(), + false, + ); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $USER changes on $BRANCH".to_string(), + description: None, + argument_hint: None, + }]); + + // Provide only one of the required args + composer.textarea.set_text("/prompts:my-prompt USER=Alice"); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!(InputResult::None, result); + assert_eq!("/prompts:my-prompt USER=Alice", composer.textarea.text()); + + let mut found_error = false; + while let Ok(event) = rx.try_recv() { + if let AppEvent::InsertHistoryCell(cell) = event { + let message = cell + .display_lines(80) + .into_iter() + .map(|line| line.to_string()) + .collect::>() + .join("\n"); + assert!(message.to_lowercase().contains("missing required args")); + assert!(message.contains("BRANCH")); + found_error = true; + break; + } + } + assert!( + found_error, + "expected missing args error history cell to be sent" + ); + } + #[test] fn selecting_custom_prompt_with_args_expands_placeholders() { // Support $1..$9 and $ARGUMENTS in prompt content. @@ -2663,6 +2892,37 @@ mod tests { assert_eq!(InputResult::Submitted(expected), result); } + #[test] + fn numeric_prompt_positional_args_does_not_error() { + // Ensure that a prompt with only numeric placeholders does not trigger + // key=value parsing errors when given positional arguments. + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + true, + sender, + false, + "Ask Codex to do anything".to_string(), + false, + ); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "elegant".to_string(), + path: "/tmp/elegant.md".to_string().into(), + content: "Echo: $ARGUMENTS".to_string(), + description: None, + argument_hint: None, + }]); + + // Type positional args; should submit with numeric expansion, no errors. + composer.textarea.set_text("/prompts:elegant hi"); + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!(InputResult::Submitted("Echo: hi".to_string()), result); + assert!(composer.textarea.is_empty()); + } + #[test] fn selecting_custom_prompt_with_no_args_inserts_template() { let prompt_text = "X:$1 Y:$2 All:[$ARGUMENTS]"; diff --git a/codex-rs/tui/src/bottom_pane/prompt_args.rs b/codex-rs/tui/src/bottom_pane/prompt_args.rs index 51c95d9998..48c3cedfab 100644 --- a/codex-rs/tui/src/bottom_pane/prompt_args.rs +++ b/codex-rs/tui/src/bottom_pane/prompt_args.rs @@ -1,6 +1,60 @@ use codex_protocol::custom_prompts::CustomPrompt; use codex_protocol::custom_prompts::PROMPTS_CMD_PREFIX; +use lazy_static::lazy_static; +use regex_lite::Regex; use shlex::Shlex; +use std::collections::HashMap; +use std::collections::HashSet; + +lazy_static! { + static ref PROMPT_ARG_REGEX: Regex = + Regex::new(r"\$[A-Z][A-Z0-9_]*").unwrap_or_else(|_| std::process::abort()); +} + +#[derive(Debug)] +pub enum PromptArgsError { + MissingAssignment { token: String }, + MissingKey { token: String }, +} + +impl PromptArgsError { + fn describe(&self, command: &str) -> String { + match self { + PromptArgsError::MissingAssignment { token } => format!( + "Could not parse {command}: expected key=value but found '{token}'. Wrap values in double quotes if they contain spaces." + ), + PromptArgsError::MissingKey { token } => { + format!("Could not parse {command}: expected a name before '=' in '{token}'.") + } + } + } +} + +#[derive(Debug)] +pub enum PromptExpansionError { + Args { + command: String, + error: PromptArgsError, + }, + MissingArgs { + command: String, + missing: Vec, + }, +} + +impl PromptExpansionError { + pub fn user_message(&self) -> String { + match self { + PromptExpansionError::Args { command, error } => error.describe(command), + PromptExpansionError::MissingArgs { command, missing } => { + let list = missing.join(", "); + format!( + "Missing required args for {command}: {list}. Provide as key=value (quote values with spaces)." + ) + } + } + } +} /// Parse a first-line slash command of the form `/name `. /// Returns `(name, rest_after_name)` if the line begins with `/` and contains @@ -27,6 +81,54 @@ pub fn parse_positional_args(rest: &str) -> Vec { Shlex::new(rest).collect() } +/// Extracts the unique placeholder variable names from a prompt template. +/// +/// A placeholder is any token that matches the pattern `$[A-Z][A-Z0-9_]*` +/// (for example `$USER`). The function returns the variable names without +/// the leading `$`, de-duplicated and in the order of first appearance. +pub fn prompt_argument_names(content: &str) -> Vec { + let mut seen = HashSet::new(); + let mut names = Vec::new(); + for m in PROMPT_ARG_REGEX.find_iter(content) { + if m.start() > 0 && content.as_bytes()[m.start() - 1] == b'$' { + continue; + } + let name = &content[m.start() + 1..m.end()]; + // Exclude special positional aggregate token from named args. + if name == "ARGUMENTS" { + continue; + } + let name = name.to_string(); + if seen.insert(name.clone()) { + names.push(name); + } + } + names +} + +/// Parses the `key=value` pairs that follow a custom prompt name. +/// +/// The input is split using shlex rules, so quoted values are supported +/// (for example `USER="Alice Smith"`). The function returns a map of parsed +/// arguments, or an error if a token is missing `=` or if the key is empty. +pub fn parse_prompt_inputs(rest: &str) -> Result, PromptArgsError> { + let mut map = HashMap::new(); + if rest.trim().is_empty() { + return Ok(map); + } + + for token in Shlex::new(rest) { + let Some((key, value)) = token.split_once('=') else { + return Err(PromptArgsError::MissingAssignment { token }); + }; + if key.is_empty() { + return Err(PromptArgsError::MissingKey { token }); + } + map.insert(key.to_string(), value.to_string()); + } + Ok(map) +} + /// Expands a message of the form `/prompts:name [value] [value] …` using a matching saved prompt. /// /// If the text does not start with `/prompts:`, or if no prompt named `name` exists, @@ -35,7 +137,7 @@ pub fn parse_positional_args(rest: &str) -> Vec { pub fn expand_custom_prompt( text: &str, custom_prompts: &[CustomPrompt], -) -> Result, ()> { +) -> Result, PromptExpansionError> { let Some((name, rest)) = parse_slash_name(text) else { return Ok(None); }; @@ -49,14 +151,45 @@ pub fn expand_custom_prompt( Some(prompt) => prompt, None => return Ok(None), }; - // Only support numeric placeholders ($1..$9) and $ARGUMENTS. - if prompt_has_numeric_placeholders(&prompt.content) { - let pos_args: Vec = Shlex::new(rest).collect(); - let expanded = expand_numeric_placeholders(&prompt.content, &pos_args); - return Ok(Some(expanded)); + // If there are named placeholders, expect key=value inputs. + let required = prompt_argument_names(&prompt.content); + if !required.is_empty() { + let inputs = parse_prompt_inputs(rest).map_err(|error| PromptExpansionError::Args { + command: format!("/{name}"), + error, + })?; + let missing: Vec = required + .into_iter() + .filter(|k| !inputs.contains_key(k)) + .collect(); + if !missing.is_empty() { + return Err(PromptExpansionError::MissingArgs { + command: format!("/{name}"), + missing, + }); + } + let content = &prompt.content; + let replaced = PROMPT_ARG_REGEX.replace_all(content, |caps: ®ex_lite::Captures<'_>| { + if let Some(matched) = caps.get(0) + && matched.start() > 0 + && content.as_bytes()[matched.start() - 1] == b'$' + { + return matched.as_str().to_string(); + } + let whole = &caps[0]; + let key = &whole[1..]; + inputs + .get(key) + .cloned() + .unwrap_or_else(|| whole.to_string()) + }); + return Ok(Some(replaced.into_owned())); } - // No recognized placeholders: return the literal content. - Ok(Some(prompt.content.clone())) + + // Otherwise, treat it as numeric/positional placeholder prompt (or none). + let pos_args: Vec = Shlex::new(rest).collect(); + let expanded = expand_numeric_placeholders(&prompt.content, &pos_args); + Ok(Some(expanded)) } /// Detect whether `content` contains numeric placeholders ($1..$9) or `$ARGUMENTS`. @@ -107,6 +240,9 @@ pub fn expand_if_numeric_with_positional_args( prompt: &CustomPrompt, first_line: &str, ) -> Option { + if !prompt_argument_names(&prompt.content).is_empty() { + return None; + } if !prompt_has_numeric_placeholders(&prompt.content) { return None; } @@ -159,3 +295,112 @@ pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { out.push_str(&content[i..]); out } + +/// Constructs a command text for a custom prompt with arguments. +/// Returns the text and the cursor position (inside the first double quote). +pub fn prompt_command_with_arg_placeholders(name: &str, args: &[String]) -> (String, usize) { + let mut text = format!("/{PROMPTS_CMD_PREFIX}:{name}"); + let mut cursor: usize = text.len(); + for (i, arg) in args.iter().enumerate() { + text.push_str(format!(" {arg}=\"\"").as_str()); + if i == 0 { + cursor = text.len() - 1; // inside first "" + } + } + (text, cursor) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn expand_arguments_basic() { + let prompts = vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $USER changes on $BRANCH".to_string(), + description: None, + argument_hint: None, + }]; + + let out = + expand_custom_prompt("/prompts:my-prompt USER=Alice BRANCH=main", &prompts).unwrap(); + assert_eq!(out, Some("Review Alice changes on main".to_string())); + } + + #[test] + fn quoted_values_ok() { + let prompts = vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Pair $USER with $BRANCH".to_string(), + description: None, + argument_hint: None, + }]; + + let out = expand_custom_prompt( + "/prompts:my-prompt USER=\"Alice Smith\" BRANCH=dev-main", + &prompts, + ) + .unwrap(); + assert_eq!(out, Some("Pair Alice Smith with dev-main".to_string())); + } + + #[test] + fn invalid_arg_token_reports_error() { + let prompts = vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $USER changes".to_string(), + description: None, + argument_hint: None, + }]; + let err = expand_custom_prompt("/prompts:my-prompt USER=Alice stray", &prompts) + .unwrap_err() + .user_message(); + assert!(err.contains("expected key=value")); + } + + #[test] + fn missing_required_args_reports_error() { + let prompts = vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $USER changes on $BRANCH".to_string(), + description: None, + argument_hint: None, + }]; + let err = expand_custom_prompt("/prompts:my-prompt USER=Alice", &prompts) + .unwrap_err() + .user_message(); + assert!(err.to_lowercase().contains("missing required args")); + assert!(err.contains("BRANCH")); + } + + #[test] + fn escaped_placeholder_is_ignored() { + assert_eq!( + prompt_argument_names("literal $$USER"), + Vec::::new() + ); + assert_eq!( + prompt_argument_names("literal $$USER and $REAL"), + vec!["REAL".to_string()] + ); + } + + #[test] + fn escaped_placeholder_remains_literal() { + let prompts = vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "literal $$USER".to_string(), + description: None, + argument_hint: None, + }]; + + let out = expand_custom_prompt("/prompts:my-prompt", &prompts).unwrap(); + assert_eq!(out, Some("literal $$USER".to_string())); + } +}