Move Guardian reviewer lifecycle into the extension (#45537)

## Why

Guardian reviews must stop when their parent shuts down or their history is reset. Shutdown must also finish reviewer cleanup before closing the parent's persistent history, including when a review is waiting to retry after a rate limit.

## What changed

- Let the Guardian extension own reviewer startup, prewarming, and shutdown through `ThreadManager`, tracking and joining outstanding work during teardown.
- Cancel reviews on history reset or parent shutdown and reject decisions returned after cancellation, including cached extension decisions.
- Keep background prewarm previews from overwriting the active turn's model metadata.

## Testing

Extend regression coverage to verify reviewer cleanup after parent shutdown, prompt shutdown during a 60-second Guardian retry with the network request denied, and preservation of active model review requirements during prewarming.

GitOrigin-RevId: 05efc369b9a1642a10365eac8a09da77f3f6e28c
This commit is contained in:
jif
2026-09-14 22:58:51 +00:00
committed by copyberry
parent 12b0164a48
commit 18d7ace221
31 changed files with 495 additions and 308 deletions
+1
View File
@@ -3596,6 +3596,7 @@ dependencies = [
"codex-extension-api",
"codex-features",
"codex-guardian-context",
"codex-guardian-reviewer",
"codex-history",
"codex-http-client",
"codex-login",
+1 -1
View File
@@ -692,7 +692,7 @@ impl CodexThread {
pub async fn guardian_trunk_rollout_path(&self) -> Option<PathBuf> {
self.session
.guardian_review_session()
.guardian_review_session()?
.trunk()
.await?
.rollout_path()
+6 -2
View File
@@ -254,12 +254,16 @@ pub(crate) fn cancel_when_either(
first: CancellationToken,
second: CancellationToken,
) -> CancellationToken {
let combined = CancellationToken::new();
let combined = first.child_token();
if combined.is_cancelled() || second.is_cancelled() {
combined.cancel();
return combined;
}
let cancel = combined.clone();
tokio::spawn(async move {
tokio::select! {
_ = first.cancelled() => {}
_ = second.cancelled() => {}
_ = cancel.cancelled() => {}
}
cancel.cancel();
});
+37 -2
View File
@@ -54,10 +54,35 @@ pub(crate) async fn decide_approval(
review_id: String,
request: impl Into<ReviewAction>,
reasons: ApprovalRequestReasons,
options: GuardianReviewOptions,
mut options: GuardianReviewOptions,
) -> Option<ReviewDecision> {
let runtime = session
.services
.thread_extension_data
.get::<codex_guardian_reviewer::ReviewerTasks>();
let _task = runtime.as_ref().map(|runtime| runtime.tasks.token());
if runtime
.as_ref()
.is_some_and(|runtime| runtime.cancellation.is_cancelled())
{
return Some(ReviewDecision::Abort);
}
let context = context.into();
let request = request.into();
let (_, history_reset) = session.history_reset().await;
let cancellation = match options.external_cancel.take() {
Some(external) => crate::exec::cancel_when_either(external, history_reset.clone()),
None => history_reset.child_token(),
};
let _cancel_on_drop = cancellation.clone().drop_guard();
let review_cancellation = match runtime.as_ref() {
Some(runtime) => {
crate::exec::cancel_when_either(runtime.cancellation.clone(), cancellation.clone())
}
None => cancellation.clone(),
};
let _review_cancel_on_drop = review_cancellation.clone().drop_guard();
options.external_cancel = Some(review_cancellation.clone());
let turn = context.turn();
let live_config = session.get_config().await;
let requirements = live_config.config_layer_stack.requirements();
@@ -109,6 +134,7 @@ pub(crate) async fn decide_approval(
};
let runtime = codex_guardian_reviewer::SynchronousReview::new(ReviewRuntime {
session: Arc::clone(&session),
history_reset: history_reset.clone(),
context: context.clone(),
review_id: review_id.clone(),
request: request.clone(),
@@ -148,7 +174,7 @@ pub(crate) async fn decide_approval(
)),
synchronous_reviewer: &runtime,
};
match session.services.extensions.decide_approval(&input).await {
let decision = match session.services.extensions.decide_approval(&input).await {
Some(ApprovalDecision::Reviewed(decision)) => Some(decision),
Some(ApprovalDecision::Allow) if !require_fresh_review => {
let request = match request.validate(&context) {
@@ -193,5 +219,14 @@ pub(crate) async fn decide_approval(
None | Some(ApprovalDecision::AskUser) => {
runtime.review(GuardianReviewReason::Policy).await
}
};
// Enforce cancellation after extension callbacks too, including cached decisions.
if history_reset.is_cancelled()
|| cancellation.is_cancelled()
|| review_cancellation.is_cancelled()
{
Some(ReviewDecision::Abort)
} else {
decision
}
}
+5 -3
View File
@@ -56,10 +56,12 @@ pub(crate) use review::new_guardian_review_id;
#[cfg(test)]
pub(crate) use review::record_guardian_denial_for_test;
pub(crate) use review::routes_approval_policy_to_guardian;
pub(crate) use review::routes_approval_to_guardian;
pub use review_session::GuardianReviewSessionHost;
pub use review_session::GuardianReviewSession;
pub(crate) use review_session::GuardianReviewSessionManager;
pub(crate) use review_session::prewarm_guardian_review_session;
pub use review_session::GuardianReviewState;
pub use review_session::PreparedGuardianContext;
pub use review_session::prepare_review_prewarm;
pub(crate) use review_session::prompt_cache_key_override_for_review_session;
pub(crate) use runtime::ReviewAction;
+9 -16
View File
@@ -102,21 +102,6 @@ pub(crate) fn new_guardian_review_id() -> String {
uuid::Uuid::new_v4().to_string()
}
/// Whether this turn should route allowed approval prompts through the guardian
/// reviewer instead of surfacing them to the user. ARC may still block actions
/// earlier in the flow.
pub(crate) fn routes_approval_to_guardian(turn: &TurnContext) -> bool {
routes_approval_to_guardian_with_reviewer(turn, turn.config.approvals_reviewer)
}
/// Whether an approval with its own reviewer selection should be routed through guardian.
pub(crate) fn routes_approval_to_guardian_with_reviewer(
turn: &TurnContext,
approvals_reviewer: ApprovalsReviewer,
) -> bool {
routes_approval_policy_to_guardian(turn.approval_policy(), approvals_reviewer)
}
/// Whether an exact approval policy and reviewer should route through Guardian.
pub(crate) fn routes_approval_policy_to_guardian(
approval_policy: AskForApproval,
@@ -333,6 +318,14 @@ async fn run_guardian_review_session_before_deadline(
external_cancel: Option<CancellationToken>,
deadline: Instant,
) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) {
let Some(pool) = session.guardian_review_session() else {
return (
GuardianReviewOutcome::Error(GuardianReviewError::prompt_build(anyhow::anyhow!(
"Guardian extension is not installed for this thread"
))),
GuardianReviewAnalyticsResult::without_session(),
);
};
let session_config = match guardian_review_session_config(session.as_ref(), &context).await {
Ok(session_config) => session_config,
Err(err) => {
@@ -344,7 +337,7 @@ async fn run_guardian_review_session_before_deadline(
};
let (session_outcome, session_analytics_result) =
Box::pin(super::review_session::run_guardian_review_session(
session.guardian_review_session(),
pool,
GuardianReviewSessionParams {
parent_session: Arc::clone(&session),
parent_context: context.clone(),
+19 -12
View File
@@ -34,6 +34,7 @@ impl ReviewHost for super::super::runtime::ReviewRuntime {
) -> Result<PreparedApproval, ReviewDecision> {
let super::super::runtime::ReviewRuntime {
session,
history_reset: _,
context,
review_id,
request,
@@ -204,20 +205,26 @@ impl ReviewHost for super::super::runtime::ReviewRuntime {
review_evidence,
} = prepared;
let session = Arc::clone(&self.session);
if session.guardian_context_mode == GuardianContextMode::ThreadOwned
&& matches!(&outcome, GuardianReviewOutcome::Completed(assessment) if assessment.outcome == GuardianAssessmentOutcome::Allow)
&& (root_authorization_version
!= session
.services
.agent_control
.root_user_authorization(session.thread_id)
.await
.map(|snapshot| snapshot.authorization_version)
|| user_message_revision
if matches!(&outcome, GuardianReviewOutcome::Completed(assessment) if assessment.outcome == GuardianAssessmentOutcome::Allow)
&& ((session.guardian_context_mode == GuardianContextMode::ThreadOwned
&& (root_authorization_version
!= session
.conversation_history_snapshot()
.services
.agent_control
.root_user_authorization(session.thread_id)
.await
.user_message_revision())
.map(|snapshot| snapshot.authorization_version)
|| user_message_revision
!= session
.conversation_history_snapshot()
.await
.user_message_revision()))
|| self.history_reset.is_cancelled()
|| self
.options
.external_cancel
.as_ref()
.is_some_and(CancellationToken::is_cancelled))
{
// A completed approval cannot outlive the owning-session or root evidence
// it evaluated, including when either changed before prompt construction.
+6 -22
View File
@@ -4,12 +4,10 @@
#[path = "review_session_setup.rs"]
mod setup;
pub(crate) use setup::prewarm_guardian_review_session;
pub use setup::PreparedGuardianContext;
pub use setup::prepare_review_prewarm;
pub(crate) use setup::run_guardian_review_session;
#[path = "review_session_threads.rs"]
mod managed_threads;
#[path = "review_session_context.rs"]
mod context_policy;
use context_policy::ReviewContextPolicy;
@@ -126,26 +124,11 @@ pub(crate) struct GuardianReviewSessionParams {
pub(crate) deadline: tokio::time::Instant,
}
/// Host capability used to spawn private reviewer runtimes for this parent.
/// The extension owns pooling; this adapter keeps the existing context and runtime paths.
#[derive(Default)]
pub struct GuardianReviewSessionHost {
managed_threads: Option<managed_threads::ManagedReviewerThreads>,
}
impl GuardianReviewSessionHost {
pub fn with_thread_manager(manager: std::sync::Weak<crate::ThreadManager>) -> Self {
Self {
managed_threads: Some(managed_threads::ManagedReviewerThreads::new(manager)),
}
}
}
pub(crate) type GuardianReviewSessionManager =
codex_guardian_reviewer::ReviewerPool<GuardianReviewSession>;
/// Opaque host session handle. Its state belongs to the existing context builder.
pub(crate) struct GuardianReviewSession {
pub struct GuardianReviewSession {
session: Arc<Session>,
io: SessionIo,
cancel_token: CancellationToken,
@@ -153,7 +136,8 @@ pub(crate) struct GuardianReviewSession {
state: Mutex<GuardianReviewState>,
}
struct GuardianReviewState {
/// Opaque conversation progress retained while ThreadManager starts a reviewer.
pub struct GuardianReviewState {
conversation: ConversationState<GuardianReviewHistory>,
last_admitted_node_repl_response_sequence: u64,
pending_node_repl_evidence_admission: Option<PendingNodeReplEvidenceAdmission>,
@@ -882,7 +866,7 @@ impl codex_guardian_reviewer::ReviewerRuntime for GuardianReviewSession {
mod tests;
impl codex_guardian_reviewer::ReviewerSession for GuardianReviewSession {
type Setup = setup::PreparedGuardianContext;
type Setup = PreparedGuardianContext;
type Context = GuardianReviewSessionReuseKey;
type Snapshot = GuardianReviewForkSnapshot;
@@ -1,4 +1,4 @@
//! Prepares opaque session inputs for the extension-owned reviewer pool.
//! Captures the context used by Guardian before and after it starts a reviewer agent.
//! Context selection and assembly stay on the existing path pending its replacement.
use super::*;
@@ -6,14 +6,14 @@ use codex_guardian_reviewer::ReviewerPool;
use codex_guardian_reviewer::ReviewerRequest;
use codex_guardian_reviewer::SessionDisposition;
pub(crate) struct PreparedGuardianContext {
pub struct PreparedGuardianContext {
parent: Arc<Session>,
context: GuardianReviewContext,
config: Config,
context_policy: ReviewContextPolicy,
key: GuardianReviewSessionReuseKey,
parent_compaction: Option<ResponseItem>,
host: Arc<GuardianReviewSessionHost>,
pub history_reset: CancellationToken,
}
impl PreparedGuardianContext {
@@ -25,6 +25,11 @@ impl PreparedGuardianContext {
node_repl_policy: &GuardianNodeReplPolicy,
compaction_model_hash: Option<&str>,
) -> anyhow::Result<Self> {
let (reset_version, history_reset) = parent.history_reset().await;
// Preparation may have raced with a history reset before capturing the reviewer context.
if reset_version != history.reset_version {
history_reset.cancel();
}
let context_policy =
ReviewContextPolicy::for_context(parent.guardian_context_mode, &config.features);
let root_authorization_version = context_policy.root_authorization_version(&parent).await;
@@ -40,10 +45,6 @@ impl PreparedGuardianContext {
.with_node_repl_policy(node_repl_policy);
key.root_authorization_version = root_authorization_version;
key.parent_reset_version = history.reset_version;
let host = parent
.services
.thread_extension_data
.get_or_init(GuardianReviewSessionHost::default);
Ok(Self {
parent,
context,
@@ -51,13 +52,16 @@ impl PreparedGuardianContext {
context_policy,
key,
parent_compaction,
host,
history_reset,
})
}
}
impl PreparedGuardianContext {
fn reuse_key(&self, previous: Option<&GuardianReviewSession>) -> GuardianReviewSessionReuseKey {
pub fn reuse_key(
&self,
previous: Option<&GuardianReviewSession>,
) -> GuardianReviewSessionReuseKey {
let mut key = self.key.clone();
if self.context_policy != ReviewContextPolicy::ThreadOwned
&& self.parent_compaction.is_none()
@@ -70,12 +74,21 @@ impl PreparedGuardianContext {
key
}
/// Returns captured parent inputs; Guardian selects the agent's identity and lifecycle.
async fn thread_options(
/// Moves fork history into startup options and retains only its context bookkeeping.
/// Guardian selects the agent's identity and lifecycle.
pub async fn thread_options(
&self,
initial_history: Option<InitialHistory>,
) -> crate::StartThreadOptions {
let initial_history = initial_history.or_else(|| {
snapshot: Option<GuardianReviewForkSnapshot>,
) -> (crate::StartThreadOptions, GuardianReviewState) {
let (conversation, history) = snapshot.map(ConversationState::fork).unzip();
let state = GuardianReviewState {
conversation: conversation.unwrap_or_default(),
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())]))
@@ -86,7 +99,7 @@ impl PreparedGuardianContext {
.services
.model_client
.responses_websocket_enabled();
crate::StartThreadOptions {
let options = crate::StartThreadOptions {
internal_parent: Some(crate::thread_manager::InternalSessionParent {
thread_id: self.parent.thread_id(),
auth_manager: Arc::clone(&self.parent.services.auth_manager),
@@ -99,24 +112,40 @@ impl PreparedGuardianContext {
inherited_environments: Some(self.context.environments().clone()),
client_mcp_extensions: self.parent.services.client_mcp_extensions.clone(),
..crate::StartThreadOptions::new(config)
}
};
(options, state)
}
/// Binds context bookkeeping to an agent that Guardian has already started.
async fn bind_session(
pub async fn bind_thread(
&self,
session: Arc<Session>,
io: SessionIo,
thread: &crate::CodexThread,
context: GuardianReviewSessionReuseKey,
state: GuardianReviewState,
cancellation: CancellationToken,
) -> GuardianReviewSession {
let session = Arc::clone(&thread.session);
let io = SessionIo {
tx_sub: thread.io.tx_sub.clone(),
rx_event: thread.io.rx_event.clone(),
agent_status: thread.io.agent_status.clone(),
session_loop_termination: thread.io.session_loop_termination.clone(),
};
let inherited = session.inherited_instructions().await;
let context = GuardianReviewSessionReuseKey {
user_instructions: inherited.user,
thread_instructions: inherited.thread,
..context
};
crate::session::emit_subagent_session_started(
&self.parent.services.analytics_events_client,
self.parent.app_server_client_metadata().await,
session.session_id(),
session.thread_id(),
Some(self.parent.thread_id()),
session.thread_config_snapshot().await,
SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_owned()),
);
GuardianReviewSession {
session,
io,
@@ -127,40 +156,6 @@ impl PreparedGuardianContext {
}
}
impl PreparedGuardianContext {
pub(crate) async fn spawn(
&self,
context: GuardianReviewSessionReuseKey,
kind: GuardianReviewSessionKind,
snapshot: Option<GuardianReviewForkSnapshot>,
cancellation: CancellationToken,
) -> anyhow::Result<GuardianReviewSession> {
let (conversation, history) = snapshot.map(ConversationState::fork).unzip();
let state = GuardianReviewState {
conversation: conversation.unwrap_or_default(),
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 mut options = self
.thread_options(history.map(|history| history.initial_history))
.await;
if matches!(kind, GuardianReviewSessionKind::EphemeralForked) {
options.config.ephemeral = true;
}
let threads = self.host.managed_threads.as_ref().ok_or_else(|| {
anyhow::anyhow!("Guardian extension is not installed for this thread")
})?;
let (session, io) = threads
.spawn(&self.parent, options, cancellation.clone())
.await?;
Ok(self
.bind_session(session, io, context, state, cancellation)
.await)
}
}
pub(super) struct PreparedReview {
context: Arc<PreparedGuardianContext>,
params: GuardianReviewSessionParams,
@@ -239,18 +234,16 @@ pub(super) async fn prepare_review(
})
}
pub(crate) fn prewarm_guardian_review_session(
parent: Arc<Session>,
turn: Arc<TurnContext>,
) -> BoxFuture<'static, anyhow::Result<()>> {
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
})
/// Captures the same startup context used by the existing prompt builder.
/// The caller owns scheduling, cancellation, and the reviewer pool.
pub async fn prepare_review_prewarm(
parent: &crate::CodexThread,
) -> anyhow::Result<PreparedGuardianContext> {
let turn = parent
.session
.new_startup_prewarm_turn_with_sub_id(crate::session::INITIAL_SUBMIT_ID.to_owned())
.await;
prepare_prewarm(Arc::clone(&parent.session), turn).await
}
pub(super) fn prepare_prewarm(
@@ -251,30 +251,6 @@ async fn test_review_params() -> GuardianReviewSessionParams {
}
}
#[tokio::test]
async fn spawned_guardian_session_preserves_windows_sandbox_proxy_settings() {
let params = test_review_params().await;
let manager = params.parent_session.guardian_review_session();
prewarm_guardian_review_session(
params.parent_session,
Arc::clone(params.parent_context.turn()),
)
.await
.expect("initialize Guardian session");
let mode = manager
.trunk()
.await
.expect("Guardian session")
.session
.windows_sandbox_proxy_settings_mode;
assert_eq!(
mode,
codex_sandboxing::WindowsSandboxProxySettingsMode::Preserve
);
manager.shutdown().await;
}
#[tokio::test]
async fn spawned_guardian_reuse_key_matches_inherited_instructions() {
let mut params = test_review_params().await;
@@ -310,17 +286,16 @@ async fn spawned_guardian_reuse_key_matches_inherited_instructions() {
thread_instructions: latest.clone(),
..stale_key.clone()
};
let manager = params
.parent_session
.guardian_review_session()
.expect("Guardian pool installed");
let prepared = setup::prepare_review(params).await.expect("prepare review");
let review = prepared
.setup()
.spawn(
stale_key,
GuardianReviewSessionKind::TrunkNew,
/*snapshot*/ None,
CancellationToken::new(),
)
manager
.prewarm(prepared.setup(), stale_key)
.await
.expect("spawn reviewer after instruction update");
let review = manager.trunk().await.expect("prewarmed reviewer");
assert_eq!(review.reuse_key, expected_key);
assert_eq!(review.session.inherited_instructions().await.thread, latest);
@@ -1,96 +0,0 @@
//! Adapts the existing reviewer lifecycle to the extension's ThreadManager.
//! Captured parent inputs support inline delegates; prompts and review outcomes stay shared.
use std::sync::Arc;
use std::sync::Weak;
use codex_async_utils::OrCancelExt;
use codex_history::InitialHistory;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::ThreadSource;
use tokio_util::sync::CancellationToken;
use super::GUARDIAN_REVIEWER_NAME;
use crate::StartThreadOptions;
use crate::ThreadManager;
use crate::codex_delegate::forward_session_io;
use crate::config::Constrained;
use crate::session::SessionIo;
use crate::session::emit_subagent_session_started;
use crate::session::session::Session;
pub(super) struct ManagedReviewerThreads {
manager: Weak<ThreadManager>,
}
impl ManagedReviewerThreads {
pub(super) fn new(manager: Weak<ThreadManager>) -> Self {
Self { manager }
}
pub(super) async fn spawn(
&self,
parent: &Arc<Session>,
mut options: StartThreadOptions,
cancel: CancellationToken,
) -> anyhow::Result<(Arc<Session>, SessionIo)> {
let manager = self
.manager
.upgrade()
.ok_or_else(|| anyhow::anyhow!("thread manager is no longer available"))?;
if options.config.permissions.approval_policy.value() != AskForApproval::Never {
anyhow::bail!("Codex delegates require approval policy `never`");
}
options.config.permissions.approval_policy = Constrained::allow_only(AskForApproval::Never);
options.session_source = Some(SessionSource::Internal(InternalSessionSource::Guardian));
options.thread_source = Some(ThreadSource::GuardianReview);
options
.thread_extension_init
.insert(codex_extension_api::SessionIsolation::Isolated);
options.initial_history = match options.initial_history {
InitialHistory::Forked(history) => InitialHistory::Forked(history),
InitialHistory::New | InitialHistory::Cleared => InitialHistory::New,
InitialHistory::Resumed(_) => {
anyhow::bail!("guardian review forks cannot resume an existing thread")
}
};
let spawned = manager
.start_thread(options)
.or_cancel(&cancel)
.await
.map_err(codex_protocol::error::CodexErr::from)??;
let thread = spawned.thread;
let session = Arc::clone(&thread.session);
let io = forward_session_io(
Arc::new(SessionIo {
tx_sub: thread.io.tx_sub.clone(),
rx_event: thread.io.rx_event.clone(),
agent_status: thread.io.agent_status.clone(),
session_loop_termination: thread.io.session_loop_termination.clone(),
}),
cancel,
);
let manager = Arc::downgrade(&manager);
drop(tokio::spawn(async move {
thread.wait_until_terminated().await;
if let Some(manager) = manager.upgrade() {
manager
.remove_thread_if_matches(&thread.session.thread_id(), &thread)
.await;
}
}));
emit_subagent_session_started(
&parent.services.analytics_events_client,
parent.app_server_client_metadata().await,
session.session_id(),
session.thread_id(),
Some(parent.thread_id()),
session.thread_config_snapshot().await,
SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_owned()),
);
Ok((session, io))
}
}
+2
View File
@@ -2,6 +2,7 @@
use codex_protocol::protocol::ReviewDecision;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use super::ApprovalRequestReasons;
use super::GuardianApprovalRequest;
@@ -74,6 +75,7 @@ impl ReviewAction {
#[derive(Clone)]
pub(super) struct ReviewRuntime {
pub(super) session: Arc<Session>,
pub(super) history_reset: CancellationToken,
pub(super) context: GuardianReviewContext,
pub(super) review_id: String,
pub(super) request: ReviewAction,
+46 -3
View File
@@ -3,10 +3,16 @@
use std::sync::Arc;
use codex_extension_api::SessionIsolation;
use codex_home::CodexHomeUserInstructionsProvider;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadSource;
use super::GuardianReviewSessionManager;
use crate::config::Config;
use crate::config::Constrained;
use crate::session::session::Session;
pub(crate) fn install(session: &Session, config: &Config) {
@@ -29,7 +35,44 @@ pub(crate) fn install(session: &Session, config: &Config) {
/*attestation_provider*/ None,
/*external_time_provider*/ None,
));
let host = super::GuardianReviewSessionHost::with_thread_manager(Arc::downgrade(&manager));
session.services.thread_extension_data.insert(host);
session.services.thread_extension_data.insert(manager);
session
.services
.thread_extension_data
.insert(GuardianReviewSessionManager::new(
move |context, key, kind, snapshot, cancel| {
let manager = Arc::clone(&manager);
Box::pin(async move {
let history_reset = context.history_reset.clone();
let (mut options, state) = context.thread_options(snapshot).await;
if matches!(
kind,
codex_analytics::GuardianReviewSessionKind::EphemeralForked
) {
options.config.ephemeral = true;
}
options.config.permissions.approval_policy =
Constrained::allow_only(AskForApproval::Never);
options.session_source =
Some(SessionSource::Internal(InternalSessionSource::Guardian));
options.thread_source = Some(ThreadSource::GuardianReview);
options
.thread_extension_init
.insert(SessionIsolation::Isolated);
let session_cancel = cancel.clone();
let until = async move {
let _cancel_on_exit = cancel.clone().drop_guard();
tokio::select! {
_ = cancel.cancelled() => {}
_ = history_reset.cancelled() => {}
}
};
let spawned = manager
.start_thread_until(options, until, &tokio_util::task::TaskTracker::new())
.await?;
Ok(context
.bind_thread(&spawned.thread, key, state, session_cancel)
.await)
})
},
));
}
+17 -1
View File
@@ -12,7 +12,6 @@ use crate::guardian::prompt::BUNDLED_GUARDIAN_POLICY;
use crate::guardian::prompt::BUNDLED_GUARDIAN_POLICY_TEMPLATE;
use crate::guardian::prompt::guardian_policy_prompt_with_config_and_template;
use crate::guardian::review::guardian_review_session_config;
use crate::guardian::review::routes_approval_to_guardian_with_reviewer;
use crate::session::session::Session;
use crate::session::tests::update_turn_settings_for_test;
use crate::session::turn_context::TurnContext;
@@ -2592,6 +2591,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow:
.await;
let committed_rollout_items = session
.guardian_review_session()
.expect("Guardian pool installed")
.trunk()
.await
.expect("reviewer")
@@ -2927,6 +2927,7 @@ async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow::
session
.guardian_review_session()
.expect("Guardian pool installed")
.trunk()
.await
.expect("reviewer")
@@ -4255,3 +4256,18 @@ async fn review_approval_request(
.await
.expect("Guardian should handle the request")
}
/// Whether this turn should route allowed approval prompts through the guardian
/// reviewer instead of surfacing them to the user. ARC may still block actions
/// earlier in the flow.
fn routes_approval_to_guardian(turn: &crate::session::turn_context::TurnContext) -> bool {
routes_approval_to_guardian_with_reviewer(turn, turn.config.approvals_reviewer)
}
/// Whether an approval with its own reviewer selection should be routed through guardian.
fn routes_approval_to_guardian_with_reviewer(
turn: &crate::session::turn_context::TurnContext,
approvals_reviewer: ApprovalsReviewer,
) -> bool {
routes_approval_policy_to_guardian(turn.approval_policy(), approvals_reviewer)
}
+6 -3
View File
@@ -1,4 +1,7 @@
//! Opaque host capabilities for the synchronous Guardian extension.
//! Core supplies context and runtime operations; the extension owns review policy and pooling.
//! Context adapter types used by the host-installed synchronous Guardian extension.
//! The extension supplies agent startup and owns the reviewer pool and its lifecycle.
pub use crate::guardian::GuardianReviewSessionHost;
pub use crate::guardian::GuardianReviewSession;
pub use crate::guardian::GuardianReviewState;
pub use crate::guardian::PreparedGuardianContext;
pub use crate::guardian::prepare_review_prewarm;
+2 -5
View File
@@ -308,12 +308,12 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc<Session>) {
sess.mcp_refresh.close();
sess.services.mcp_runtime.shutdown().await;
}
sess.guardian_review_session().shutdown().await;
crate::hook_runtime::run_session_end_hooks(sess).await;
emit_thread_stop_lifecycle(sess).await;
}
pub(super) async fn emit_thread_stop_lifecycle(sess: &Session) {
async fn emit_thread_stop_lifecycle(sess: &Session) {
for contributor in sess.services.extensions.thread_lifecycle_contributors() {
contributor
.on_thread_stop(codex_extension_api::ThreadStopInput {
@@ -338,8 +338,6 @@ pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
&[],
);
emit_thread_stop_lifecycle(sess.as_ref()).await;
// Gracefully flush and shutdown thread persistence on session end so tests
// that inspect durable state do not race with the background writer.
if let Some(live_thread) = sess.live_thread()
@@ -607,7 +605,6 @@ pub(super) async fn submission_loop(
// explicit shutdown op, still run session teardown.
if !shutdown_received {
shutdown_session_runtime(&sess).await;
emit_thread_stop_lifecycle(sess.as_ref()).await;
if let Some(live_thread) = sess.live_thread()
&& let Err(err) = live_thread.shutdown().await
{
+14 -7
View File
@@ -2245,13 +2245,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::new(|context, key, kind, snapshot, cancel| {
Box::pin(async move { context.spawn(key, kind, snapshot, cancel).await })
})
})
/// Returns the reviewer pool installed by the Guardian extension.
pub(crate) fn guardian_review_session(&self) -> Option<Arc<GuardianReviewSessionManager>> {
self.services
.thread_extension_data
.get::<GuardianReviewSessionManager>()
}
pub(crate) async fn emit_turn_started(&self, turn_context: &TurnContext) {
@@ -4393,6 +4391,15 @@ impl Session {
state.clone_history()
}
/// Captures the history lifetime without exposing mutable session state.
pub(crate) async fn history_reset(&self) -> (u64, CancellationToken) {
let state = self.state.lock().await;
(
state.history.reset_version,
state.history_reset.child_token(),
)
}
pub(crate) async fn conversation_history_snapshot(
&self,
) -> Arc<dyn ConversationHistorySnapshot> {
+4 -4
View File
@@ -1021,12 +1021,12 @@ impl Session {
self.features.enabled(Feature::Personality),
)
.await;
self.services
.thread_extension_data
.insert(model_info.clone());
let multi_agent_version = match multi_agent_runtime {
TurnMultiAgentRuntime::ResolveAndStore => {
// A background preview must not overwrite a newer turn's model metadata.
self.services
.thread_extension_data
.insert(model_info.clone());
self.resolve_multi_agent_version_for_model(&model_info, &per_turn_config)
}
TurnMultiAgentRuntime::Preview => per_turn_config.multi_agent_version_for_model(
+2 -3
View File
@@ -107,9 +107,8 @@ pub(super) async fn suspend_turn_and_shutdown(
live_thread.shutdown().await.map_err(|error| {
CodexErr::Fatal(format!("close suspended root turn writer failed: {error}"))
})?;
// Announce thread shutdown only after its writer closes so a replacement worker
// cannot write the same thread concurrently.
handlers::emit_thread_stop_lifecycle(session.as_ref()).await;
// Announce completion only after extension cleanup and writer closure so a
// replacement worker cannot write the same thread concurrently.
session
.deliver_event_raw(Event {
id: submission_id,
@@ -12,7 +12,6 @@ use tracing::trace_span;
use tracing::warn;
use crate::client::ModelClientSession;
use crate::guardian::routes_approval_to_guardian;
use crate::responses_metadata::CodexResponsesRequestKind;
use crate::session::INITIAL_SUBMIT_ID;
use crate::session::RequestEffortUsage;
@@ -273,20 +272,6 @@ async fn schedule_startup_prewarm_inner(
prewarm_started_at.elapsed(),
/*status*/ None,
);
if routes_approval_to_guardian(&startup_turn_context) {
let guardian_session = Arc::clone(&session);
let guardian_parent_turn = Arc::clone(&startup_turn_context);
drop(tokio::spawn(async move {
if let Err(err) = crate::guardian::prewarm_guardian_review_session(
guardian_session,
guardian_parent_turn,
)
.await
{
warn!("failed to initialize guardian review session: {err:#}");
}
}));
}
let startup_cancellation_token = CancellationToken::new();
let built_tools_started_at = Instant::now();
// Startup prewarm runs before run_turn and needs its own tool-building snapshot.
+13 -5
View File
@@ -28,6 +28,7 @@ use codex_protocol::protocol::TokenUsageInfo;
use codex_protocol::protocol::TokenUsageRecord;
use codex_protocol::protocol::TurnContextItem;
use codex_utils_output_truncation::TruncationPolicy;
use tokio_util::sync::CancellationToken;
use tokio_util::task::AbortOnDropHandle;
/// Runtime request effort, initially unset and established by prewarm or sampling.
@@ -72,6 +73,8 @@ pub(crate) struct SessionState {
/// Persisted origin of the session base instructions, when known.
pub(crate) base_instructions_provenance: Option<BaseInstructionsProvenance>,
pub(crate) history: ContextManager,
/// Cancels work bound to discarded history. Appends and compaction preserve it.
pub(crate) history_reset: CancellationToken,
pub(crate) latest_rate_limits: Option<RateLimitSnapshot>,
pub(crate) latest_token_usage_record: Option<TokenUsageRecord>,
pub(crate) server_reasoning_included: bool,
@@ -117,6 +120,7 @@ impl SessionState {
session_configuration,
base_instructions_provenance: None,
history,
history_reset: CancellationToken::new(),
latest_rate_limits: None,
latest_token_usage_record: None,
server_reasoning_included: false,
@@ -174,10 +178,11 @@ impl SessionState {
items: Vec<ResponseItem>,
reference_context_item: Option<TurnContextItem>,
) {
self.history.replace(items);
self.history
.set_reference_context_item(reference_context_item);
self.auto_compact_window.clear_prefill();
self.replace_annotated_history(
items.into_iter().map(ResponseItemEnvelope::new).collect(),
reference_context_item,
HistoryReplacement::Reset,
);
}
pub(crate) fn replace_annotated_history(
@@ -188,7 +193,10 @@ impl SessionState {
) {
match replacement {
HistoryReplacement::Compaction => self.history.replace_compacted(items),
HistoryReplacement::Reset => self.history.replace_annotated(items),
HistoryReplacement::Reset => {
self.history.replace_annotated(items);
std::mem::take(&mut self.history_reset).cancel();
}
}
self.history
.set_reference_context_item(reference_context_item);
@@ -1052,6 +1052,19 @@ async fn guardian_session_prewarms_and_is_reused_for_first_review(
.await
.expect("guardian trunk rollout path");
test.codex.shutdown_and_wait().await?;
// Parent stop joins ThreadManager cleanup through the real extension registration.
assert!(
test.codex
.thread_extension_data()
.get::<codex_guardian_reviewer::ReviewerTasks>()
.expect("Guardian tasks")
.tasks
.is_empty()
);
assert!(matches!(
test.thread_store.flush_thread(guardian_thread_id).await,
Err(codex_thread_store::ThreadStoreError::ThreadNotFound { .. })
));
let guardian_rollout = fs::read_to_string(guardian_rollout_path)?
.lines()
.map(codex_rollout::parse_rollout_line)
+66 -2
View File
@@ -15,6 +15,7 @@ use codex_exec_server::LOCAL_ENVIRONMENT_ID;
use codex_exec_server::REMOTE_ENVIRONMENT_ID;
use codex_exec_server::RemoveOptions;
use codex_features::Feature;
use codex_history::RolloutItem;
use codex_network_proxy::NetworkProxyConfig;
use codex_protocol::approvals::NetworkApprovalContext;
use codex_protocol::approvals::NetworkApprovalProtocol;
@@ -45,6 +46,7 @@ use codex_protocol::request_permissions::PermissionGrantScope;
use codex_protocol::request_permissions::RequestPermissionProfile;
use codex_protocol::request_permissions::RequestPermissionsResponse;
use codex_protocol::user_input::UserInput;
use codex_thread_store::LoadThreadHistoryParams;
use codex_utils_path_uri::PathUri;
use core_test_support::PathBufExt;
use core_test_support::PathExt;
@@ -1582,12 +1584,16 @@ async fn failed_network_policy_amendment_denies_request_and_does_not_approve_hos
Ok(())
}
#[test_case(ApprovalsReviewer::User; "user_approval")]
#[test_case(ApprovalsReviewer::AutoReview; "shutdown_during_guardian_retry")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[cfg_attr(
not(target_os = "linux"),
ignore = "requires the trusted Linux proxy bridge"
)]
async fn unattributed_network_request_uses_active_turn_environment_fallback() -> Result<()> {
async fn unattributed_network_request_uses_active_turn_environment_fallback(
approvals_reviewer: ApprovalsReviewer,
) -> Result<()> {
skip_if_target_windows!(Ok(()), "uses a raw TCP proxy fixture");
skip_if_host_windows!(Ok(()));
skip_if_no_network!(Ok(()));
@@ -1610,7 +1616,7 @@ async fn unattributed_network_request_uses_active_turn_environment_fallback() ->
&test,
"hold the active turn",
vec![local(test.config.cwd.clone())],
ApprovalsReviewer::User,
approvals_reviewer,
AskForApproval::OnRequest,
)
.await?;
@@ -1623,6 +1629,64 @@ async fn unattributed_network_request_uses_active_turn_environment_fallback() ->
.context("expected managed network proxy")?
.http_addr
.clone();
if approvals_reviewer == ApprovalsReviewer::AutoReview {
mount_sse_once_match(
&server,
is_guardian_request,
core_test_support::responses::sse_failed(
"guardian-first-failure",
"rate_limit_exceeded",
"Rate limit exceeded. Please try again in 0s.",
),
)
.await;
let retrying = mount_sse_once_match(
&server,
is_guardian_request,
core_test_support::responses::sse_failed(
"guardian-rate-limited",
"rate_limit_exceeded",
"Rate limit exceeded. Please try again in 60s.",
),
)
.await;
let proxy_request = tokio::spawn(raw_http_proxy_request(proxy_addr, NETWORK_TEST_HOST));
wait_for_response_request(&retrying).await;
let request = retrying.single_request().body_json();
let thread_id = request["client_metadata"]["thread_id"]
.as_str()
.context("Guardian thread id")?;
let thread_id = codex_protocol::ThreadId::from_string(thread_id)?;
// Internal reviewers are hidden from get_thread. Their persisted terminal event
// confirms the reviewer's own retries finished before testing Guardian's retry wait.
tokio::time::timeout(Duration::from_secs(5), async {
loop {
test.thread_store.flush_thread(thread_id).await?;
let history = test
.thread_store
.load_latest_model_context(LoadThreadHistoryParams {
thread_id,
include_archived: false,
})
.await?;
if history.items.iter().any(|item| {
matches!(item, RolloutItem::EventMsg(EventMsg::TurnComplete(event)) if event.error.is_some())
}) {
return Ok::<(), anyhow::Error>(());
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.context("Guardian reviewer did not finish its rate-limited turn")??;
tokio::time::timeout(Duration::from_secs(5), test.codex.shutdown_and_wait())
.await
.context("parent shutdown waited for Guardian's 60-second retry")??;
let response = tokio::time::timeout(Duration::from_secs(5), proxy_request).await???;
assert!(response.starts_with("HTTP/1.1 403"), "{response}");
assert_eq!(retrying.requests().len(), 1);
return Ok(());
}
let proxy_request = tokio::spawn(raw_http_proxy_request(proxy_addr, NETWORK_TEST_HOST));
let approval = expect_network_approval(&test, LOCAL_ENVIRONMENT_ID).await?;
assert_eq!(approval.command, ["network-access", NETWORK_TEST_TARGET]);
@@ -170,7 +170,8 @@ pub trait ThreadLifecycleContributor<C: Sync>: Send + Sync {
})
}
/// Called before the host drops the thread runtime and thread-scoped store.
/// Called during runtime teardown, before the host closes persistent history.
/// Contributors must cancel and join their background work before returning.
fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
let _self = self;
@@ -75,7 +75,7 @@ pub struct ThreadIdleInput<'a> {
pub thread_store: &'a ExtensionData,
}
/// Input supplied when the host stops a thread runtime.
/// Input supplied during runtime teardown, before persistent history closes.
pub struct ThreadStopInput<'a> {
/// Store scoped to the host session runtime.
pub session_store: &'a ExtensionData,
+1 -1
View File
@@ -26,7 +26,7 @@ rand = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true }
tokio-util = { workspace = true }
tokio-util = { workspace = true, features = ["rt"] }
tracing = { workspace = true }
[dev-dependencies]
@@ -44,6 +44,7 @@ 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::ReviewerTasks;
pub use pool::SessionDisposition;
pub use review::ReviewHost;
@@ -18,6 +18,14 @@ 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.
#[derive(Default)]
pub struct ReviewerTasks {
pub tasks: tokio_util::task::TaskTracker,
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.
pub trait ReviewerSession: Send + Sync + 'static {
+2 -1
View File
@@ -13,12 +13,14 @@ doctest = false
workspace = true
[dependencies]
anyhow = { workspace = true }
codex-analytics = { workspace = true }
codex-api = { workspace = true }
codex-core = { workspace = true }
codex-extension-api = { workspace = true }
codex-features = { workspace = true }
codex-guardian-context = { workspace = true }
codex-guardian-reviewer = { workspace = true }
codex-history = { workspace = true }
codex-http-client = { workspace = true }
codex-login = { workspace = true }
@@ -33,7 +35,6 @@ tracing = { workspace = true }
uuid = { workspace = true }
[dev-dependencies]
anyhow = { workspace = true }
codex-model-provider-info = { workspace = true }
core_test_support = { workspace = true }
pretty_assertions = { workspace = true }
@@ -2203,6 +2203,8 @@ async fn contributor_skips_required_models_in_standard_scope() -> Result<()> {
.await;
model_info.slug = "protected-model".to_owned();
thread_store.insert(model_info);
// A late prewarm preview must leave the active model's review requirements intact.
let _ = codex_core::guardian_review::prepare_review_prewarm(&test.codex).await?;
let authorization = ScoreAuthorization::current(&test.codex).await;
let progress = thread_store
.get::<GuardianV2ScoreProgress>()
@@ -1,18 +1,31 @@
//! Installs the synchronous reviewer independently of async scorer startup.
//! Reviewer policy and pool state live in the extension; core supplies context and runtime handles.
//! Runs Guardian agents through ThreadManager and owns their background tasks.
//! Context stays in the temporary core adapter. Parent stop joins all reviewer
//! work before the parent closes its history, including partial startup.
use std::sync::Arc;
use std::sync::Weak;
use codex_core::ThreadManager;
use codex_core::config::Config;
use codex_core::guardian_review::GuardianReviewSessionHost;
use codex_core::config::Constrained;
use codex_core::guardian_review::GuardianReviewSession;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::SessionIsolation;
use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::ThreadReadyInput;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ThreadStopInput;
use codex_guardian_reviewer::ReviewerPool;
use codex_guardian_reviewer::ReviewerTasks;
use codex_protocol::ThreadId;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadSource;
/// Owns reviewer state through the same thread manager as the parent conversation.
/// Owns reviewer agents through the same thread manager as the parent conversation.
#[derive(Clone, Debug)]
pub struct GuardianExtension {
thread_manager: Weak<ThreadManager>,
@@ -33,11 +46,137 @@ impl ThreadLifecycleContributor<Config> for GuardianExtension {
if input.session_source.is_internal() {
return;
}
let manager = self.thread_manager.clone();
let runtime = input.thread_store.get_or_init(ReviewerTasks::default);
input.thread_store.get_or_init(|| {
GuardianReviewSessionHost::with_thread_manager(self.thread_manager.clone())
ReviewerPool::<GuardianReviewSession>::new(
move |context, key, kind, snapshot, cancel| {
let manager = manager.clone();
let runtime = Arc::clone(&runtime);
Box::pin(async move {
let history_reset = context.history_reset.clone();
// Register before checking cancellation so parent stop also joins
// a spawn racing with shutdown.
let _task = runtime.tasks.token();
anyhow::ensure!(
!runtime.cancellation.is_cancelled()
&& !cancel.is_cancelled()
&& !history_reset.is_cancelled(),
"Guardian is stopping"
);
let manager = manager.upgrade().ok_or_else(|| {
anyhow::anyhow!("thread manager is no longer available")
})?;
let (mut options, state) = context.thread_options(snapshot).await;
if matches!(
kind,
codex_analytics::GuardianReviewSessionKind::EphemeralForked
) {
options.config.ephemeral = true;
}
options.config.permissions.approval_policy =
Constrained::allow_only(AskForApproval::Never);
options.session_source =
Some(SessionSource::Internal(InternalSessionSource::Guardian));
options.thread_source = Some(ThreadSource::GuardianReview);
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
.start_thread_until(options, until, &runtime.tasks)
.await?;
Ok(context
.bind_thread(&spawned.thread, key, state, session_cancel)
.await)
})
},
)
});
})
}
fn on_thread_ready<'a>(
&'a self,
input: ThreadReadyInput<'a, Config>,
) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
if input.session_source.is_internal()
|| !matches!(
input.config.permissions.approval_policy.value(),
AskForApproval::OnRequest | AskForApproval::Granular(_)
)
|| input.config.approvals_reviewer != ApprovalsReviewer::AutoReview
{
return;
}
let Some(runtime) = input.thread_store.get::<ReviewerTasks>() else {
return;
};
let task = runtime.tasks.token();
if runtime.cancellation.is_cancelled() {
return;
}
let Some(manager) = self.thread_manager.upgrade() else {
return;
};
let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else {
return;
};
let Ok(parent) = manager.get_thread(thread_id).await else {
return;
};
let Some(pool) = input
.thread_store
.get::<ReviewerPool<GuardianReviewSession>>()
else {
return;
};
let cancel = runtime.cancellation.clone();
tokio::spawn(async move {
let _task = task;
let prepare = async {
let context =
codex_core::guardian_review::prepare_review_prewarm(&parent).await?;
let key = context.reuse_key(/*previous*/ None);
pool.prewarm(Arc::new(context), key).await
};
tokio::select! {
_ = cancel.cancelled() => {}
result = prepare => {
if let Err(error) = result {
tracing::warn!("failed to prewarm Guardian reviewer: {error:#}");
}
}
}
});
})
}
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;
}
})
}
}
pub fn install(