Custom prompt args (numeric) (#4470)
[Cherry picked from /pull/3565] Adds $1, $2, $3, $ARGUMENTS param parsing for custom prompts.
This commit is contained in:
151
codex-rs/tui/src/bottom_pane/prompt_args.rs
Normal file
151
codex-rs/tui/src/bottom_pane/prompt_args.rs
Normal file
@@ -0,0 +1,151 @@
|
||||
use codex_protocol::custom_prompts::CustomPrompt;
|
||||
use shlex::Shlex;
|
||||
|
||||
/// Parse a first-line slash command of the form `/name <rest>`.
|
||||
/// 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)> {
|
||||
let stripped = line.strip_prefix('/')?;
|
||||
let mut name_end = stripped.len();
|
||||
for (idx, ch) in stripped.char_indices() {
|
||||
if ch.is_whitespace() {
|
||||
name_end = idx;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let name = &stripped[..name_end];
|
||||
if name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let rest = stripped[name_end..].trim_start();
|
||||
Some((name, rest))
|
||||
}
|
||||
|
||||
/// Parse positional arguments using shlex semantics (supports quoted tokens).
|
||||
pub fn parse_positional_args(rest: &str) -> Vec<String> {
|
||||
Shlex::new(rest).collect()
|
||||
}
|
||||
|
||||
/// Expands a message of the form `/name key=value …` using a matching saved prompt.
|
||||
///
|
||||
/// If the text does not start with `/`, 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(
|
||||
text: &str,
|
||||
custom_prompts: &[CustomPrompt],
|
||||
) -> Result<Option<String>, ()> {
|
||||
let Some((name, rest)) = parse_slash_name(text) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let prompt = match custom_prompts.iter().find(|p| p.name == name) {
|
||||
Some(prompt) => prompt,
|
||||
None => return Ok(None),
|
||||
};
|
||||
// Only support numeric placeholders ($1..$9) and $ARGUMENTS.
|
||||
if prompt_has_numeric_placeholders(&prompt.content) {
|
||||
let pos_args: Vec<String> = Shlex::new(rest).collect();
|
||||
let expanded = expand_numeric_placeholders(&prompt.content, &pos_args);
|
||||
return Ok(Some(expanded));
|
||||
}
|
||||
// No recognized placeholders: return the literal content.
|
||||
Ok(Some(prompt.content.clone()))
|
||||
}
|
||||
|
||||
/// Detect whether `content` contains numeric placeholders ($1..$9) or `$ARGUMENTS`.
|
||||
pub fn prompt_has_numeric_placeholders(content: &str) -> bool {
|
||||
if content.contains("$ARGUMENTS") {
|
||||
return true;
|
||||
}
|
||||
let bytes = content.as_bytes();
|
||||
let mut i = 0;
|
||||
while i + 1 < bytes.len() {
|
||||
if bytes[i] == b'$' {
|
||||
let b1 = bytes[i + 1];
|
||||
if (b'1'..=b'9').contains(&b1) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// 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<String> {
|
||||
let trimmed = line.trim_start();
|
||||
let Some(rest) = trimmed.strip_prefix('/') else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut parts = rest.splitn(2, char::is_whitespace);
|
||||
let cmd = parts.next().unwrap_or("");
|
||||
if cmd != prompt_name {
|
||||
return Vec::new();
|
||||
}
|
||||
let args_str = parts.next().unwrap_or("").trim();
|
||||
if args_str.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
parse_positional_args(args_str)
|
||||
}
|
||||
|
||||
/// If the prompt only uses numeric placeholders and the first line contains
|
||||
/// positional args for it, expand and return Some(expanded); otherwise None.
|
||||
pub fn expand_if_numeric_with_positional_args(
|
||||
prompt: &CustomPrompt,
|
||||
first_line: &str,
|
||||
) -> Option<String> {
|
||||
if !prompt_has_numeric_placeholders(&prompt.content) {
|
||||
return None;
|
||||
}
|
||||
let args = extract_positional_args_for_prompt_line(first_line, &prompt.name);
|
||||
if args.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(expand_numeric_placeholders(&prompt.content, &args))
|
||||
}
|
||||
|
||||
/// Expand `$1..$9` and `$ARGUMENTS` in `content` with values from `args`.
|
||||
pub fn expand_numeric_placeholders(content: &str, args: &[String]) -> String {
|
||||
let mut out = String::with_capacity(content.len());
|
||||
let mut i = 0;
|
||||
let mut cached_joined_args: Option<String> = None;
|
||||
while let Some(off) = content[i..].find('$') {
|
||||
let j = i + off;
|
||||
out.push_str(&content[i..j]);
|
||||
let rest = &content[j..];
|
||||
let bytes = rest.as_bytes();
|
||||
if bytes.len() >= 2 {
|
||||
match bytes[1] {
|
||||
b'$' => {
|
||||
out.push_str("$$");
|
||||
i = j + 2;
|
||||
continue;
|
||||
}
|
||||
b'1'..=b'9' => {
|
||||
let idx = (bytes[1] - b'1') as usize;
|
||||
if let Some(val) = args.get(idx) {
|
||||
out.push_str(val);
|
||||
}
|
||||
i = j + 2;
|
||||
continue;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
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);
|
||||
}
|
||||
i = j + 1 + "ARGUMENTS".len();
|
||||
continue;
|
||||
}
|
||||
out.push('$');
|
||||
i = j + 1;
|
||||
}
|
||||
out.push_str(&content[i..]);
|
||||
out
|
||||
}
|
||||
Reference in New Issue
Block a user