mirror of
https://github.com/openai/codex.git
synced 2026-11-06 20:20:34 +00:00
Extract Guardian action rendering into its own module (#41151)
## What changed - Move bounded planned-action JSON rendering from the Guardian V2 extension implementation into a dedicated `async_scorer::action` module. - Keep the structural size-limit tests alongside the extracted implementation. GitOrigin-RevId: 9292b164588ba0119b4c5ff81056bf7416352d69
This commit is contained in:
jif
@@ -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<RenderedAction> {
|
||||
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::<Vec<_>>();
|
||||
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<String, serde_json::Value>,
|
||||
max_action_bytes: usize,
|
||||
max_action_tokens: usize,
|
||||
) -> serde_json::Result<Option<String>> {
|
||||
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<String, serde_json::Value>,
|
||||
max_tokens: usize,
|
||||
) -> serde_json::Result<String> {
|
||||
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;
|
||||
@@ -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::<Vec<_>>(),
|
||||
})
|
||||
.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::<serde_json::Value>(&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::<serde_json::Value>(&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(())
|
||||
}
|
||||
@@ -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<RenderedAction> {
|
||||
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::<Vec<_>>();
|
||||
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<String, serde_json::Value>,
|
||||
max_action_bytes: usize,
|
||||
max_action_tokens: usize,
|
||||
) -> serde_json::Result<Option<String>> {
|
||||
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<String, serde_json::Value>,
|
||||
max_tokens: usize,
|
||||
) -> serde_json::Result<String> {
|
||||
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 {
|
||||
|
||||
@@ -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::<Vec<_>>(),
|
||||
})
|
||||
.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::<serde_json::Value>(&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::<serde_json::Value>(&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(())
|
||||
}
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod action;
|
||||
mod config;
|
||||
mod extension;
|
||||
mod metrics;
|
||||
|
||||
Reference in New Issue
Block a user