mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
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:
Generated
+1
@@ -3596,6 +3596,7 @@ dependencies = [
|
||||
"codex-extension-api",
|
||||
"codex-features",
|
||||
"codex-guardian-context",
|
||||
"codex-guardian-reviewer",
|
||||
"codex-history",
|
||||
"codex-http-client",
|
||||
"codex-login",
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,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,
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
},
|
||||
));
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
{
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user