From d09a5cd8cfbbdaeb956c1de5ac1a9d733ef475bd Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 28 Oct 2026 10:52:06 +0000 Subject: [PATCH] Move Guardian review state into the reviewer extension (#57697) ## What changed Move review state, fork history, checkpoint recovery bookkeeping, and Node REPL evidence admission into `codex-guardian-reviewer`. Expose helpers for session initialization and snapshot publication, and update core to delegate to them while preserving the existing review logic. GitOrigin-RevId: 56d6a0124301a3c70908614d4bc3bfe970ac1c85 --- codex-rs/Cargo.lock | 1 + codex-rs/core/src/guardian/input_budget.rs | 10 +- codex-rs/core/src/guardian/prompt.rs | 2 +- codex-rs/core/src/guardian/review_session.rs | 76 +------- .../core/src/guardian/review_session_setup.rs | 72 ++------ codex-rs/ext/guardian-reviewer/Cargo.toml | 1 + codex-rs/ext/guardian-reviewer/src/lib.rs | 10 ++ codex-rs/ext/guardian-reviewer/src/prompt.rs | 2 + .../ext/guardian-reviewer/src/review_state.rs | 168 ++++++++++++++++++ 9 files changed, 205 insertions(+), 137 deletions(-) create mode 100644 codex-rs/ext/guardian-reviewer/src/review_state.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 389c138446..b516ac6ec0 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3652,6 +3652,7 @@ dependencies = [ "codex-extension-api", "codex-feedback", "codex-guardian-context", + "codex-history", "codex-mcp", "codex-otel", "codex-prompts", diff --git a/codex-rs/core/src/guardian/input_budget.rs b/codex-rs/core/src/guardian/input_budget.rs index 28af6f9ad2..e801a0974f 100644 --- a/codex-rs/core/src/guardian/input_budget.rs +++ b/codex-rs/core/src/guardian/input_budget.rs @@ -6,8 +6,6 @@ //! Loss telemetry describes accepted inputs, excluding feasibility checks and failed retries. use std::collections::BTreeMap; -use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use codex_features::Feature; @@ -35,13 +33,7 @@ use crate::session::turn_context::TurnContext; #[derive(Clone)] pub(crate) struct PendingReviewContext(pub ComposedContext); -/// Requests one bounded restart from the captured parent checkpoint when -/// reviewer compaction invalidates the pending input. The flag belongs to this attempt. -pub(super) struct CheckpointRecovery { - pub(super) requested: Arc, - pub(super) history_version: u64, - pub(super) fresh_parent_checkpoint: bool, -} +pub(super) use codex_guardian_reviewer::CheckpointRecovery; /// Restart a reused reviewer before paying for a summary we would discard. /// A fresh checkpoint stays intact while final input selection makes the review fit. diff --git a/codex-rs/core/src/guardian/prompt.rs b/codex-rs/core/src/guardian/prompt.rs index 77ae4a197b..4c341ae1de 100644 --- a/codex-rs/core/src/guardian/prompt.rs +++ b/codex-rs/core/src/guardian/prompt.rs @@ -34,7 +34,7 @@ use super::GuardianReviewContext; // Bound both JSON and permission evidence without restricting manual approvals. const MAX_GUARDIAN_ENVIRONMENT_ID_BYTES: usize = 256; -pub(super) const GUARDIAN_TRANSCRIPT_START: &str = ">>> TRANSCRIPT START\n"; +pub(super) use codex_guardian_reviewer::GUARDIAN_TRANSCRIPT_START; /// Selects transcript evidence using the checkpoint already loaded by the reviewer. pub(crate) enum GuardianTranscriptHistory<'a> { diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 23bbea47fa..7f61145f8a 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -25,13 +25,18 @@ use codex_analytics::GuardianReviewAnalyticsResult; use codex_analytics::GuardianReviewSessionAnalyticsParams; use codex_analytics::GuardianReviewSessionKind; use codex_extension_api::Instructions; -use codex_guardian_reviewer::ConversationCheckpoint; +#[cfg(test)] use codex_guardian_reviewer::ConversationState; +use codex_guardian_reviewer::GuardianReviewForkSnapshot; +pub use codex_guardian_reviewer::GuardianReviewState; +use codex_guardian_reviewer::GuardianTranscriptSource; +use codex_guardian_reviewer::PendingNodeReplEvidenceAdmission; use codex_guardian_reviewer::ReviewModel; use codex_guardian_reviewer::ReviewSessionResult; use codex_guardian_reviewer::SessionDisposition; use codex_guardian_reviewer::token_usage_delta; use codex_history::InitialHistory; +#[cfg(test)] use codex_history::RolloutItem; use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::config_types::Personality; @@ -45,6 +50,7 @@ use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::Event; +#[cfg(test)] use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::SubAgentSource; @@ -130,44 +136,10 @@ pub struct GuardianReviewSession { state: Mutex, } -/// Opaque conversation progress retained while ThreadManager starts a reviewer. -pub struct GuardianReviewState { - conversation: ConversationState, - // Only the first attempt may use the checkpoint loaded without prior reviewer history. - fresh_parent_checkpoint: bool, - // The cursor describes evidence admitted before any later reviewer compaction. - transcript_history_version: u64, - transcript_source: Option, - last_admitted_node_repl_response_sequence: u64, - pending_node_repl_evidence_admission: Option, -} - -// These histories have independent version counters; a cursor must identify its source. -#[derive(Clone, Copy, PartialEq, Eq)] -enum GuardianTranscriptSource { - Retained, - LoadedParentCheckpoint, -} - -struct PendingNodeReplEvidenceAdmission { - turn_id: String, - response_sequence: u64, -} - fn had_prior_review_context(prompt_mode: &GuardianPromptMode) -> bool { matches!(prompt_mode, GuardianPromptMode::Delta { .. }) } -type GuardianReviewForkSnapshot = ConversationCheckpoint; - -/// Host-owned history and admitted evidence used to seed a private reviewer fork. -#[derive(Clone)] -pub struct GuardianReviewHistory { - initial_history: InitialHistory, - transcript_source: Option, - last_admitted_node_repl_response_sequence: u64, -} - /// Opaque compatibility key derived by the existing context builder. #[derive(Debug, Clone, PartialEq)] pub struct GuardianReviewSessionReuseKey { @@ -265,38 +237,6 @@ impl GuardianReviewSessionReuseKey { } } -impl GuardianReviewSession { - async fn admit_node_repl_evidence(&self, event: &Event) { - // Annotated review inputs are recorded as response items without UI turn - // items. Both input paths emit this event after history admission. - let EventMsg::RawResponseItem(recorded) = &event.msg else { - return; - }; - let ResponseItem::Message { role, content, .. } = &recorded.item else { - return; - }; - if role != "user" - || !content.iter().any(|item| { - matches!(item, ContentItem::InputText { text } - if text == GUARDIAN_TRANSCRIPT_START || text == ">>> TRANSCRIPT DELTA START\n") - }) - { - return; - } - - let mut state = self.state.lock().await; - let Some(pending) = state.pending_node_repl_evidence_admission.as_ref() else { - return; - }; - if event.id == pending.turn_id { - state.last_admitted_node_repl_response_sequence = state - .last_admitted_node_repl_response_sequence - .max(pending.response_sequence); - state.pending_node_repl_evidence_admission = None; - } - } -} - async fn run_review_on_session( review_session: &GuardianReviewSession, params: &GuardianReviewSessionParams, @@ -920,7 +860,7 @@ impl codex_guardian_reviewer::ReviewerRuntime for GuardianReviewSession { } async fn admit_context(&self, event: &Event) { - self.admit_node_repl_evidence(event).await; + codex_guardian_reviewer::admit_node_repl_evidence(&self.state, event).await; } async fn interrupt(&self) -> anyhow::Result<()> { diff --git a/codex-rs/core/src/guardian/review_session_setup.rs b/codex-rs/core/src/guardian/review_session_setup.rs index 163f6b60e4..15279a0848 100644 --- a/codex-rs/core/src/guardian/review_session_setup.rs +++ b/codex-rs/core/src/guardian/review_session_setup.rs @@ -12,15 +12,9 @@ use crate::guardian::input_budget::CheckpointRecovery; use codex_features::Feature; use codex_guardian_reviewer::ReviewerPool; use codex_guardian_reviewer::ReviewerRequest; +use codex_guardian_reviewer::ReviewerSelection; use codex_protocol::protocol::TurnEnvironmentSelection; -/// Controls whether selection may reuse a session or must start from the parent checkpoint. -#[derive(Clone, Copy, PartialEq, Eq)] -enum ReviewerSelection { - ReuseIfAvailable, - FreshParentCheckpoint, -} - #[derive(Clone)] pub struct PreparedGuardianContext { parent: Arc, @@ -102,41 +96,12 @@ impl PreparedGuardianContext { kind: GuardianReviewSessionKind, snapshot: Option, ) -> (crate::StartThreadOptions, GuardianReviewState) { - 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 = 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, - fresh_parent_checkpoint: history.is_none() && self.parent_compaction.is_some(), - transcript_history_version: 0, - transcript_source: history - .as_ref() - .and_then(|history| history.transcript_source), - last_admitted_node_repl_response_sequence: history.as_ref().map_or(0, |history| { - history.last_admitted_node_repl_response_sequence - }), - pending_node_repl_evidence_admission: None, - }; - let initial_history = history.map(|history| history.initial_history).or_else(|| { - self.parent_compaction - .clone() - .map(|item| InitialHistory::Forked(vec![RolloutItem::ResponseItem(item.into())])) - }); + let (state, initial_history) = GuardianReviewState::for_session( + kind, + snapshot, + self.reviewer_selection, + self.parent_compaction.as_ref(), + ); let mut config = self.config.clone(); config.model_provider.supports_websockets &= self .parent @@ -343,27 +308,16 @@ impl ReviewerRequest for PreparedReview { return result; } if let Some((items, history_version, mut state)) = checkpoint { - // Forks must not inherit a cursor for evidence the completed turn compacted away. - if state.transcript_history_version != history_version - && session + state.commit_snapshot( + items, + history_version, + session .session .services .thread_extension_data .get::() - .is_some() - { - state.conversation.reset_transcript(); - state.transcript_source = None; - state.last_admitted_node_repl_response_sequence = 0; - } - let last_admitted_node_repl_response_sequence = - state.last_admitted_node_repl_response_sequence; - let transcript_source = state.transcript_source; - state.conversation.commit_snapshot(GuardianReviewHistory { - initial_history: InitialHistory::Forked(items), - transcript_source, - last_admitted_node_repl_response_sequence, - }); + .as_deref(), + ); } result } diff --git a/codex-rs/ext/guardian-reviewer/Cargo.toml b/codex-rs/ext/guardian-reviewer/Cargo.toml index 04e7bc388e..b13427fdfd 100644 --- a/codex-rs/ext/guardian-reviewer/Cargo.toml +++ b/codex-rs/ext/guardian-reviewer/Cargo.toml @@ -18,6 +18,7 @@ codex-analytics = { workspace = true } codex-extension-api = { workspace = true } codex-feedback = { workspace = true } codex-guardian-context = { workspace = true } +codex-history = { workspace = true } codex-mcp = { workspace = true } codex-otel = { workspace = true } codex-prompts = { workspace = true } diff --git a/codex-rs/ext/guardian-reviewer/src/lib.rs b/codex-rs/ext/guardian-reviewer/src/lib.rs index 24f6da0adc..a521c4f04c 100644 --- a/codex-rs/ext/guardian-reviewer/src/lib.rs +++ b/codex-rs/ext/guardian-reviewer/src/lib.rs @@ -25,6 +25,15 @@ mod prompt; mod reporting; mod retry; mod review; +mod review_state; +pub use review_state::CheckpointRecovery; +pub use review_state::GuardianReviewForkSnapshot; +pub use review_state::GuardianReviewHistory; +pub use review_state::GuardianReviewState; +pub use review_state::GuardianTranscriptSource; +pub use review_state::PendingNodeReplEvidenceAdmission; +pub use review_state::ReviewerSelection; +pub use review_state::admit_node_repl_evidence; mod routing; mod session; mod settings; @@ -68,6 +77,7 @@ pub use pool::ReviewerRequest; pub use pool::ReviewerSession; pub use pool::ReviewerTasks; pub use pool::SessionDisposition; +pub use prompt::GUARDIAN_TRANSCRIPT_START; pub use prompt::GuardianPromptItems; pub use prompt::compose_guardian_prompt; diff --git a/codex-rs/ext/guardian-reviewer/src/prompt.rs b/codex-rs/ext/guardian-reviewer/src/prompt.rs index 788951ccb3..3cc0f99d76 100644 --- a/codex-rs/ext/guardian-reviewer/src/prompt.rs +++ b/codex-rs/ext/guardian-reviewer/src/prompt.rs @@ -9,6 +9,8 @@ use codex_guardian_context::TranscriptCursor as GuardianTranscriptCursor; use codex_guardian_context::TranscriptMode as GuardianPromptMode; use codex_guardian_context::TranscriptSelection; +pub const GUARDIAN_TRANSCRIPT_START: &str = ">>> TRANSCRIPT START\n"; + pub struct GuardianPromptItems { pub context: ComposedContext, pub transcript_cursor: GuardianTranscriptCursor, diff --git a/codex-rs/ext/guardian-reviewer/src/review_state.rs b/codex-rs/ext/guardian-reviewer/src/review_state.rs new file mode 100644 index 0000000000..a497be54ff --- /dev/null +++ b/codex-rs/ext/guardian-reviewer/src/review_state.rs @@ -0,0 +1,168 @@ +//! Owns reviewer transcript progress, fork history, and checkpoint recovery bookkeeping. +//! The host captures history and validates authorization before publishing a checkpoint. +//! Evidence advances only after the matching review turn records its transcript input. + +use std::sync::Arc; +use std::sync::atomic::AtomicBool; + +use codex_analytics::GuardianReviewSessionKind; +use codex_history::InitialHistory; +use codex_history::RolloutItem; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use tokio::sync::Mutex; + +use crate::ConversationCheckpoint; +use crate::ConversationState; +use crate::GUARDIAN_TRANSCRIPT_START; + +/// Opaque conversation progress retained while ThreadManager starts a reviewer. +pub struct GuardianReviewState { + pub conversation: ConversationState, + // Only the first attempt may use the checkpoint loaded without prior reviewer history. + pub fresh_parent_checkpoint: bool, + // The cursor describes evidence admitted before any later reviewer compaction. + pub transcript_history_version: u64, + pub transcript_source: Option, + pub last_admitted_node_repl_response_sequence: u64, + pub pending_node_repl_evidence_admission: Option, +} + +// These histories have independent version counters; a cursor must identify its source. +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum GuardianTranscriptSource { + Retained, + LoadedParentCheckpoint, +} + +pub struct PendingNodeReplEvidenceAdmission { + pub turn_id: String, + pub response_sequence: u64, +} + +pub type GuardianReviewForkSnapshot = ConversationCheckpoint; + +/// Host-owned history and admitted evidence used to seed a private reviewer fork. +#[derive(Clone)] +pub struct GuardianReviewHistory { + pub initial_history: InitialHistory, + transcript_source: Option, + last_admitted_node_repl_response_sequence: u64, +} + +/// Controls whether selection may reuse a session or must start from the parent checkpoint. +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum ReviewerSelection { + ReuseIfAvailable, + FreshParentCheckpoint, +} + +/// Requests one bounded restart from the captured parent checkpoint when +/// reviewer compaction invalidates the pending input. The flag belongs to this attempt. +pub struct CheckpointRecovery { + pub requested: Arc, + pub history_version: u64, + pub fresh_parent_checkpoint: bool, +} + +impl GuardianReviewState { + /// Restores fork progress or starts from the captured parent checkpoint. + pub fn for_session( + kind: GuardianReviewSessionKind, + snapshot: Option, + reviewer_selection: ReviewerSelection, + parent_compaction: Option<&ResponseItem>, + ) -> (Self, Option) { + let mut conversation = snapshot + .filter(|_| reviewer_selection == ReviewerSelection::ReuseIfAvailable) + .map(ConversationState::from) + .unwrap_or_default(); + if 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 = 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 = Self { + conversation, + fresh_parent_checkpoint: history.is_none() && parent_compaction.is_some(), + transcript_history_version: 0, + transcript_source: history + .as_ref() + .and_then(|history| history.transcript_source), + last_admitted_node_repl_response_sequence: history.as_ref().map_or(0, |history| { + history.last_admitted_node_repl_response_sequence + }), + pending_node_repl_evidence_admission: None, + }; + let initial_history = history.map(|history| history.initial_history).or_else(|| { + parent_compaction + .cloned() + .map(|item| InitialHistory::Forked(vec![RolloutItem::ResponseItem(item.into())])) + }); + (state, initial_history) + } + + /// Publishes captured history after the host's final authorization validation. + pub fn commit_snapshot( + &mut self, + items: Vec, + history_version: u64, + checkpoint_recovery: Option<&CheckpointRecovery>, + ) { + // Forks must not inherit a cursor for evidence the completed turn compacted away. + if self.transcript_history_version != history_version && checkpoint_recovery.is_some() { + self.conversation.reset_transcript(); + self.transcript_source = None; + self.last_admitted_node_repl_response_sequence = 0; + } + let last_admitted_node_repl_response_sequence = + self.last_admitted_node_repl_response_sequence; + let transcript_source = self.transcript_source; + self.conversation.commit_snapshot(GuardianReviewHistory { + initial_history: InitialHistory::Forked(items), + transcript_source, + last_admitted_node_repl_response_sequence, + }); + } +} + +pub async fn admit_node_repl_evidence(state: &Mutex, event: &Event) { + // Annotated review inputs are recorded as response items without UI turn + // items. Both input paths emit this event after history admission. + let EventMsg::RawResponseItem(recorded) = &event.msg else { + return; + }; + let ResponseItem::Message { role, content, .. } = &recorded.item else { + return; + }; + if role != "user" + || !content.iter().any(|item| { + matches!(item, ContentItem::InputText { text } + if text == GUARDIAN_TRANSCRIPT_START || text == ">>> TRANSCRIPT DELTA START\n") + }) + { + return; + } + + let mut state = state.lock().await; + let Some(pending) = state.pending_node_repl_evidence_admission.as_ref() else { + return; + }; + if event.id == pending.turn_id { + state.last_admitted_node_repl_response_sequence = state + .last_admitted_node_repl_response_sequence + .max(pending.response_sequence); + state.pending_node_repl_evidence_admission = None; + } +}