diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 29b02c006..4acc1ef4e 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -51,6 +51,7 @@ use codex_protocol::items::PlanItem; use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; use codex_protocol::models::BaseInstructions; +use codex_protocol::models::format_allow_prefixes; use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::FileChange; use codex_protocol::protocol::HasLegacyEvent; @@ -212,7 +213,6 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; -use codex_protocol::models::render_command_prefix_list; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::InitialHistory; use codex_protocol::user_input::UserInput; @@ -1522,7 +1522,7 @@ impl Session { sub_id: &str, amendment: &ExecPolicyAmendment, ) { - let Some(prefixes) = render_command_prefix_list([amendment.command.as_slice()]) else { + let Some(prefixes) = format_allow_prefixes(vec![amendment.command.clone()]) else { warn!("execpolicy amendment for {sub_id} had no command prefix"); return; }; diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 8e67fc40f..9dffb0256 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -233,10 +233,13 @@ impl DeveloperInstructions { if !request_rule_enabled { APPROVAL_POLICY_ON_REQUEST.to_string() } else { - let command_prefixes = format_allow_prefixes(exec_policy); + let command_prefixes = + format_allow_prefixes(exec_policy.get_allowed_prefixes()); match command_prefixes { Some(prefixes) => { - format!("{APPROVAL_POLICY_ON_REQUEST_RULE}\n{prefixes}") + format!( + "{APPROVAL_POLICY_ON_REQUEST_RULE}\nApproved command prefixes:\n{prefixes}" + ) } None => APPROVAL_POLICY_ON_REQUEST_RULE.to_string(), } @@ -371,20 +374,51 @@ impl DeveloperInstructions { } } -pub fn render_command_prefix_list(prefixes: I) -> Option -where - I: IntoIterator, - P: AsRef<[String]>, -{ - let lines = prefixes - .into_iter() - .map(|prefix| format!("- {}", render_command_prefix(prefix.as_ref()))) - .collect::>(); - if lines.is_empty() { - return None; +const MAX_RENDERED_PREFIXES: usize = 100; +const MAX_ALLOW_PREFIX_TEXT_BYTES: usize = 5000; +const TRUNCATED_MARKER: &str = "...\n[Some commands were truncated]"; + +pub fn format_allow_prefixes(prefixes: Vec>) -> Option { + let mut truncated = false; + if prefixes.len() > MAX_RENDERED_PREFIXES { + truncated = true; } - Some(lines.join("\n")) + let mut prefixes = prefixes; + prefixes.sort_by(|a, b| { + a.len() + .cmp(&b.len()) + .then_with(|| prefix_combined_str_len(a).cmp(&prefix_combined_str_len(b))) + .then_with(|| a.cmp(b)) + }); + + let full_text = prefixes + .into_iter() + .take(MAX_RENDERED_PREFIXES) + .map(|prefix| format!("- {}", render_command_prefix(&prefix))) + .collect::>() + .join("\n"); + + // truncate to last UTF8 char + let mut output = full_text; + let byte_idx = output + .char_indices() + .nth(MAX_ALLOW_PREFIX_TEXT_BYTES) + .map(|(i, _)| i); + if let Some(byte_idx) = byte_idx { + truncated = true; + output = output[..byte_idx].to_string(); + } + + if truncated { + Some(format!("{output}{TRUNCATED_MARKER}")) + } else { + Some(output) + } +} + +fn prefix_combined_str_len(prefix: &[String]) -> usize { + prefix.iter().map(String::len).sum() } fn render_command_prefix(prefix: &[String]) -> String { @@ -396,12 +430,6 @@ fn render_command_prefix(prefix: &[String]) -> String { format!("[{tokens}]") } -fn format_allow_prefixes(exec_policy: &Policy) -> Option { - let prefixes = exec_policy.get_allowed_prefixes(); - let lines = render_command_prefix_list(prefixes)?; - Some(format!("Approved command prefixes:\n{lines}")) -} - impl From for ResponseItem { fn from(di: DeveloperInstructions) -> Self { ResponseItem::Message { @@ -1000,6 +1028,62 @@ mod tests { assert!(text.contains(r#"["git", "pull"]"#)); } + #[test] + fn render_command_prefix_list_sorts_by_len_then_total_len_then_alphabetical() { + let prefixes = vec![ + vec!["b".to_string(), "zz".to_string()], + vec!["aa".to_string()], + vec!["b".to_string()], + vec!["a".to_string(), "b".to_string(), "c".to_string()], + vec!["a".to_string()], + vec!["b".to_string(), "a".to_string()], + ]; + + let output = format_allow_prefixes(prefixes).expect("rendered list"); + assert_eq!( + output, + r#"- ["a"] +- ["b"] +- ["aa"] +- ["b", "a"] +- ["b", "zz"] +- ["a", "b", "c"]"# + .to_string(), + ); + } + + #[test] + fn render_command_prefix_list_limits_output_to_max_prefixes() { + let prefixes = (0..(MAX_RENDERED_PREFIXES + 5)) + .map(|i| vec![format!("{i:03}")]) + .collect::>(); + + let output = format_allow_prefixes(prefixes).expect("rendered list"); + assert_eq!(output.ends_with(TRUNCATED_MARKER), true); + eprintln!("output: {output}"); + assert_eq!(output.lines().count(), MAX_RENDERED_PREFIXES + 1); + } + + #[test] + fn format_allow_prefixes_limits_output() { + let mut exec_policy = Policy::empty(); + for i in 0..200 { + exec_policy + .add_prefix_rule( + &[format!("tool-{i:03}"), "x".repeat(500)], + codex_execpolicy::Decision::Allow, + ) + .expect("add rule"); + } + + let output = + format_allow_prefixes(exec_policy.get_allowed_prefixes()).expect("formatted prefixes"); + assert!( + output.len() <= MAX_ALLOW_PREFIX_TEXT_BYTES + TRUNCATED_MARKER.len(), + "output length exceeds expected limit: {output}", + ); + } + #[test] fn serializes_success_as_plain_string() -> Result<()> { let item = ResponseInputItem::FunctionCallOutput {