clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-04-24 14:45:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Generalize context snapshot rendering helpers

Browse Source
This commit is contained in:
Charles Cunningham
2026-02-11 18:02:41 -08:00
parent 6264c63223
commit 0bd8ea1997
3 changed files with 213 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

245
codex-rs/core/tests/common/context_snapshot.rs
View File

@@ -2,45 +2,154 @@ use serde_json::Value;
use crate::responses::ResponsesRequest;
#[derive(Debug, Clone, Default)]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ContextSnapshotRenderMode {
#[default]
RedactedText,
FullText,
KindOnly,
}
#[derive(Debug, Clone)]
struct ContextSnapshotPrefixCapture {
prefix: String,
marker_prefix: String,
marker_suffix: String,
}
#[derive(Debug, Clone)]
pub struct ContextSnapshotOptions {
summary_prefix: Option<String>,
summarization_prompt: Option<String>,
cwd_marker: Option<String>,
tool_call_name: Option<String>,
render_mode: ContextSnapshotRenderMode,
normalize_environment_context: bool,
text_exact_replacements: Vec<(String, String)>,
text_prefix_replacements: Vec<(String, String)>,
text_contains_replacements: Vec<(String, String)>,
text_prefix_captures: Vec<ContextSnapshotPrefixCapture>,
cwd_contains_replacements: Vec<(String, String)>,
tool_name_replacements: Vec<(String, String)>,
}
impl Default for ContextSnapshotOptions {
fn default() -> Self {
Self {
render_mode: ContextSnapshotRenderMode::RedactedText,
normalize_environment_context: true,
text_exact_replacements: Vec::new(),
text_prefix_replacements: vec![(
"# AGENTS.md instructions for ".to_string(),
"<AGENTS_MD>".to_string(),
)],
text_contains_replacements: vec![(
"<permissions instructions>".to_string(),
"<PERMISSIONS_INSTRUCTIONS>".to_string(),
)],
text_prefix_captures: Vec::new(),
cwd_contains_replacements: Vec::new(),
tool_name_replacements: Vec::new(),
}
}
}
impl ContextSnapshotOptions {
pub fn summary_prefix(mut self, summary_prefix: impl Into<String>) -> Self {
self.summary_prefix = Some(summary_prefix.into());
pub fn render_mode(mut self, render_mode: ContextSnapshotRenderMode) -> Self {
self.render_mode = render_mode;
self
}
pub fn summarization_prompt(mut self, summarization_prompt: impl Into<String>) -> Self {
self.summarization_prompt = Some(summarization_prompt.into());
pub fn normalize_environment_context(mut self, enabled: bool) -> Self {
self.normalize_environment_context = enabled;
self
}
pub fn cwd_marker(mut self, cwd_marker: impl Into<String>) -> Self {
self.cwd_marker = Some(cwd_marker.into());
pub fn replace_exact_text(
mut self,
text: impl Into<String>,
replacement: impl Into<String>,
) -> Self {
self.text_exact_replacements
.push((text.into(), replacement.into()));
self
}
pub fn tool_call_name(mut self, tool_call_name: impl Into<String>) -> Self {
self.tool_call_name = Some(tool_call_name.into());
pub fn replace_text_with_prefix(
mut self,
prefix: impl Into<String>,
replacement: impl Into<String>,
) -> Self {
self.text_prefix_replacements
.push((prefix.into(), replacement.into()));
self
}
pub fn replace_text_containing(
mut self,
needle: impl Into<String>,
replacement: impl Into<String>,
) -> Self {
self.text_contains_replacements
.push((needle.into(), replacement.into()));
self
}
pub fn capture_text_suffix_after_prefix(
mut self,
prefix: impl Into<String>,
marker_prefix: impl Into<String>,
marker_suffix: impl Into<String>,
) -> Self {
self.text_prefix_captures
.push(ContextSnapshotPrefixCapture {
prefix: prefix.into(),
marker_prefix: marker_prefix.into(),
marker_suffix: marker_suffix.into(),
});
self
}
pub fn replace_cwd_when_contains(
mut self,
cwd_substring: impl Into<String>,
replacement: impl Into<String>,
) -> Self {
self.cwd_contains_replacements
.push((cwd_substring.into(), replacement.into()));
self
}
pub fn replace_tool_name(
mut self,
original_name: impl Into<String>,
replacement: impl Into<String>,
) -> Self {
self.tool_name_replacements
.push((original_name.into(), replacement.into()));
self
}
}
pub fn request_input_shape(request: &ResponsesRequest, options: &ContextSnapshotOptions) -> String {
request
.input()
.into_iter()
let items = request.input();
response_items_shape(items.as_slice(), options)
}
pub fn response_items_shape(items: &[Value], options: &ContextSnapshotOptions) -> String {
items
.iter()
.enumerate()
.map(|(idx, item)| {
let Some(item_type) = item.get("type").and_then(Value::as_str) else {
return format!("{idx:02}:<MISSING_TYPE>");
};
if options.render_mode == ContextSnapshotRenderMode::KindOnly {
return if item_type == "message" {
let role = item.get("role").and_then(Value::as_str).unwrap_or("unknown");
format!("{idx:02}:message/{role}")
} else {
format!("{idx:02}:{item_type}")
};
}
match item_type {
"message" => {
let role = item.get("role").and_then(Value::as_str).unwrap_or("unknown");
@@ -61,25 +170,31 @@ pub fn request_input_shape(request: &ResponsesRequest, options: &ContextSnapshot
}
"function_call" => {
let name = item.get("name").and_then(Value::as_str).unwrap_or("unknown");
let normalized_name = if options.tool_call_name.as_deref() == Some(name) {
"<TOOL_CALL>"
} else {
name
};
let normalized_name = options
.tool_name_replacements
.iter()
.find_map(|(original_name, replacement)| {
(original_name == name).then_some(replacement.as_str())
})
.unwrap_or(name);
format!("{idx:02}:function_call/{normalized_name}")
}
"function_call_output" => {
let output = item
.get("output")
.and_then(Value::as_str)
.map(|output| {
if output.starts_with("unsupported call: ")
|| output.starts_with("unsupported custom tool call: ")
{
"<TOOL_ERROR_OUTPUT>".to_string()
} else {
normalize_shape_text(output, options)
.map(|output| match options.render_mode {
ContextSnapshotRenderMode::RedactedText => {
if output.starts_with("unsupported call: ")
|| output.starts_with("unsupported custom tool call: ")
{
"<TOOL_ERROR_OUTPUT>".to_string()
} else {
normalize_shape_text(output, options)
}
}
ContextSnapshotRenderMode::FullText => output.replace('\n', "\\n"),
ContextSnapshotRenderMode::KindOnly => unreachable!(),
})
.unwrap_or_else(|| "<NON_STRING_OUTPUT>".to_string());
format!("{idx:02}:function_call_output:{output}")
@@ -144,29 +259,64 @@ pub fn sectioned_request_shapes(
format!("Scenario: {scenario}\n\n{sections}")
}
pub fn sectioned_item_shapes(
scenario: &str,
sections: &[(&str, &[Value])],
options: &ContextSnapshotOptions,
) -> String {
let sections = sections
.iter()
.map(|(title, items)| format!("## {title}\n{}", response_items_shape(items, options)))
.collect::<Vec<String>>()
.join("\n\n");
format!("Scenario: {scenario}\n\n{sections}")
}
fn normalize_shape_text(text: &str, options: &ContextSnapshotOptions) -> String {
if options.summarization_prompt.as_deref() == Some(text) {
return "<SUMMARIZATION_PROMPT>".to_string();
if options.render_mode == ContextSnapshotRenderMode::FullText {
return text.replace('\n', "\\n");
}
if let Some(summary_prefix) = options.summary_prefix.as_deref() {
let summary_prefix_line = format!("{summary_prefix}\n");
if let Some(summary) = text.strip_prefix(summary_prefix_line.as_str()) {
return format!("<SUMMARY:{summary}>");
}
if let Some((_, replacement)) = options
.text_exact_replacements
.iter()
.find(|(target, _)| target == text)
{
return replacement.clone();
}
if text.starts_with("# AGENTS.md instructions for ") {
return "<AGENTS_MD>".to_string();
if let Some((_, replacement)) = options
.text_prefix_replacements
.iter()
.find(|(prefix, _)| text.starts_with(prefix.as_str()))
{
return replacement.clone();
}
if text.starts_with("<environment_context>") {
if let Some(capture) = options
.text_prefix_captures
.iter()
.find(|capture| text.starts_with(capture.prefix.as_str()))
{
let suffix = text
.strip_prefix(capture.prefix.as_str())
.unwrap_or_default();
return format!(
"{}{}{}",
capture.marker_prefix, suffix, capture.marker_suffix
);
}
if options.normalize_environment_context && text.starts_with("<environment_context>") {
let cwd = text.lines().find_map(|line| {
let trimmed = line.trim();
let cwd = trimmed.strip_prefix("<cwd>")?.strip_suffix("</cwd>")?;
if options
.cwd_marker
.as_deref()
.is_some_and(|marker| cwd.contains(marker))
if let Some((_, replacement)) = options
.cwd_contains_replacements
.iter()
.find(|(needle, _)| cwd.contains(needle.as_str()))
{
return options.cwd_marker.clone();
return Some(replacement.clone());
}
Some("<CWD>".to_string())
});
@@ -175,8 +325,13 @@ fn normalize_shape_text(text: &str, options: &ContextSnapshotOptions) -> String
None => "<ENVIRONMENT_CONTEXT:cwd=<NONE>>".to_string(),
};
}
if text.contains("<permissions instructions>") {
return "<PERMISSIONS_INSTRUCTIONS>".to_string();
if let Some((_, replacement)) = options
.text_contains_replacements
.iter()
.find(|(needle, _)| text.contains(needle.as_str()))
{
return replacement.clone();
}
text.replace('\n', "\\n")

11
codex-rs/core/tests/suite/compact.rs
View File

@@ -193,10 +193,13 @@ async fn assert_compaction_uses_turn_lifecycle_id(codex: &std::sync::Arc<codex_c
}
fn context_snapshot_options() -> ContextSnapshotOptions {
ContextSnapshotOptions::default()
.summary_prefix(SUMMARY_PREFIX)
.summarization_prompt(SUMMARIZATION_PROMPT)
.cwd_marker(PRETURN_CONTEXT_DIFF_CWD_MARKER)
.tool_call_name(DUMMY_FUNCTION_NAME)
.replace_exact_text(SUMMARIZATION_PROMPT, "<SUMMARIZATION_PROMPT>")
.capture_text_suffix_after_prefix(format!("{SUMMARY_PREFIX}\n"), "<SUMMARY:", ">")
.replace_cwd_when_contains(
PRETURN_CONTEXT_DIFF_CWD_MARKER,
PRETURN_CONTEXT_DIFF_CWD_MARKER,
)
.replace_tool_name(DUMMY_FUNCTION_NAME, "<TOOL_CALL>")
}
fn request_input_shape(request: &core_test_support::responses::ResponsesRequest) -> String {

9
codex-rs/core/tests/suite/compact_remote.rs
View File

@@ -66,9 +66,12 @@ fn user_message_item(text: &str) -> ResponseItem {
fn context_snapshot_options() -> ContextSnapshotOptions {
ContextSnapshotOptions::default()
.summary_prefix(SUMMARY_PREFIX)
.cwd_marker(PRETURN_CONTEXT_DIFF_CWD_MARKER)
.tool_call_name(DUMMY_FUNCTION_NAME)
.capture_text_suffix_after_prefix(format!("{SUMMARY_PREFIX}\n"), "<SUMMARY:", ">")
.replace_cwd_when_contains(
PRETURN_CONTEXT_DIFF_CWD_MARKER,
PRETURN_CONTEXT_DIFF_CWD_MARKER,
)
.replace_tool_name(DUMMY_FUNCTION_NAME, "<TOOL_CALL>")
}
fn request_input_shape(request: &responses::ResponsesRequest) -> String {