diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index f387a3d57e..bb47793ce8 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -911,23 +911,6 @@ impl ChatComposer { return (InputResult::None, true); } } - // If we have pending placeholder pastes, submit immediately to expand them. - if !self.pending_pastes.is_empty() { - let mut text = self.textarea.text().to_string(); - self.textarea.set_text(""); - for (placeholder, actual) in &self.pending_pastes { - if text.contains(placeholder) { - text = text.replace(placeholder, actual); - } - } - self.pending_pastes.clear(); - if text.is_empty() { - return (InputResult::None, true); - } - self.history.record_local_submission(&text); - return (InputResult::Submitted(text), true); - } - // During a paste-like burst, treat Enter as a newline instead of submit. let now = Instant::now(); if self @@ -943,31 +926,33 @@ impl ChatComposer { let original_input = text.clone(); self.textarea.set_text(""); - // Replace all pending pastes in the text - for (placeholder, actual) in &self.pending_pastes { - if text.contains(placeholder) { - text = text.replace(placeholder, actual); - } - } - self.pending_pastes.clear(); + // Attempt token-aware prompt expansion that also resolves any + // placeholder pastes; otherwise fall back to raw string replacement. + let pending_pastes = self.pending_pastes.clone(); // 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 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); - } - }; + let expanded_prompt = + match expand_custom_prompt(&text, &self.custom_prompts, &pending_pastes) { + 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; + } else { + for (placeholder, actual) in &pending_pastes { + text = text.replace(placeholder, actual); + } } + self.pending_pastes.clear(); if text.is_empty() && !has_attachments { return (InputResult::None, true); } @@ -2804,6 +2789,58 @@ mod tests { assert!(composer.textarea.is_empty()); } + #[test] + fn custom_prompt_large_paste_argument_expands_before_submit() { + 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: "echo".to_string(), + path: "/tmp/echo.md".to_string().into(), + content: "Echo back: $1".to_string(), + description: None, + argument_hint: None, + }]); + + type_chars_humanlike( + &mut composer, + &[ + '/', 'p', 'r', 'o', 'm', 'p', 't', 's', ':', 'e', 'c', 'h', 'o', ' ', '"', + ], + ); + + let mut large = String::new(); + while large.chars().count() <= LARGE_PASTE_CHAR_THRESHOLD + 5 { + large.push_str("Segment with \"double\" quotes and 'single' ticks. "); + } + let count = large.chars().count(); + let placeholder = format!("[Pasted Content {count} chars]"); + assert!(composer.handle_paste(large.clone())); + assert!(composer.textarea.text().contains(&placeholder)); + assert_eq!(composer.pending_pastes.len(), 1); + type_chars_humanlike(&mut composer, &['"']); + + let (result, _needs_redraw) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!( + InputResult::Submitted(format!("Echo back: {large}")), + result + ); + assert!(composer.pending_pastes.is_empty()); + } + #[test] fn custom_prompt_invalid_args_reports_error() { let (tx, mut rx) = unbounded_channel::(); diff --git a/codex-rs/tui/src/bottom_pane/prompt_args.rs b/codex-rs/tui/src/bottom_pane/prompt_args.rs index 48c3cedfab..721a5a2786 100644 --- a/codex-rs/tui/src/bottom_pane/prompt_args.rs +++ b/codex-rs/tui/src/bottom_pane/prompt_args.rs @@ -106,17 +106,17 @@ pub fn prompt_argument_names(content: &str) -> Vec { names } -/// Parses the `key=value` pairs that follow a custom prompt name. -/// -/// The input is split using shlex rules, so quoted values are supported -/// (for example `USER="Alice Smith"`). The function returns a map of parsed -/// arguments, or an error if a token is missing `=` or if the key is empty. -pub fn parse_prompt_inputs(rest: &str) -> Result, PromptArgsError> { +/// Parses `key=value` pairs using shlex. If `pastes` are provided, any value +/// token that exactly equals a paste placeholder is replaced with its original +/// pasted content. +fn parse_prompt_inputs( + rest: &str, + pastes: &[(String, String)], +) -> 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 }); @@ -124,19 +124,49 @@ pub fn parse_prompt_inputs(rest: &str) -> Result, Prompt if key.is_empty() { return Err(PromptArgsError::MissingKey { token }); } - map.insert(key.to_string(), value.to_string()); + let replaced = if let Some((_, actual)) = pastes.iter().find(|(ph, _)| ph == value) { + actual.clone() + } else { + value.to_string() + }; + map.insert(key.to_string(), replaced); } Ok(map) } +/// Replace paste placeholders in `rest` with synthetic, space-free tokens and +/// return both the rewritten string and the mapping from token -> original +/// pasted content. +fn tokenize_pastes(rest: &str, pastes: &[(String, String)]) -> (String, Vec<(String, String)>) { + let mut rest_rewritten = rest.to_string(); + let mut tokenized: Vec<(String, String)> = Vec::with_capacity(pastes.len()); + for (i, (placeholder, actual)) in pastes.iter().enumerate() { + let token = format!("__codex_paste_{i}__"); + // `replace` is a no-op if `placeholder` is absent; no need to pre-check. + rest_rewritten = rest_rewritten.replace(placeholder, &token); + tokenized.push((token, actual.clone())); + } + (rest_rewritten, tokenized) +} + +/// Given a token (possibly a synthetic paste token), return the original +/// pasted content if it matches; otherwise return the token unchanged. +fn resolve_paste_tokens(tok: &str, tokenized_pastes: &[(String, String)]) -> String { + if let Some((_, actual)) = tokenized_pastes.iter().find(|(t, _)| t == tok) { + return actual.clone(); + } + tok.to_string() +} + /// Expands a message of the form `/prompts:name [value] [value] …` using a matching saved prompt. /// /// If the text does not start with `/prompts:`, or if no prompt named `name` exists, /// the function returns `Ok(None)`. On success it returns /// `Ok(Some(expanded))`; otherwise it returns a descriptive error. -pub fn expand_custom_prompt( +pub(crate) fn expand_custom_prompt( text: &str, custom_prompts: &[CustomPrompt], + pastes: &[(String, String)], ) -> Result, PromptExpansionError> { let Some((name, rest)) = parse_slash_name(text) else { return Ok(None); @@ -151,12 +181,20 @@ pub fn expand_custom_prompt( Some(prompt) => prompt, None => return Ok(None), }; + + // Protect paste placeholders from shlex splitting by rewriting them to + // unique, space-free tokens, then map those tokens back to the original + // content during parsing. + let (rest_rewritten, tokenized_pastes) = tokenize_pastes(rest, pastes); + // If there are named placeholders, expect key=value inputs. let required = prompt_argument_names(&prompt.content); if !required.is_empty() { - let inputs = parse_prompt_inputs(rest).map_err(|error| PromptExpansionError::Args { - command: format!("/{name}"), - error, + let inputs = parse_prompt_inputs(&rest_rewritten, &tokenized_pastes).map_err(|error| { + PromptExpansionError::Args { + command: format!("/{name}"), + error, + } })?; let missing: Vec = required .into_iter() @@ -186,8 +224,11 @@ pub fn expand_custom_prompt( return Ok(Some(replaced.into_owned())); } - // Otherwise, treat it as numeric/positional placeholder prompt (or none). - let pos_args: Vec = Shlex::new(rest).collect(); + // Otherwise positional placeholders: tokenize args and replace any + // placeholder-matching token before expansion. + let pos_args: Vec = Shlex::new(&rest_rewritten) + .map(|t| resolve_paste_tokens(&t, &tokenized_pastes)) + .collect(); let expanded = expand_numeric_placeholders(&prompt.content, &pos_args); Ok(Some(expanded)) } @@ -324,8 +365,8 @@ mod tests { argument_hint: None, }]; - let out = - expand_custom_prompt("/prompts:my-prompt USER=Alice BRANCH=main", &prompts).unwrap(); + 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())); } @@ -342,6 +383,7 @@ 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())); @@ -356,7 +398,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 +413,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 +442,7 @@ mod tests { argument_hint: None, }]; - let out = expand_custom_prompt("/prompts:my-prompt", &prompts).unwrap(); + let out = expand_custom_prompt("/prompts:my-prompt", &prompts, &[]).unwrap(); assert_eq!(out, Some("literal $$USER".to_string())); } }