Fix large paste issue in custom prompt args

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::<AppEvent>();
+        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::<AppEvent>();

codex-rs/tui/src/bottom_pane/prompt_args.rs

@@ -106,17 +106,17 @@ pub fn prompt_argument_names(content: &str) -> Vec<String> {
     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<HashMap<String, String>, 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<HashMap<String, String>, 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<HashMap<String, String>, 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<Option<String>, 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<String> = 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<String> = Shlex::new(rest).collect();
+    // Otherwise positional placeholders: tokenize args and replace any
+    // placeholder-matching token before expansion.
+    let pos_args: Vec<String> = 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()));
     }
 }