diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/action.rs b/codex-rs/ext/guardian-v2/src/async_scorer/action.rs new file mode 100644 index 0000000000..eea41dfe0c --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/action.rs @@ -0,0 +1,170 @@ +//! Renders bounded planned-action JSON while preserving tool and call identity. +//! String truncation precedes omission of optional fields. + +use codex_extension_api::ToolName; +use codex_extension_api::ToolPayload; +use codex_protocol::protocol::TruncationPolicy; +use serde_json::json; + +use super::transcript::truncate_entry; + +pub(super) struct GuardianAction { + pub(super) tool_name: ToolName, + pub(super) payload: ToolPayload, +} + +pub(super) struct RenderedAction { + pub(super) text: String, + pub(super) original_bytes: usize, +} + +impl GuardianAction { + pub(super) fn render(self, max_action_tokens: usize) -> serde_json::Result { + let arguments = match self.payload { + ToolPayload::Function { arguments } => { + serde_json::from_str(&arguments).unwrap_or(serde_json::Value::String(arguments)) + } + ToolPayload::Custom { input } => serde_json::Value::String(input), + ToolPayload::ToolSearch { arguments } => json!(arguments), + }; + let mut action = match arguments { + serde_json::Value::Object(arguments) => arguments, + arguments => serde_json::Map::from_iter([("arguments".to_owned(), arguments)]), + }; + action.insert( + "tool".to_owned(), + serde_json::Value::String(self.tool_name.to_string()), + ); + + action.sort_keys(); + action + .values_mut() + .for_each(serde_json::Value::sort_all_objects); + let max_action_bytes = TruncationPolicy::Tokens(max_action_tokens).byte_budget(); + let rendered = serde_json::to_string_pretty(&action)?; + let original_bytes = rendered.len(); + if rendered.len().saturating_add(1) <= max_action_bytes { + return Ok(RenderedAction { + text: rendered, + original_bytes, + }); + } + + if let Some(rendered) = fit_action_to_budget(&action, max_action_bytes, max_action_tokens)? + { + return Ok(RenderedAction { + text: rendered, + original_bytes, + }); + } + + let mut omission_key = "_guardian_omitted_fields".to_owned(); + while action.contains_key(&omission_key) { + omission_key.push('_'); + } + let mut retained = serde_json::Map::new(); + for key in ["tool", "call_id"] { + if let Some(value) = action.get(key) { + retained.insert(key.to_owned(), value.clone()); + } + } + let mut omitted = action.len().saturating_sub(retained.len()); + retained.insert(omission_key.clone(), json!(omitted)); + + let mut optional_fields = action + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "tool" | "call_id")) + .collect::>(); + optional_fields.sort_by_key(|(key, _)| { + !matches!( + key.as_str(), + "arguments" | "cmd" | "command" | "input" | "patch" | "path" | "url" + ) + }); + for (key, value) in optional_fields { + let mut candidate = retained.clone(); + candidate.insert(key.clone(), value.clone()); + candidate.insert(omission_key.clone(), json!(omitted.saturating_sub(1))); + candidate.sort_keys(); + let minimized = render_action_with_limit(&candidate, /*max_tokens*/ 0)?; + if minimized.len().saturating_add(1) <= max_action_bytes { + retained = candidate; + omitted = omitted.saturating_sub(1); + } + } + + retained.sort_keys(); + let rendered = fit_action_to_budget(&retained, max_action_bytes, max_action_tokens)? + .ok_or_else(|| { + serde_json::Error::io(std::io::Error::other(format!( + "Guardian action identity exceeds the {max_action_tokens}-token limit" + ))) + })?; + Ok(RenderedAction { + text: rendered, + original_bytes, + }) + } +} + +fn fit_action_to_budget( + action: &serde_json::Map, + max_action_bytes: usize, + max_action_tokens: usize, +) -> serde_json::Result> { + let mut low = 0usize; + let mut high = max_action_tokens.saturating_add(1); + let mut best = None; + + while low < high { + let max_tokens = low + (high - low) / 2; + let rendered = render_action_with_limit(action, max_tokens)?; + if rendered.len().saturating_add(1) <= max_action_bytes { + best = Some(rendered); + low = max_tokens.saturating_add(1); + } else { + high = max_tokens; + } + } + + Ok(best) +} + +fn render_action_with_limit( + action: &serde_json::Map, + max_tokens: usize, +) -> serde_json::Result { + let mut truncated = action.clone(); + for (key, value) in &mut truncated { + if !matches!(key.as_str(), "tool" | "call_id") { + truncate_action_value(value, max_tokens); + } + } + serde_json::to_string_pretty(&truncated) +} + +fn truncate_action_value(value: &mut serde_json::Value, max_tokens: usize) { + match value { + serde_json::Value::String(text) => { + let truncated = truncate_entry(text, max_tokens); + if truncated.len() < text.len() { + *text = truncated; + } + } + serde_json::Value::Array(values) => { + for value in values { + truncate_action_value(value, max_tokens); + } + } + serde_json::Value::Object(values) => { + for value in values.values_mut() { + truncate_action_value(value, max_tokens); + } + } + serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => {} + } +} + +#[cfg(test)] +#[path = "action_tests.rs"] +mod tests; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/action_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/action_tests.rs new file mode 100644 index 0000000000..cfa52a59e2 --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/action_tests.rs @@ -0,0 +1,96 @@ +//! Tests for bounded action rendering and field retention. + +use anyhow::Result; +use codex_extension_api::ToolName; +use codex_extension_api::ToolPayload; +use codex_protocol::protocol::TruncationPolicy; +use pretty_assertions::assert_eq; +use serde_json::json; + +use super::GuardianAction; +use crate::async_scorer::config::DEFAULT_MODEL_CONTEXT_ITEM_TOKENS; + +#[test] +fn guardian_action_bounds_structurally_oversized_arrays() -> Result<()> { + let action = GuardianAction { + tool_name: ToolName::plain("inspect_values"), + payload: ToolPayload::Function { + arguments: json!({ + "call_id": "genuine-call", + "tool": "spoofed-tool", + "values": (0..6_000).collect::>(), + }) + .to_string(), + }, + }; + + let rendered = action.render(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS)?.text; + assert!( + rendered.len().saturating_add(1) + <= TruncationPolicy::Tokens(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS).byte_budget() + ); + let action = serde_json::from_str::(&rendered)?; + assert_eq!( + action, + json!({ + "_guardian_omitted_fields": 1, + "call_id": "genuine-call", + "tool": "inspect_values", + }) + ); + + Ok(()) +} + +#[test] +fn guardian_action_bounds_structurally_oversized_object_keys() -> Result<()> { + let oversized_key = "oversized_key_".to_owned() + + &"k".repeat(TruncationPolicy::Tokens(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS).byte_budget()); + let mut arguments = serde_json::Map::from_iter([ + ( + "_guardian_omitted_fields".to_owned(), + json!("actual-tool-argument"), + ), + ("call_id".to_owned(), json!("genuine-call")), + ("cmd".to_owned(), json!("remove-important-file")), + ("tool".to_owned(), json!("spoofed-tool")), + (oversized_key.clone(), json!(true)), + ]); + for index in 0..600 { + arguments.insert(format!("a_{index:04}_{}", "k".repeat(64)), json!(index)); + } + let original_field_count = arguments.len(); + let action = GuardianAction { + tool_name: ToolName::plain("inspect_fields"), + payload: ToolPayload::Function { + arguments: serde_json::Value::Object(arguments).to_string(), + }, + }; + + let rendered = action.render(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS)?.text; + assert!( + rendered.len().saturating_add(1) + <= TruncationPolicy::Tokens(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS).byte_budget() + ); + let action = serde_json::from_str::(&rendered)?; + let fields = action + .as_object() + .expect("the bounded action must remain a JSON object"); + assert_eq!(fields.get("tool"), Some(&json!("inspect_fields"))); + assert_eq!(fields.get("call_id"), Some(&json!("genuine-call"))); + assert_eq!(fields.get("cmd"), Some(&json!("remove-important-file"))); + assert_eq!( + fields.get("_guardian_omitted_fields"), + Some(&json!("actual-tool-argument")) + ); + assert!(fields.len() < original_field_count); + assert!(!fields.contains_key(&oversized_key)); + assert!( + fields + .get("_guardian_omitted_fields_") + .and_then(serde_json::Value::as_u64) + .is_some_and(|omitted| omitted > 0) + ); + + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs index eeb4b5b3e6..330a139af2 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -41,8 +41,9 @@ use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::TruncationPolicy; use codex_protocol::security_risk::SecurityRiskScore; -use serde_json::json; +use super::action::GuardianAction; +use super::action::RenderedAction; use super::config::GuardianV2Config; use super::config::GuardianV2ReviewScope; use super::metrics::REVIEW_FALLBACK_METRIC; @@ -61,16 +62,6 @@ use super::trusted_skills::TrustedSkillInvocations; use super::trusted_skills::TrustedSkillRoots; use super::trusted_tools::trusted_tool_context; -struct GuardianAction { - tool_name: ToolName, - payload: ToolPayload, -} - -struct RenderedAction { - text: String, - original_bytes: usize, -} - fn should_classify_tool( tool_name: &ToolName, payload: &ToolPayload, @@ -103,152 +94,6 @@ fn should_classify_tool( ) } -impl GuardianAction { - fn render(self, max_action_tokens: usize) -> serde_json::Result { - let arguments = match self.payload { - ToolPayload::Function { arguments } => { - serde_json::from_str(&arguments).unwrap_or(serde_json::Value::String(arguments)) - } - ToolPayload::Custom { input } => serde_json::Value::String(input), - ToolPayload::ToolSearch { arguments } => json!(arguments), - }; - let mut action = match arguments { - serde_json::Value::Object(arguments) => arguments, - arguments => serde_json::Map::from_iter([("arguments".to_owned(), arguments)]), - }; - action.insert( - "tool".to_owned(), - serde_json::Value::String(self.tool_name.to_string()), - ); - - action.sort_keys(); - action - .values_mut() - .for_each(serde_json::Value::sort_all_objects); - let max_action_bytes = TruncationPolicy::Tokens(max_action_tokens).byte_budget(); - let rendered = serde_json::to_string_pretty(&action)?; - let original_bytes = rendered.len(); - if rendered.len().saturating_add(1) <= max_action_bytes { - return Ok(RenderedAction { - text: rendered, - original_bytes, - }); - } - - if let Some(rendered) = fit_action_to_budget(&action, max_action_bytes, max_action_tokens)? - { - return Ok(RenderedAction { - text: rendered, - original_bytes, - }); - } - - let mut omission_key = "_guardian_omitted_fields".to_owned(); - while action.contains_key(&omission_key) { - omission_key.push('_'); - } - let mut retained = serde_json::Map::new(); - for key in ["tool", "call_id"] { - if let Some(value) = action.get(key) { - retained.insert(key.to_owned(), value.clone()); - } - } - let mut omitted = action.len().saturating_sub(retained.len()); - retained.insert(omission_key.clone(), json!(omitted)); - - let mut optional_fields = action - .iter() - .filter(|(key, _)| !matches!(key.as_str(), "tool" | "call_id")) - .collect::>(); - optional_fields.sort_by_key(|(key, _)| { - !matches!( - key.as_str(), - "arguments" | "cmd" | "command" | "input" | "patch" | "path" | "url" - ) - }); - for (key, value) in optional_fields { - let mut candidate = retained.clone(); - candidate.insert(key.clone(), value.clone()); - candidate.insert(omission_key.clone(), json!(omitted.saturating_sub(1))); - candidate.sort_keys(); - let minimized = render_action_with_limit(&candidate, /*max_tokens*/ 0)?; - if minimized.len().saturating_add(1) <= max_action_bytes { - retained = candidate; - omitted = omitted.saturating_sub(1); - } - } - - retained.sort_keys(); - let rendered = fit_action_to_budget(&retained, max_action_bytes, max_action_tokens)? - .ok_or_else(|| { - serde_json::Error::io(std::io::Error::other(format!( - "Guardian action identity exceeds the {max_action_tokens}-token limit" - ))) - })?; - Ok(RenderedAction { - text: rendered, - original_bytes, - }) - } -} - -fn fit_action_to_budget( - action: &serde_json::Map, - max_action_bytes: usize, - max_action_tokens: usize, -) -> serde_json::Result> { - let mut low = 0usize; - let mut high = max_action_tokens.saturating_add(1); - let mut best = None; - - while low < high { - let max_tokens = low + (high - low) / 2; - let rendered = render_action_with_limit(action, max_tokens)?; - if rendered.len().saturating_add(1) <= max_action_bytes { - best = Some(rendered); - low = max_tokens.saturating_add(1); - } else { - high = max_tokens; - } - } - - Ok(best) -} - -fn render_action_with_limit( - action: &serde_json::Map, - max_tokens: usize, -) -> serde_json::Result { - let mut truncated = action.clone(); - for (key, value) in &mut truncated { - if !matches!(key.as_str(), "tool" | "call_id") { - truncate_action_value(value, max_tokens); - } - } - serde_json::to_string_pretty(&truncated) -} - -fn truncate_action_value(value: &mut serde_json::Value, max_tokens: usize) { - match value { - serde_json::Value::String(text) => { - let truncated = super::transcript::truncate_entry(text, max_tokens); - if truncated.len() < text.len() { - *text = truncated; - } - } - serde_json::Value::Array(values) => { - for value in values { - truncate_action_value(value, max_tokens); - } - } - serde_json::Value::Object(values) => { - for value in values.values_mut() { - truncate_action_value(value, max_tokens); - } - } - serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => {} - } -} /// Explains why Guardian v2 requires synchronous approval review. #[derive(Debug, Eq, PartialEq)] pub enum StrictReviewReason { diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index 48bca87e2c..75b360c0be 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs @@ -2889,88 +2889,3 @@ async fn contributor_bounds_oversized_actions_and_fairly_truncates_nested_fields Ok(()) } - -#[test] -fn guardian_action_bounds_structurally_oversized_arrays() -> Result<()> { - let action = super::GuardianAction { - tool_name: ToolName::plain("inspect_values"), - payload: ToolPayload::Function { - arguments: json!({ - "call_id": "genuine-call", - "tool": "spoofed-tool", - "values": (0..6_000).collect::>(), - }) - .to_string(), - }, - }; - - let rendered = action.render(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS)?.text; - assert!( - rendered.len().saturating_add(1) - <= TruncationPolicy::Tokens(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS).byte_budget() - ); - let action = serde_json::from_str::(&rendered)?; - assert_eq!( - action, - json!({ - "_guardian_omitted_fields": 1, - "call_id": "genuine-call", - "tool": "inspect_values", - }) - ); - - Ok(()) -} - -#[test] -fn guardian_action_bounds_structurally_oversized_object_keys() -> Result<()> { - let oversized_key = "oversized_key_".to_owned() - + &"k".repeat(TruncationPolicy::Tokens(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS).byte_budget()); - let mut arguments = serde_json::Map::from_iter([ - ( - "_guardian_omitted_fields".to_owned(), - json!("actual-tool-argument"), - ), - ("call_id".to_owned(), json!("genuine-call")), - ("cmd".to_owned(), json!("remove-important-file")), - ("tool".to_owned(), json!("spoofed-tool")), - (oversized_key.clone(), json!(true)), - ]); - for index in 0..600 { - arguments.insert(format!("a_{index:04}_{}", "k".repeat(64)), json!(index)); - } - let original_field_count = arguments.len(); - let action = super::GuardianAction { - tool_name: ToolName::plain("inspect_fields"), - payload: ToolPayload::Function { - arguments: serde_json::Value::Object(arguments).to_string(), - }, - }; - - let rendered = action.render(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS)?.text; - assert!( - rendered.len().saturating_add(1) - <= TruncationPolicy::Tokens(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS).byte_budget() - ); - let action = serde_json::from_str::(&rendered)?; - let fields = action - .as_object() - .expect("the bounded action must remain a JSON object"); - assert_eq!(fields.get("tool"), Some(&json!("inspect_fields"))); - assert_eq!(fields.get("call_id"), Some(&json!("genuine-call"))); - assert_eq!(fields.get("cmd"), Some(&json!("remove-important-file"))); - assert_eq!( - fields.get("_guardian_omitted_fields"), - Some(&json!("actual-tool-argument")) - ); - assert!(fields.len() < original_field_count); - assert!(!fields.contains_key(&oversized_key)); - assert!( - fields - .get("_guardian_omitted_fields_") - .and_then(serde_json::Value::as_u64) - .is_some_and(|omitted| omitted > 0) - ); - - Ok(()) -} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs b/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs index b552b68abb..4ce01e4500 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs @@ -1,3 +1,4 @@ +mod action; mod config; mod extension; mod metrics;