mirror of
https://github.com/openai/codex.git
synced 2026-10-31 19:28:49 +00:00
560 lines
20 KiB
Rust
560 lines
20 KiB
Rust
use serde_json::Value;
|
|
use std::collections::HashMap;
|
|
|
|
use crate::responses::ResponsesRequest;
|
|
|
|
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
|
pub enum ContextSnapshotRenderMode {
|
|
#[default]
|
|
RedactedText,
|
|
FullText,
|
|
KindOnly,
|
|
KindWithTextPrefix {
|
|
max_chars: usize,
|
|
},
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct ContextSnapshotOptions {
|
|
render_mode: ContextSnapshotRenderMode,
|
|
}
|
|
|
|
impl Default for ContextSnapshotOptions {
|
|
fn default() -> Self {
|
|
Self {
|
|
render_mode: ContextSnapshotRenderMode::RedactedText,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ContextSnapshotOptions {
|
|
pub fn render_mode(mut self, render_mode: ContextSnapshotRenderMode) -> Self {
|
|
self.render_mode = render_mode;
|
|
self
|
|
}
|
|
}
|
|
|
|
pub fn format_request_input_snapshot(
|
|
request: &ResponsesRequest,
|
|
options: &ContextSnapshotOptions,
|
|
) -> String {
|
|
let items = request.input();
|
|
format_response_items_snapshot(items.as_slice(), options)
|
|
}
|
|
|
|
pub fn format_response_items_snapshot(items: &[Value], options: &ContextSnapshotOptions) -> String {
|
|
let mut canonicalizer = SnapshotCanonicalizer::default();
|
|
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");
|
|
let rendered_parts = item
|
|
.get("content")
|
|
.and_then(Value::as_array)
|
|
.map(|content| {
|
|
content
|
|
.iter()
|
|
.map(|entry| {
|
|
if let Some(text) = entry.get("text").and_then(Value::as_str) {
|
|
return format_snapshot_text(
|
|
text,
|
|
options,
|
|
&mut canonicalizer,
|
|
);
|
|
}
|
|
let Some(content_type) =
|
|
entry.get("type").and_then(Value::as_str)
|
|
else {
|
|
return "<UNKNOWN_CONTENT_ITEM>".to_string();
|
|
};
|
|
let Some(content_object) = entry.as_object() else {
|
|
return format!("<{content_type}>");
|
|
};
|
|
let mut extra_keys = content_object
|
|
.keys()
|
|
.filter(|key| *key != "type" && *key != "text")
|
|
.cloned()
|
|
.collect::<Vec<String>>();
|
|
extra_keys.sort();
|
|
if extra_keys.is_empty() {
|
|
format!("<{content_type}>")
|
|
} else {
|
|
format!("<{content_type}:{}>", extra_keys.join(","))
|
|
}
|
|
})
|
|
.collect::<Vec<String>>()
|
|
})
|
|
.unwrap_or_default();
|
|
let role = if rendered_parts.len() > 1 {
|
|
format!("{role}[{}]", rendered_parts.len())
|
|
} else {
|
|
role.to_string()
|
|
};
|
|
if rendered_parts.is_empty() {
|
|
return format!("{idx:02}:message/{role}:<NO_TEXT>");
|
|
}
|
|
if rendered_parts.len() == 1 {
|
|
return format!("{idx:02}:message/{role}:{}", rendered_parts[0]);
|
|
}
|
|
|
|
let parts = rendered_parts
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(part_idx, part)| format!(" [{:02}] {part}", part_idx + 1))
|
|
.collect::<Vec<String>>()
|
|
.join("\n");
|
|
format!("{idx:02}:message/{role}:\n{parts}")
|
|
}
|
|
"function_call" => {
|
|
let name = item.get("name").and_then(Value::as_str).unwrap_or("unknown");
|
|
format!("{idx:02}:function_call/{name}")
|
|
}
|
|
"function_call_output" => {
|
|
let output = item
|
|
.get("output")
|
|
.and_then(Value::as_str)
|
|
.map(|output| format_snapshot_text(output, options, &mut canonicalizer))
|
|
.unwrap_or_else(|| "<NON_STRING_OUTPUT>".to_string());
|
|
format!("{idx:02}:function_call_output:{output}")
|
|
}
|
|
"local_shell_call" => {
|
|
let command = item
|
|
.get("action")
|
|
.and_then(|action| action.get("command"))
|
|
.and_then(Value::as_array)
|
|
.map(|parts| {
|
|
parts
|
|
.iter()
|
|
.filter_map(Value::as_str)
|
|
.collect::<Vec<&str>>()
|
|
.join(" ")
|
|
})
|
|
.map(|command| {
|
|
format_snapshot_text(&command, options, &mut canonicalizer)
|
|
})
|
|
.filter(|cmd| !cmd.is_empty())
|
|
.unwrap_or_else(|| "<NO_COMMAND>".to_string());
|
|
format!("{idx:02}:local_shell_call:{command}")
|
|
}
|
|
"reasoning" => {
|
|
let summary_text = item
|
|
.get("summary")
|
|
.and_then(Value::as_array)
|
|
.and_then(|summary| summary.first())
|
|
.and_then(|entry| entry.get("text"))
|
|
.and_then(Value::as_str)
|
|
.map(|text| format_snapshot_text(text, options, &mut canonicalizer))
|
|
.unwrap_or_else(|| "<NO_SUMMARY>".to_string());
|
|
let has_encrypted_content = item
|
|
.get("encrypted_content")
|
|
.and_then(Value::as_str)
|
|
.is_some_and(|value| !value.is_empty());
|
|
format!(
|
|
"{idx:02}:reasoning:summary={summary_text}:encrypted={has_encrypted_content}"
|
|
)
|
|
}
|
|
"compaction" => {
|
|
let has_encrypted_content = item
|
|
.get("encrypted_content")
|
|
.and_then(Value::as_str)
|
|
.is_some_and(|value| !value.is_empty());
|
|
format!("{idx:02}:compaction:encrypted={has_encrypted_content}")
|
|
}
|
|
other => format!("{idx:02}:{other}"),
|
|
}
|
|
})
|
|
.collect::<Vec<String>>()
|
|
.join("\n")
|
|
}
|
|
|
|
pub fn format_labeled_requests_snapshot(
|
|
scenario: &str,
|
|
sections: &[(&str, &ResponsesRequest)],
|
|
options: &ContextSnapshotOptions,
|
|
) -> String {
|
|
let sections = sections
|
|
.iter()
|
|
.map(|(title, request)| {
|
|
format!(
|
|
"## {title}\n{}",
|
|
format_request_input_snapshot(request, options)
|
|
)
|
|
})
|
|
.collect::<Vec<String>>()
|
|
.join("\n\n");
|
|
format!("Scenario: {scenario}\n\n{sections}")
|
|
}
|
|
|
|
pub fn format_labeled_items_snapshot(
|
|
scenario: &str,
|
|
sections: &[(&str, &[Value])],
|
|
options: &ContextSnapshotOptions,
|
|
) -> String {
|
|
let sections = sections
|
|
.iter()
|
|
.map(|(title, items)| {
|
|
format!(
|
|
"## {title}\n{}",
|
|
format_response_items_snapshot(items, options)
|
|
)
|
|
})
|
|
.collect::<Vec<String>>()
|
|
.join("\n\n");
|
|
format!("Scenario: {scenario}\n\n{sections}")
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct SnapshotCanonicalizer {
|
|
cwd_placeholders: HashMap<String, usize>,
|
|
}
|
|
|
|
impl SnapshotCanonicalizer {
|
|
fn canonicalize_text(&mut self, text: &str) -> String {
|
|
if text.starts_with("<permissions instructions>") {
|
|
return "<PERMISSIONS_INSTRUCTIONS>".to_string();
|
|
}
|
|
if text.starts_with("# AGENTS.md instructions for ") {
|
|
return "<AGENTS_MD>".to_string();
|
|
}
|
|
if text.starts_with("<user_instructions>") {
|
|
return "<USER_INSTRUCTIONS>".to_string();
|
|
}
|
|
if text.starts_with("<skills_section>") {
|
|
return "<SKILLS_SECTION>".to_string();
|
|
}
|
|
if text.starts_with("<plugins>") {
|
|
return "<PLUGINS>".to_string();
|
|
}
|
|
if text.starts_with("<js_repl_instructions>") {
|
|
return "<JS_REPL_INSTRUCTIONS>".to_string();
|
|
}
|
|
if text.starts_with("<child_agents_instructions>") {
|
|
return "<CHILD_AGENTS_INSTRUCTIONS>".to_string();
|
|
}
|
|
if text.starts_with("<environment_context>") {
|
|
if let (Some(cwd_start), Some(cwd_end)) = (text.find("<cwd>"), text.find("</cwd>")) {
|
|
let cwd = &text[cwd_start + "<cwd>".len()..cwd_end];
|
|
return if cwd.ends_with("PRETURN_CONTEXT_DIFF_CWD") {
|
|
"<ENVIRONMENT_CONTEXT:cwd=PRETURN_CONTEXT_DIFF_CWD>".to_string()
|
|
} else {
|
|
let next_idx = self.cwd_placeholders.len() + 1;
|
|
let idx = *self
|
|
.cwd_placeholders
|
|
.entry(cwd.to_string())
|
|
.or_insert(next_idx);
|
|
format!("<ENVIRONMENT_CONTEXT:cwd=<CWD#{idx}>>")
|
|
};
|
|
}
|
|
return "<ENVIRONMENT_CONTEXT>".to_string();
|
|
}
|
|
if text.starts_with("<subagents>") {
|
|
let subagent_count = text
|
|
.lines()
|
|
.filter(|line| line.trim_start().starts_with("- "))
|
|
.count();
|
|
return format!("<SUBAGENTS:count={subagent_count}>");
|
|
}
|
|
if text.starts_with("You are performing a CONTEXT CHECKPOINT COMPACTION.") {
|
|
return "<SUMMARIZATION_PROMPT>".to_string();
|
|
}
|
|
if text.starts_with("Another language model started to solve this problem")
|
|
&& let Some((_, summary)) = text.split_once('\n')
|
|
{
|
|
return format!("<COMPACTION_SUMMARY>\n{summary}");
|
|
}
|
|
text.to_string()
|
|
}
|
|
}
|
|
|
|
fn format_snapshot_text(
|
|
text: &str,
|
|
options: &ContextSnapshotOptions,
|
|
canonicalizer: &mut SnapshotCanonicalizer,
|
|
) -> String {
|
|
match options.render_mode {
|
|
ContextSnapshotRenderMode::RedactedText => {
|
|
normalize_snapshot_line_endings(&canonicalizer.canonicalize_text(text))
|
|
.replace('\n', "\\n")
|
|
}
|
|
ContextSnapshotRenderMode::FullText => {
|
|
normalize_snapshot_line_endings(text).replace('\n', "\\n")
|
|
}
|
|
ContextSnapshotRenderMode::KindWithTextPrefix { max_chars } => {
|
|
let normalized =
|
|
normalize_snapshot_line_endings(&canonicalizer.canonicalize_text(text))
|
|
.replace('\n', "\\n");
|
|
if normalized.chars().count() <= max_chars {
|
|
normalized
|
|
} else {
|
|
let prefix = normalized.chars().take(max_chars).collect::<String>();
|
|
format!("{prefix}...")
|
|
}
|
|
}
|
|
ContextSnapshotRenderMode::KindOnly => unreachable!(),
|
|
}
|
|
}
|
|
|
|
fn normalize_snapshot_line_endings(text: &str) -> String {
|
|
text.replace("\r\n", "\n").replace('\r', "\n")
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::ContextSnapshotOptions;
|
|
use super::ContextSnapshotRenderMode;
|
|
use super::format_response_items_snapshot;
|
|
use pretty_assertions::assert_eq;
|
|
use serde_json::json;
|
|
|
|
#[test]
|
|
fn full_text_mode_preserves_unredacted_text() {
|
|
let items = vec![json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "# AGENTS.md instructions for /tmp/example\n\n<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>"
|
|
}]
|
|
})];
|
|
|
|
let rendered = format_response_items_snapshot(
|
|
&items,
|
|
&ContextSnapshotOptions::default().render_mode(ContextSnapshotRenderMode::FullText),
|
|
);
|
|
|
|
assert_eq!(
|
|
rendered,
|
|
r"00:message/user:# AGENTS.md instructions for /tmp/example\n\n<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn full_text_mode_normalizes_crlf_line_endings() {
|
|
let items = vec![json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "line one\r\n\r\nline two"
|
|
}]
|
|
})];
|
|
|
|
let rendered = format_response_items_snapshot(
|
|
&items,
|
|
&ContextSnapshotOptions::default().render_mode(ContextSnapshotRenderMode::FullText),
|
|
);
|
|
|
|
assert_eq!(rendered, r"00:message/user:line one\n\nline two");
|
|
}
|
|
|
|
#[test]
|
|
fn redacted_text_mode_keeps_canonical_placeholders() {
|
|
let items = vec![json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "# AGENTS.md instructions for /tmp/example\n\n<INSTRUCTIONS>\nbody\n</INSTRUCTIONS>"
|
|
}]
|
|
})];
|
|
|
|
let rendered = format_response_items_snapshot(
|
|
&items,
|
|
&ContextSnapshotOptions::default().render_mode(ContextSnapshotRenderMode::RedactedText),
|
|
);
|
|
|
|
assert_eq!(rendered, "00:message/user:<AGENTS_MD>");
|
|
}
|
|
|
|
#[test]
|
|
fn redacted_text_mode_normalizes_subagents_fragment() {
|
|
let items = vec![json!({
|
|
"type": "message",
|
|
"role": "developer",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "<subagents>\n - agent-1: atlas\n - agent-2\n</subagents>"
|
|
}]
|
|
})];
|
|
|
|
let rendered = format_response_items_snapshot(
|
|
&items,
|
|
&ContextSnapshotOptions::default().render_mode(ContextSnapshotRenderMode::RedactedText),
|
|
);
|
|
|
|
assert_eq!(rendered, "00:message/developer:<SUBAGENTS:count=2>");
|
|
}
|
|
|
|
#[test]
|
|
fn redacted_text_mode_canonicalizes_split_contextual_user_fragments() {
|
|
let items = vec![
|
|
json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "<user_instructions>\nbe nice\n</user_instructions>"
|
|
}]
|
|
}),
|
|
json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "<skills_section>\n## Skills\n- foo\n</skills_section>"
|
|
}]
|
|
}),
|
|
json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "<plugins>\n## Plugins\n- bar\n</plugins>"
|
|
}]
|
|
}),
|
|
];
|
|
|
|
let rendered = format_response_items_snapshot(
|
|
&items,
|
|
&ContextSnapshotOptions::default().render_mode(ContextSnapshotRenderMode::RedactedText),
|
|
);
|
|
|
|
assert_eq!(
|
|
rendered,
|
|
concat!(
|
|
"00:message/user:<USER_INSTRUCTIONS>\n",
|
|
"01:message/user:<SKILLS_SECTION>\n",
|
|
"02:message/user:<PLUGINS>"
|
|
)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn redacted_text_mode_assigns_distinct_cwd_placeholders_within_one_snapshot() {
|
|
let items = vec![
|
|
json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "<environment_context>\n <cwd>/tmp/one</cwd>\n</environment_context>"
|
|
}]
|
|
}),
|
|
json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "<environment_context>\n <cwd>/tmp/two</cwd>\n</environment_context>"
|
|
}]
|
|
}),
|
|
json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "<environment_context>\n <cwd>/tmp/one</cwd>\n</environment_context>"
|
|
}]
|
|
}),
|
|
];
|
|
|
|
let rendered = format_response_items_snapshot(
|
|
&items,
|
|
&ContextSnapshotOptions::default().render_mode(ContextSnapshotRenderMode::RedactedText),
|
|
);
|
|
|
|
assert_eq!(
|
|
rendered,
|
|
concat!(
|
|
"00:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD#1>>\n",
|
|
"01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD#2>>\n",
|
|
"02:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD#1>>"
|
|
)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn kind_with_text_prefix_mode_normalizes_crlf_line_endings() {
|
|
let items = vec![json!({
|
|
"type": "message",
|
|
"role": "developer",
|
|
"content": [{
|
|
"type": "input_text",
|
|
"text": "<realtime_conversation>\r\nRealtime conversation started.\r\n\r\nYou are..."
|
|
}]
|
|
})];
|
|
|
|
let rendered = format_response_items_snapshot(
|
|
&items,
|
|
&ContextSnapshotOptions::default()
|
|
.render_mode(ContextSnapshotRenderMode::KindWithTextPrefix { max_chars: 64 }),
|
|
);
|
|
|
|
assert_eq!(
|
|
rendered,
|
|
r"00:message/developer:<realtime_conversation>\nRealtime conversation started.\n\nYou a..."
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn image_only_message_is_rendered_as_non_text_span() {
|
|
let items = vec![json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [{
|
|
"type": "input_image",
|
|
"image_url": "data:image/png;base64,AAAA"
|
|
}]
|
|
})];
|
|
|
|
let rendered = format_response_items_snapshot(&items, &ContextSnapshotOptions::default());
|
|
|
|
assert_eq!(rendered, "00:message/user:<input_image:image_url>");
|
|
}
|
|
|
|
#[test]
|
|
fn mixed_text_and_image_message_keeps_image_span() {
|
|
let items = vec![json!({
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "input_text",
|
|
"text": "<image>"
|
|
},
|
|
{
|
|
"type": "input_image",
|
|
"image_url": "data:image/png;base64,AAAA"
|
|
},
|
|
{
|
|
"type": "input_text",
|
|
"text": "</image>"
|
|
}
|
|
]
|
|
})];
|
|
|
|
let rendered = format_response_items_snapshot(&items, &ContextSnapshotOptions::default());
|
|
|
|
assert_eq!(
|
|
rendered,
|
|
"00:message/user[3]:\n [01] <image>\n [02] <input_image:image_url>\n [03] </image>"
|
|
);
|
|
}
|
|
}
|