From 84f3343b2d9d2906a41c0d044d3868044ac503c6 Mon Sep 17 00:00:00 2001 From: Dylan Hurd Date: Fri, 12 Jun 2026 00:19:13 -0700 Subject: [PATCH] guardian: store review session manager in extension state --- codex-rs/core/src/codex_thread.rs | 8 ++--- codex-rs/core/src/guardian/mod.rs | 8 ++++- codex-rs/core/src/guardian/review.rs | 36 ++++++++++---------- codex-rs/core/src/guardian/review_session.rs | 18 ++++++++++ codex-rs/core/src/guardian/tests.rs | 6 ++-- codex-rs/core/src/session/handlers.rs | 4 ++- codex-rs/core/src/session/mod.rs | 1 - codex-rs/core/src/session/session.rs | 2 -- codex-rs/core/src/session/tests.rs | 23 +++++-------- 9 files changed, 60 insertions(+), 46 deletions(-) diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index c4cb1e8e12..e83e039364 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -480,11 +480,9 @@ impl CodexThread { } pub async fn guardian_trunk_rollout_path(&self) -> Option { - self.codex - .session - .guardian_review_session - .trunk_rollout_path() - .await + let manager = + crate::guardian::existing_guardian_review_session_manager(&self.codex.session)?; + manager.trunk_rollout_path().await } pub async fn load_history( diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index ebe8c67698..3f16faec9b 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -50,7 +50,7 @@ pub(crate) use review::review_approval_request_with_cancel; pub(crate) use review::routes_approval_to_guardian; pub(crate) use review::routes_approval_to_guardian_with_reviewer; pub(crate) use review::spawn_approval_request_review; -pub(crate) use review_session::GuardianReviewSessionManager; +pub(crate) use review_session::existing_guardian_review_session_manager; pub(crate) use review_session::prompt_cache_key_override_for_review_session; pub(crate) const GUARDIAN_REVIEW_TIMEOUT: Duration = Duration::from_secs(90); @@ -63,6 +63,12 @@ const GUARDIAN_MAX_MESSAGE_ENTRY_TOKENS: usize = 2_000; const GUARDIAN_MAX_TOOL_ENTRY_TOKENS: usize = 1_000; const GUARDIAN_RECENT_ENTRY_LIMIT: usize = 40; +pub(crate) fn guardian_review_session_manager( + session: &crate::session::session::Session, +) -> std::sync::Arc { + review_session::guardian_review_session_manager(session) +} + /// Structured output contract that the guardian reviewer must satisfy. #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub(crate) struct GuardianAssessment { diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index b1b4691571..b5e29fb55b 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -42,6 +42,7 @@ use super::approval_request::guardian_assessment_action; use super::approval_request::guardian_request_target_item_id; use super::approval_request::guardian_request_turn_id; use super::approval_request::guardian_reviewed_action; +use super::guardian_review_session_manager; use super::metrics::emit_guardian_review_metrics; use super::prompt::guardian_output_schema; use super::prompt::parse_guardian_assessment; @@ -769,24 +770,23 @@ async fn run_guardian_review_session_before_deadline( } }; - let (session_outcome, session_analytics_result) = Box::pin( - session - .guardian_review_session - .run_review(GuardianReviewSessionParams { - parent_session: Arc::clone(&session), - parent_turn: turn.clone(), - spawn_config: guardian_config, - request, - retry_reason, - schema, - model: guardian_model, - reasoning_effort: guardian_reasoning_effort, - reasoning_summary: turn.reasoning_summary, - personality: turn.personality, - external_cancel, - deadline, - }), - ) + let review_session_manager = guardian_review_session_manager(&session); + let (session_outcome, session_analytics_result) = Box::pin(review_session_manager.run_review( + GuardianReviewSessionParams { + parent_session: Arc::clone(&session), + parent_turn: turn.clone(), + spawn_config: guardian_config, + request, + retry_reason, + schema, + model: guardian_model, + reasoning_effort: guardian_reasoning_effort, + reasoning_summary: turn.reasoning_summary, + personality: turn.personality, + external_cancel, + deadline, + }, + )) .await; match session_outcome { diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 63b778838e..1b476fe970 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -89,6 +89,24 @@ pub(crate) struct GuardianReviewSessionManager { state: Arc>, } +pub(crate) fn guardian_review_session_manager( + session: &Session, +) -> Arc { + session + .services + .thread_extension_data + .get_or_init(GuardianReviewSessionManager::default) +} + +pub(crate) fn existing_guardian_review_session_manager( + session: &Session, +) -> Option> { + session + .services + .thread_extension_data + .get::() +} + #[derive(Default)] struct GuardianReviewSessionState { trunk: Option>, diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 68d2a72dc4..d278ad1682 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -1772,8 +1772,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: 1, "later follow-up guardian requests should not append the reminder again" ); - let committed_rollout_items = session - .guardian_review_session + let committed_rollout_items = guardian_review_session_manager(&session) .committed_fork_rollout_items_for_test() .await .expect("committed guardian fork snapshot"); @@ -1882,8 +1881,7 @@ async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow:: Some(codex_analytics::GuardianReviewSessionKind::TrunkNew) )); - session - .guardian_review_session + guardian_review_session_manager(&session) .send_trunk_event_raw_for_test(Event { id: "stale-turn".to_string(), msg: EventMsg::TurnComplete(TurnCompleteEvent { diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index bfb87a0ff3..927bff883a 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -597,7 +597,9 @@ async fn shutdown_session_runtime(sess: &Arc) { .load_full() .shutdown() .await; - sess.guardian_review_session.shutdown().await; + if let Some(manager) = crate::guardian::existing_guardian_review_session_manager(sess) { + manager.shutdown().await; + } } async fn emit_thread_stop_lifecycle(sess: &Session) { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 12bc8d7f40..7e8e492663 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -296,7 +296,6 @@ use crate::SkillsManager; use crate::agents_md::load_project_instructions; use crate::context::UserInstructions; use crate::exec_policy::ExecPolicyUpdateError; -use crate::guardian::GuardianReviewSessionManager; use crate::mcp::McpManager; use crate::network_policy_decision::execpolicy_network_rule_amendment; use crate::rollout::map_session_init_error; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 09d361c5f3..bec0a7e866 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -38,7 +38,6 @@ pub(crate) struct Session { pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, pub(crate) input_queue: InputQueue, - pub(crate) guardian_review_session: GuardianReviewSessionManager, pub(crate) services: SessionServices, pub(super) next_internal_sub_id: AtomicU64, } @@ -1048,7 +1047,6 @@ impl Session { conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: InputQueue::new(), - guardian_review_session: GuardianReviewSessionManager::default(), services, next_internal_sub_id: AtomicU64::new(0), }); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 61f52e772f..55bfb1e636 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5054,7 +5054,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: super::input_queue::InputQueue::new(), - guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, next_internal_sub_id: AtomicU64::new(0), }; @@ -6783,8 +6782,7 @@ async fn shutdown_and_wait_shuts_down_cached_guardian_subagent() { session: Arc::new(child_session), session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), }; - parent_session - .guardian_review_session + crate::guardian::guardian_review_session_manager(&parent_session) .cache_for_test(child_codex) .await; @@ -6816,16 +6814,15 @@ async fn cached_guardian_subagent_exposes_its_rollout_path() { session: Arc::new(child_session), session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), }; - parent_session - .guardian_review_session - .cache_for_test(child_codex) - .await; + let review_session_manager = crate::guardian::guardian_review_session_manager(&parent_session); + assert!(Arc::ptr_eq( + &review_session_manager, + &crate::guardian::guardian_review_session_manager(&parent_session) + )); + review_session_manager.cache_for_test(child_codex).await; assert_eq!( - parent_session - .guardian_review_session - .trunk_rollout_path() - .await, + review_session_manager.trunk_rollout_path().await, Some(child_rollout_path) ); } @@ -6872,8 +6869,7 @@ async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() { session: Arc::new(child_session), session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), }; - parent_session - .guardian_review_session + crate::guardian::guardian_review_session_manager(&parent_session) .register_ephemeral_for_test(child_codex) .await; @@ -7127,7 +7123,6 @@ where conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: super::input_queue::InputQueue::new(), - guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, next_internal_sub_id: AtomicU64::new(0), });