From 531748a080a5772fe3dc55face2261a3cfdea9da Mon Sep 17 00:00:00 2001 From: charley-oai Date: Tue, 20 Jan 2026 18:30:20 -0800 Subject: [PATCH] Prompt Expansion: Preserve Text Elements (#9518) Summary - Preserve `text_elements` through custom prompt argument parsing and expansion (named and numeric placeholders). - Translate text element ranges through Shlex parsing using sentinel substitution, and rehydrate text + element ranges per arg. - Drop image attachments when their placeholder does not survive prompt expansion, keeping attachments consistent with rendered elements. - Mirror changes in TUI2 and expand tests for prompt parsing/expansion edge cases. Tests - placeholders with spaces as single tokens (positional + key=value, quoted + unquoted), - prompt expansion with image placeholders, - large paste + image arg combinations, - unused image arg dropped after expansion. --- codex-rs/protocol/src/user_input.rs | 20 +- codex-rs/tui/src/bottom_pane/chat_composer.rs | 580 ++++++++++++++++-- codex-rs/tui/src/bottom_pane/prompt_args.rs | 578 +++++++++++++++-- .../tui2/src/bottom_pane/chat_composer.rs | 580 ++++++++++++++++-- codex-rs/tui2/src/bottom_pane/prompt_args.rs | 578 +++++++++++++++-- docs/tui-chat-composer.md | 28 + 6 files changed, 2148 insertions(+), 216 deletions(-) diff --git a/codex-rs/protocol/src/user_input.rs b/codex-rs/protocol/src/user_input.rs index 0f32f5a33..e9ca09580 100644 --- a/codex-rs/protocol/src/user_input.rs +++ b/codex-rs/protocol/src/user_input.rs @@ -47,14 +47,30 @@ impl TextElement { } } + /// Returns a copy of this element with a remapped byte range. + /// + /// The placeholder is preserved as-is; callers must ensure the new range + /// still refers to the same logical element (and same placeholder) + /// within the new text. + pub fn map_range(&self, map: F) -> Self + where + F: FnOnce(ByteRange) -> ByteRange, + { + Self { + byte_range: map(self.byte_range), + placeholder: self.placeholder.clone(), + } + } + pub fn set_placeholder(&mut self, placeholder: Option) { self.placeholder = placeholder; } /// Returns the stored placeholder without falling back to the text buffer. /// - /// This is intended only for protocol conversions where the full text is not - /// available; prefer `placeholder(text)` for UI logic. + /// This must only be used inside `From` implementations on equivalent + /// protocol types where the source text is unavailable. Prefer `placeholder(text)` + /// everywhere else. #[doc(hidden)] pub fn _placeholder_for_conversion_only(&self) -> Option<&str> { self.placeholder.as_deref() diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index ed0872c47..724c267a5 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -15,6 +15,18 @@ //! [`ChatComposer::handle_key_event_without_popup`]. After every handled key, we call //! [`ChatComposer::sync_popups`] so UI state follows the latest buffer/cursor. //! +//! # Submission and Prompt Expansion +//! +//! On submit/queue paths, the composer: +//! +//! - Expands pending paste placeholders so element ranges align with the final text. +//! - Trims whitespace and rebases text elements accordingly. +//! - Expands `/prompts:` custom prompts (named or numeric args), preserving text elements. +//! - Prunes attached images so only placeholders that survive expansion are sent. +//! +//! The numeric auto-submit path used by the slash popup performs the same pending-paste expansion +//! and attachment pruning, and clears pending paste state on success. +//! //! # Non-bracketed Paste Bursts //! //! On some terminals (especially on Windows), pastes arrive as a rapid sequence of @@ -171,8 +183,14 @@ enum PromptSelectionMode { } enum PromptSelectionAction { - Insert { text: String, cursor: Option }, - Submit { text: String }, + Insert { + text: String, + cursor: Option, + }, + Submit { + text: String, + text_elements: Vec, + }, } pub(crate) struct ChatComposer { @@ -609,6 +627,18 @@ impl ChatComposer { .collect() } + fn prune_attached_images_for_submission(&mut self, text: &str, text_elements: &[TextElement]) { + if self.attached_images.is_empty() { + return; + } + let image_placeholders: HashSet<&str> = text_elements + .iter() + .filter_map(|elem| elem.placeholder(text)) + .collect(); + self.attached_images + .retain(|img| image_placeholders.contains(img.placeholder.as_str())); + } + /// Insert an attachment placeholder and track it for the next submission. pub fn attach_image(&mut self, path: PathBuf) { let image_number = self.attached_images.len() + 1; @@ -836,6 +866,7 @@ impl ChatComposer { prompt, first_line, PromptSelectionMode::Completion, + &self.textarea.text_elements(), ) { PromptSelectionAction::Insert { text, cursor } => { let target = cursor.unwrap_or(text.len()); @@ -862,19 +893,31 @@ impl ChatComposer { // If the current line starts with a custom prompt name and includes // positional args for a numeric-style template, expand and submit // immediately regardless of the popup selection. - let first_line = self.textarea.text().lines().next().unwrap_or(""); - if let Some((name, _rest)) = parse_slash_name(first_line) + let mut text = self.textarea.text().to_string(); + let mut text_elements = self.textarea.text_elements(); + if !self.pending_pastes.is_empty() { + let (expanded, expanded_elements) = + Self::expand_pending_pastes(&text, text_elements, &self.pending_pastes); + text = expanded; + text_elements = expanded_elements; + } + let first_line = text.lines().next().unwrap_or(""); + if let Some((name, _rest, _rest_offset)) = parse_slash_name(first_line) && let Some(prompt_name) = name.strip_prefix(&format!("{PROMPTS_CMD_PREFIX}:")) && let Some(prompt) = self.custom_prompts.iter().find(|p| p.name == prompt_name) && let Some(expanded) = - expand_if_numeric_with_positional_args(prompt, first_line) + expand_if_numeric_with_positional_args(prompt, first_line, &text_elements) { + self.prune_attached_images_for_submission( + &expanded.text, + &expanded.text_elements, + ); + self.pending_pastes.clear(); self.textarea.set_text_clearing_elements(""); return ( InputResult::Submitted { - text: expanded, - // Expanded prompt is plain text; no UI element ranges to preserve. - text_elements: Vec::new(), + text: expanded.text, + text_elements: expanded.text_elements, }, true, ); @@ -892,14 +935,21 @@ impl ChatComposer { prompt, first_line, PromptSelectionMode::Submit, + &self.textarea.text_elements(), ) { - PromptSelectionAction::Submit { text } => { + PromptSelectionAction::Submit { + text, + text_elements, + } => { + self.prune_attached_images_for_submission( + &text, + &text_elements, + ); self.textarea.set_text_clearing_elements(""); return ( InputResult::Submitted { text, - // Submitting a slash/custom prompt generates plain text, so there are no UI element ranges. - text_elements: Vec::new(), + text_elements, }, true, ); @@ -1558,11 +1608,10 @@ impl ChatComposer { let expanded_input = text.clone(); // If there is neither text nor attachments, suppress submission entirely. - let has_attachments = !self.attached_images.is_empty(); text = text.trim().to_string(); text_elements = Self::trim_text_elements(&expanded_input, &text, text_elements); - if let Some((name, _rest)) = parse_slash_name(&text) { + if let Some((name, _rest, _rest_offset)) = parse_slash_name(&text) { let treat_as_plain_text = input_starts_with_space || name.contains('/'); if !treat_as_plain_text { let is_builtin = @@ -1596,29 +1645,31 @@ impl ChatComposer { } } - 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.set_text_content( - original_input.clone(), - original_text_elements, - original_local_image_paths, - ); - self.pending_pastes.clone_from(&original_pending_pastes); - self.textarea.set_cursor(original_input.len()); - return None; - } - }; + let expanded_prompt = + match expand_custom_prompt(&text, &text_elements, &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.set_text_content( + original_input.clone(), + original_text_elements, + original_local_image_paths, + ); + self.pending_pastes.clone_from(&original_pending_pastes); + self.textarea.set_cursor(original_input.len()); + return None; + } + }; if let Some(expanded) = expanded_prompt { - text = expanded; - // Expanded prompt (e.g. custom prompt) is plain text; text elements not supported yet. - // TODO: Preserve UI element ranges through prompt expansion in a follow-up PR. - text_elements = Vec::new(); + text = expanded.text; + text_elements = expanded.text_elements; } - if text.is_empty() && !has_attachments { + // Custom prompt expansion can remove or rewrite image placeholders, so prune any + // attachments that no longer have a corresponding placeholder in the expanded text. + self.prune_attached_images_for_submission(&text, &text_elements); + if text.is_empty() && self.attached_images.is_empty() { return None; } if !text.is_empty() { @@ -1720,7 +1771,7 @@ impl ChatComposer { /// Returns Some(InputResult) if a command was dispatched, None otherwise. fn try_dispatch_bare_slash_command(&mut self) -> Option { let first_line = self.textarea.text().lines().next().unwrap_or(""); - if let Some((name, rest)) = parse_slash_name(first_line) + if let Some((name, rest, _rest_offset)) = parse_slash_name(first_line) && rest.is_empty() && let Some((_n, cmd)) = Self::built_in_slash_commands_for_input(self.collaboration_modes_enabled) @@ -1741,7 +1792,7 @@ impl ChatComposer { if !input_starts_with_space { let text = self.textarea.text().to_string(); - if let Some((name, rest)) = parse_slash_name(&text) + if let Some((name, rest, _rest_offset)) = parse_slash_name(&text) && !rest.is_empty() && !name.contains('/') && let Some((_n, cmd)) = @@ -2495,6 +2546,7 @@ fn prompt_selection_action( prompt: &CustomPrompt, first_line: &str, mode: PromptSelectionMode, + text_elements: &[TextElement], ) -> PromptSelectionAction { let named_args = prompt_argument_names(&prompt.content); let has_numeric = prompt_has_numeric_placeholders(&prompt.content); @@ -2526,14 +2578,21 @@ fn prompt_selection_action( }; } if has_numeric { - if let Some(expanded) = expand_if_numeric_with_positional_args(prompt, first_line) { - return PromptSelectionAction::Submit { text: expanded }; + if let Some(expanded) = + expand_if_numeric_with_positional_args(prompt, first_line, text_elements) + { + return PromptSelectionAction::Submit { + text: expanded.text, + text_elements: expanded.text_elements, + }; } let text = format!("/{PROMPTS_CMD_PREFIX}:{} ", prompt.name); return PromptSelectionAction::Insert { text, cursor: None }; } PromptSelectionAction::Submit { text: prompt.content.clone(), + // By now we know this custom prompt has no args, so no text elements to preserve. + text_elements: Vec::new(), } } } @@ -2554,6 +2613,7 @@ 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::prompt_args::PromptArg; use crate::bottom_pane::prompt_args::extract_positional_args_for_prompt_line; use crate::bottom_pane::textarea::TextArea; use tokio::sync::mpsc::unbounded_channel; @@ -3801,15 +3861,37 @@ mod tests { let args = extract_positional_args_for_prompt_line( "/prompts:review \"docs/My File.md\"", "review", + &[], + ); + assert_eq!( + args, + vec![PromptArg { + text: "docs/My File.md".to_string(), + text_elements: Vec::new(), + }] ); - assert_eq!(args, vec!["docs/My File.md".to_string()]); } #[test] fn extract_args_supports_mixed_quoted_and_unquoted() { - let args = - extract_positional_args_for_prompt_line("/prompts:cmd \"with spaces\" simple", "cmd"); - assert_eq!(args, vec!["with spaces".to_string(), "simple".to_string()]); + let args = extract_positional_args_for_prompt_line( + "/prompts:cmd \"with spaces\" simple", + "cmd", + &[], + ); + assert_eq!( + args, + vec![ + PromptArg { + text: "with spaces".to_string(), + text_elements: Vec::new(), + }, + PromptArg { + text: "simple".to_string(), + text_elements: Vec::new(), + } + ] + ); } #[test] @@ -4868,6 +4950,166 @@ mod tests { assert!(composer.textarea.is_empty()); } + #[test] + fn custom_prompt_submission_preserves_image_placeholder_unquoted() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $IMG".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt IMG="); + composer.textarea.set_cursor(composer.textarea.text().len()); + let path = PathBuf::from("/tmp/image_prompt.png"); + composer.attach_image(path); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + match result { + InputResult::Submitted { + text, + text_elements, + } => { + let placeholder = local_image_label_text(1); + assert_eq!(text, format!("Review {placeholder}")); + assert_eq!( + text_elements, + vec![TextElement::new( + ByteRange { + start: "Review ".len(), + end: "Review ".len() + placeholder.len(), + }, + Some(placeholder), + )] + ); + } + _ => panic!("expected Submitted"), + } + } + + #[test] + fn custom_prompt_submission_preserves_image_placeholder_quoted() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $IMG".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt IMG=\""); + composer.textarea.set_cursor(composer.textarea.text().len()); + let path = PathBuf::from("/tmp/image_prompt_quoted.png"); + composer.attach_image(path); + composer.handle_paste("\"".to_string()); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + match result { + InputResult::Submitted { + text, + text_elements, + } => { + let placeholder = local_image_label_text(1); + assert_eq!(text, format!("Review {placeholder}")); + assert_eq!( + text_elements, + vec![TextElement::new( + ByteRange { + start: "Review ".len(), + end: "Review ".len() + placeholder.len(), + }, + Some(placeholder), + )] + ); + } + _ => panic!("expected Submitted"), + } + } + + #[test] + fn custom_prompt_submission_drops_unused_image_arg() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review changes".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt IMG="); + composer.textarea.set_cursor(composer.textarea.text().len()); + let path = PathBuf::from("/tmp/unused_image.png"); + composer.attach_image(path); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + match result { + InputResult::Submitted { + text, + text_elements, + } => { + assert_eq!(text, "Review changes"); + assert!(text_elements.is_empty()); + } + _ => panic!("expected Submitted"), + } + assert!(composer.take_recent_submission_images().is_empty()); + } + /// Behavior: selecting a custom prompt that includes a large paste placeholder should expand /// to the full pasted content before submission. #[test] @@ -4935,6 +5177,65 @@ mod tests { assert!(composer.pending_pastes.is_empty()); } + #[test] + fn custom_prompt_with_large_paste_and_image_preserves_elements() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $IMG\n\n$CODE".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt IMG="); + composer.textarea.set_cursor(composer.textarea.text().len()); + let path = PathBuf::from("/tmp/image_prompt_combo.png"); + composer.attach_image(path); + composer.handle_paste(" CODE=".to_string()); + let large_content = "x".repeat(LARGE_PASTE_CHAR_THRESHOLD + 5); + composer.handle_paste(large_content.clone()); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + match result { + InputResult::Submitted { + text, + text_elements, + } => { + let placeholder = local_image_label_text(1); + assert_eq!(text, format!("Review {placeholder}\n\n{large_content}")); + assert_eq!( + text_elements, + vec![TextElement::new( + ByteRange { + start: "Review ".len(), + end: "Review ".len() + placeholder.len(), + }, + Some(placeholder), + )] + ); + } + _ => panic!("expected Submitted"), + } + } + #[test] fn slash_path_input_submits_without_command_error() { use crossterm::event::KeyCode; @@ -5153,6 +5454,201 @@ mod tests { )); } + #[test] + fn popup_prompt_submission_prunes_unused_image_attachments() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Hello".to_string(), + description: None, + argument_hint: None, + }]); + + composer.attach_image(PathBuf::from("/tmp/unused.png")); + composer.textarea.set_cursor(0); + composer.handle_paste(format!("/{PROMPTS_CMD_PREFIX}:my-prompt ")); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert!(matches!( + result, + InputResult::Submitted { text, .. } if text == "Hello" + )); + assert!( + composer + .take_recent_submission_images_with_placeholders() + .is_empty() + ); + } + + #[test] + fn numeric_prompt_auto_submit_prunes_unused_image_attachments() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Hello $1".to_string(), + description: None, + argument_hint: None, + }]); + + type_chars_humanlike( + &mut composer, + &[ + '/', 'p', 'r', 'o', 'm', 'p', 't', 's', ':', 'm', 'y', '-', 'p', 'r', 'o', 'm', + 'p', 't', ' ', 'f', 'o', 'o', ' ', + ], + ); + composer.attach_image(PathBuf::from("/tmp/unused.png")); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert!(matches!( + result, + InputResult::Submitted { text, .. } if text == "Hello foo" + )); + assert!( + composer + .take_recent_submission_images_with_placeholders() + .is_empty() + ); + } + + #[test] + fn numeric_prompt_auto_submit_expands_pending_pastes() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Echo: $1".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt "); + composer.textarea.set_cursor(composer.textarea.text().len()); + let large_content = "x".repeat(LARGE_PASTE_CHAR_THRESHOLD + 5); + composer.handle_paste(large_content.clone()); + + assert_eq!(composer.pending_pastes.len(), 1); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + let expected = format!("Echo: {large_content}"); + assert!(matches!( + result, + InputResult::Submitted { text, .. } if text == expected + )); + assert!(composer.pending_pastes.is_empty()); + } + + #[test] + fn queued_prompt_submission_prunes_unused_image_attachments() { + 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_steer_enabled(false); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Hello $1".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt foo "); + composer.textarea.set_cursor(composer.textarea.text().len()); + composer.attach_image(PathBuf::from("/tmp/unused.png")); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert!(matches!( + result, + InputResult::Queued { text, .. } if text == "Hello foo" + )); + assert!( + composer + .take_recent_submission_images_with_placeholders() + .is_empty() + ); + } + + #[test] + fn selecting_custom_prompt_with_positional_args_submits_numeric_expansion() { + let prompt_text = "Header: $1\nArgs: $ARGUMENTS\n"; + + let prompt = CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: prompt_text.to_string(), + description: None, + argument_hint: None, + }; + + let action = prompt_selection_action( + &prompt, + "/prompts:my-prompt foo bar", + PromptSelectionMode::Submit, + &[], + ); + match action { + PromptSelectionAction::Submit { + text, + text_elements, + } => { + assert_eq!(text, "Header: foo\nArgs: foo bar\n"); + assert!(text_elements.is_empty()); + } + _ => panic!("expected Submit action"), + } + } + #[test] fn numeric_prompt_positional_args_does_not_error() { // Ensure that a prompt with only numeric placeholders does not trigger diff --git a/codex-rs/tui/src/bottom_pane/prompt_args.rs b/codex-rs/tui/src/bottom_pane/prompt_args.rs index 48c3cedfa..efe0a0071 100644 --- a/codex-rs/tui/src/bottom_pane/prompt_args.rs +++ b/codex-rs/tui/src/bottom_pane/prompt_args.rs @@ -1,5 +1,7 @@ use codex_protocol::custom_prompts::CustomPrompt; use codex_protocol::custom_prompts::PROMPTS_CMD_PREFIX; +use codex_protocol::user_input::ByteRange; +use codex_protocol::user_input::TextElement; use lazy_static::lazy_static; use regex_lite::Regex; use shlex::Shlex; @@ -57,28 +59,49 @@ impl PromptExpansionError { } /// Parse a first-line slash command of the form `/name `. -/// Returns `(name, rest_after_name)` if the line begins with `/` and contains -/// a non-empty name; otherwise returns `None`. -pub fn parse_slash_name(line: &str) -> Option<(&str, &str)> { +/// Returns `(name, rest_after_name, rest_offset)` if the line begins with `/` +/// and contains a non-empty name; otherwise returns `None`. +/// +/// `rest_offset` is the byte index into the original line where `rest_after_name` +/// starts after trimming leading whitespace (so `line[rest_offset..] == rest_after_name`). +pub fn parse_slash_name(line: &str) -> Option<(&str, &str, usize)> { let stripped = line.strip_prefix('/')?; - let mut name_end = stripped.len(); + let mut name_end_in_stripped = stripped.len(); for (idx, ch) in stripped.char_indices() { if ch.is_whitespace() { - name_end = idx; + name_end_in_stripped = idx; break; } } - let name = &stripped[..name_end]; + let name = &stripped[..name_end_in_stripped]; if name.is_empty() { return None; } - let rest = stripped[name_end..].trim_start(); - Some((name, rest)) + let rest_untrimmed = &stripped[name_end_in_stripped..]; + let rest = rest_untrimmed.trim_start(); + let rest_start_in_stripped = name_end_in_stripped + (rest_untrimmed.len() - rest.len()); + // `stripped` is `line` without the leading '/', so add 1 to get the original offset. + let rest_offset = rest_start_in_stripped + 1; + Some((name, rest, rest_offset)) +} + +#[derive(Debug, Clone, PartialEq)] +pub struct PromptArg { + pub text: String, + pub text_elements: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct PromptExpansion { + pub text: String, + pub text_elements: Vec, } /// Parse positional arguments using shlex semantics (supports quoted tokens). -pub fn parse_positional_args(rest: &str) -> Vec { - Shlex::new(rest).collect() +/// +/// `text_elements` must be relative to `rest`. +pub fn parse_positional_args(rest: &str, text_elements: &[TextElement]) -> Vec { + parse_tokens_with_elements(rest, text_elements) } /// Extracts the unique placeholder variable names from a prompt template. @@ -106,25 +129,56 @@ pub fn prompt_argument_names(content: &str) -> Vec { names } +/// Shift a text element's byte range left by `offset`, returning `None` if empty. +/// +/// `offset` is the byte length of the prefix removed from the original text. +fn shift_text_element_left(elem: &TextElement, offset: usize) -> Option { + if elem.byte_range.end <= offset { + return None; + } + let start = elem.byte_range.start.saturating_sub(offset); + let end = elem.byte_range.end.saturating_sub(offset); + (start < end).then_some(elem.map_range(|_| ByteRange { start, end })) +} + /// 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> { +pub fn parse_prompt_inputs( + rest: &str, + text_elements: &[TextElement], +) -> 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 }); + // Tokenize the rest of the command using shlex rules, but keep text element + // ranges relative to each emitted token. + for token in parse_tokens_with_elements(rest, text_elements) { + let Some((key, value)) = token.text.split_once('=') else { + return Err(PromptArgsError::MissingAssignment { token: token.text }); }; if key.is_empty() { - return Err(PromptArgsError::MissingKey { token }); + return Err(PromptArgsError::MissingKey { token: token.text }); } - map.insert(key.to_string(), value.to_string()); + // The token is `key=value`; translate element ranges into the value-only + // coordinate space by subtracting the `key=` prefix length. + let value_start = key.len() + 1; + let value_elements = token + .text_elements + .iter() + .filter_map(|elem| shift_text_element_left(elem, value_start)) + .collect(); + map.insert( + key.to_string(), + PromptArg { + text: value.to_string(), + text_elements: value_elements, + }, + ); } Ok(map) } @@ -136,9 +190,10 @@ pub fn parse_prompt_inputs(rest: &str) -> Result, Prompt /// `Ok(Some(expanded))`; otherwise it returns a descriptive error. pub fn expand_custom_prompt( text: &str, + text_elements: &[TextElement], custom_prompts: &[CustomPrompt], -) -> Result, PromptExpansionError> { - let Some((name, rest)) = parse_slash_name(text) else { +) -> Result, PromptExpansionError> { + let Some((name, rest, rest_offset)) = parse_slash_name(text) else { return Ok(None); }; @@ -153,10 +208,24 @@ pub fn expand_custom_prompt( }; // If there are named placeholders, expect key=value inputs. let required = prompt_argument_names(&prompt.content); + let local_elements: Vec = text_elements + .iter() + .filter_map(|elem| { + let mut shifted = shift_text_element_left(elem, rest_offset)?; + if shifted.byte_range.start >= rest.len() { + return None; + } + let end = shifted.byte_range.end.min(rest.len()); + shifted.byte_range.end = end; + (shifted.byte_range.start < shifted.byte_range.end).then_some(shifted) + }) + .collect(); if !required.is_empty() { - let inputs = parse_prompt_inputs(rest).map_err(|error| PromptExpansionError::Args { - command: format!("/{name}"), - error, + let inputs = parse_prompt_inputs(rest, &local_elements).map_err(|error| { + PromptExpansionError::Args { + command: format!("/{name}"), + error, + } })?; let missing: Vec = required .into_iter() @@ -168,28 +237,19 @@ pub fn expand_custom_prompt( 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())); + let (text, elements) = expand_named_placeholders_with_elements(&prompt.content, &inputs); + return Ok(Some(PromptExpansion { + text, + text_elements: elements, + })); } // 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)) + let pos_args = parse_positional_args(rest, &local_elements); + Ok(Some(expand_numeric_placeholders( + &prompt.content, + &pos_args, + ))) } /// Detect whether `content` contains numeric placeholders ($1..$9) or `$ARGUMENTS`. @@ -213,25 +273,42 @@ pub fn prompt_has_numeric_placeholders(content: &str) -> bool { /// Extract positional arguments from a composer first line like "/name a b" for a given prompt name. /// Returns empty when the command name does not match or when there are no args. -pub fn extract_positional_args_for_prompt_line(line: &str, prompt_name: &str) -> Vec { +pub fn extract_positional_args_for_prompt_line( + line: &str, + prompt_name: &str, + text_elements: &[TextElement], +) -> Vec { let trimmed = line.trim_start(); - let Some(rest) = trimmed.strip_prefix('/') else { + let trim_offset = line.len() - trimmed.len(); + let Some((name, rest, rest_offset)) = parse_slash_name(trimmed) else { return Vec::new(); }; // Require the explicit prompts prefix for custom prompt invocations. - let Some(after_prefix) = rest.strip_prefix(&format!("{PROMPTS_CMD_PREFIX}:")) else { + let Some(after_prefix) = name.strip_prefix(&format!("{PROMPTS_CMD_PREFIX}:")) else { return Vec::new(); }; - let mut parts = after_prefix.splitn(2, char::is_whitespace); - let cmd = parts.next().unwrap_or(""); - if cmd != prompt_name { + if after_prefix != prompt_name { return Vec::new(); } - let args_str = parts.next().unwrap_or("").trim(); + let rest_trimmed_start = rest.trim_start(); + let args_str = rest_trimmed_start.trim_end(); if args_str.is_empty() { return Vec::new(); } - parse_positional_args(args_str) + let args_offset = trim_offset + rest_offset + (rest.len() - rest_trimmed_start.len()); + let local_elements: Vec = text_elements + .iter() + .filter_map(|elem| { + let mut shifted = shift_text_element_left(elem, args_offset)?; + if shifted.byte_range.start >= args_str.len() { + return None; + } + let end = shifted.byte_range.end.min(args_str.len()); + shifted.byte_range.end = end; + (shifted.byte_range.start < shifted.byte_range.end).then_some(shifted) + }) + .collect(); + parse_positional_args(args_str, &local_elements) } /// If the prompt only uses numeric placeholders and the first line contains @@ -239,14 +316,15 @@ pub fn extract_positional_args_for_prompt_line(line: &str, prompt_name: &str) -> pub fn expand_if_numeric_with_positional_args( prompt: &CustomPrompt, first_line: &str, -) -> Option { + text_elements: &[TextElement], +) -> Option { if !prompt_argument_names(&prompt.content).is_empty() { return None; } if !prompt_has_numeric_placeholders(&prompt.content) { return None; } - let args = extract_positional_args_for_prompt_line(first_line, &prompt.name); + let args = extract_positional_args_for_prompt_line(first_line, &prompt.name, text_elements); if args.is_empty() { return None; } @@ -254,10 +332,10 @@ pub fn expand_if_numeric_with_positional_args( } /// Expand `$1..$9` and `$ARGUMENTS` in `content` with values from `args`. -pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { +pub fn expand_numeric_placeholders(content: &str, args: &[PromptArg]) -> PromptExpansion { let mut out = String::with_capacity(content.len()); + let mut out_elements = Vec::new(); let mut i = 0; - let mut cached_joined_args: Option = None; while let Some(off) = content[i..].find('$') { let j = i + off; out.push_str(&content[i..j]); @@ -272,8 +350,8 @@ pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { } b'1'..=b'9' => { let idx = (bytes[1] - b'1') as usize; - if let Some(val) = args.get(idx) { - out.push_str(val); + if let Some(arg) = args.get(idx) { + append_arg_with_elements(&mut out, &mut out_elements, arg); } i = j + 2; continue; @@ -283,8 +361,7 @@ pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { } if rest.len() > "ARGUMENTS".len() && rest[1..].starts_with("ARGUMENTS") { if !args.is_empty() { - let joined = cached_joined_args.get_or_insert_with(|| args.join(" ")); - out.push_str(joined); + append_joined_args_with_elements(&mut out, &mut out_elements, args); } i = j + 1 + "ARGUMENTS".len(); continue; @@ -293,7 +370,179 @@ pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { i = j + 1; } out.push_str(&content[i..]); - out + PromptExpansion { + text: out, + text_elements: out_elements, + } +} + +fn parse_tokens_with_elements(rest: &str, text_elements: &[TextElement]) -> Vec { + let mut elements = text_elements.to_vec(); + elements.sort_by_key(|elem| elem.byte_range.start); + // Keep element placeholders intact across shlex splitting by replacing + // each element range with a unique sentinel token first. + let (rest_for_shlex, replacements) = replace_text_elements_with_sentinels(rest, &elements); + Shlex::new(&rest_for_shlex) + .map(|token| apply_replacements_to_token(token, &replacements)) + .collect() +} + +#[derive(Debug, Clone)] +struct ElementReplacement { + sentinel: String, + text: String, + placeholder: Option, +} + +/// Replace each text element range with a unique sentinel token. +/// +/// The sentinel is chosen so it will survive shlex tokenization as a single word. +fn replace_text_elements_with_sentinels( + rest: &str, + elements: &[TextElement], +) -> (String, Vec) { + let mut out = String::with_capacity(rest.len()); + let mut replacements = Vec::new(); + let mut cursor = 0; + + for (idx, elem) in elements.iter().enumerate() { + let start = elem.byte_range.start; + let end = elem.byte_range.end; + out.push_str(&rest[cursor..start]); + let mut sentinel = format!("__CODEX_ELEM_{idx}__"); + // Ensure we never collide with user content so a sentinel can't be mistaken for text. + while rest.contains(&sentinel) { + sentinel.push('_'); + } + out.push_str(&sentinel); + replacements.push(ElementReplacement { + sentinel, + text: rest[start..end].to_string(), + placeholder: elem.placeholder(rest).map(str::to_string), + }); + cursor = end; + } + + out.push_str(&rest[cursor..]); + (out, replacements) +} + +/// Rehydrate a shlex token by swapping sentinels back to the original text +/// and rebuilding text element ranges relative to the resulting token. +fn apply_replacements_to_token(token: String, replacements: &[ElementReplacement]) -> PromptArg { + if replacements.is_empty() { + return PromptArg { + text: token, + text_elements: Vec::new(), + }; + } + + let mut out = String::with_capacity(token.len()); + let mut out_elements = Vec::new(); + let mut cursor = 0; + + while cursor < token.len() { + let Some((offset, replacement)) = next_replacement(&token, cursor, replacements) else { + out.push_str(&token[cursor..]); + break; + }; + let start_in_token = cursor + offset; + out.push_str(&token[cursor..start_in_token]); + let start = out.len(); + out.push_str(&replacement.text); + let end = out.len(); + if start < end { + out_elements.push(TextElement::new( + ByteRange { start, end }, + replacement.placeholder.clone(), + )); + } + cursor = start_in_token + replacement.sentinel.len(); + } + + PromptArg { + text: out, + text_elements: out_elements, + } +} + +/// Find the earliest sentinel occurrence at or after `cursor`. +fn next_replacement<'a>( + token: &str, + cursor: usize, + replacements: &'a [ElementReplacement], +) -> Option<(usize, &'a ElementReplacement)> { + let slice = &token[cursor..]; + let mut best: Option<(usize, &'a ElementReplacement)> = None; + for replacement in replacements { + if let Some(pos) = slice.find(&replacement.sentinel) { + match best { + Some((best_pos, _)) if best_pos <= pos => {} + _ => best = Some((pos, replacement)), + } + } + } + best +} + +fn expand_named_placeholders_with_elements( + content: &str, + args: &HashMap, +) -> (String, Vec) { + let mut out = String::with_capacity(content.len()); + let mut out_elements = Vec::new(); + let mut cursor = 0; + for m in PROMPT_ARG_REGEX.find_iter(content) { + let start = m.start(); + let end = m.end(); + if start > 0 && content.as_bytes()[start - 1] == b'$' { + out.push_str(&content[cursor..end]); + cursor = end; + continue; + } + out.push_str(&content[cursor..start]); + cursor = end; + let key = &content[start + 1..end]; + if let Some(arg) = args.get(key) { + append_arg_with_elements(&mut out, &mut out_elements, arg); + } else { + out.push_str(&content[start..end]); + } + } + out.push_str(&content[cursor..]); + (out, out_elements) +} + +fn append_arg_with_elements( + out: &mut String, + out_elements: &mut Vec, + arg: &PromptArg, +) { + let start = out.len(); + out.push_str(&arg.text); + if arg.text_elements.is_empty() { + return; + } + out_elements.extend(arg.text_elements.iter().map(|elem| { + elem.map_range(|range| ByteRange { + start: start + range.start, + end: start + range.end, + }) + })); +} + +fn append_joined_args_with_elements( + out: &mut String, + out_elements: &mut Vec, + args: &[PromptArg], +) { + // `$ARGUMENTS` joins args with single spaces while preserving element ranges. + for (idx, arg) in args.iter().enumerate() { + if idx > 0 { + out.push(' '); + } + append_arg_with_elements(out, out_elements, arg); + } } /// Constructs a command text for a custom prompt with arguments. @@ -313,6 +562,7 @@ pub fn prompt_command_with_arg_placeholders(name: &str, args: &[String]) -> (Str #[cfg(test)] mod tests { use super::*; + use pretty_assertions::assert_eq; #[test] fn expand_arguments_basic() { @@ -324,9 +574,15 @@ mod tests { 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())); + let out = expand_custom_prompt("/prompts:my-prompt USER=Alice BRANCH=main", &[], &prompts) + .unwrap(); + assert_eq!( + out, + Some(PromptExpansion { + text: "Review Alice changes on main".to_string(), + text_elements: Vec::new(), + }) + ); } #[test] @@ -341,10 +597,17 @@ mod tests { 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())); + assert_eq!( + out, + Some(PromptExpansion { + text: "Pair Alice Smith with dev-main".to_string(), + text_elements: Vec::new(), + }) + ); } #[test] @@ -356,7 +619,7 @@ mod tests { description: None, argument_hint: None, }]; - let err = expand_custom_prompt("/prompts:my-prompt USER=Alice stray", &prompts) + let err = expand_custom_prompt("/prompts:my-prompt USER=Alice stray", &[], &prompts) .unwrap_err() .user_message(); assert!(err.contains("expected key=value")); @@ -371,7 +634,7 @@ mod tests { description: None, argument_hint: None, }]; - let err = expand_custom_prompt("/prompts:my-prompt USER=Alice", &prompts) + let err = expand_custom_prompt("/prompts:my-prompt USER=Alice", &[], &prompts) .unwrap_err() .user_message(); assert!(err.to_lowercase().contains("missing required args")); @@ -400,7 +663,192 @@ mod tests { argument_hint: None, }]; - let out = expand_custom_prompt("/prompts:my-prompt", &prompts).unwrap(); - assert_eq!(out, Some("literal $$USER".to_string())); + let out = expand_custom_prompt("/prompts:my-prompt", &[], &prompts).unwrap(); + assert_eq!( + out, + Some(PromptExpansion { + text: "literal $$USER".to_string(), + text_elements: Vec::new(), + }) + ); + } + + #[test] + fn positional_args_treat_placeholder_with_spaces_as_single_token() { + let placeholder = "[Image #1]"; + let rest = format!("alpha {placeholder} beta"); + let start = rest.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = parse_positional_args(&rest, &text_elements); + assert_eq!( + args, + vec![ + PromptArg { + text: "alpha".to_string(), + text_elements: Vec::new(), + }, + PromptArg { + text: placeholder.to_string(), + text_elements: vec![TextElement::new( + ByteRange { + start: 0, + end: placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }, + PromptArg { + text: "beta".to_string(), + text_elements: Vec::new(), + } + ] + ); + } + + #[test] + fn extract_positional_args_shifts_element_offsets_into_args_str() { + let placeholder = "[Image #1]"; + let line = format!(" /{PROMPTS_CMD_PREFIX}:my-prompt alpha {placeholder} beta "); + let start = line.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = extract_positional_args_for_prompt_line(&line, "my-prompt", &text_elements); + assert_eq!( + args, + vec![ + PromptArg { + text: "alpha".to_string(), + text_elements: Vec::new(), + }, + PromptArg { + text: placeholder.to_string(), + text_elements: vec![TextElement::new( + ByteRange { + start: 0, + end: placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }, + PromptArg { + text: "beta".to_string(), + text_elements: Vec::new(), + } + ] + ); + } + + #[test] + fn key_value_args_treat_placeholder_with_spaces_as_single_token() { + let placeholder = "[Image #1]"; + let rest = format!("IMG={placeholder} NOTE=hello"); + let start = rest.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = parse_prompt_inputs(&rest, &text_elements).expect("inputs"); + assert_eq!( + args.get("IMG"), + Some(&PromptArg { + text: placeholder.to_string(), + text_elements: vec![TextElement::new( + ByteRange { + start: 0, + end: placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }) + ); + assert_eq!( + args.get("NOTE"), + Some(&PromptArg { + text: "hello".to_string(), + text_elements: Vec::new(), + }) + ); + } + + #[test] + fn positional_args_allow_placeholder_inside_quotes() { + let placeholder = "[Image #1]"; + let rest = format!("alpha \"see {placeholder} here\" beta"); + let start = rest.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = parse_positional_args(&rest, &text_elements); + assert_eq!( + args, + vec![ + PromptArg { + text: "alpha".to_string(), + text_elements: Vec::new(), + }, + PromptArg { + text: format!("see {placeholder} here"), + text_elements: vec![TextElement::new( + ByteRange { + start: "see ".len(), + end: "see ".len() + placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }, + PromptArg { + text: "beta".to_string(), + text_elements: Vec::new(), + } + ] + ); + } + + #[test] + fn key_value_args_allow_placeholder_inside_quotes() { + let placeholder = "[Image #1]"; + let rest = format!("IMG=\"see {placeholder} here\" NOTE=ok"); + let start = rest.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = parse_prompt_inputs(&rest, &text_elements).expect("inputs"); + assert_eq!( + args.get("IMG"), + Some(&PromptArg { + text: format!("see {placeholder} here"), + text_elements: vec![TextElement::new( + ByteRange { + start: "see ".len(), + end: "see ".len() + placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }) + ); + assert_eq!( + args.get("NOTE"), + Some(&PromptArg { + text: "ok".to_string(), + text_elements: Vec::new(), + }) + ); } } diff --git a/codex-rs/tui2/src/bottom_pane/chat_composer.rs b/codex-rs/tui2/src/bottom_pane/chat_composer.rs index a17fd59f0..c8bd88d41 100644 --- a/codex-rs/tui2/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui2/src/bottom_pane/chat_composer.rs @@ -15,6 +15,18 @@ //! [`ChatComposer::handle_key_event_without_popup`]. After every handled key, we call //! [`ChatComposer::sync_popups`] so UI state follows the latest buffer/cursor. //! +//! # Submission and Prompt Expansion +//! +//! On submit/queue paths, the composer: +//! +//! - Expands pending paste placeholders so element ranges align with the final text. +//! - Trims whitespace and rebases text elements accordingly. +//! - Expands `/prompts:` custom prompts (named or numeric args), preserving text elements. +//! - Prunes attached images so only placeholders that survive expansion are sent. +//! +//! The numeric auto-submit path used by the slash popup performs the same pending-paste expansion +//! and attachment pruning, and clears pending paste state on success. +//! //! # Non-bracketed Paste Bursts //! //! On some terminals (especially on Windows), pastes arrive as a rapid sequence of @@ -172,8 +184,14 @@ enum PromptSelectionMode { } enum PromptSelectionAction { - Insert { text: String, cursor: Option }, - Submit { text: String }, + Insert { + text: String, + cursor: Option, + }, + Submit { + text: String, + text_elements: Vec, + }, } pub(crate) struct ChatComposer { @@ -541,6 +559,18 @@ impl ChatComposer { .collect() } + fn prune_attached_images_for_submission(&mut self, text: &str, text_elements: &[TextElement]) { + if self.attached_images.is_empty() { + return; + } + let image_placeholders: HashSet<&str> = text_elements + .iter() + .filter_map(|elem| elem.placeholder(text)) + .collect(); + self.attached_images + .retain(|img| image_placeholders.contains(img.placeholder.as_str())); + } + /// Insert an attachment placeholder and track it for the next submission. pub fn attach_image(&mut self, path: PathBuf) { let image_number = self.attached_images.len() + 1; @@ -769,6 +799,7 @@ impl ChatComposer { prompt, first_line, PromptSelectionMode::Completion, + &self.textarea.text_elements(), ) { PromptSelectionAction::Insert { text, cursor } => { let target = cursor.unwrap_or(text.len()); @@ -795,19 +826,31 @@ impl ChatComposer { // If the current line starts with a custom prompt name and includes // positional args for a numeric-style template, expand and submit // immediately regardless of the popup selection. - let first_line = self.textarea.text().lines().next().unwrap_or(""); - if let Some((name, _rest)) = parse_slash_name(first_line) + let mut text = self.textarea.text().to_string(); + let mut text_elements = self.textarea.text_elements(); + if !self.pending_pastes.is_empty() { + let (expanded, expanded_elements) = + Self::expand_pending_pastes(&text, text_elements, &self.pending_pastes); + text = expanded; + text_elements = expanded_elements; + } + let first_line = text.lines().next().unwrap_or(""); + if let Some((name, _rest, _rest_offset)) = parse_slash_name(first_line) && let Some(prompt_name) = name.strip_prefix(&format!("{PROMPTS_CMD_PREFIX}:")) && let Some(prompt) = self.custom_prompts.iter().find(|p| p.name == prompt_name) && let Some(expanded) = - expand_if_numeric_with_positional_args(prompt, first_line) + expand_if_numeric_with_positional_args(prompt, first_line, &text_elements) { + self.prune_attached_images_for_submission( + &expanded.text, + &expanded.text_elements, + ); + self.pending_pastes.clear(); self.textarea.set_text_clearing_elements(""); return ( InputResult::Submitted { - text: expanded, - // Expanded prompt is plain text; no UI element ranges to preserve. - text_elements: Vec::new(), + text: expanded.text, + text_elements: expanded.text_elements, }, true, ); @@ -825,14 +868,21 @@ impl ChatComposer { prompt, first_line, PromptSelectionMode::Submit, + &self.textarea.text_elements(), ) { - PromptSelectionAction::Submit { text } => { + PromptSelectionAction::Submit { + text, + text_elements, + } => { + self.prune_attached_images_for_submission( + &text, + &text_elements, + ); self.textarea.set_text_clearing_elements(""); return ( InputResult::Submitted { text, - // Submitting a slash/custom prompt generates plain text, so there are no UI element ranges. - text_elements: Vec::new(), + text_elements, }, true, ); @@ -1487,11 +1537,10 @@ impl ChatComposer { let expanded_input = text.clone(); // If there is neither text nor attachments, suppress submission entirely. - let has_attachments = !self.attached_images.is_empty(); text = text.trim().to_string(); text_elements = Self::trim_text_elements(&expanded_input, &text, text_elements); - if let Some((name, _rest)) = parse_slash_name(&text) { + if let Some((name, _rest, _rest_offset)) = parse_slash_name(&text) { let treat_as_plain_text = input_starts_with_space || name.contains('/'); if !treat_as_plain_text { let is_builtin = @@ -1525,29 +1574,31 @@ impl ChatComposer { } } - 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.set_text_content( - original_input.clone(), - original_text_elements, - original_local_image_paths, - ); - self.pending_pastes.clone_from(&original_pending_pastes); - self.textarea.set_cursor(original_input.len()); - return None; - } - }; + let expanded_prompt = + match expand_custom_prompt(&text, &text_elements, &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.set_text_content( + original_input.clone(), + original_text_elements, + original_local_image_paths, + ); + self.pending_pastes.clone_from(&original_pending_pastes); + self.textarea.set_cursor(original_input.len()); + return None; + } + }; if let Some(expanded) = expanded_prompt { - text = expanded; - // Expanded prompt (e.g. custom prompt) is plain text; text elements not supported yet. - // TODO: Preserve UI element ranges through prompt expansion in a follow-up PR. - text_elements = Vec::new(); + text = expanded.text; + text_elements = expanded.text_elements; } - if text.is_empty() && !has_attachments { + // Custom prompt expansion can remove or rewrite image placeholders, so prune any + // attachments that no longer have a corresponding placeholder in the expanded text. + self.prune_attached_images_for_submission(&text, &text_elements); + if text.is_empty() && self.attached_images.is_empty() { return None; } if !text.is_empty() { @@ -1650,7 +1701,7 @@ impl ChatComposer { /// Returns Some(InputResult) if a command was dispatched, None otherwise. fn try_dispatch_bare_slash_command(&mut self) -> Option { let first_line = self.textarea.text().lines().next().unwrap_or(""); - if let Some((name, rest)) = parse_slash_name(first_line) + if let Some((name, rest, _rest_offset)) = parse_slash_name(first_line) && rest.is_empty() && let Some((_n, cmd)) = Self::built_in_slash_commands_for_input(self.collaboration_modes_enabled) @@ -1671,7 +1722,7 @@ impl ChatComposer { if !input_starts_with_space { let text = self.textarea.text().to_string(); - if let Some((name, rest)) = parse_slash_name(&text) + if let Some((name, rest, _rest_offset)) = parse_slash_name(&text) && !rest.is_empty() && !name.contains('/') && let Some((_n, cmd)) = @@ -2463,6 +2514,7 @@ fn prompt_selection_action( prompt: &CustomPrompt, first_line: &str, mode: PromptSelectionMode, + text_elements: &[TextElement], ) -> PromptSelectionAction { let named_args = prompt_argument_names(&prompt.content); let has_numeric = prompt_has_numeric_placeholders(&prompt.content); @@ -2494,14 +2546,21 @@ fn prompt_selection_action( }; } if has_numeric { - if let Some(expanded) = expand_if_numeric_with_positional_args(prompt, first_line) { - return PromptSelectionAction::Submit { text: expanded }; + if let Some(expanded) = + expand_if_numeric_with_positional_args(prompt, first_line, text_elements) + { + return PromptSelectionAction::Submit { + text: expanded.text, + text_elements: expanded.text_elements, + }; } let text = format!("/{PROMPTS_CMD_PREFIX}:{} ", prompt.name); return PromptSelectionAction::Insert { text, cursor: None }; } PromptSelectionAction::Submit { text: prompt.content.clone(), + // By now we know this custom prompt has no args, so no text elements to preserve. + text_elements: Vec::new(), } } } @@ -2522,6 +2581,7 @@ 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::prompt_args::PromptArg; use crate::bottom_pane::prompt_args::extract_positional_args_for_prompt_line; use crate::bottom_pane::textarea::TextArea; use tokio::sync::mpsc::unbounded_channel; @@ -3792,15 +3852,37 @@ mod tests { let args = extract_positional_args_for_prompt_line( "/prompts:review \"docs/My File.md\"", "review", + &[], + ); + assert_eq!( + args, + vec![PromptArg { + text: "docs/My File.md".to_string(), + text_elements: Vec::new(), + }] ); - assert_eq!(args, vec!["docs/My File.md".to_string()]); } #[test] fn extract_args_supports_mixed_quoted_and_unquoted() { - let args = - extract_positional_args_for_prompt_line("/prompts:cmd \"with spaces\" simple", "cmd"); - assert_eq!(args, vec!["with spaces".to_string(), "simple".to_string()]); + let args = extract_positional_args_for_prompt_line( + "/prompts:cmd \"with spaces\" simple", + "cmd", + &[], + ); + assert_eq!( + args, + vec![ + PromptArg { + text: "with spaces".to_string(), + text_elements: Vec::new(), + }, + PromptArg { + text: "simple".to_string(), + text_elements: Vec::new(), + } + ] + ); } #[test] @@ -4831,6 +4913,166 @@ mod tests { assert!(composer.textarea.is_empty()); } + #[test] + fn custom_prompt_submission_preserves_image_placeholder_unquoted() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $IMG".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt IMG="); + composer.textarea.set_cursor(composer.textarea.text().len()); + let path = PathBuf::from("/tmp/image_prompt.png"); + composer.attach_image(path); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + match result { + InputResult::Submitted { + text, + text_elements, + } => { + let placeholder = local_image_label_text(1); + assert_eq!(text, format!("Review {placeholder}")); + assert_eq!( + text_elements, + vec![TextElement::new( + ByteRange { + start: "Review ".len(), + end: "Review ".len() + placeholder.len(), + }, + Some(placeholder), + )] + ); + } + _ => panic!("expected Submitted"), + } + } + + #[test] + fn custom_prompt_submission_preserves_image_placeholder_quoted() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $IMG".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt IMG=\""); + composer.textarea.set_cursor(composer.textarea.text().len()); + let path = PathBuf::from("/tmp/image_prompt_quoted.png"); + composer.attach_image(path); + composer.handle_paste("\"".to_string()); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + match result { + InputResult::Submitted { + text, + text_elements, + } => { + let placeholder = local_image_label_text(1); + assert_eq!(text, format!("Review {placeholder}")); + assert_eq!( + text_elements, + vec![TextElement::new( + ByteRange { + start: "Review ".len(), + end: "Review ".len() + placeholder.len(), + }, + Some(placeholder), + )] + ); + } + _ => panic!("expected Submitted"), + } + } + + #[test] + fn custom_prompt_submission_drops_unused_image_arg() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review changes".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt IMG="); + composer.textarea.set_cursor(composer.textarea.text().len()); + let path = PathBuf::from("/tmp/unused_image.png"); + composer.attach_image(path); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + match result { + InputResult::Submitted { + text, + text_elements, + } => { + assert_eq!(text, "Review changes"); + assert!(text_elements.is_empty()); + } + _ => panic!("expected Submitted"), + } + assert!(composer.take_recent_submission_images().is_empty()); + } + /// Behavior: selecting a custom prompt that includes a large paste placeholder should expand /// to the full pasted content before submission. #[test] @@ -4898,6 +5140,65 @@ mod tests { assert!(composer.pending_pastes.is_empty()); } + #[test] + fn custom_prompt_with_large_paste_and_image_preserves_elements() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Review $IMG\n\n$CODE".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt IMG="); + composer.textarea.set_cursor(composer.textarea.text().len()); + let path = PathBuf::from("/tmp/image_prompt_combo.png"); + composer.attach_image(path); + composer.handle_paste(" CODE=".to_string()); + let large_content = "x".repeat(LARGE_PASTE_CHAR_THRESHOLD + 5); + composer.handle_paste(large_content.clone()); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + match result { + InputResult::Submitted { + text, + text_elements, + } => { + let placeholder = local_image_label_text(1); + assert_eq!(text, format!("Review {placeholder}\n\n{large_content}")); + assert_eq!( + text_elements, + vec![TextElement::new( + ByteRange { + start: "Review ".len(), + end: "Review ".len() + placeholder.len(), + }, + Some(placeholder), + )] + ); + } + _ => panic!("expected Submitted"), + } + } + #[test] fn slash_path_input_submits_without_command_error() { use crossterm::event::KeyCode; @@ -5116,6 +5417,201 @@ mod tests { )); } + #[test] + fn popup_prompt_submission_prunes_unused_image_attachments() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Hello".to_string(), + description: None, + argument_hint: None, + }]); + + composer.attach_image(PathBuf::from("/tmp/unused.png")); + composer.textarea.set_cursor(0); + composer.handle_paste(format!("/{PROMPTS_CMD_PREFIX}:my-prompt ")); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert!(matches!( + result, + InputResult::Submitted { text, .. } if text == "Hello" + )); + assert!( + composer + .take_recent_submission_images_with_placeholders() + .is_empty() + ); + } + + #[test] + fn numeric_prompt_auto_submit_prunes_unused_image_attachments() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Hello $1".to_string(), + description: None, + argument_hint: None, + }]); + + type_chars_humanlike( + &mut composer, + &[ + '/', 'p', 'r', 'o', 'm', 'p', 't', 's', ':', 'm', 'y', '-', 'p', 'r', 'o', 'm', + 'p', 't', ' ', 'f', 'o', 'o', ' ', + ], + ); + composer.attach_image(PathBuf::from("/tmp/unused.png")); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert!(matches!( + result, + InputResult::Submitted { text, .. } if text == "Hello foo" + )); + assert!( + composer + .take_recent_submission_images_with_placeholders() + .is_empty() + ); + } + + #[test] + fn numeric_prompt_auto_submit_expands_pending_pastes() { + 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_steer_enabled(true); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Echo: $1".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt "); + composer.textarea.set_cursor(composer.textarea.text().len()); + let large_content = "x".repeat(LARGE_PASTE_CHAR_THRESHOLD + 5); + composer.handle_paste(large_content.clone()); + + assert_eq!(composer.pending_pastes.len(), 1); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + let expected = format!("Echo: {large_content}"); + assert!(matches!( + result, + InputResult::Submitted { text, .. } if text == expected + )); + assert!(composer.pending_pastes.is_empty()); + } + + #[test] + fn queued_prompt_submission_prunes_unused_image_attachments() { + 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_steer_enabled(false); + + composer.set_custom_prompts(vec![CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: "Hello $1".to_string(), + description: None, + argument_hint: None, + }]); + + composer + .textarea + .set_text_clearing_elements("/prompts:my-prompt foo "); + composer.textarea.set_cursor(composer.textarea.text().len()); + composer.attach_image(PathBuf::from("/tmp/unused.png")); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert!(matches!( + result, + InputResult::Queued { text, .. } if text == "Hello foo" + )); + assert!( + composer + .take_recent_submission_images_with_placeholders() + .is_empty() + ); + } + + #[test] + fn selecting_custom_prompt_with_positional_args_submits_numeric_expansion() { + let prompt_text = "Header: $1\nArgs: $ARGUMENTS\n"; + + let prompt = CustomPrompt { + name: "my-prompt".to_string(), + path: "/tmp/my-prompt.md".to_string().into(), + content: prompt_text.to_string(), + description: None, + argument_hint: None, + }; + + let action = prompt_selection_action( + &prompt, + "/prompts:my-prompt foo bar", + PromptSelectionMode::Submit, + &[], + ); + match action { + PromptSelectionAction::Submit { + text, + text_elements, + } => { + assert_eq!(text, "Header: foo\nArgs: foo bar\n"); + assert!(text_elements.is_empty()); + } + _ => panic!("expected Submit action"), + } + } + #[test] fn numeric_prompt_positional_args_does_not_error() { // Ensure that a prompt with only numeric placeholders does not trigger diff --git a/codex-rs/tui2/src/bottom_pane/prompt_args.rs b/codex-rs/tui2/src/bottom_pane/prompt_args.rs index 48c3cedfa..ade750f67 100644 --- a/codex-rs/tui2/src/bottom_pane/prompt_args.rs +++ b/codex-rs/tui2/src/bottom_pane/prompt_args.rs @@ -1,5 +1,7 @@ use codex_protocol::custom_prompts::CustomPrompt; use codex_protocol::custom_prompts::PROMPTS_CMD_PREFIX; +use codex_protocol::user_input::ByteRange; +use codex_protocol::user_input::TextElement; use lazy_static::lazy_static; use regex_lite::Regex; use shlex::Shlex; @@ -57,28 +59,49 @@ impl PromptExpansionError { } /// Parse a first-line slash command of the form `/name `. -/// Returns `(name, rest_after_name)` if the line begins with `/` and contains -/// a non-empty name; otherwise returns `None`. -pub fn parse_slash_name(line: &str) -> Option<(&str, &str)> { +/// Returns `(name, rest_after_name, rest_offset)` if the line begins with `/` +/// and contains a non-empty name; otherwise returns `None`. +/// +/// `rest_offset` is the byte index into the original line where `rest_after_name` +/// starts after trimming leading whitespace (so `line[rest_offset..] == rest_after_name`). +pub fn parse_slash_name(line: &str) -> Option<(&str, &str, usize)> { let stripped = line.strip_prefix('/')?; - let mut name_end = stripped.len(); + let mut name_end_in_stripped = stripped.len(); for (idx, ch) in stripped.char_indices() { if ch.is_whitespace() { - name_end = idx; + name_end_in_stripped = idx; break; } } - let name = &stripped[..name_end]; + let name = &stripped[..name_end_in_stripped]; if name.is_empty() { return None; } - let rest = stripped[name_end..].trim_start(); - Some((name, rest)) + let rest_untrimmed = &stripped[name_end_in_stripped..]; + let rest = rest_untrimmed.trim_start(); + let rest_start_in_stripped = name_end_in_stripped + (rest_untrimmed.len() - rest.len()); + // `stripped` is `line` without the leading '/', so add 1 to get the original offset. + let rest_offset = rest_start_in_stripped + 1; + Some((name, rest, rest_offset)) +} + +#[derive(Debug, Clone, PartialEq)] +pub struct PromptArg { + pub text: String, + pub text_elements: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct PromptExpansion { + pub text: String, + pub text_elements: Vec, } /// Parse positional arguments using shlex semantics (supports quoted tokens). -pub fn parse_positional_args(rest: &str) -> Vec { - Shlex::new(rest).collect() +/// +/// `text_elements` must be relative to `rest`. +pub fn parse_positional_args(rest: &str, text_elements: &[TextElement]) -> Vec { + parse_tokens_with_elements(rest, text_elements) } /// Extracts the unique placeholder variable names from a prompt template. @@ -111,20 +134,39 @@ pub fn prompt_argument_names(content: &str) -> Vec { /// 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> { +pub fn parse_prompt_inputs( + rest: &str, + text_elements: &[TextElement], +) -> 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 }); + // Tokenize the rest of the command using shlex rules, but keep text element + // ranges relative to each emitted token. + for token in parse_tokens_with_elements(rest, text_elements) { + let Some((key, value)) = token.text.split_once('=') else { + return Err(PromptArgsError::MissingAssignment { token: token.text }); }; if key.is_empty() { - return Err(PromptArgsError::MissingKey { token }); + return Err(PromptArgsError::MissingKey { token: token.text }); } - map.insert(key.to_string(), value.to_string()); + // The token is `key=value`; translate element ranges into the value-only + // coordinate space by subtracting the `key=` prefix length. + let value_start = key.len() + 1; + let value_elements = token + .text_elements + .iter() + .filter_map(|elem| shift_text_element_left(elem, value_start)) + .collect(); + map.insert( + key.to_string(), + PromptArg { + text: value.to_string(), + text_elements: value_elements, + }, + ); } Ok(map) } @@ -136,9 +178,10 @@ pub fn parse_prompt_inputs(rest: &str) -> Result, Prompt /// `Ok(Some(expanded))`; otherwise it returns a descriptive error. pub fn expand_custom_prompt( text: &str, + text_elements: &[TextElement], custom_prompts: &[CustomPrompt], -) -> Result, PromptExpansionError> { - let Some((name, rest)) = parse_slash_name(text) else { +) -> Result, PromptExpansionError> { + let Some((name, rest, rest_offset)) = parse_slash_name(text) else { return Ok(None); }; @@ -153,10 +196,24 @@ pub fn expand_custom_prompt( }; // If there are named placeholders, expect key=value inputs. let required = prompt_argument_names(&prompt.content); + let local_elements: Vec = text_elements + .iter() + .filter_map(|elem| { + let mut shifted = shift_text_element_left(elem, rest_offset)?; + if shifted.byte_range.start >= rest.len() { + return None; + } + let end = shifted.byte_range.end.min(rest.len()); + shifted.byte_range.end = end; + (shifted.byte_range.start < shifted.byte_range.end).then_some(shifted) + }) + .collect(); if !required.is_empty() { - let inputs = parse_prompt_inputs(rest).map_err(|error| PromptExpansionError::Args { - command: format!("/{name}"), - error, + let inputs = parse_prompt_inputs(rest, &local_elements).map_err(|error| { + PromptExpansionError::Args { + command: format!("/{name}"), + error, + } })?; let missing: Vec = required .into_iter() @@ -168,28 +225,19 @@ pub fn expand_custom_prompt( 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())); + let (text, elements) = expand_named_placeholders_with_elements(&prompt.content, &inputs); + return Ok(Some(PromptExpansion { + text, + text_elements: elements, + })); } // 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)) + let pos_args = parse_positional_args(rest, &local_elements); + Ok(Some(expand_numeric_placeholders( + &prompt.content, + &pos_args, + ))) } /// Detect whether `content` contains numeric placeholders ($1..$9) or `$ARGUMENTS`. @@ -213,25 +261,42 @@ pub fn prompt_has_numeric_placeholders(content: &str) -> bool { /// Extract positional arguments from a composer first line like "/name a b" for a given prompt name. /// Returns empty when the command name does not match or when there are no args. -pub fn extract_positional_args_for_prompt_line(line: &str, prompt_name: &str) -> Vec { +pub fn extract_positional_args_for_prompt_line( + line: &str, + prompt_name: &str, + text_elements: &[TextElement], +) -> Vec { let trimmed = line.trim_start(); - let Some(rest) = trimmed.strip_prefix('/') else { + let trim_offset = line.len() - trimmed.len(); + let Some((name, rest, rest_offset)) = parse_slash_name(trimmed) else { return Vec::new(); }; // Require the explicit prompts prefix for custom prompt invocations. - let Some(after_prefix) = rest.strip_prefix(&format!("{PROMPTS_CMD_PREFIX}:")) else { + let Some(after_prefix) = name.strip_prefix(&format!("{PROMPTS_CMD_PREFIX}:")) else { return Vec::new(); }; - let mut parts = after_prefix.splitn(2, char::is_whitespace); - let cmd = parts.next().unwrap_or(""); - if cmd != prompt_name { + if after_prefix != prompt_name { return Vec::new(); } - let args_str = parts.next().unwrap_or("").trim(); + let rest_trimmed_start = rest.trim_start(); + let args_str = rest_trimmed_start.trim_end(); if args_str.is_empty() { return Vec::new(); } - parse_positional_args(args_str) + let args_offset = trim_offset + rest_offset + (rest.len() - rest_trimmed_start.len()); + let local_elements: Vec = text_elements + .iter() + .filter_map(|elem| { + let mut shifted = shift_text_element_left(elem, args_offset)?; + if shifted.byte_range.start >= args_str.len() { + return None; + } + let end = shifted.byte_range.end.min(args_str.len()); + shifted.byte_range.end = end; + (shifted.byte_range.start < shifted.byte_range.end).then_some(shifted) + }) + .collect(); + parse_positional_args(args_str, &local_elements) } /// If the prompt only uses numeric placeholders and the first line contains @@ -239,14 +304,15 @@ pub fn extract_positional_args_for_prompt_line(line: &str, prompt_name: &str) -> pub fn expand_if_numeric_with_positional_args( prompt: &CustomPrompt, first_line: &str, -) -> Option { + text_elements: &[TextElement], +) -> Option { if !prompt_argument_names(&prompt.content).is_empty() { return None; } if !prompt_has_numeric_placeholders(&prompt.content) { return None; } - let args = extract_positional_args_for_prompt_line(first_line, &prompt.name); + let args = extract_positional_args_for_prompt_line(first_line, &prompt.name, text_elements); if args.is_empty() { return None; } @@ -254,10 +320,10 @@ pub fn expand_if_numeric_with_positional_args( } /// Expand `$1..$9` and `$ARGUMENTS` in `content` with values from `args`. -pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { +pub fn expand_numeric_placeholders(content: &str, args: &[PromptArg]) -> PromptExpansion { let mut out = String::with_capacity(content.len()); + let mut out_elements = Vec::new(); let mut i = 0; - let mut cached_joined_args: Option = None; while let Some(off) = content[i..].find('$') { let j = i + off; out.push_str(&content[i..j]); @@ -272,8 +338,8 @@ pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { } b'1'..=b'9' => { let idx = (bytes[1] - b'1') as usize; - if let Some(val) = args.get(idx) { - out.push_str(val); + if let Some(arg) = args.get(idx) { + append_arg_with_elements(&mut out, &mut out_elements, arg); } i = j + 2; continue; @@ -283,8 +349,7 @@ pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { } if rest.len() > "ARGUMENTS".len() && rest[1..].starts_with("ARGUMENTS") { if !args.is_empty() { - let joined = cached_joined_args.get_or_insert_with(|| args.join(" ")); - out.push_str(joined); + append_joined_args_with_elements(&mut out, &mut out_elements, args); } i = j + 1 + "ARGUMENTS".len(); continue; @@ -293,7 +358,191 @@ pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String { i = j + 1; } out.push_str(&content[i..]); - out + PromptExpansion { + text: out, + text_elements: out_elements, + } +} + +fn parse_tokens_with_elements(rest: &str, text_elements: &[TextElement]) -> Vec { + let mut elements = text_elements.to_vec(); + elements.sort_by_key(|elem| elem.byte_range.start); + // Keep element placeholders intact across shlex splitting by replacing + // each element range with a unique sentinel token first. + let (rest_for_shlex, replacements) = replace_text_elements_with_sentinels(rest, &elements); + Shlex::new(&rest_for_shlex) + .map(|token| apply_replacements_to_token(token, &replacements)) + .collect() +} + +#[derive(Debug, Clone)] +struct ElementReplacement { + sentinel: String, + text: String, + placeholder: Option, +} + +/// Replace each text element range with a unique sentinel token. +/// +/// The sentinel is chosen so it will survive shlex tokenization as a single word. +fn replace_text_elements_with_sentinels( + rest: &str, + elements: &[TextElement], +) -> (String, Vec) { + let mut out = String::with_capacity(rest.len()); + let mut replacements = Vec::new(); + let mut cursor = 0; + + for (idx, elem) in elements.iter().enumerate() { + let start = elem.byte_range.start; + let end = elem.byte_range.end; + out.push_str(&rest[cursor..start]); + let mut sentinel = format!("__CODEX_ELEM_{idx}__"); + // Ensure we never collide with user content so a sentinel can't be mistaken for text. + while rest.contains(&sentinel) { + sentinel.push('_'); + } + out.push_str(&sentinel); + replacements.push(ElementReplacement { + sentinel, + text: rest[start..end].to_string(), + placeholder: elem.placeholder(rest).map(str::to_string), + }); + cursor = end; + } + + out.push_str(&rest[cursor..]); + (out, replacements) +} + +/// Rehydrate a shlex token by swapping sentinels back to the original text +/// and rebuilding text element ranges relative to the resulting token. +fn apply_replacements_to_token(token: String, replacements: &[ElementReplacement]) -> PromptArg { + if replacements.is_empty() { + return PromptArg { + text: token, + text_elements: Vec::new(), + }; + } + + let mut out = String::with_capacity(token.len()); + let mut out_elements = Vec::new(); + let mut cursor = 0; + + while cursor < token.len() { + let Some((offset, replacement)) = next_replacement(&token, cursor, replacements) else { + out.push_str(&token[cursor..]); + break; + }; + let start_in_token = cursor + offset; + out.push_str(&token[cursor..start_in_token]); + let start = out.len(); + out.push_str(&replacement.text); + let end = out.len(); + if start < end { + out_elements.push(TextElement::new( + ByteRange { start, end }, + replacement.placeholder.clone(), + )); + } + cursor = start_in_token + replacement.sentinel.len(); + } + + PromptArg { + text: out, + text_elements: out_elements, + } +} + +/// Find the earliest sentinel occurrence at or after `cursor`. +fn next_replacement<'a>( + token: &str, + cursor: usize, + replacements: &'a [ElementReplacement], +) -> Option<(usize, &'a ElementReplacement)> { + let slice = &token[cursor..]; + let mut best: Option<(usize, &'a ElementReplacement)> = None; + for replacement in replacements { + if let Some(pos) = slice.find(&replacement.sentinel) { + match best { + Some((best_pos, _)) if best_pos <= pos => {} + _ => best = Some((pos, replacement)), + } + } + } + best +} + +/// Shift a text element's byte range left by `offset`, returning `None` if empty. +/// +/// `offset` is the byte length of the prefix removed from the original text. +fn shift_text_element_left(elem: &TextElement, offset: usize) -> Option { + if elem.byte_range.end <= offset { + return None; + } + let start = elem.byte_range.start.saturating_sub(offset); + let end = elem.byte_range.end.saturating_sub(offset); + (start < end).then_some(elem.map_range(|_| ByteRange { start, end })) +} + +fn expand_named_placeholders_with_elements( + content: &str, + args: &HashMap, +) -> (String, Vec) { + let mut out = String::with_capacity(content.len()); + let mut out_elements = Vec::new(); + let mut cursor = 0; + for m in PROMPT_ARG_REGEX.find_iter(content) { + let start = m.start(); + let end = m.end(); + if start > 0 && content.as_bytes()[start - 1] == b'$' { + out.push_str(&content[cursor..end]); + cursor = end; + continue; + } + out.push_str(&content[cursor..start]); + cursor = end; + let key = &content[start + 1..end]; + if let Some(arg) = args.get(key) { + append_arg_with_elements(&mut out, &mut out_elements, arg); + } else { + out.push_str(&content[start..end]); + } + } + out.push_str(&content[cursor..]); + (out, out_elements) +} + +fn append_arg_with_elements( + out: &mut String, + out_elements: &mut Vec, + arg: &PromptArg, +) { + let start = out.len(); + out.push_str(&arg.text); + if arg.text_elements.is_empty() { + return; + } + out_elements.extend(arg.text_elements.iter().map(|elem| { + elem.map_range(|range| ByteRange { + start: start + range.start, + end: start + range.end, + }) + })); +} + +fn append_joined_args_with_elements( + out: &mut String, + out_elements: &mut Vec, + args: &[PromptArg], +) { + // `$ARGUMENTS` joins args with single spaces while preserving element ranges. + for (idx, arg) in args.iter().enumerate() { + if idx > 0 { + out.push(' '); + } + append_arg_with_elements(out, out_elements, arg); + } } /// Constructs a command text for a custom prompt with arguments. @@ -313,6 +562,7 @@ pub fn prompt_command_with_arg_placeholders(name: &str, args: &[String]) -> (Str #[cfg(test)] mod tests { use super::*; + use pretty_assertions::assert_eq; #[test] fn expand_arguments_basic() { @@ -324,9 +574,15 @@ mod tests { 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())); + let out = expand_custom_prompt("/prompts:my-prompt USER=Alice BRANCH=main", &[], &prompts) + .unwrap(); + assert_eq!( + out, + Some(PromptExpansion { + text: "Review Alice changes on main".to_string(), + text_elements: Vec::new(), + }) + ); } #[test] @@ -341,10 +597,17 @@ mod tests { 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())); + assert_eq!( + out, + Some(PromptExpansion { + text: "Pair Alice Smith with dev-main".to_string(), + text_elements: Vec::new(), + }) + ); } #[test] @@ -356,7 +619,7 @@ mod tests { description: None, argument_hint: None, }]; - let err = expand_custom_prompt("/prompts:my-prompt USER=Alice stray", &prompts) + let err = expand_custom_prompt("/prompts:my-prompt USER=Alice stray", &[], &prompts) .unwrap_err() .user_message(); assert!(err.contains("expected key=value")); @@ -371,7 +634,7 @@ mod tests { description: None, argument_hint: None, }]; - let err = expand_custom_prompt("/prompts:my-prompt USER=Alice", &prompts) + let err = expand_custom_prompt("/prompts:my-prompt USER=Alice", &[], &prompts) .unwrap_err() .user_message(); assert!(err.to_lowercase().contains("missing required args")); @@ -400,7 +663,192 @@ mod tests { argument_hint: None, }]; - let out = expand_custom_prompt("/prompts:my-prompt", &prompts).unwrap(); - assert_eq!(out, Some("literal $$USER".to_string())); + let out = expand_custom_prompt("/prompts:my-prompt", &[], &prompts).unwrap(); + assert_eq!( + out, + Some(PromptExpansion { + text: "literal $$USER".to_string(), + text_elements: Vec::new(), + }) + ); + } + + #[test] + fn positional_args_treat_placeholder_with_spaces_as_single_token() { + let placeholder = "[Image #1]"; + let rest = format!("alpha {placeholder} beta"); + let start = rest.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = parse_positional_args(&rest, &text_elements); + assert_eq!( + args, + vec![ + PromptArg { + text: "alpha".to_string(), + text_elements: Vec::new(), + }, + PromptArg { + text: placeholder.to_string(), + text_elements: vec![TextElement::new( + ByteRange { + start: 0, + end: placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }, + PromptArg { + text: "beta".to_string(), + text_elements: Vec::new(), + } + ] + ); + } + + #[test] + fn extract_positional_args_shifts_element_offsets_into_args_str() { + let placeholder = "[Image #1]"; + let line = format!(" /{PROMPTS_CMD_PREFIX}:my-prompt alpha {placeholder} beta "); + let start = line.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = extract_positional_args_for_prompt_line(&line, "my-prompt", &text_elements); + assert_eq!( + args, + vec![ + PromptArg { + text: "alpha".to_string(), + text_elements: Vec::new(), + }, + PromptArg { + text: placeholder.to_string(), + text_elements: vec![TextElement::new( + ByteRange { + start: 0, + end: placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }, + PromptArg { + text: "beta".to_string(), + text_elements: Vec::new(), + } + ] + ); + } + + #[test] + fn key_value_args_treat_placeholder_with_spaces_as_single_token() { + let placeholder = "[Image #1]"; + let rest = format!("IMG={placeholder} NOTE=hello"); + let start = rest.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = parse_prompt_inputs(&rest, &text_elements).expect("inputs"); + assert_eq!( + args.get("IMG"), + Some(&PromptArg { + text: placeholder.to_string(), + text_elements: vec![TextElement::new( + ByteRange { + start: 0, + end: placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }) + ); + assert_eq!( + args.get("NOTE"), + Some(&PromptArg { + text: "hello".to_string(), + text_elements: Vec::new(), + }) + ); + } + + #[test] + fn positional_args_allow_placeholder_inside_quotes() { + let placeholder = "[Image #1]"; + let rest = format!("alpha \"see {placeholder} here\" beta"); + let start = rest.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = parse_positional_args(&rest, &text_elements); + assert_eq!( + args, + vec![ + PromptArg { + text: "alpha".to_string(), + text_elements: Vec::new(), + }, + PromptArg { + text: format!("see {placeholder} here"), + text_elements: vec![TextElement::new( + ByteRange { + start: "see ".len(), + end: "see ".len() + placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }, + PromptArg { + text: "beta".to_string(), + text_elements: Vec::new(), + } + ] + ); + } + + #[test] + fn key_value_args_allow_placeholder_inside_quotes() { + let placeholder = "[Image #1]"; + let rest = format!("IMG=\"see {placeholder} here\" NOTE=ok"); + let start = rest.find(placeholder).expect("placeholder"); + let end = start + placeholder.len(); + let text_elements = vec![TextElement::new( + ByteRange { start, end }, + Some(placeholder.to_string()), + )]; + + let args = parse_prompt_inputs(&rest, &text_elements).expect("inputs"); + assert_eq!( + args.get("IMG"), + Some(&PromptArg { + text: format!("see {placeholder} here"), + text_elements: vec![TextElement::new( + ByteRange { + start: "see ".len(), + end: "see ".len() + placeholder.len(), + }, + Some(placeholder.to_string()), + )], + }) + ); + assert_eq!( + args.get("NOTE"), + Some(&PromptArg { + text: "ok".to_string(), + text_elements: Vec::new(), + }) + ); } } diff --git a/docs/tui-chat-composer.md b/docs/tui-chat-composer.md index bda9b7ff7..76408e098 100644 --- a/docs/tui-chat-composer.md +++ b/docs/tui-chat-composer.md @@ -51,6 +51,34 @@ The solution is to detect paste-like _bursts_ and buffer them into a single expl - After handling the key, `sync_popups()` runs so popup visibility/filters stay consistent with the latest text + cursor. +## Submission flow (Enter/Tab) + +There are multiple submission paths, but they share the same core rules: + +### Normal submit/queue path + +`handle_submission` calls `prepare_submission_text` for both submit and queue. That method: + +1. Expands any pending paste placeholders so element ranges align with the final text. +2. Trims whitespace and rebases element ranges to the trimmed buffer. +3. Expands `/prompts:` custom prompts: + - Named args use key=value parsing. + - Numeric args use positional parsing for `$1..$9` and `$ARGUMENTS`. + The expansion preserves text elements and yields the final submission payload. +4. Prunes attachments so only placeholders that survive expansion are sent. +5. Clears pending pastes on success and suppresses submission if the final text is empty and there + are no attachments. + +### Numeric auto-submit path + +When the slash popup is open and the first line matches a numeric-only custom prompt with +positional args, Enter auto-submits without calling `prepare_submission_text`. That path still: + +- Expands pending pastes before parsing positional args. +- Uses expanded text elements for prompt expansion. +- Prunes attachments based on expanded placeholders. +- Clears pending pastes after a successful auto-submit. + ## Paste burst: concepts and assumptions The burst detector is intentionally conservative: it only processes “plain” character input