From 0f7c8608bdce4874acba3f3a4467f1bf8275fecc Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 18 Sep 2026 12:03:48 +0000 Subject: [PATCH] Capture Guardian review checkpoints from live context (#46512) ## Why Committing a Guardian review snapshot currently flushes and reloads the transcript from storage, even though the completed model context is already available in memory. ## What changed - Build forkable checkpoints directly from live session state, preserving model history, retained Guardian context, world state, turn context, and token usage. - Preserve compaction window metadata and MCP resource origins in the checkpoint. - Recognize prior user turns inside compacted replacement history. ## Testing Add round-trip checkpoint tests for legacy and thread-owned Guardian context without storage, including preservation after live history is mutated. Extend the review compaction test to assert that checkpoints do not load history from storage. GitOrigin-RevId: cb2638ca93bba6ab4472de50c7e2913e1a2f25a8 --- codex-rs/core/src/context_manager/history.rs | 6 + codex-rs/core/src/guardian/review_session.rs | 38 ++---- codex-rs/core/src/guardian/tests.rs | 17 +-- .../core/src/session/guardian_checkpoint.rs | 46 +++++++ .../src/session/guardian_checkpoint_tests.rs | 112 ++++++++++++++++++ codex-rs/core/src/session/mod.rs | 1 + .../core/src/thread_rollout_truncation.rs | 5 + codex-rs/core/tests/suite/guardian_review.rs | 9 ++ 8 files changed, 198 insertions(+), 36 deletions(-) create mode 100644 codex-rs/core/src/session/guardian_checkpoint.rs create mode 100644 codex-rs/core/src/session/guardian_checkpoint_tests.rs diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index abb7aecc7b..194be3a8f9 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -381,6 +381,12 @@ impl ContextManager { self.world_state_baseline = Some(snapshot); } + pub(crate) fn world_state_checkpoint(&self) -> Option { + self.world_state_baseline + .clone() + .map(|snapshot| WorldStateItem::full(snapshot.into_object())) + } + pub(crate) fn set_token_usage_full(&mut self, context_window: i64) { match &mut self.token_info { Some(info) => info.fill_to_context_window(context_window), diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index d176191972..b186804af7 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -855,18 +855,6 @@ async fn ensure_guardian_node_repl_policy( Ok(()) } -async fn load_rollout_items_for_fork( - session: &Session, -) -> anyhow::Result>> { - session - .try_ensure_rollout_materialized(PersistContext::Standard) - .await?; - session.flush_rollout().await?; - let live_thread = session.live_thread_for_persistence("guardian review fork")?; - let history = live_thread.load_history(/*include_archived*/ true).await?; - Ok(Some(history.items)) -} - impl codex_guardian_reviewer::ReviewerRuntime for GuardianReviewSession { async fn submit_turn(&self, request: TurnInputRequest) -> anyhow::Result { Ok(self @@ -915,22 +903,16 @@ impl codex_guardian_reviewer::ReviewerSession for GuardianReviewSession { } async fn commit_snapshot(&self) { - match load_rollout_items_for_fork(&self.session).await { - Ok(Some(items)) if !items.is_empty() => { - let mut state = self.state.lock().await; - let last_admitted_node_repl_response_sequence = - state.last_admitted_node_repl_response_sequence; - state.conversation.commit_snapshot(GuardianReviewHistory { - initial_history: InitialHistory::Forked(items), - last_admitted_node_repl_response_sequence, - }); - } - Ok(Some(_)) => {} - Ok(None) => {} - Err(err) => { - warn!("failed to refresh guardian trunk rollout snapshot: {err}"); - } - } + // The pool holds the review lock until this checkpoint is published. Capture the + // completed model context directly; saving and reloading the transcript adds no state. + let items = self.session.guardian_fork_history().await; + let mut state = self.state.lock().await; + let last_admitted_node_repl_response_sequence = + state.last_admitted_node_repl_response_sequence; + state.conversation.commit_snapshot(GuardianReviewHistory { + initial_history: InitialHistory::Forked(items), + last_admitted_node_repl_response_sequence, + }); } } diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 45287045b9..960fdcf15f 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -295,13 +295,6 @@ async fn seed_guardian_parent_history(session: &Arc, turn: &Arc bool { - let RolloutItem::ResponseItem(response_item) = item else { - return false; - }; - response_item_contains_message_text(response_item, needle) -} - fn response_item_contains_message_text(item: &ResponseItem, needle: &str) -> bool { let ResponseItem::Message { content, .. } = item else { return false; @@ -2335,7 +2328,15 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: assert_eq!( committed_rollout_items .iter() - .filter(|item| rollout_item_contains_message_text( + .flat_map(|item| match item { + RolloutItem::ResponseItem(item) => std::slice::from_ref(item), + RolloutItem::Compacted(checkpoint) => checkpoint + .replacement_history + .as_deref() + .unwrap_or_default(), + _ => &[], + }) + .filter(|item| response_item_contains_message_text( item, "Use prior reviews as context, not binding precedent." )) diff --git a/codex-rs/core/src/session/guardian_checkpoint.rs b/codex-rs/core/src/session/guardian_checkpoint.rs new file mode 100644 index 0000000000..934aa7b8c4 --- /dev/null +++ b/codex-rs/core/src/session/guardian_checkpoint.rs @@ -0,0 +1,46 @@ +//! Captures a completed Guardian model context using the existing fork replay format. +//! This snapshot is in memory only; transcript durability remains part of turn completion. + +use super::session::Session; +use codex_history::CompactedItem; +use codex_history::RolloutItem; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::TokenCountEvent; + +impl Session { + pub(crate) async fn guardian_fork_history(&self) -> Vec { + let state = self.state.lock().await; + let history = &state.history; + let window_ids = state.auto_compact_window_ids(); + let mut items = vec![RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: Some(history.annotated_items().to_vec()), + guardian_history: history.guardian_history_checkpoint(), + retained_context: Some(history.retained_context().clone()), + mcp_resource_origins: self.services.mcp_runtime.resource_origin_checkpoint(), + window_number: Some(state.auto_compact_window_number()), + first_window_id: Some(window_ids.first_window_id.to_string()), + previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()), + window_id: Some(window_ids.window_id.to_string()), + compaction_response_id: None, + latest_token_usage_record: state.latest_token_usage_record.clone(), + })]; + if let Some(world_state) = history.world_state_checkpoint() { + items.push(RolloutItem::WorldState(world_state)); + } + if let Some(context) = history.reference_context_item() { + items.push(RolloutItem::TurnContext(context)); + } + items.push(RolloutItem::EventMsg(EventMsg::TokenCount( + TokenCountEvent { + info: history.token_info(), + rate_limits: None, + }, + ))); + items + } +} + +#[cfg(test)] +#[path = "guardian_checkpoint_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/session/guardian_checkpoint_tests.rs b/codex-rs/core/src/session/guardian_checkpoint_tests.rs new file mode 100644 index 0000000000..f6ba9dc340 --- /dev/null +++ b/codex-rs/core/src/session/guardian_checkpoint_tests.rs @@ -0,0 +1,112 @@ +use crate::context::GuardianContextMode; +use crate::context::world_state::WorldStateSnapshot; +use crate::context_manager::ContextManager; +use crate::session::tests::make_session_and_context; +use codex_history::InitialHistory; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TokenUsage; +use pretty_assertions::assert_eq; +use serde_json::json; +use test_case::test_case; + +#[test_case(GuardianContextMode::Legacy; "legacy")] +#[test_case(GuardianContextMode::ThreadOwned; "thread_owned")] +#[tokio::test] +async fn guardian_checkpoint_preserves_live_context_without_storage(mode: GuardianContextMode) { + let (mut session, turn) = make_session_and_context().await; + session.guardian_context_mode = mode; + // This session has no live store. A checkpoint must still capture the complete context. + assert!(session.live_thread().is_none()); + let instruction: ResponseItem = serde_json::from_value(json!({ + "type": "message", "id": "user-restriction", "role": "user", "content": [ + {"type": "input_text", "text": "Do not publish the private report."} + ], + "internal_chat_message_metadata_passthrough": {"content_item_kinds": ["unknown"]} + })) + .unwrap(); + let context = turn.to_turn_context_item(); + let baseline = json!({"test": "baseline"}); + let world_state = WorldStateSnapshot::from(baseline.as_object().unwrap()); + { + let mut state = session.state.lock().await; + state.history = ContextManager::with_guardian_context_mode(mode, &SessionSource::default()); + state + .history + .record_items([&instruction], turn.model_info().truncation_policy.into()); + // The model window no longer contains the old restriction. Legacy review must + // preserve it in its retained transcript when this checkpoint is forked. + if mode == GuardianContextMode::Legacy { + let compacted: ResponseItem = serde_json::from_value(json!({ + "type": "compaction", "id": "compaction-1", "encrypted_content": "summary", + })) + .unwrap(); + state.history.replace_compacted( + vec![compacted.into()], + /*reviewer_compaction_hash*/ None, + ); + } + let followup: ResponseItem = serde_json::from_value(json!({ + "type": "message", "id": "user-followup", "role": "user", "content": [ + {"type": "input_text", "text": "Review the next action."} + ], + "internal_chat_message_metadata_passthrough": {"content_item_kinds": ["unknown"]} + })) + .unwrap(); + state + .history + .record_items([&followup], turn.model_info().truncation_policy.into()); + state.history.set_reference_context_item(Some(context)); + state.history.set_world_state_baseline(world_state.clone()); + state.history.update_token_info( + &TokenUsage { + input_tokens: 123, + total_tokens: 123, + ..Default::default() + }, + Some(32_000), + ); + } + let expected = session.clone_history().await; + if mode == GuardianContextMode::Legacy { + assert!( + expected + .guardian_history_checkpoint() + .unwrap() + .0 + .contains(&instruction) + ); + } + let items = session.guardian_fork_history().await; + assert!( + crate::thread_rollout_truncation::initial_history_has_prior_user_turns( + &InitialHistory::Forked(items.clone()), + ) + ); + // Mutating the live reviewer must not change the previously completed checkpoint. + session + .replace_history(Vec::new(), /*reference_context_item*/ None) + .await; + session + .record_initial_history(InitialHistory::Forked(items)) + .await; + let restored = session.clone_history().await; + assert_eq!(restored.annotated_items(), expected.annotated_items()); + assert_eq!(restored.retained_context(), expected.retained_context()); + assert_eq!( + restored.guardian_history_checkpoint(), + expected.guardian_history_checkpoint() + ); + assert_eq!( + restored.world_state_checkpoint(), + expected.world_state_checkpoint() + ); + assert_eq!( + serde_json::to_value(restored.reference_context_item()).unwrap(), + serde_json::to_value(expected.reference_context_item()).unwrap() + ); + assert_eq!( + serde_json::to_value(restored.token_info()).unwrap(), + serde_json::to_value(expected.token_info()).unwrap() + ); +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 2875c5f9ad..d09a612b16 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -235,6 +235,7 @@ mod daemon_recovery; mod environment; mod extension_interruption; pub(crate) mod extension_metrics; +mod guardian_checkpoint; mod handlers; mod inject; mod reasoning_effort; diff --git a/codex-rs/core/src/thread_rollout_truncation.rs b/codex-rs/core/src/thread_rollout_truncation.rs index cdbde89414..16f67a47b8 100644 --- a/codex-rs/core/src/thread_rollout_truncation.rs +++ b/codex-rs/core/src/thread_rollout_truncation.rs @@ -23,6 +23,11 @@ pub(crate) fn initial_history_has_prior_user_turns(conversation_history: &Initia fn rollout_item_is_user_turn_boundary(item: &RolloutItem) -> bool { match item { RolloutItem::ResponseItem(item) => is_user_turn_boundary(item), + RolloutItem::Compacted(checkpoint) => checkpoint + .replacement_history + .iter() + .flatten() + .any(|item| is_user_turn_boundary(item)), RolloutItem::InterAgentCommunication(_) => true, _ => false, } diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index 5102e7d0cb..35445ea966 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/codex-rs/core/tests/suite/guardian_review.rs @@ -510,7 +510,10 @@ async fn guardian_review_compacts_with_summary_despite_parent_token_budget( let server = start_mock_server().await; let summary = "Guardian retained the user's standing authorization."; + let store = Arc::new(codex_thread_store::InMemoryThreadStore::default()); let mut builder = test_codex() + .with_thread_store(store.clone()) + .with_history_mode(codex_protocol::protocol::ThreadHistoryMode::Legacy) .with_model_info_override("gpt-5.5", |model| { model.auto_review_model_override = Some(model.slug.clone()); model.supports_experimental_context = true; @@ -592,6 +595,12 @@ async fn guardian_review_compacts_with_summary_despite_parent_token_budget( let user_authorization = "Read the internal evaluation samples I have authorized."; test.submit_text_turn(user_authorization).await?; + assert_eq!( + store.calls().await.load_history, + 0, + "review checkpoints use live context" + ); + let requests = responses.requests(); let guardian_requests = requests .iter()