openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-09 20:50:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Move Guardian reviewer startup into the pool (#45521)

Browse Source 
## What changed

Replace `ReviewerSessionFactory` with a startup callback installed through
`ReviewerPool::new`. Review requests supply shared setup data and reuse context,
while the pool uses its callback to create both reusable and forked reviewers.

Update Guardian session setup, prewarming, and existing test fixtures to use the
new pool API.

GitOrigin-RevId: fa171503afcee8bdbdf6822573bea8ded50ce117
This commit is contained in:
jif
2026-09-14 21:34:33 +00:00
committed by copyberry
parent 78dfc1349e
commit e84a594636
6 changed files with 112 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/guardian/review_session.rs
+1
View File
@@ -882,6 +882,7 @@ impl codex_guardian_reviewer::ReviewerRuntime for GuardianReviewSession {
mod tests;
impl codex_guardian_reviewer::ReviewerSession for GuardianReviewSession {
type Setup = setup::PreparedGuardianContext;
type Context = GuardianReviewSessionReuseKey;
type Snapshot = GuardianReviewForkSnapshot;
codex-rs/core/src/guardian/review_session_setup.rs
+31 -20
View File
@@ -4,10 +4,9 @@
use super::*;
use codex_guardian_reviewer::ReviewerPool;
use codex_guardian_reviewer::ReviewerRequest;
use codex_guardian_reviewer::ReviewerSessionFactory;
use codex_guardian_reviewer::SessionDisposition;
pub(super) struct PreparedGuardianContext {
pub(crate) struct PreparedGuardianContext {
parent: Arc<Session>,
context: GuardianReviewContext,
config: Config,
@@ -128,14 +127,8 @@ impl PreparedGuardianContext {
}
}
impl ReviewerSessionFactory for PreparedGuardianContext {
type Session = GuardianReviewSession;
fn context(&self, previous: Option<&GuardianReviewSession>) -> GuardianReviewSessionReuseKey {
self.reuse_key(previous)
}
async fn spawn(
impl PreparedGuardianContext {
pub(crate) async fn spawn(
&self,
context: GuardianReviewSessionReuseKey,
kind: GuardianReviewSessionKind,
@@ -169,15 +162,18 @@ impl ReviewerSessionFactory for PreparedGuardianContext {
}
pub(super) struct PreparedReview {
factory: PreparedGuardianContext,
context: Arc<PreparedGuardianContext>,
params: GuardianReviewSessionParams,
}
impl ReviewerRequest for PreparedReview {
type Factory = PreparedGuardianContext;
type Session = GuardianReviewSession;
fn factory(&self) -> &PreparedGuardianContext {
&self.factory
fn setup(&self) -> Arc<PreparedGuardianContext> {
Arc::clone(&self.context)
}
fn context(&self, previous: Option<&GuardianReviewSession>) -> GuardianReviewSessionReuseKey {
self.context.reuse_key(previous)
}
fn deadline(&self) -> tokio::time::Instant {
self.params.deadline
@@ -228,7 +224,7 @@ pub(crate) async fn run_guardian_review_session(
pub(super) async fn prepare_review(
params: GuardianReviewSessionParams,
) -> anyhow::Result<PreparedReview> {
let factory = PreparedGuardianContext::prepare(
let context = PreparedGuardianContext::prepare(
Arc::clone(&params.parent_session),
params.parent_context.clone(),
params.spawn_config.clone(),
@@ -237,19 +233,35 @@ pub(super) async fn prepare_review(
params.compaction_model_hash.as_deref(),
)
.await?;
Ok(PreparedReview { factory, params })
Ok(PreparedReview {
context: Arc::new(context),
params,
})
}
pub(crate) fn prewarm_guardian_review_session(
parent: Arc<Session>,
turn: Arc<TurnContext>,
) -> BoxFuture<'static, anyhow::Result<()>> {
// Keep the Session -> Guardian -> Session startup future on the heap.
Box::pin(async move {
let context = prepare_prewarm(Arc::clone(&parent), turn).await?;
let key = context.reuse_key(/*previous*/ None);
parent
.guardian_review_session()
.prewarm(Arc::new(context), key)
.await
})
}
pub(super) fn prepare_prewarm(
parent: Arc<Session>,
turn: Arc<TurnContext>,
) -> BoxFuture<'static, anyhow::Result<PreparedGuardianContext>> {
Box::pin(async move {
let context = GuardianReviewContext::from(turn);
let config = guardian_review_session_config(&parent, &context).await?;
let history = parent.clone_history().await;
let factory = PreparedGuardianContext::prepare(
PreparedGuardianContext::prepare(
Arc::clone(&parent),
context,
config.spawn_config,
@@ -257,7 +269,6 @@ pub(crate) fn prewarm_guardian_review_session(
&config.node_repl_policy,
config.compaction_model_hash.as_deref(),
)
.await?;
parent.guardian_review_session().prewarm(&factory).await
.await
})
}
codex-rs/core/src/guardian/review_session_tests.rs
+21 -34
View File
@@ -5,7 +5,6 @@ use crate::agents_md_manager::AgentsMdManager;
use crate::context_manager::ContextManager;
use codex_guardian_reviewer::ReviewerRequest;
use codex_guardian_reviewer::ReviewerSession;
use codex_guardian_reviewer::ReviewerSessionFactory;
use codex_history::CodexHarnessMetadata;
use codex_history::ResponseItemEnvelope;
use codex_protocol::openai_models::AutoReviewMessages;
@@ -73,8 +72,7 @@ async fn run_review_preserves_evidence_during_parent_compaction() {
.node_repl_auto_review_required,
)
.with_node_repl_policy(&params.node_repl_policy);
let manager = GuardianReviewSessionManager::default();
prewarm_test_session(&manager, reviewer).await;
let manager = prewarm_test_session(&params, reviewer).await;
// Capture the review context, then compact the parent before the reviewer builds its prompt.
let prepared = setup::prepare_review(params).await.unwrap();
let checkpoint: ResponseItem = serde_json::from_value(serde_json::json!({
@@ -314,7 +312,7 @@ async fn spawned_guardian_reuse_key_matches_inherited_instructions() {
};
let prepared = setup::prepare_review(params).await.expect("prepare review");
let review = prepared
.factory()
.setup()
.spawn(
stale_key,
GuardianReviewSessionKind::TrunkNew,
@@ -901,8 +899,7 @@ async fn run_review_removes_trunk_when_event_stream_is_broken() {
)
.with_environments(params.parent_context.environments())
.with_node_repl_policy(&params.node_repl_policy);
let manager = Arc::new(GuardianReviewSessionManager::default());
prewarm_test_session(&manager, review_session).await;
let manager = Arc::new(prewarm_test_session(&params, review_session).await);
let manager_for_review = Arc::clone(&manager);
let review =
tokio::spawn(async move { run_guardian_review_session(manager_for_review, params).await });
@@ -1174,32 +1171,22 @@ async fn interrupt_and_drain_turn_ignores_prior_turn_completion() {
}
// Reuse the existing in-memory reviewer fixture through the production prewarm path.
async fn prewarm_test_session(pool: &GuardianReviewSessionManager, session: GuardianReviewSession) {
struct ReadySession {
context: GuardianReviewSessionReuseKey,
session: Mutex<Option<GuardianReviewSession>>,
}
impl codex_guardian_reviewer::ReviewerSessionFactory for ReadySession {
type Session = GuardianReviewSession;
fn context(
&self,
_previous: Option<&GuardianReviewSession>,
) -> GuardianReviewSessionReuseKey {
self.context.clone()
}
async fn spawn(
&self,
_context: GuardianReviewSessionReuseKey,
_kind: GuardianReviewSessionKind,
_snapshot: Option<GuardianReviewForkSnapshot>,
_cancellation: CancellationToken,
) -> anyhow::Result<GuardianReviewSession> {
Ok(self.session.lock().await.take().expect("one fixture spawn"))
}
}
let factory = ReadySession {
context: session.reuse_key.clone(),
session: Mutex::new(Some(session)),
};
pool.prewarm(&factory).await.unwrap();
async fn prewarm_test_session(
params: &GuardianReviewSessionParams,
session: GuardianReviewSession,
) -> GuardianReviewSessionManager {
let key = session.reuse_key.clone();
let session = Arc::new(Mutex::new(Some(session)));
let pool = GuardianReviewSessionManager::new(move |_, _, _, _, _| {
let session = Arc::clone(&session);
Box::pin(async move { Ok(session.lock().await.take().expect("one fixture spawn")) })
});
let context = setup::prepare_prewarm(
Arc::clone(&params.parent_session),
Arc::clone(params.parent_context.turn()),
)
.await
.unwrap();
pool.prewarm(Arc::new(context), key).await.unwrap();
pool
}
codex-rs/core/src/session/mod.rs
+5 -3
View File
@@ -2247,9 +2247,11 @@ impl Session {
/// Uses the extension-owned reviewer, or the same manager for standalone hosts.
pub(crate) fn guardian_review_session(&self) -> Arc<GuardianReviewSessionManager> {
self.services
.thread_extension_data
.get_or_init(GuardianReviewSessionManager::default)
self.services.thread_extension_data.get_or_init(|| {
GuardianReviewSessionManager::new(|context, key, kind, snapshot, cancel| {
Box::pin(async move { context.spawn(key, kind, snapshot, cancel).await })
})
})
}
pub(crate) async fn emit_turn_started(&self, turn_context: &TurnContext) {
codex-rs/ext/guardian-reviewer/src/lib.rs
-1
View File
@@ -44,7 +44,6 @@ pub(crate) use deadline::run_before_review_deadline_with_cancel;
pub use pool::ReviewerPool;
pub use pool::ReviewerRequest;
pub use pool::ReviewerSession;
pub use pool::ReviewerSessionFactory;
pub use pool::SessionDisposition;
pub use review::ReviewHost;
codex-rs/ext/guardian-reviewer/src/pool.rs
+54 -46
View File
@@ -1,5 +1,5 @@
//! Owns the reusable reviewer and temporary forks for one parent thread.
//! The host owns session execution and context construction. Selection stays serialized;
//! Guardian supplies agent startup; the host supplies captured context. Selection stays serialized;
//! concurrent reviews fork committed context and shutdown joins every tracked session.
//! Startup and fork futures stay boxed to bound the orchestration stack frames.
@@ -8,6 +8,7 @@ use std::sync::Arc;
use codex_analytics::GuardianReviewAnalyticsResult;
use codex_analytics::GuardianReviewSessionKind;
use codex_extension_api::ExtensionFuture;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio::time::Instant;
@@ -20,6 +21,7 @@ use crate::run_before_review_deadline_with_cancel;
/// A host-owned reviewer session. Context and snapshots remain opaque to the pool.
/// Shutdown must cancel the runtime and await its termination.
pub trait ReviewerSession: Send + Sync + 'static {
type Setup: Send + Sync + 'static;
type Context: Clone + PartialEq + Send + Sync;
type Snapshot: Send + Sync;
@@ -30,36 +32,21 @@ pub trait ReviewerSession: Send + Sync + 'static {
fn shutdown(&self) -> impl Future<Output = ()> + Send;
}
/// Builds a session from one captured parent context. The host must preserve that
/// context's authority and honor cancellation during spawning, including partial startup.
pub trait ReviewerSessionFactory: Send + Sync {
/// Executes one approval on a selected session. The host must drain the submitted
/// turn before returning Reusable, and must keep the issuing action and permissions bound.
pub trait ReviewerRequest: Send + Sync {
type Session: ReviewerSession;
fn setup(&self) -> Arc<<Self::Session as ReviewerSession>::Setup>;
fn context(
&self,
previous: Option<&Self::Session>,
) -> <Self::Session as ReviewerSession>::Context;
fn spawn(
&self,
context: <Self::Session as ReviewerSession>::Context,
kind: GuardianReviewSessionKind,
snapshot: Option<<Self::Session as ReviewerSession>::Snapshot>,
cancellation: CancellationToken,
) -> impl Future<Output = anyhow::Result<Self::Session>> + Send;
}
/// Executes one approval on a selected session. The host must drain the submitted
/// turn before returning Reusable, and must keep the issuing action and permissions bound.
pub trait ReviewerRequest: Send + Sync {
type Factory: ReviewerSessionFactory;
fn factory(&self) -> &Self::Factory;
fn deadline(&self) -> Instant;
fn cancellation(&self) -> Option<&CancellationToken>;
fn run(
&self,
session: &<Self::Factory as ReviewerSessionFactory>::Session,
session: &Self::Session,
kind: GuardianReviewSessionKind,
) -> impl Future<
Output = (
@@ -81,8 +68,19 @@ pub enum SessionDisposition {
pub struct ReviewerPool<S: ReviewerSession> {
state: Arc<Mutex<PoolState<S>>>,
cancellation: CancellationToken,
spawn: Box<SpawnReviewer<S>>,
}
type SpawnReviewer<S> = dyn Fn(
Arc<<S as ReviewerSession>::Setup>,
<S as ReviewerSession>::Context,
GuardianReviewSessionKind,
Option<<S as ReviewerSession>::Snapshot>,
CancellationToken,
) -> ExtensionFuture<'static, anyhow::Result<S>>
+ Send
+ Sync;
struct PoolState<S: ReviewerSession> {
trunk: Option<Arc<Trunk<S>>>,
ephemeral_reviews: Vec<Arc<S>>,
@@ -93,14 +91,28 @@ struct Trunk<S: ReviewerSession> {
review_lock: Semaphore,
}
impl<S: ReviewerSession> Default for ReviewerPool<S> {
fn default() -> Self {
impl<S: ReviewerSession> ReviewerPool<S> {
/// Installs Guardian's startup function once. It must finish or clean up partial startup
/// even when the caller drops its future, and preserve the supplied cancellation token.
pub fn new(
spawn: impl Fn(
Arc<S::Setup>,
S::Context,
GuardianReviewSessionKind,
Option<S::Snapshot>,
CancellationToken,
) -> ExtensionFuture<'static, anyhow::Result<S>>
+ Send
+ Sync
+ 'static,
) -> Self {
Self {
state: Arc::new(Mutex::new(PoolState {
trunk: None,
ephemeral_reviews: Vec::new(),
})),
cancellation: CancellationToken::new(),
spawn: Box::new(spawn),
}
}
}
@@ -117,20 +129,17 @@ impl<S: ReviewerSession> ReviewerPool<S> {
}
/// Prepares the first reviewer without replacing a review that won the startup race.
pub async fn prewarm(
&self,
factory: &impl ReviewerSessionFactory<Session = S>,
) -> anyhow::Result<()> {
pub async fn prewarm(&self, setup: Arc<S::Setup>, context: S::Context) -> anyhow::Result<()> {
let cancellation = self.cancellation.child_token();
let guard = cancellation.clone().drop_guard();
let session = factory
.spawn(
factory.context(/*previous*/ None),
GuardianReviewSessionKind::TrunkNew,
/*snapshot*/ None,
cancellation.clone(),
)
.await?;
let session = (self.spawn)(
setup,
context,
GuardianReviewSessionKind::TrunkNew,
/*snapshot*/ None,
cancellation.clone(),
)
.await?;
let mut state = self.state.lock().await;
if !cancellation.is_cancelled() && state.trunk.is_none() {
state.trunk = Some(Arc::new(Trunk {
@@ -181,8 +190,7 @@ impl<S: ReviewerSession> ReviewerPool<S> {
request: R,
) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult)
where
R: ReviewerRequest,
R::Factory: ReviewerSessionFactory<Session = S>,
R: ReviewerRequest<Session = S>,
{
let mut spawned_trunk = false;
let (trunk, context) = match run_before_review_deadline(
@@ -193,9 +201,8 @@ impl<S: ReviewerSession> ReviewerPool<S> {
.await
{
Ok(mut state) => {
let context = request
.factory()
.context(state.trunk.as_ref().map(|trunk| trunk.session.as_ref()));
let context =
request.context(state.trunk.as_ref().map(|trunk| trunk.session.as_ref()));
if let Some(trunk) = state.trunk.as_ref()
&& trunk.session.context() != &context
&& trunk.review_lock.try_acquire().is_ok()
@@ -209,12 +216,13 @@ impl<S: ReviewerSession> ReviewerPool<S> {
request.deadline(),
request.cancellation(),
&cancellation,
Box::pin(request.factory().spawn(
(self.spawn)(
request.setup(),
context.clone(),
GuardianReviewSessionKind::TrunkNew,
/*snapshot*/ None,
cancellation.clone(),
)),
),
)
.await
{
@@ -294,20 +302,20 @@ impl<S: ReviewerSession> ReviewerPool<S> {
snapshot: Option<S::Snapshot>,
) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult)
where
R: ReviewerRequest,
R::Factory: ReviewerSessionFactory<Session = S>,
R: ReviewerRequest<Session = S>,
{
let cancellation = self.cancellation.child_token();
let session = match run_before_review_deadline_with_cancel(
request.deadline(),
request.cancellation(),
&cancellation,
Box::pin(request.factory().spawn(
(self.spawn)(
request.setup(),
context,
GuardianReviewSessionKind::EphemeralForked,
snapshot,
cancellation.clone(),
)),
),
)
.await
{