From 885bf8e0988a010d7dddd2b1cdd3111fe21e4260 Mon Sep 17 00:00:00 2001 From: Daniel Edrisian Date: Tue, 23 Sep 2025 12:28:56 -0700 Subject: [PATCH] Custom prompt args --- codex-rs/tui/src/bottom_pane/chat_composer.rs | 299 ++++++++++++++++-- codex-rs/tui/src/bottom_pane/command_popup.rs | 13 +- codex-rs/tui/src/bottom_pane/mod.rs | 1 + codex-rs/tui/src/bottom_pane/prompt_args.rs | 200 ++++++++++++ 4 files changed, 471 insertions(+), 42 deletions(-) create mode 100644 codex-rs/tui/src/bottom_pane/prompt_args.rs diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index 65203b5193..575147dcaf 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -31,12 +31,14 @@ use crate::bottom_pane::paste_burst::FlushResult; use crate::slash_command::SlashCommand; use codex_protocol::custom_prompts::CustomPrompt; +use super::prompt_args; use crate::app_event::AppEvent; use crate::app_event_sender::AppEventSender; 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::key_hint; use crate::ui_consts::LIVE_PREFIX_COLS; use codex_file_search::FileMatch; @@ -383,6 +385,7 @@ impl ChatComposer { let first_line = self.textarea.text().lines().next().unwrap_or(""); popup.on_composer_text_change(first_line.to_string()); if let Some(sel) = popup.selected_item() { + let mut cursor_target: Option = None; match sel { CommandItem::Builtin(cmd) => { let starts_with_cmd = first_line @@ -391,21 +394,24 @@ impl ChatComposer { if !starts_with_cmd { self.textarea.set_text(&format!("/{} ", cmd.command())); } + if !self.textarea.text().is_empty() { + cursor_target = Some(self.textarea.text().len()); + } } CommandItem::UserPrompt(idx) => { - if let Some(name) = popup.prompt_name(idx) { - let starts_with_cmd = - first_line.trim_start().starts_with(&format!("/{name}")); - if !starts_with_cmd { - self.textarea.set_text(&format!("/{name} ")); + if let Some(prompt) = popup.prompt(idx) { + let args = prompt_args::prompt_argument_names(&prompt.content); + let (text, cursor) = Self::prompt_command_text(&prompt.name, &args); + self.textarea.set_text(&text); + cursor_target = cursor; + if cursor_target.is_none() && !self.textarea.text().is_empty() { + cursor_target = Some(self.textarea.text().len()); } } } } - // After completing the command, move cursor to the end. - if !self.textarea.text().is_empty() { - let end = self.textarea.text().len(); - self.textarea.set_cursor(end); + if let Some(pos) = cursor_target { + self.textarea.set_cursor(pos); } } (InputResult::None, true) @@ -416,25 +422,29 @@ impl ChatComposer { .. } => { if let Some(sel) = popup.selected_item() { - // Clear textarea so no residual text remains. - self.textarea.set_text(""); - // Capture any needed data from popup before clearing it. - let prompt_content = match sel { - CommandItem::UserPrompt(idx) => { - popup.prompt_content(idx).map(str::to_string) - } - _ => None, - }; - // Hide popup since an action has been dispatched. - self.active_popup = ActivePopup::None; - match sel { CommandItem::Builtin(cmd) => { + self.textarea.set_text(""); return (InputResult::Command(cmd), true); } - CommandItem::UserPrompt(_) => { - if let Some(contents) = prompt_content { - return (InputResult::Submitted(contents), true); + CommandItem::UserPrompt(idx) => { + if let Some(prompt) = popup.prompt(idx) { + // If the prompt has no $ARGS, auto-submit its content (legacy behavior). + let args = prompt_args::prompt_argument_names(&prompt.content); + if args.is_empty() { + self.textarea.set_text(""); + return (InputResult::Submitted(prompt.content.clone()), true); + } + // Otherwise, insert the template and place the cursor. + let (text, cursor) = Self::prompt_command_text(&prompt.name, &args); + self.textarea.set_text(&text); + let target = cursor.or_else(|| { + (!self.textarea.text().is_empty()) + .then_some(self.textarea.text().len()) + }); + if let Some(pos) = target { + self.textarea.set_cursor(pos); + } } return (InputResult::None, true); } @@ -723,6 +733,21 @@ impl ChatComposer { self.textarea.set_cursor(new_cursor); } + fn prompt_command_text(name: &str, args: &[String]) -> (String, Option) { + let mut text = format!("/{name}"); + let mut first_cursor: Option = None; + for arg in args { + text.push_str(format!(" {arg}=\"\"").as_str()); + if first_cursor.is_none() { + first_cursor = Some(text.len() - 1); + } + } + if first_cursor.is_none() && !text.is_empty() { + first_cursor = Some(text.len()); + } + (text, first_cursor) + } + /// Handle key event when no popup is visible. fn handle_key_event_without_popup(&mut self, key_event: KeyEvent) -> (InputResult, bool) { match key_event { @@ -811,6 +836,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 @@ -824,6 +850,21 @@ 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(); + let expanded_prompt = + match prompt_args::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); } @@ -1145,18 +1186,34 @@ impl ChatComposer { /// textarea. This must be called after every modification that can change /// the text so the popup is shown/updated/hidden as appropriate. fn sync_command_popup(&mut self) { - let first_line = self.textarea.text().lines().next().unwrap_or(""); - let input_starts_with_slash = first_line.starts_with('/'); + // Determine whether the caret is inside the initial '/name' token on the first line. + let text = self.textarea.text(); + let first_line_end = text.find('\n').unwrap_or(text.len()); + let first_line = &text[..first_line_end]; + let cursor = self.textarea.cursor(); + let caret_on_first_line = cursor <= first_line_end; + + let in_slash_name_context = if first_line.starts_with('/') && caret_on_first_line { + let token_end = first_line + .char_indices() + .find(|(_, c)| c.is_whitespace()) + .map(|(i, _)| i) + .unwrap_or(first_line.len()); + cursor <= token_end + } else { + false + }; + match &mut self.active_popup { ActivePopup::Command(popup) => { - if input_starts_with_slash { + if in_slash_name_context { popup.on_composer_text_change(first_line.to_string()); } else { self.active_popup = ActivePopup::None; } } _ => { - if input_starts_with_slash { + if in_slash_name_context { let mut command_popup = CommandPopup::new(self.custom_prompts.clone()); command_popup.on_composer_text_change(first_line.to_string()); self.active_popup = ActivePopup::Command(command_popup); @@ -2315,9 +2372,7 @@ mod tests { } #[test] - fn selecting_custom_prompt_submits_file_contents() { - let prompt_text = "Hello from saved prompt"; - + fn selecting_custom_prompt_without_args_submits_content() { let (tx, _rx) = unbounded_channel::(); let sender = AppEventSender::new(tx); let mut composer = ChatComposer::new( @@ -2332,7 +2387,7 @@ mod tests { composer.set_custom_prompts(vec![CustomPrompt { name: "my-prompt".to_string(), path: "/tmp/my-prompt.md".to_string().into(), - content: prompt_text.to_string(), + content: "Hello from saved prompt".to_string(), }]); type_chars_humanlike( @@ -2343,7 +2398,185 @@ mod tests { let (result, _needs_redraw) = composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); - assert_eq!(InputResult::Submitted(prompt_text.to_string()), result); + assert_eq!( + InputResult::Submitted("Hello from saved prompt".to_string()), + result + ); + assert!(composer.textarea.is_empty()); + } + + #[test] + fn custom_prompt_submission_expands_arguments() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + 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(), + }]); + + composer + .textarea + .set_text("/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()); + } + + // Note: manual UI flow test omitted for Enter-with-popup; core logic + // is covered by other prompt expansion tests. + + #[test] + fn custom_prompt_submission_accepts_quoted_values() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + 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(), + }]); + + composer + .textarea + .set_text("/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() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + 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(), + }]); + + composer.textarea.set_text("/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!("/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() { + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + 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(), + }]); + + // Provide only one of the required args + composer.textarea.set_text("/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!("/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] diff --git a/codex-rs/tui/src/bottom_pane/command_popup.rs b/codex-rs/tui/src/bottom_pane/command_popup.rs index a0933f4ae6..e630b1ce1e 100644 --- a/codex-rs/tui/src/bottom_pane/command_popup.rs +++ b/codex-rs/tui/src/bottom_pane/command_popup.rs @@ -53,12 +53,8 @@ impl CommandPopup { self.prompts = prompts; } - pub(crate) fn prompt_name(&self, idx: usize) -> Option<&str> { - self.prompts.get(idx).map(|p| p.name.as_str()) - } - - pub(crate) fn prompt_content(&self, idx: usize) -> Option<&str> { - self.prompts.get(idx).map(|p| p.content.as_str()) + pub(crate) fn prompt(&self, idx: usize) -> Option<&CustomPrompt> { + self.prompts.get(idx) } /// Update the filter string based on the current composer text. The text @@ -218,7 +214,6 @@ impl WidgetRef for CommandPopup { #[cfg(test)] mod tests { use super::*; - use std::string::ToString; #[test] fn filter_includes_init_when_typing_prefix() { @@ -288,7 +283,7 @@ mod tests { let mut prompt_names: Vec = items .into_iter() .filter_map(|it| match it { - CommandItem::UserPrompt(i) => popup.prompt_name(i).map(ToString::to_string), + CommandItem::UserPrompt(i) => popup.prompt(i).map(|p| p.name.clone()), _ => None, }) .collect(); @@ -306,7 +301,7 @@ mod tests { }]); let items = popup.filtered_items(); let has_collision_prompt = items.into_iter().any(|it| match it { - CommandItem::UserPrompt(i) => popup.prompt_name(i) == Some("init"), + CommandItem::UserPrompt(i) => popup.prompt(i).is_some_and(|p| p.name == "init"), _ => false, }); assert!( diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index fc9f06c39d..5067e5000d 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -24,6 +24,7 @@ mod command_popup; pub mod custom_prompt_view; mod file_search_popup; mod list_selection_view; +mod prompt_args; pub(crate) use list_selection_view::SelectionViewParams; mod paste_burst; pub mod popup_consts; diff --git a/codex-rs/tui/src/bottom_pane/prompt_args.rs b/codex-rs/tui/src/bottom_pane/prompt_args.rs new file mode 100644 index 0000000000..519466da37 --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/prompt_args.rs @@ -0,0 +1,200 @@ +use codex_protocol::custom_prompts::CustomPrompt; +use once_cell::sync::Lazy; +use regex_lite::Regex; +use shlex::Shlex; +use std::collections::HashMap; +use std::collections::HashSet; + +static PROMPT_ARG_REGEX: Lazy = Lazy::new(|| { + // Regex is a hard-coded literal; abort if it ever fails to compile. + 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)." + ) + } + } + } +} + +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) { + let name = &content[m.start() + 1..m.end()]; + let name = name.to_string(); + if seen.insert(name.clone()) { + names.push(name); + } + } + names +} + +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) +} + +pub fn expand_custom_prompt( + text: &str, + custom_prompts: &[CustomPrompt], +) -> Result, PromptExpansionError> { + let Some(stripped) = text.strip_prefix('/') else { + return Ok(None); + }; + let mut name_end = stripped.len(); + for (idx, ch) in stripped.char_indices() { + if ch.is_whitespace() { + name_end = idx; + break; + } + } + + let name = &stripped[..name_end]; + if name.is_empty() { + return Ok(None); + } + + let prompt = match custom_prompts.iter().find(|p| p.name == name) { + Some(prompt) => prompt, + None => return Ok(None), + }; + let rest = stripped[name_end..].trim(); + let inputs = parse_prompt_inputs(rest).map_err(|error| PromptExpansionError::Args { + command: format!("/{name}"), + error, + })?; + + // Ensure that all required variables are provided. + let required = prompt_argument_names(&prompt.content); + 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 replaced = + PROMPT_ARG_REGEX.replace_all(&prompt.content, |caps: ®ex_lite::Captures<'_>| { + let whole = caps.get(0).map(|m| m.as_str()).unwrap_or(""); + let key = &whole[1..]; + inputs + .get(key) + .cloned() + .unwrap_or_else(|| whole.to_string()) + }); + + Ok(Some(replaced.into_owned())) +} + +#[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(), + }]; + + let out = expand_custom_prompt("/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(), + }]; + + let out = expand_custom_prompt("/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(), + }]; + let err = expand_custom_prompt("/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(), + }]; + let err = expand_custom_prompt("/my-prompt USER=Alice", &prompts) + .unwrap_err() + .user_message(); + assert!(err.to_lowercase().contains("missing required args")); + assert!(err.contains("BRANCH")); + } +}