openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-15 12:08:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Consolidate Guardian reviewer lifecycle ownership (#45672)

Browse Source 
## What changed

- Use cancellation guards to tie reusable reviewers and temporary forks to their lifetimes, including when a review future is dropped. Replace cancelled reusable reviewers before reuse.
- Share `ReviewerTasks` between the pool and `ThreadManager`, and wait for tracked cleanup during pool shutdown. Remove the separate session shutdown protocol.
- Move denial cleanup from core task handling into Guardian's turn start, stop, and abort hooks. Expose `install_reviewer` to register both thread and turn lifecycle contributors together.

## Testing

Add an integration test that exercises stale denial cleanup at turn start, completion, and interruption, including starting another turn after interruption.

GitOrigin-RevId: 85badb235309c1e1576547d1746f7ee9fc66d545
This commit is contained in:
jif
2026-09-15 10:27:12 +00:00
committed by copyberry
parent 4415f985dc
commit a113f3e063
15 changed files with 231 additions and 315 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
codex-rs/Cargo.lock generated
View File

@@ -3609,6 +3609,7 @@ dependencies = [
"http 1.4.0",
"pretty_assertions",
"serde_json",
"tempfile",
"thiserror 2.0.18",
"tokio",
"tracing",

9
codex-rs/core/src/guardian/review_session.rs
View File

@@ -874,15 +874,6 @@ impl codex_guardian_reviewer::ReviewerSession for GuardianReviewSession {
fn context(&self) -> &Self::Context {
&self.reuse_key
}
fn cancel(&self) {
self.cancel_token.cancel();
}
async fn shutdown(&self) {
self.cancel_token.cancel();
let _ = self.io.shutdown_and_wait().await;
}
async fn snapshot(&self) -> Option<GuardianReviewForkSnapshot> {
self.state.lock().await.conversation.snapshot().cloned()
}

14
codex-rs/core/src/guardian/review_session_tests.rs
View File

@@ -4,7 +4,6 @@ use super::*;
use crate::agents_md_manager::AgentsMdManager;
use crate::context_manager::ContextManager;
use codex_guardian_reviewer::ReviewerRequest;
use codex_guardian_reviewer::ReviewerSession;
use codex_history::CodexHarnessMetadata;
use codex_history::ResponseItemEnvelope;
use codex_protocol::openai_models::AutoReviewMessages;
@@ -299,7 +298,7 @@ async fn spawned_guardian_reuse_key_matches_inherited_instructions() {
assert_eq!(review.reuse_key, expected_key);
assert_eq!(review.session.inherited_instructions().await.thread, latest);
review.shutdown().await;
manager.shutdown().await;
}
#[tokio::test]
@@ -1152,10 +1151,13 @@ async fn prewarm_test_session(
) -> 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 pool = GuardianReviewSessionManager::new(
Arc::new(codex_guardian_reviewer::ReviewerTasks::default()),
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()),

5
codex-rs/core/src/guardian/test_host.rs
View File

@@ -35,12 +35,15 @@ pub(crate) fn install(session: &Session, config: &Config) {
/*attestation_provider*/ None,
/*external_time_provider*/ None,
));
let runtime = Arc::new(codex_guardian_reviewer::ReviewerTasks::default());
session
.services
.thread_extension_data
.insert(GuardianReviewSessionManager::new(
Arc::clone(&runtime),
move |context, key, kind, snapshot, cancel| {
let manager = Arc::clone(&manager);
let runtime = Arc::clone(&runtime);
Box::pin(async move {
let history_reset = context.history_reset.clone();
let (mut options, state) = context.thread_options(snapshot).await;
@@ -67,7 +70,7 @@ pub(crate) fn install(session: &Session, config: &Config) {
}
};
let spawned = manager
.start_thread_until(options, until, &tokio_util::task::TaskTracker::new())
.start_thread_until(options, until, &runtime.tasks)
.await?;
Ok(context
.bind_thread(&spawned.thread, key, state, session_cancel)

45
codex-rs/core/src/session/tests/guardian_tests.rs
View File

@@ -1330,48 +1330,3 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
);
drop(io);
}
#[test_case(TerminalEventKind::TurnComplete; "completion")]
#[test_case(TerminalEventKind::TurnAborted; "interruption")]
#[tokio::test]
async fn terminal_turn_clears_extension_owned_denials(terminal: TerminalEventKind) {
let (session, turn, events) = make_session_and_context_with_rx().await;
let finish = Arc::new(tokio::sync::Notify::new());
session
.spawn_task(
Arc::clone(&turn),
Vec::new(),
HeldStepTask {
kind: TaskKind::Regular,
finish: Arc::clone(&finish),
},
)
.await;
let denials =
codex_guardian_reviewer::ReviewDenials::for_thread(&session.services.thread_extension_data);
for _ in 0..2 {
assert_eq!(
denials.record_denial(&turn.sub_id, turn.model_info()).await,
None
);
}
match terminal {
TerminalEventKind::TurnComplete => finish.notify_one(),
TerminalEventKind::TurnAborted => {
session.abort_all_tasks(TurnAbortReason::Interrupted).await
}
}
recv_terminal_event(&events, terminal).await;
// Delivery precedes accounting cleanup. Wait for the runtime to finish the turn.
timeout(Duration::from_secs(/*secs*/ 5), async {
while session.active_turn.lock().await.is_some() {
tokio::task::yield_now().await;
}
})
.await
.expect("turn becomes idle");
assert_eq!(
denials.record_denial(&turn.sub_id, turn.model_info()).await,
None
);
}

16
codex-rs/core/src/tasks/mod.rs
View File

@@ -306,12 +306,6 @@ impl Session {
let cancellation_token = CancellationToken::new();
let done = Arc::new(Notify::new());
codex_guardian_reviewer::ReviewDenials::clear_turn(
&self.services.thread_extension_data,
&turn_context.sub_id,
)
.await;
let (pending_items, _) = self.input_queue.drain_mailbox_input_items().await;
let turn_state = {
let mut active = self.active_turn.lock().await;
@@ -812,11 +806,6 @@ impl Session {
})
};
self.send_event(turn_context.as_ref(), event).await;
codex_guardian_reviewer::ReviewDenials::clear_turn(
&self.services.thread_extension_data,
&turn_context.sub_id,
)
.await;
let cleared_active_turn = {
let mut active = self.active_turn.lock().await;
@@ -968,11 +957,6 @@ impl Session {
duration_ms,
});
self.send_event(task.turn_context.as_ref(), event).await;
codex_guardian_reviewer::ReviewDenials::clear_turn(
&self.services.thread_extension_data,
&task.turn_context.sub_id,
)
.await;
// Regular items were flushed before this terminal event was appended; buffering
// thread writers may not flush it without another explicit barrier.
if let Err(err) = self.flush_rollout().await {

5
codex-rs/core/tests/common/test_codex.rs
View File

@@ -732,10 +732,7 @@ impl TestCodexBuilder {
.or_else(|| codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").ok());
let thread_manager = Arc::new_cyclic(|manager| {
let mut extensions = self.extensions.to_builder();
let guardian = Arc::new(codex_guardian_v2::GuardianReviewerExtension::new(
manager.clone(),
));
extensions.thread_lifecycle_contributor(guardian);
codex_guardian_v2::install_reviewer(&mut extensions, manager.clone());
let thread_manager = ThreadManager::new(
&config,
auth_manager.clone(),

17
codex-rs/ext/guardian-reviewer/src/deadline.rs
View File

@@ -21,20 +21,3 @@ pub async fn run_before_review_deadline<T>(
} => Err(GuardianReviewSessionOutcome::Aborted),
}
}
pub(crate) async fn run_before_review_deadline_with_cancel<T>(
deadline: tokio::time::Instant,
external_cancel: Option<&CancellationToken>,
cancel_token: &CancellationToken,
future: impl Future<Output = T>,
) -> Result<T, GuardianReviewSessionOutcome> {
let result = run_before_review_deadline(deadline, external_cancel, future).await;
if result.is_err() {
cancel_token.cancel();
}
result
}
#[cfg(test)]
#[path = "deadline_tests.rs"]
mod tests;

69
codex-rs/ext/guardian-reviewer/src/deadline_tests.rs
View File

@@ -1,69 +0,0 @@
use std::time::Duration;
use pretty_assertions::assert_eq;
use tokio_util::sync::CancellationToken;
use super::run_before_review_deadline_with_cancel;
use crate::GuardianReviewSessionOutcome;
#[tokio::test(flavor = "current_thread")]
async fn run_before_review_deadline_with_cancel_cancels_token_on_timeout() {
let cancel_token = CancellationToken::new();
let outcome = run_before_review_deadline_with_cancel(
tokio::time::Instant::now() + Duration::from_millis(10),
/*external_cancel*/ None,
&cancel_token,
async {
tokio::time::sleep(Duration::from_millis(50)).await;
},
)
.await;
assert!(matches!(
outcome,
Err(GuardianReviewSessionOutcome::TimedOut)
));
assert!(cancel_token.is_cancelled());
}
#[tokio::test(flavor = "current_thread")]
async fn run_before_review_deadline_with_cancel_cancels_token_on_abort() {
let external_cancel = CancellationToken::new();
let external_canceller = external_cancel.clone();
let cancel_token = CancellationToken::new();
drop(tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await;
external_canceller.cancel();
}));
let outcome = run_before_review_deadline_with_cancel(
tokio::time::Instant::now() + Duration::from_secs(1),
Some(&external_cancel),
&cancel_token,
std::future::pending::<()>(),
)
.await;
assert!(matches!(
outcome,
Err(GuardianReviewSessionOutcome::Aborted)
));
assert!(cancel_token.is_cancelled());
}
#[tokio::test(flavor = "current_thread")]
async fn run_before_review_deadline_with_cancel_preserves_token_on_success() {
let cancel_token = CancellationToken::new();
let outcome = run_before_review_deadline_with_cancel(
tokio::time::Instant::now() + Duration::from_secs(1),
/*external_cancel*/ None,
&cancel_token,
async { 42usize },
)
.await;
assert_eq!(outcome.unwrap(), 42);
assert!(!cancel_token.is_cancelled());
}

1
codex-rs/ext/guardian-reviewer/src/lib.rs
View File

@@ -40,7 +40,6 @@ pub const MAX_REVIEW_ATTEMPTS: i64 = 3;
pub const REVIEW_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90);
pub use deadline::run_before_review_deadline;
pub(crate) use deadline::run_before_review_deadline_with_cancel;
pub use pool::ReviewerPool;
pub use pool::ReviewerRequest;
pub use pool::ReviewerSession;

174
codex-rs/ext/guardian-reviewer/src/pool.rs
View File

@@ -1,6 +1,7 @@
//! Owns the reusable reviewer and temporary forks for one parent thread.
//! Guardian supplies agent startup; the host supplies captured context. Selection stays serialized;
//! concurrent reviews fork committed context and shutdown joins every tracked session.
//! concurrent reviews fork committed context. Lifetime guards cancel agents; ThreadManager
//! performs cleanup and tracks its completion. The pool never runs a second shutdown protocol.
//! Startup and fork futures stay boxed to bound the orchestration stack frames.
use std::future::Future;
@@ -16,7 +17,6 @@ use tokio_util::sync::CancellationToken;
use crate::GuardianReviewSessionOutcome;
use crate::run_before_review_deadline;
use crate::run_before_review_deadline_with_cancel;
/// Background work owned by Guardian for one parent runtime.
/// Stop cancels the work and joins this tracker before the parent closes its history.
@@ -26,18 +26,16 @@ pub struct ReviewerTasks {
pub cancellation: CancellationToken,
}
/// A host-owned reviewer session. Context and snapshots remain opaque to the pool.
/// Shutdown must cancel the runtime and await its termination.
/// Context bookkeeping for a reviewer. Lifetime and cleanup belong to ThreadManager.
/// Context and snapshots remain opaque while the host context builder is being A/B tested.
pub trait ReviewerSession: Send + Sync + 'static {
type Setup: Send + Sync + 'static;
type Context: Clone + PartialEq + Send + Sync;
type Snapshot: Send + Sync;
fn context(&self) -> &Self::Context;
fn cancel(&self);
fn snapshot(&self) -> impl Future<Output = Option<Self::Snapshot>> + Send;
fn commit_snapshot(&self) -> impl Future<Output = ()> + Send;
fn shutdown(&self) -> impl Future<Output = ()> + Send;
}
/// Executes one approval on a selected session. The host must drain the submitted
@@ -74,8 +72,8 @@ pub enum SessionDisposition {
/// Per-parent reviewer state. The same pool serves prewarm, review, invalidation and shutdown.
pub struct ReviewerPool<S: ReviewerSession> {
state: Arc<Mutex<PoolState<S>>>,
cancellation: CancellationToken,
trunk: Mutex<Option<Arc<Trunk<S>>>>,
runtime: Arc<ReviewerTasks>,
spawn: Box<SpawnReviewer<S>>,
}
@@ -89,20 +87,23 @@ type SpawnReviewer<S> = dyn Fn(
+ Send
+ Sync;
struct PoolState<S: ReviewerSession> {
trunk: Option<Arc<Trunk<S>>>,
ephemeral_reviews: Vec<Arc<S>>,
}
struct Trunk<S: ReviewerSession> {
session: Arc<S>,
review_lock: Semaphore,
cancellation: CancellationToken,
}
impl<S: ReviewerSession> Drop for Trunk<S> {
fn drop(&mut self) {
self.cancellation.cancel();
}
}
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(
runtime: Arc<ReviewerTasks>,
spawn: impl Fn(
Arc<S::Setup>,
S::Context,
@@ -115,11 +116,8 @@ impl<S: ReviewerSession> ReviewerPool<S> {
+ 'static,
) -> Self {
Self {
state: Arc::new(Mutex::new(PoolState {
trunk: None,
ephemeral_reviews: Vec::new(),
})),
cancellation: CancellationToken::new(),
trunk: Mutex::new(None),
runtime,
spawn: Box::new(spawn),
}
}
@@ -128,17 +126,16 @@ impl<S: ReviewerSession> ReviewerPool<S> {
impl<S: ReviewerSession> ReviewerPool<S> {
/// Returns the current reviewer handle for host inspection and feedback collection.
pub async fn trunk(&self) -> Option<Arc<S>> {
self.state
self.trunk
.lock()
.await
.trunk
.as_ref()
.map(|trunk| Arc::clone(&trunk.session))
}
/// Prepares the first reviewer without replacing a review that won the startup race.
pub async fn prewarm(&self, setup: Arc<S::Setup>, context: S::Context) -> anyhow::Result<()> {
let cancellation = self.cancellation.child_token();
let cancellation = self.runtime.cancellation.child_token();
let guard = cancellation.clone().drop_guard();
let session = (self.spawn)(
setup,
@@ -148,44 +145,23 @@ impl<S: ReviewerSession> ReviewerPool<S> {
cancellation.clone(),
)
.await?;
let mut state = self.state.lock().await;
if !cancellation.is_cancelled() && state.trunk.is_none() {
state.trunk = Some(Arc::new(Trunk {
let mut trunk = self.trunk.lock().await;
if !cancellation.is_cancelled() && trunk.is_none() {
*trunk = Some(Arc::new(Trunk {
session: Arc::new(session),
review_lock: Semaphore::new(/*permits*/ 1),
cancellation: guard.disarm(),
}));
drop(guard.disarm());
}
Ok(())
}
/// Permanently stops this parent's reviewer pool and waits for tracked runtimes.
pub async fn shutdown(&self) {
self.cancellation.cancel();
self.invalidate().await;
}
/// Drops reusable context after parent history rollback or another host invalidation.
pub async fn invalidate(&self) {
let (trunk, ephemeral) = {
let mut state = self.state.lock().await;
(
state.trunk.take(),
std::mem::take(&mut state.ephemeral_reviews),
)
};
for session in trunk
.into_iter()
.map(|trunk| Arc::clone(&trunk.session))
.chain(ephemeral)
{
if self.cancellation.is_cancelled() {
session.shutdown().await;
} else {
session.cancel();
shutdown_in_background(session);
}
}
self.runtime.cancellation.cancel();
self.trunk.lock().await.take();
self.runtime.tasks.close();
self.runtime.tasks.wait().await;
}
/// Selects one reviewer; busy or incompatible trunks use an isolated temporary session.
@@ -204,26 +180,24 @@ impl<S: ReviewerSession> ReviewerPool<S> {
let (trunk, context) = match run_before_review_deadline(
request.deadline(),
request.cancellation(),
self.state.lock(),
self.trunk.lock(),
)
.await
{
Ok(mut state) => {
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
let context = request.context(state.as_ref().map(|trunk| trunk.session.as_ref()));
if let Some(trunk) = state.as_ref()
&& (trunk.cancellation.is_cancelled() || trunk.session.context() != &context)
&& trunk.review_lock.try_acquire().is_ok()
&& let Some(stale) = state.trunk.take()
{
shutdown_in_background(Arc::clone(&stale.session));
state.take();
}
if state.trunk.is_none() {
let cancellation = self.cancellation.child_token();
let session = match run_before_review_deadline_with_cancel(
if state.is_none() {
let cancellation = self.runtime.cancellation.child_token();
let lifetime = cancellation.clone().drop_guard();
let session = match run_before_review_deadline(
request.deadline(),
request.cancellation(),
&cancellation,
(self.spawn)(
request.setup(),
context.clone(),
@@ -245,13 +219,14 @@ impl<S: ReviewerSession> ReviewerPool<S> {
return (outcome, GuardianReviewAnalyticsResult::without_session());
}
};
state.trunk = Some(Arc::new(Trunk {
*state = Some(Arc::new(Trunk {
session,
review_lock: Semaphore::new(/*permits*/ 1),
cancellation: lifetime.disarm(),
}));
spawned_trunk = true;
}
(state.trunk.as_ref().cloned(), context)
(state.as_ref().cloned(), context)
}
Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()),
};
@@ -282,22 +257,27 @@ impl<S: ReviewerSession> ReviewerPool<S> {
} else {
GuardianReviewSessionKind::TrunkReused
};
// Dropping a review before it drains its turn must not leave a reusable agent.
let review_lifetime = trunk.cancellation.clone().drop_guard();
let (outcome, disposition, analytics) = request.run(&trunk.session, kind).await;
if disposition == SessionDisposition::Reusable
&& matches!(outcome, GuardianReviewSessionOutcome::Completed(_))
{
trunk.session.commit_snapshot().await;
}
if disposition == SessionDisposition::Reusable {
review_lifetime.disarm();
} else {
drop(review_lifetime);
}
drop(guard);
if disposition == SessionDisposition::Discard {
let mut state = self.state.lock().await;
let mut state = self.trunk.lock().await;
if state
.trunk
.as_ref()
.is_some_and(|current| Arc::ptr_eq(current, &trunk))
&& let Some(removed) = state.trunk.take()
{
shutdown_in_background(Arc::clone(&removed.session));
state.take();
}
}
(outcome, analytics)
@@ -312,11 +292,11 @@ impl<S: ReviewerSession> ReviewerPool<S> {
where
R: ReviewerRequest<Session = S>,
{
let cancellation = self.cancellation.child_token();
let session = match run_before_review_deadline_with_cancel(
let cancellation = self.runtime.cancellation.child_token();
let _lifetime = cancellation.clone().drop_guard();
let session = match run_before_review_deadline(
request.deadline(),
request.cancellation(),
&cancellation,
(self.spawn)(
request.setup(),
context,
@@ -336,63 +316,9 @@ impl<S: ReviewerSession> ReviewerPool<S> {
}
Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()),
};
self.state
.lock()
.await
.ephemeral_reviews
.push(Arc::clone(&session));
let mut cleanup = EphemeralCleanup {
state: Arc::clone(&self.state),
session: Some(Arc::clone(&session)),
};
let (outcome, _, analytics) = request
.run(&session, GuardianReviewSessionKind::EphemeralForked)
.await;
let removed = {
let mut state = self.state.lock().await;
state
.ephemeral_reviews
.iter()
.position(|active| Arc::ptr_eq(active, &session))
.map(|index| state.ephemeral_reviews.swap_remove(index))
};
if let Some(removed) = removed {
cleanup.session = None;
shutdown_in_background(removed);
}
(outcome, analytics)
}
}
fn shutdown_in_background<S: ReviewerSession>(session: Arc<S>) {
drop(tokio::spawn(async move {
session.shutdown().await;
}));
}
struct EphemeralCleanup<S: ReviewerSession> {
state: Arc<Mutex<PoolState<S>>>,
session: Option<Arc<S>>,
}
impl<S: ReviewerSession> Drop for EphemeralCleanup<S> {
fn drop(&mut self) {
let Some(session) = self.session.take() else {
return;
};
let state = Arc::clone(&self.state);
drop(tokio::spawn(async move {
let removed = {
let mut state = state.lock().await;
state
.ephemeral_reviews
.iter()
.position(|active| Arc::ptr_eq(active, &session))
.map(|index| state.ephemeral_reviews.swap_remove(index))
};
if let Some(removed) = removed {
removed.shutdown().await;
}
}));
}
}

1
codex-rs/ext/guardian-v2/Cargo.toml
View File

@@ -38,4 +38,5 @@ uuid = { workspace = true }
codex-model-provider-info = { workspace = true }
core_test_support = { workspace = true }
pretty_assertions = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

4
codex-rs/ext/guardian-v2/src/lib.rs
View File

@@ -9,7 +9,7 @@ use codex_login::AuthManager;
mod async_scorer;
mod sync_reviewer;
pub use sync_reviewer::GuardianExtension as GuardianReviewerExtension;
pub use sync_reviewer::install as install_reviewer;
/// Installs the guardian contributors into the extension registry.
pub fn install(
@@ -18,5 +18,5 @@ pub fn install(
thread_manager: Weak<ThreadManager>,
) {
async_scorer::install(registry, auth_manager, thread_manager.clone());
sync_reviewer::install(registry, thread_manager);
install_reviewer(registry, thread_manager);
}

125
codex-rs/ext/guardian-v2/src/sync_reviewer/lifecycle_tests.rs Normal file
View File

@@ -0,0 +1,125 @@
//! Exercises Guardian registration and denial cleanup through real turns.
use super::*;
use codex_core::TurnInputRequest;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::user_input::UserInput;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::sse;
use core_test_support::responses::start_mock_server;
use core_test_support::streaming_sse::StreamingSseChunk;
use core_test_support::streaming_sse::start_streaming_sse_server;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use core_test_support::wait_for_event_match;
use pretty_assertions::assert_eq;
use std::time::Duration;
struct StaleDenials(ModelInfo);
impl TurnLifecycleContributor for StaleDenials {
fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
let denials = ReviewDenials::for_thread(input.thread_store);
for _ in 0..2 {
assert_eq!(denials.record_denial(input.turn_id, &self.0).await, None);
}
})
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn terminal_turn_clears_extension_owned_denials() -> anyhow::Result<()> {
let home = tempfile::tempdir()?;
let config = core_test_support::load_default_config_for_test(&home).await;
let model = codex_core::test_support::construct_model_info_offline("gpt-5", &config);
let mut extensions = ExtensionRegistryBuilder::<Config>::new();
// Seed stale state before the production Guardian contributor runs its start hook.
extensions.turn_lifecycle_contributor(Arc::new(StaleDenials(model.clone())));
let mut releases = Vec::new();
let streams = (0..3)
.map(|_| {
let (release, gate) = tokio::sync::oneshot::channel();
releases.push(release);
vec![
StreamingSseChunk {
gate: None,
body: sse(vec![ev_response_created("response")]),
},
StreamingSseChunk {
gate: Some(gate),
body: sse(vec![ev_completed("response")]),
},
]
})
.collect();
let (streaming, _) = start_streaming_sse_server(streams).await;
let server = start_mock_server().await;
let base_url = format!("{}/v1", streaming.uri());
let test = test_codex()
.with_extensions(Arc::new(extensions.build()))
.with_config(move |config| {
config.model_provider.base_url = Some(base_url);
config.model_provider.supports_websockets = false;
})
.build_with_auto_env(&server)
.await?;
let denials = ReviewDenials::for_thread(test.codex.thread_extension_data());
let mut previous_turn = None;
// The third turn also checks that an interrupted turn leaves the next turn usable.
for (index, release) in releases.into_iter().enumerate() {
test.codex
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "Continue.".into(),
text_elements: Vec::new(),
}]))
.await?;
let turn_id = wait_for_event_match(&test.codex, |event| match event {
EventMsg::TurnStarted(event) => Some(event.turn_id.clone()),
_ => None,
})
.await;
// An interruption event can precede its cleanup. The next turn starts after it.
if let Some(previous) = previous_turn.replace(turn_id.clone()) {
assert_eq!(denials.record_denial(&previous, &model).await, None);
}
tokio::time::timeout(
Duration::from_secs(/*secs*/ 10),
streaming.wait_for_request_count(index + 1),
)
.await?;
// Without start cleanup, these denials reach the three-denial limit.
for _ in 0..2 {
assert_eq!(denials.record_denial(&turn_id, &model).await, None);
}
if index == 1 {
test.codex.submit(Op::Interrupt).await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnAborted(event)
if event.reason == TurnAbortReason::Interrupted
&& event.turn_id.as_deref() == Some(turn_id.as_str()))
})
.await;
drop(release);
} else {
release.send(()).unwrap();
wait_for_event(
&test.codex,
|event| matches!(event, EventMsg::TurnComplete(event) if event.turn_id == turn_id),
)
.await;
}
}
test.codex.shutdown_and_wait().await?;
assert_eq!(
denials.record_denial(&previous_turn.unwrap(), &model).await,
None
);
streaming.shutdown().await;
Ok(())
}

60
codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs
View File

@@ -16,6 +16,11 @@ use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::ThreadReadyInput;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ThreadStopInput;
use codex_extension_api::TurnAbortInput;
use codex_extension_api::TurnLifecycleContributor;
use codex_extension_api::TurnStartInput;
use codex_extension_api::TurnStopInput;
use codex_guardian_reviewer::ReviewDenials;
use codex_guardian_reviewer::ReviewerPool;
use codex_guardian_reviewer::ReviewerTasks;
use codex_protocol::ThreadId;
@@ -26,17 +31,11 @@ use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadSource;
/// Owns reviewer agents through the same thread manager as the parent conversation.
#[derive(Clone, Debug)]
pub struct GuardianExtension {
#[derive(Debug)]
struct GuardianExtension {
thread_manager: Weak<ThreadManager>,
}
impl GuardianExtension {
pub fn new(thread_manager: Weak<ThreadManager>) -> Self {
Self { thread_manager }
}
}
impl ThreadLifecycleContributor<Config> for GuardianExtension {
fn on_thread_start<'a>(
&'a self,
@@ -50,6 +49,7 @@ impl ThreadLifecycleContributor<Config> for GuardianExtension {
let runtime = input.thread_store.get_or_init(ReviewerTasks::default);
input.thread_store.get_or_init(|| {
ReviewerPool::<GuardianReviewSession>::new(
Arc::clone(&runtime),
move |context, key, kind, snapshot, cancel| {
let manager = manager.clone();
let runtime = Arc::clone(&runtime);
@@ -82,14 +82,12 @@ impl ThreadLifecycleContributor<Config> for GuardianExtension {
options
.thread_extension_init
.insert(SessionIsolation::Isolated);
let parent_cancel = runtime.cancellation.clone();
let session_cancel = cancel.clone();
let until = async move {
let _cancel_on_exit = cancel.clone().drop_guard();
tokio::select! {
_ = cancel.cancelled() => {}
_ = history_reset.cancelled() => {}
_ = parent_cancel.cancelled() => {}
}
};
let spawned = manager
@@ -164,24 +162,44 @@ impl ThreadLifecycleContributor<Config> for GuardianExtension {
fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
if let Some(runtime) = input.thread_store.get::<ReviewerTasks>() {
runtime.cancellation.cancel();
if let Some(pool) = input
.thread_store
.get::<ReviewerPool<GuardianReviewSession>>()
{
pool.shutdown().await;
}
runtime.tasks.close();
runtime.tasks.wait().await;
if let Some(pool) = input
.thread_store
.get::<ReviewerPool<GuardianReviewSession>>()
{
pool.shutdown().await;
}
})
}
}
impl TurnLifecycleContributor for GuardianExtension {
fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(ReviewDenials::clear_turn(input.thread_store, input.turn_id))
}
fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(ReviewDenials::clear_turn(
input.thread_store,
input.turn_store.level_id(),
))
}
fn on_turn_abort<'a>(&'a self, input: TurnAbortInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(ReviewDenials::clear_turn(
input.thread_store,
input.turn_store.level_id(),
))
}
}
/// Registers the synchronous reviewer and its thread and turn cleanup.
pub fn install(
registry: &mut ExtensionRegistryBuilder<Config>,
thread_manager: Weak<ThreadManager>,
) {
registry.thread_lifecycle_contributor(Arc::new(GuardianExtension::new(thread_manager)));
let extension = Arc::new(GuardianExtension { thread_manager });
registry.thread_lifecycle_contributor(extension.clone());
registry.turn_lifecycle_contributor(extension);
}
#[cfg(test)]
#[path = "lifecycle_tests.rs"]
mod tests;