github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 16:17:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Share Guardian image deduplication and checkpoint restoration (#54278)

Browse Source 
## What changed

- Add `ComposedContext::retain_new_images` and use it in Guardian review sessions and the async scorer. Compare host-prepared image references against admitted history, ignoring image detail, and skip scanning history when no images need checking.
- Restore `ConversationState` directly from a checkpoint, retaining committed history for rollback or transferring it with `take_history` for ephemeral forks.

## Testing

Extend tests for prepared-image matching across differing detail levels, image-free continuations that avoid scanning history, and checkpoint restoration and history transfer.

GitOrigin-RevId: bafeb6de9babc721182be43d0eb1a653955a0ebb
This commit is contained in:
felixxia-oai
2026-10-15 12:42:15 +00:00
committed by copyberry
parent 3bae578af3
commit 5fff6378d6
8 changed files with 127 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/guardian/review_session.rs
+9 -32
View File
@@ -14,7 +14,6 @@ use context_policy::ReviewContextPolicy;
use std::borrow::Cow;
use std::collections::HashMap;
use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::Ordering;
@@ -581,26 +580,6 @@ async fn run_review_on_session(
.any(|(_, cost)| cost.image_count > 0)
{
let reviewer_history = review_session.session.clone_history().await;
let mut reviewer_image_urls = HashSet::new();
let mut reviewer_file_ids = HashSet::new();
for item in reviewer_history.raw_items() {
let ResponseItem::Message { content, .. } = item else {
continue;
};
for item in content {
let ContentItem::InputImage { image, .. } = item else {
continue;
};
match image {
ImageReference::Inline { image_url } => {
reviewer_image_urls.insert(image_url.as_str());
}
ImageReference::File { file_id } => {
reviewer_file_ids.insert(file_id.as_str());
}
}
}
}
let context_window = model_info.resolved_context_window().map(|supported| {
params
.spawn_config
@@ -621,7 +600,7 @@ async fn run_review_on_session(
} else {
ImagePreparationMode::DetailBased
};
prompt_items.context.retain_images(|image, detail| {
prompt_items.context.retain_new_images(reviewer_history.raw_items(), |image, detail| {
*detail = match normalize_output_image_detail(&model_info, *detail) {
_ if mode == ImagePreparationMode::UnifiedBudget => {
Some(ImageDetail::Original)
@@ -631,20 +610,18 @@ async fn run_review_on_session(
};
match image {
ImageReference::Inline { image_url } => {
let prepared_image_url = match resize_image(image_url, detail, mode)
{
Ok(Some(prepared)) => Cow::Owned(prepared.into_data_url()),
Ok(None) => Cow::Borrowed(image_url),
match resize_image(image_url, detail, mode) {
Ok(Some(prepared)) => Some(Cow::Owned(ImageReference::Inline {
image_url: prepared.into_data_url(),
})),
Ok(None) => Some(Cow::Borrowed(image)),
Err(error) => {
warn!(%error, "failed to prepare guardian review image");
return false;
None
}
};
!reviewer_image_urls.contains(prepared_image_url.as_str())
}
ImageReference::File { file_id } => {
!reviewer_file_ids.contains(file_id.as_str())
}
}
ImageReference::File { .. } => Some(Cow::Borrowed(image)),
}
});
let prompt_tokens = prompt_items
codex-rs/core/src/guardian/review_session_setup.rs
+16 -17
View File
@@ -99,27 +99,26 @@ impl PreparedGuardianContext {
kind: GuardianReviewSessionKind,
snapshot: Option<GuardianReviewForkSnapshot>,
) -> (crate::StartThreadOptions, GuardianReviewState) {
let snapshot =
snapshot.filter(|_| self.reviewer_selection == ReviewerSelection::ReuseIfAvailable);
let (mut conversation, mut history) = snapshot.map(ConversationState::fork).unzip();
if self.parent_compaction.is_some()
&& conversation
.as_ref()
.is_some_and(|state| state.cursor().is_none())
{
let mut conversation = snapshot
.filter(|_| self.reviewer_selection == ReviewerSelection::ReuseIfAvailable)
.map(ConversationState::from)
.unwrap_or_default();
if self.parent_compaction.is_some() && conversation.cursor().is_none() {
// A compacted fork has no valid transcript progress. Prefer the
// captured parent checkpoint to its lossy reviewer summary.
conversation = None;
history = None;
}
if !matches!(kind, GuardianReviewSessionKind::EphemeralForked)
&& let Some((conversation, history)) = conversation.as_mut().zip(history.as_ref())
{
// Replacement trunks need a rollback point if their first review is invalidated.
conversation.commit_snapshot(history.clone());
conversation = ConversationState::default();
}
let history = match kind {
GuardianReviewSessionKind::EphemeralForked => conversation.take_history(),
// Replacement trunks keep the checkpoint if their first review is invalidated.
GuardianReviewSessionKind::TrunkNew | GuardianReviewSessionKind::TrunkReused => {
conversation
.snapshot()
.map(|checkpoint| checkpoint.history().clone())
}
};
let state = GuardianReviewState {
conversation: conversation.unwrap_or_default(),
conversation,
fresh_parent_checkpoint: history.is_none() && self.parent_compaction.is_some(),
transcript_history_version: 0,
transcript_source: history
codex-rs/ext/guardian-reviewer/src/conversation.rs
+11 -11
View File
@@ -35,19 +35,19 @@ impl<H> Default for ConversationState<H> {
}
}
impl<H> ConversationState<H> {
/// Starts an independent conversation from exactly the committed history and cursor.
pub fn fork(checkpoint: ConversationCheckpoint<H>) -> (Self, H) {
(
Self {
cursor: checkpoint.cursor,
completed_review_count: checkpoint.completed_review_count,
..Self::default()
},
checkpoint.history,
)
impl<H> From<ConversationCheckpoint<H>> for ConversationState<H> {
/// Restores committed progress and keeps its history available for rollback or a fork.
/// An exclusive request can transfer the history with `take_history` after admission.
fn from(checkpoint: ConversationCheckpoint<H>) -> Self {
Self {
cursor: checkpoint.cursor,
completed_review_count: checkpoint.completed_review_count,
checkpoint: Some(checkpoint),
}
}
}
impl<H> ConversationState<H> {
pub fn cursor(&self) -> Option<TranscriptCursor> {
self.cursor
}
codex-rs/ext/guardian-reviewer/src/conversation_tests.rs
+14 -8
View File
@@ -19,31 +19,37 @@ fn fork_keeps_committed_history_and_cursor_while_next_review_is_uncommitted() {
// A failed review can retain new input without completing or replacing the
// last successful checkpoint used by concurrent forks.
state.admit_transcript(second);
let (fork, history) = ConversationState::fork(state.snapshot().unwrap().clone());
let fork = ConversationState::from(state.snapshot().unwrap().clone());
assert_eq!(
(
state.cursor(),
state.completed_review_count(),
history,
fork.snapshot().unwrap().history(),
fork.cursor(),
fork.completed_review_count(),
),
(Some(second), 1, vec!["first review"], Some(first), 1),
(Some(second), 1, &vec!["first review"], Some(first), 1),
);
state.complete_review(second);
let (fork, history) = ConversationState::fork(state.snapshot().unwrap().clone());
let fork = ConversationState::from(state.snapshot().unwrap().clone());
assert_eq!(
(history, fork.cursor(), fork.completed_review_count()),
(vec!["first review"], Some(first), 1),
(
fork.snapshot().unwrap().history(),
fork.cursor(),
fork.completed_review_count(),
),
(&vec!["first review"], Some(first), 1),
);
state.commit_snapshot(vec!["first review", "second review"]);
let (fork, history) = ConversationState::fork(state.snapshot().unwrap().clone());
let mut fork = ConversationState::from(state.snapshot().unwrap().clone());
let history = fork.take_history();
assert_eq!(
(history, fork.cursor(), fork.completed_review_count()),
(vec!["first review", "second review"], Some(second), 2),
(Some(vec!["first review", "second review"]), Some(second), 2),
);
assert!(fork.snapshot().is_none());
assert_eq!(
(
state.take_history(),
codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs
+1 -5
View File
@@ -130,11 +130,7 @@ impl Reservation {
// Delayed preparation must never import a later action's evidence.
.filter(|previous| previous.index < self.index && previous.key == key)
.and_then(|previous| previous.state.snapshot().cloned())
.map(|checkpoint| {
let (mut state, history) = ConversationState::fork(checkpoint);
state.commit_snapshot(history);
state
})
.map(ConversationState::from)
.unwrap_or_default();
let had_history = state.snapshot().is_some();
let prepared = request
codex-rs/ext/guardian-v2/src/async_scorer/request.rs
+4 -12
View File
@@ -9,8 +9,7 @@ use codex_guardian_context::estimate_input_tokens;
use codex_guardian_reviewer::ConversationCheckpoint;
use codex_guardian_reviewer::ConversationState;
use codex_history::ResponseItemEnvelope;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use std::borrow::Cow;
use super::sampler::LunaSampler;
use super::sampler::LunaSamplingRequest;
@@ -41,16 +40,9 @@ impl LunaSamplingRequest {
.unwrap_or_default();
let (mut context, _) = evidence.clone().compose(cursor, history).ok()?;
context.retain_new_instructions(history);
if let Some(snapshot) = conversation.snapshot() {
context.retain_images(|image, _| {
!snapshot.history().iter().any(|item| match &item.item {
ResponseItem::Message { content, .. } => content.iter().any(|item| {
matches!(item, ContentItem::InputImage { image: retained, .. } if retained == image)
}),
_ => false,
})
});
}
context.retain_new_images(history.iter().map(|item| &item.item), |image, _| {
Some(Cow::Borrowed(image))
});
let truncations = std::mem::take(&mut context.truncations);
let section_costs = context.section_costs().collect();
let history = conversation.snapshot().map(ConversationCheckpoint::history);
codex-rs/guardian-context/src/enforcement_tests.rs
+31 -1
View File
@@ -7,6 +7,7 @@ use crate::composition::user_message;
use codex_protocol::models::ImageReference;
use codex_protocol::models::ResponseItem;
use pretty_assertions::assert_eq;
use std::borrow::Cow;
fn text(value: &str) -> ContentItem {
ContentItem::InputText {
@@ -355,7 +356,36 @@ fn image_accounting_preserves_later_eviction_policy() {
text("evidence omitted")
])]
);
context.retain_images(|image, _| matches!(image, ImageReference::File { .. }));
let delivered = [user_message(vec![ContentItem::InputImage {
image: ImageReference::Inline {
image_url: "prepared-image".to_owned(),
},
detail: Some(codex_protocol::models::ImageDetail::High),
}])];
// A host-prepared reference matches delivered history even when source detail differs.
context.retain_new_images(&delivered, |image, _| match image {
ImageReference::Inline { .. } => Some(Cow::Owned(ImageReference::Inline {
image_url: "prepared-image".to_owned(),
})),
ImageReference::File { .. } => Some(Cow::Borrowed(image)),
});
let mut continued = context.clone();
let completed = context.clone().into_messages();
continued.retain_new_images(&completed, |image, _| Some(Cow::Borrowed(image)));
// Once all images were delivered, deduplication must not inspect retained history.
continued.retain_new_images(
completed
.iter()
.inspect(|_| panic!("image-free context must not scan retained history")),
|image, _| Some(Cow::Borrowed(image)),
);
assert_eq!(
continued.into_messages(),
vec![user_message(vec![
evidence.clone(),
text("user restriction")
])]
);
let available = context.estimated_tokens();
let retained = context
.clone()
codex-rs/guardian-context/src/images.rs
+41
View File
@@ -2,7 +2,9 @@
//! Keeps source order and evicts oldest images using the existing count/byte caps.
//! Byte caps cover inline URLs and file ID strings, not referenced file contents.
//! Consumers still choose source visibility and model-specific image detail.
//! Continuations compare prepared image references against their own admitted history.
use crate::ComposedContext;
use crate::ContextSection;
use crate::SectionContributor;
use crate::SectionError;
@@ -13,11 +15,50 @@ use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::ImageDetail;
use codex_protocol::models::ImageReference;
use codex_protocol::models::ResponseItem;
use std::borrow::Cow;
use std::cell::LazyCell;
use std::collections::VecDeque;
const MAX_TRANSCRIPT_IMAGES: usize = 4;
const MAX_TRANSCRIPT_IMAGE_BYTES: usize = 8 * 1024 * 1024;
impl ComposedContext {
/// Omits images already present in this conversation, ignoring image detail.
/// The host supplies the reference its model will receive (for example, after resizing)
/// and may normalize detail or reject an image. The original payload stays unchanged.
/// Only admitted history proves delivery; images kept here are not committed yet.
pub fn retain_new_images<'a>(
&mut self,
history: impl IntoIterator<Item = &'a ResponseItem>,
mut prepare: impl for<'image> FnMut(
&'image ImageReference,
&mut Option<ImageDetail>,
) -> Option<Cow<'image, ImageReference>>,
) {
// Collect references only when an image needs checking; avoid hashing screenshot bytes.
let retained_images = LazyCell::new(|| {
let mut images = Vec::new();
for item in history {
let ResponseItem::Message { content, .. } = item else {
continue;
};
for item in content {
if let ContentItem::InputImage { image, .. } = item {
images.push(image);
}
}
}
images
});
self.retain_images(|image, detail| {
let Some(prepared) = prepare(image, detail) else {
return false;
};
!retained_images.contains(&prepared.as_ref())
});
}
}
/// Request-local image policy and a frozen view of captured REPL screenshots.
#[derive(Clone, Copy)]
pub struct TranscriptImageInput<'a> {