diff --git a/codex-rs/analytics/src/events.rs b/codex-rs/analytics/src/events.rs index 7ea459b0a3..cb0aacff8d 100644 --- a/codex-rs/analytics/src/events.rs +++ b/codex-rs/analytics/src/events.rs @@ -425,6 +425,7 @@ pub struct GuardianReviewEventParams { pub guardian_review_model_override: Option, pub guardian_model_provider_id: Option, pub had_prior_review_context: Option, + /// Configured overall review budget in milliseconds; zero means unlimited. pub review_timeout_ms: u64, pub tool_call_count: Option, pub time_to_first_token_ms: Option, diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index b31342230e..c3ac3e0b4b 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -581,6 +581,10 @@ pub enum ThreadStoreToml { #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] pub struct AutoReviewToml { + /// Overall Guardian review budget in milliseconds, including setup and retries. + /// Defaults to 90,000. Set to 0 to wait without a review deadline; cancellation + /// and independent transport timeouts still apply. + pub timeout_ms: Option, /// Controls whether circuit-breaker interruptions include a structured error. /// Strict mode writes structured errors that older clients may not /// recognize when reading shared history. Defaults to `default`. diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 7c59145359..26df04230e 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -453,6 +453,12 @@ "policy": { "description": "Additional policy instructions inserted into the guardian prompt.", "type": "string" + }, + "timeout_ms": { + "description": "Overall Guardian review budget in milliseconds, including setup and retries. Defaults to 90,000. Set to 0 to wait without a review deadline; cancellation and independent transport timeouts still apply.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" } }, "type": "object" diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 751dd48ad1..d253342031 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -8742,6 +8742,7 @@ async fn load_config_uses_auto_review_guardian_policy_config_and_template( let cfg = ConfigToml { features: Some(toml::from_str(features).unwrap()), auto_review: Some(AutoReviewToml { + timeout_ms: None, circuit_break_action: None, policy: Some(" Use the user-configured guardian policy. ".to_string()), extra_policy: Some(" Use the user-configured additional policy. ".to_string()), @@ -8805,6 +8806,7 @@ async fn requirements_guardian_policy_beats_auto_review() -> std::io::Result<()> .map_err(std::io::Error::other)?; let cfg = ConfigToml { auto_review: Some(AutoReviewToml { + timeout_ms: None, circuit_break_action: None, policy: Some("Use the user-configured guardian policy.".to_string()), extra_policy: Some("Use the user-configured additional policy.".to_string()), @@ -8846,6 +8848,7 @@ async fn load_config_ignores_empty_auto_review_guardian_policy_config() -> std:: let codex_home = TempDir::new()?; let cfg = ConfigToml { auto_review: Some(AutoReviewToml { + timeout_ms: None, circuit_break_action: None, policy: Some(" ".to_string()), extra_policy: Some(" ".to_string()), diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 076983dfc5..38832d1f16 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -742,6 +742,9 @@ pub struct Config { /// Include a structured error when Guardian's circuit breaker interrupts a turn. pub guardian_circuit_break_action: CircuitBreakAction, + /// Overall Guardian review budget. None means unlimited; cancellation still applies. + pub guardian_review_timeout: Option, + /// Whether to inject the `` developer block. pub include_permissions_instructions: bool, @@ -4614,6 +4617,13 @@ impl Config { guardian_transcript_mode, guardian_conversation_history_prompt, guardian_conversation_history_max_output_tokens, + guardian_review_timeout: cfg + .auto_review + .as_ref() + .and_then(|auto_review| auto_review.timeout_ms) + .map_or(Some(codex_guardian_reviewer::REVIEW_TIMEOUT), |timeout_ms| { + (timeout_ms != 0).then(|| Duration::from_millis(u64::from(timeout_ms))) + }), guardian_circuit_break_action: cfg .auto_review .as_ref() diff --git a/codex-rs/core/src/guardian/decision.rs b/codex-rs/core/src/guardian/decision.rs index c44fedf597..79981f57e0 100644 --- a/codex-rs/core/src/guardian/decision.rs +++ b/codex-rs/core/src/guardian/decision.rs @@ -101,6 +101,7 @@ pub(crate) async fn decide_approval( escalated_exec: matches!(&request.request, Ok(GuardianApprovalRequest::ExecCommand { sandbox_permissions, .. }) if sandbox_permissions.requires_escalated_permissions()), full_access, cancellation: cancellation.clone(), + timeout: turn.config.guardian_review_timeout, // Keep the existing turn-level reporting policy during this ownership move. model: turn.model_info(), telemetry: &turn.session_telemetry, diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index d07623cdee..81aef2fb4f 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -13,6 +13,7 @@ use codex_guardian_reviewer::GuardianReviewError; use codex_guardian_reviewer::GuardianReviewOutcome; #[cfg(test)] use codex_guardian_reviewer::GuardianReviewSessionLimits; +use codex_guardian_reviewer::ReviewDeadline; use codex_guardian_reviewer::ReviewModel; #[cfg(test)] use codex_guardian_reviewer::run_with_retry; @@ -164,7 +165,7 @@ pub(super) async fn guardian_review_session_config( struct ReviewAttemptOptions { authorization: Option, external_cancel: Option, - deadline: Instant, + deadline: ReviewDeadline, } /// Runs the guardian in a locked-down reusable review session. @@ -249,7 +250,7 @@ pub(super) async fn run_guardian_review_session_with_retry( external_cancel, GuardianReviewSessionLimits { max_attempts, - deadline: Instant::now() + GUARDIAN_REVIEW_TIMEOUT, + deadline: ReviewDeadline::from_timeout(Some(GUARDIAN_REVIEW_TIMEOUT)), }, ) .await diff --git a/codex-rs/core/src/guardian/review_request.rs b/codex-rs/core/src/guardian/review_request.rs index 613b7efb61..e7ea989cc1 100644 --- a/codex-rs/core/src/guardian/review_request.rs +++ b/codex-rs/core/src/guardian/review_request.rs @@ -4,6 +4,7 @@ use super::*; use crate::codex_thread::GuardianAuthorizationVersion; use codex_config::config_toml::CircuitBreakAction; +use codex_guardian_reviewer::ReviewDeadline; use codex_guardian_reviewer::ReviewHost; use codex_protocol::approvals::GuardianReviewReason; use codex_protocol::protocol::ErrorEvent; @@ -118,7 +119,7 @@ impl ReviewHost for super::super::runtime::ReviewRuntime { &self, review_id: &str, review_reason: GuardianReviewReason, - deadline: Instant, + deadline: ReviewDeadline, cancellation: &CancellationToken, ) -> Result<(PreparedApproval, codex_guardian_reviewer::ReviewReport), ReviewDecision> { let super::super::runtime::ReviewRuntime { @@ -148,10 +149,8 @@ impl ReviewHost for super::super::runtime::ReviewRuntime { let plugin_attribution = match plugin_attribution_override { Some(attribution) => Some(attribution), None if matches!(&request, GuardianApprovalRequest::ExecCommand { .. }) => { - let attribution_deadline = std::cmp::min( - deadline, - Instant::now() + GUARDIAN_PLUGIN_ATTRIBUTION_TIMEOUT, - ); + let attribution_deadline = + deadline.cap(Instant::now() + GUARDIAN_PLUGIN_ATTRIBUTION_TIMEOUT); let attribution = tokio::select! { biased; _ = cancellation.cancelled() => return Err(ReviewDecision::Abort), @@ -190,6 +189,7 @@ impl ReviewHost for super::super::runtime::ReviewRuntime { action: guardian_assessment_action(&request), review_reason, model_context, + timeout: turn.config.guardian_review_timeout, }); let root_history = session .services @@ -211,7 +211,7 @@ impl ReviewHost for super::super::runtime::ReviewRuntime { async fn attempt( &self, prepared: &PreparedApproval, - deadline: Instant, + deadline: ReviewDeadline, cancellation: &CancellationToken, ) -> ( GuardianReviewOutcome, diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 3f2fff3f2f..095233c642 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -10,6 +10,7 @@ pub(crate) use setup::run_guardian_review_session; #[path = "review_session_context.rs"] mod context_policy; +use codex_guardian_reviewer::ReviewDeadline; use context_policy::ReviewContextPolicy; use std::borrow::Cow; @@ -117,7 +118,7 @@ pub(crate) struct GuardianReviewSessionParams { pub(crate) reasoning_summary: ReasoningSummaryConfig, pub(crate) personality: Option, pub(crate) external_cancel: Option, - pub(crate) deadline: tokio::time::Instant, + pub(crate) deadline: ReviewDeadline, } pub(crate) type GuardianReviewSessionManager = @@ -331,7 +332,7 @@ async fn run_review_on_session( review_session: &GuardianReviewSession, params: &GuardianReviewSessionParams, guardian_session_kind: GuardianReviewSessionKind, - deadline: tokio::time::Instant, + deadline: ReviewDeadline, ) -> ReviewSessionResult { let review_model = ¶ms.review_model; let model_info = params diff --git a/codex-rs/core/src/guardian/review_session_setup.rs b/codex-rs/core/src/guardian/review_session_setup.rs index 39052af7cd..163f6b60e4 100644 --- a/codex-rs/core/src/guardian/review_session_setup.rs +++ b/codex-rs/core/src/guardian/review_session_setup.rs @@ -4,6 +4,7 @@ //! Checkpoints are captured before final authorization validation and published without yielding. //! Recorded reviewer items and interruptions are forwarded to the parent's raw-event subscription. +use codex_guardian_reviewer::ReviewDeadline; use std::sync::atomic::AtomicBool; use super::*; @@ -254,7 +255,7 @@ impl ReviewerRequest for PreparedReview { fn requires_fresh_session(&self) -> bool { self.context.reviewer_selection == ReviewerSelection::FreshParentCheckpoint } - fn deadline(&self) -> tokio::time::Instant { + fn deadline(&self) -> ReviewDeadline { self.params.deadline } fn cancellation(&self) -> Option<&CancellationToken> { diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index a9dfc53949..6e71ad6b91 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -261,7 +261,9 @@ async fn test_review_params() -> GuardianReviewSessionParams { reasoning_summary, personality, external_cancel: None, - deadline: tokio::time::Instant::now() + Duration::from_secs(30), + deadline: codex_guardian_reviewer::ReviewDeadline::At( + tokio::time::Instant::now() + Duration::from_secs(30), + ), } } @@ -727,7 +729,9 @@ async fn run_review_on_reused_session_waits_for_submitted_turn() { &review_session, ¶ms, GuardianReviewSessionKind::TrunkReused, - tokio::time::Instant::now() + Duration::from_secs(1), + codex_guardian_reviewer::ReviewDeadline::At( + tokio::time::Instant::now() + Duration::from_secs(1), + ), ) .await }); @@ -849,7 +853,9 @@ async fn wait_for_guardian_review_ignores_prior_turn_errors() { } = wait_for_guardian_review( &review_session, "current-turn", - tokio::time::Instant::now() + Duration::from_secs(1), + codex_guardian_reviewer::ReviewDeadline::At( + tokio::time::Instant::now() + Duration::from_secs(1), + ), /*external_cancel*/ None, &mut analytics_result, ) @@ -895,7 +901,9 @@ async fn wait_for_guardian_review_preserves_structured_session_error() { } = wait_for_guardian_review( &review_session, "current-turn", - tokio::time::Instant::now() + Duration::from_secs(1), + codex_guardian_reviewer::ReviewDeadline::At( + tokio::time::Instant::now() + Duration::from_secs(1), + ), /*external_cancel*/ None, &mut analytics_result, ) @@ -933,7 +941,9 @@ async fn wait_for_guardian_review_ignores_prior_turn_aborts() { } = wait_for_guardian_review( &review_session, "current-turn", - tokio::time::Instant::now() + Duration::from_secs(1), + codex_guardian_reviewer::ReviewDeadline::At( + tokio::time::Instant::now() + Duration::from_secs(1), + ), /*external_cancel*/ None, &mut analytics_result, ) @@ -973,7 +983,9 @@ async fn wait_for_guardian_review_timeout_drains_expected_turn_after_stale_termi } = wait_for_guardian_review( &review_session, "current-turn", - tokio::time::Instant::now() + Duration::from_millis(10), + codex_guardian_reviewer::ReviewDeadline::At( + tokio::time::Instant::now() + Duration::from_millis(10), + ), /*external_cancel*/ None, &mut analytics_result, ) @@ -1014,7 +1026,9 @@ async fn wait_for_guardian_review_cancel_drains_expected_turn_after_stale_termin } = wait_for_guardian_review( &review_session, "current-turn", - tokio::time::Instant::now() + Duration::from_secs(1), + codex_guardian_reviewer::ReviewDeadline::At( + tokio::time::Instant::now() + Duration::from_secs(1), + ), Some(&external_cancel), &mut analytics_result, ) @@ -1045,7 +1059,7 @@ async fn interrupt_and_drain_turn_ignores_prior_turn_completion() { let result = wait_for_guardian_review( &review_session, "current-turn", - tokio::time::Instant::now(), + codex_guardian_reviewer::ReviewDeadline::At(tokio::time::Instant::now()), Some(&cancellation), &mut GuardianReviewAnalyticsResult::without_session(), ) diff --git a/codex-rs/core/tests/suite/guardian_retry.rs b/codex-rs/core/tests/suite/guardian_retry.rs index 09af5467ed..e6eeac038b 100644 --- a/codex-rs/core/tests/suite/guardian_retry.rs +++ b/codex-rs/core/tests/suite/guardian_retry.rs @@ -31,11 +31,14 @@ enum ReviewInterruption { Timeout, } -#[test_case(ReviewInterruption::Cancellation; "cancellation")] -#[test_case(ReviewInterruption::Timeout; "timeout")] +#[test_case(ReviewInterruption::Cancellation, None; "cancellation")] +#[test_case(ReviewInterruption::Cancellation, Some(0); "unlimited_cancellation")] +#[test_case(ReviewInterruption::Timeout, None; "timeout")] +#[test_case(ReviewInterruption::Timeout, Some(180_000); "configured_timeout")] #[tokio::test] async fn interrupted_guardian_review_preserves_evidence_once( interruption: ReviewInterruption, + timeout_ms: Option, ) -> Result<()> { skip_if_no_network!(Ok(())); skip_if_wine_exec!( @@ -93,9 +96,18 @@ async fn interrupted_guardian_review_preserves_evidence_once( let server = start_mock_server().await; let base_url = format!("{}/v1", model.uri()); let test = test_codex() + .with_pre_build_hook(move |home| { + if let Some(timeout_ms) = timeout_ms { + std::fs::write( + home.join("config.toml"), + format!("[auto_review]\ntimeout_ms = {timeout_ms}\n"), + ) + .expect("write timeout config"); + } + }) .with_config(move |config| { config.model_provider.base_url = Some(base_url); - config.model_provider.stream_idle_timeout_ms = Some(180_000); + config.model_provider.stream_idle_timeout_ms = Some(600_000); config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); config.approvals_reviewer = ApprovalsReviewer::AutoReview; config @@ -136,12 +148,20 @@ async fn interrupted_guardian_review_preserves_evidence_once( .context("wait for pending Guardian review")?; match interruption { ReviewInterruption::Cancellation => { + if timeout_ms == Some(0) { + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(120)).await; + tokio::time::resume(); + } test.codex.submit(Op::Interrupt).await?; } ReviewInterruption::Timeout => { tokio::time::pause(); - tokio::time::advance(codex_guardian_reviewer::REVIEW_TIMEOUT + Duration::from_secs(1)) - .await; + tokio::time::advance( + Duration::from_millis(u64::from(timeout_ms.unwrap_or(90_000))) + + Duration::from_secs(1), + ) + .await; tokio::time::resume(); } } diff --git a/codex-rs/ext/guardian-reviewer/src/deadline.rs b/codex-rs/ext/guardian-reviewer/src/deadline.rs index c6f63f7180..898f435e94 100644 --- a/codex-rs/ext/guardian-reviewer/src/deadline.rs +++ b/codex-rs/ext/guardian-reviewer/src/deadline.rs @@ -1,16 +1,57 @@ -//! Bounds session startup and context work by the review's existing deadline and cancellation. +//! Bounds session startup and context work when a review deadline is configured. +//! Unlimited reviews still honor cancellation at every stage. use crate::GuardianReviewSessionOutcome; use std::future::Future; +use std::time::Duration; +use tokio::time::Instant; use tokio_util::sync::CancellationToken; +/// One shared review deadline, optionally disabled without disabling cancellation. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReviewDeadline { + At(Instant), + Unlimited, +} + +impl ReviewDeadline { + pub fn from_timeout(timeout: Option) -> Self { + match timeout { + Some(timeout) => Self::At(Instant::now() + timeout), + None => Self::Unlimited, + } + } + + pub(crate) fn is_elapsed(self) -> bool { + match self { + Self::At(deadline) => Instant::now() >= deadline, + Self::Unlimited => false, + } + } + + /// Retain local bounds, such as retry backoff and attribution lookup timeouts. + pub fn cap(self, local_deadline: Instant) -> Instant { + match self { + Self::At(deadline) => deadline.min(local_deadline), + Self::Unlimited => local_deadline, + } + } + + pub async fn wait(self) { + match self { + Self::At(deadline) => tokio::time::sleep_until(deadline).await, + Self::Unlimited => std::future::pending::<()>().await, + } + } +} + pub async fn run_before_review_deadline( - deadline: tokio::time::Instant, + deadline: ReviewDeadline, external_cancel: Option<&CancellationToken>, future: impl Future, ) -> Result { tokio::select! { - _ = tokio::time::sleep_until(deadline) => Err(GuardianReviewSessionOutcome::TimedOut), + _ = deadline.wait() => Err(GuardianReviewSessionOutcome::TimedOut), result = future => Ok(result), _ = async { if let Some(cancel_token) = external_cancel { @@ -21,3 +62,7 @@ pub async fn run_before_review_deadline( } => Err(GuardianReviewSessionOutcome::Aborted), } } + +#[cfg(test)] +#[path = "deadline_tests.rs"] +mod tests; diff --git a/codex-rs/ext/guardian-reviewer/src/deadline_tests.rs b/codex-rs/ext/guardian-reviewer/src/deadline_tests.rs new file mode 100644 index 0000000000..934c6610dc --- /dev/null +++ b/codex-rs/ext/guardian-reviewer/src/deadline_tests.rs @@ -0,0 +1,38 @@ +//! Checks overall review budgets without sleeping in wall-clock time. + +use super::*; + +#[tokio::test(start_paused = true)] +async fn setup_uses_remaining_budget() { + let deadline = ReviewDeadline::from_timeout(Some(Duration::from_secs(10))); + run_before_review_deadline(deadline, /*external_cancel*/ None, async { + tokio::time::sleep(Duration::from_secs(6)).await; + }) + .await + .expect("first stage fits"); + let result = run_before_review_deadline(deadline, /*external_cancel*/ None, async { + tokio::time::sleep(Duration::from_secs(6)).await; + }) + .await; + assert!(matches!( + result, + Err(GuardianReviewSessionOutcome::TimedOut) + )); +} + +#[tokio::test(start_paused = true)] +async fn unlimited_setup_can_be_cancelled() { + let cancel = CancellationToken::new(); + let (result, ()) = tokio::join!( + run_before_review_deadline( + ReviewDeadline::Unlimited, + Some(&cancel), + std::future::pending::<()>() + ), + async { + tokio::time::sleep(Duration::from_secs(120)).await; + cancel.cancel(); + }, + ); + assert!(matches!(result, Err(GuardianReviewSessionOutcome::Aborted))); +} diff --git a/codex-rs/ext/guardian-reviewer/src/execution.rs b/codex-rs/ext/guardian-reviewer/src/execution.rs index 7d1653d426..5e1f553e46 100644 --- a/codex-rs/ext/guardian-reviewer/src/execution.rs +++ b/codex-rs/ext/guardian-reviewer/src/execution.rs @@ -2,6 +2,7 @@ //! A session is reusable only after its submitted turn has completed or drained. //! Interrupted turns roll back admitted context; completed inference errors can retain it. +use crate::ReviewDeadline; use std::future::Future; use std::time::Duration; @@ -38,7 +39,7 @@ pub trait ReviewerRuntime: Sync { pub async fn start_review_turn( runtime: &impl ReviewerRuntime, request: TurnInputRequest, - deadline: Instant, + deadline: ReviewDeadline, external_cancel: Option<&CancellationToken>, ) -> Result { match crate::run_before_review_deadline(deadline, external_cancel, runtime.submit_turn(request)) @@ -67,11 +68,11 @@ pub struct ReviewTurnResult { pub async fn wait_for_guardian_review( runtime: &impl ReviewerRuntime, expected_turn_id: &str, - deadline: tokio::time::Instant, + deadline: ReviewDeadline, external_cancel: Option<&CancellationToken>, analytics_result: &mut GuardianReviewAnalyticsResult, ) -> ReviewTurnResult { - let timeout = tokio::time::sleep_until(deadline); + let timeout = deadline.wait(); tokio::pin!(timeout); let mut last_error: Option = None; diff --git a/codex-rs/ext/guardian-reviewer/src/lib.rs b/codex-rs/ext/guardian-reviewer/src/lib.rs index 94b078f738..50e8c8fcce 100644 --- a/codex-rs/ext/guardian-reviewer/src/lib.rs +++ b/codex-rs/ext/guardian-reviewer/src/lib.rs @@ -41,8 +41,10 @@ pub use retry::GuardianReviewSessionLimits; pub use retry::run_with_retry; pub const MAX_REVIEW_ATTEMPTS: i64 = 3; +/// Default overall review budget when `auto_review.timeout_ms` is omitted. pub const REVIEW_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90); +pub use deadline::ReviewDeadline; pub use deadline::run_before_review_deadline; pub use pool::ReviewSessionResult; pub use pool::ReviewerPool; diff --git a/codex-rs/ext/guardian-reviewer/src/metrics.rs b/codex-rs/ext/guardian-reviewer/src/metrics.rs index fdb857508d..54dbc52024 100644 --- a/codex-rs/ext/guardian-reviewer/src/metrics.rs +++ b/codex-rs/ext/guardian-reviewer/src/metrics.rs @@ -554,7 +554,9 @@ mod tests { let (_, result, _) = crate::run_with_retry( crate::GuardianReviewSessionLimits { max_attempts: 2, - deadline: start + Duration::from_secs(if outcome == "deadline" { 1 } else { 10 }), + deadline: crate::ReviewDeadline::At( + start + Duration::from_secs(if outcome == "deadline" { 1 } else { 10 }), + ), }, Some(&cancellation), Some(&telemetry), diff --git a/codex-rs/ext/guardian-reviewer/src/pool.rs b/codex-rs/ext/guardian-reviewer/src/pool.rs index dc31a7238e..6956f06b68 100644 --- a/codex-rs/ext/guardian-reviewer/src/pool.rs +++ b/codex-rs/ext/guardian-reviewer/src/pool.rs @@ -4,6 +4,7 @@ //! performs cleanup and tracks its completion. The pool never runs a second shutdown protocol. //! Startup and fork futures stay boxed to bound the orchestration stack frames. +use crate::ReviewDeadline; use std::future::Future; use std::sync::Arc; use std::sync::atomic::AtomicBool; @@ -14,7 +15,6 @@ use codex_analytics::GuardianReviewSessionKind; use codex_extension_api::ExtensionFuture; use tokio::sync::Mutex; use tokio::sync::Semaphore; -use tokio::time::Instant; use tokio_util::sync::CancellationToken; use crate::GuardianReviewSessionOutcome; @@ -55,7 +55,7 @@ pub trait ReviewerRequest: Send + Sync { fn requires_fresh_session(&self) -> bool { false } - fn deadline(&self) -> Instant; + fn deadline(&self) -> ReviewDeadline; fn cancellation(&self) -> Option<&CancellationToken>; fn run( &self, diff --git a/codex-rs/ext/guardian-reviewer/src/pool_tests.rs b/codex-rs/ext/guardian-reviewer/src/pool_tests.rs index c392e2301f..7c1c6305eb 100644 --- a/codex-rs/ext/guardian-reviewer/src/pool_tests.rs +++ b/codex-rs/ext/guardian-reviewer/src/pool_tests.rs @@ -2,6 +2,7 @@ use super::*; use pretty_assertions::assert_eq; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use tokio::time::Instant; #[tokio::test] async fn tracked_review_stays_reserved_until_cleanup_finishes() { @@ -109,8 +110,8 @@ impl ReviewerRequest for Request { fn requires_fresh_session(&self) -> bool { self.fresh } - fn deadline(&self) -> Instant { - Instant::now() + std::time::Duration::from_secs(5) + fn deadline(&self) -> crate::ReviewDeadline { + crate::ReviewDeadline::At(Instant::now() + std::time::Duration::from_secs(5)) } fn cancellation(&self) -> Option<&CancellationToken> { None diff --git a/codex-rs/ext/guardian-reviewer/src/reporting.rs b/codex-rs/ext/guardian-reviewer/src/reporting.rs index 30fc49c9ba..d2af080dad 100644 --- a/codex-rs/ext/guardian-reviewer/src/reporting.rs +++ b/codex-rs/ext/guardian-reviewer/src/reporting.rs @@ -42,6 +42,8 @@ pub struct ReviewMetadata { pub action: GuardianAssessmentAction, pub review_reason: GuardianReviewReason, pub model_context: ModelInvocationContext, + /// Configured review budget; None is reported as zero (unlimited). + pub timeout: Option, } pub struct ReviewReport { @@ -62,7 +64,9 @@ impl ReviewReport { metadata.target_item_id.clone(), metadata.approval_request_source, metadata.reviewed_action.clone(), - crate::REVIEW_TIMEOUT.as_millis() as u64, + metadata + .timeout + .map_or(0, |timeout| timeout.as_millis() as u64), ); let started = GuardianAssessmentEvent { review_reason: Some(metadata.review_reason), diff --git a/codex-rs/ext/guardian-reviewer/src/reporting_tests.rs b/codex-rs/ext/guardian-reviewer/src/reporting_tests.rs index 5bb364eab3..5b3eece1c1 100644 --- a/codex-rs/ext/guardian-reviewer/src/reporting_tests.rs +++ b/codex-rs/ext/guardian-reviewer/src/reporting_tests.rs @@ -55,6 +55,7 @@ fn review_completion_uses_one_monotonic_duration() { port: 443, }, review_reason: GuardianReviewReason::Policy, + timeout: Some(crate::REVIEW_TIMEOUT), model_context: ModelInvocationContext { model_slug: "test-model".to_string(), reasoning_effort: None, diff --git a/codex-rs/ext/guardian-reviewer/src/retry.rs b/codex-rs/ext/guardian-reviewer/src/retry.rs index 5740cda985..579f3349fe 100644 --- a/codex-rs/ext/guardian-reviewer/src/retry.rs +++ b/codex-rs/ext/guardian-reviewer/src/retry.rs @@ -4,6 +4,7 @@ use crate::GuardianReviewError; use crate::GuardianReviewOutcome; +use crate::ReviewDeadline; use codex_analytics::GuardianReviewAnalyticsResult; use codex_protocol::protocol::CodexErrorInfo; use rand::Rng; @@ -15,7 +16,7 @@ use tokio_util::sync::CancellationToken; pub struct GuardianReviewSessionLimits { pub max_attempts: i64, - pub deadline: Instant, + pub deadline: ReviewDeadline, } /// Retries only recoverable review failures within one shared deadline. @@ -31,7 +32,7 @@ pub async fn run_with_retry( Option, ) where - F: FnMut(Instant) -> Attempt, + F: FnMut(ReviewDeadline) -> Attempt, Attempt: Future< Output = ( GuardianReviewOutcome, @@ -107,7 +108,7 @@ async fn wait_before_guardian_retry( max_attempts: i64, outcome: &GuardianReviewOutcome, retry_not_before: Option, - deadline: Instant, + deadline: ReviewDeadline, external_cancel: Option<&CancellationToken>, ) -> Option { if matches!( @@ -119,7 +120,7 @@ async fn wait_before_guardian_retry( tokio::task::yield_now().await; return if external_cancel.is_some_and(CancellationToken::is_cancelled) { Some(GuardianReviewError::Cancelled) - } else if Instant::now() >= deadline { + } else if deadline.is_elapsed() { Some(GuardianReviewError::Timeout) } else { None @@ -130,7 +131,7 @@ async fn wait_before_guardian_retry( let retry_delay = Duration::from_millis((exponential_delay * jitter) as u64); let now = Instant::now(); let requested_retry_at = (now + retry_delay).max(retry_not_before.unwrap_or(now)); - let retry_at = requested_retry_at.min(deadline); + let retry_at = deadline.cap(requested_retry_at); let (failure_reason, error_info) = match outcome { GuardianReviewOutcome::Error(error) => ( Some(error.failure_reason()), @@ -152,13 +153,16 @@ async fn wait_before_guardian_retry( backoff_ms = retry_delay.as_millis() as u64, retry_not_before_ms = retry_not_before .map(|instant| instant.saturating_duration_since(now).as_millis() as u64), - remaining_budget_ms = deadline.saturating_duration_since(now).as_millis() as u64, - deadline_limited = requested_retry_at >= deadline, + remaining_budget_ms = match deadline { + ReviewDeadline::At(at) => Some(at.saturating_duration_since(now).as_millis() as u64), + ReviewDeadline::Unlimited => None, + }, + deadline_limited = matches!(deadline, ReviewDeadline::At(at) if requested_retry_at >= at), "Guardian review retry scheduled" ); tokio::select! { _ = sleep_until(retry_at) => { - (Instant::now() >= deadline).then_some(GuardianReviewError::Timeout) + deadline.is_elapsed().then_some(GuardianReviewError::Timeout) } _ = async { if let Some(cancel_token) = external_cancel { diff --git a/codex-rs/ext/guardian-reviewer/src/retry_tests.rs b/codex-rs/ext/guardian-reviewer/src/retry_tests.rs index c33aaff5ea..785b7e9e7a 100644 --- a/codex-rs/ext/guardian-reviewer/src/retry_tests.rs +++ b/codex-rs/ext/guardian-reviewer/src/retry_tests.rs @@ -130,8 +130,10 @@ fn guardian_review_retry_only_retries_recoverable_errors() { } } +#[test_case::test_case(Some(Duration::from_secs(1)); "finite")] +#[test_case::test_case(None; "unlimited")] #[tokio::test] -async fn guardian_review_retry_wait_honors_cancellation() { +async fn guardian_review_retry_wait_honors_cancellation(timeout: Option) { let cancel_token = CancellationToken::new(); cancel_token.cancel(); @@ -140,7 +142,7 @@ async fn guardian_review_retry_wait_honors_cancellation() { /*max_attempts*/ 3, &GuardianReviewOutcome::Error(GuardianReviewError::StaleAuthorization), /*retry_not_before*/ None, - Instant::now() + Duration::from_secs(/*secs*/ 1), + ReviewDeadline::from_timeout(timeout), Some(&cancel_token), ) .await; @@ -155,7 +157,7 @@ async fn guardian_review_retry_wait_honors_deadline() { /*max_attempts*/ 3, &GuardianReviewOutcome::Error(GuardianReviewError::StaleAuthorization), /*retry_not_before*/ None, - Instant::now(), + ReviewDeadline::At(Instant::now()), /*external_cancel*/ None, ) .await; @@ -163,13 +165,17 @@ async fn guardian_review_retry_wait_honors_deadline() { assert!(matches!(error, Some(GuardianReviewError::Timeout))); } -#[test_case::test_case(2; "recovers with fresh evidence")] -#[test_case::test_case(4; "exhausts the shared budget")] +#[test_case::test_case(2, Some(Duration::from_secs(90)); "recovers with fresh evidence")] +#[test_case::test_case(4, Some(Duration::from_secs(90)); "exhausts the shared budget")] +#[test_case::test_case(4, None; "unlimited exhausts attempt budget")] #[tokio::test(start_paused = true)] -async fn stale_authorization_uses_the_shared_attempt_budget(success_on_attempt: i64) { +async fn stale_authorization_uses_the_shared_attempt_budget( + success_on_attempt: i64, + timeout: Option, +) { let mut attempts = 0; let started_at = Instant::now(); - let deadline = started_at + Duration::from_secs(/*secs*/ 90); + let deadline = ReviewDeadline::from_timeout(timeout); let (outcome, analytics, evidence) = run_with_retry( GuardianReviewSessionLimits { max_attempts: 3, @@ -221,7 +227,7 @@ async fn stale_authorization_uses_the_shared_attempt_budget(success_on_attempt: #[tokio::test] async fn stale_authorization_does_not_extend_the_deadline() { - let deadline = Instant::now() + Duration::from_secs(/*secs*/ 1); + let deadline = ReviewDeadline::At(Instant::now() + Duration::from_secs(/*secs*/ 1)); let mut attempts = 0; let (outcome, analytics, _) = run_with_retry( GuardianReviewSessionLimits { @@ -234,7 +240,7 @@ async fn stale_authorization_does_not_extend_the_deadline() { attempts += 1; assert_eq!(attempt_deadline, deadline); async move { - tokio::time::sleep_until(deadline).await; + deadline.wait().await; ( GuardianReviewOutcome::Error(GuardianReviewError::StaleAuthorization), GuardianReviewAnalyticsResult::without_session(), diff --git a/codex-rs/ext/guardian-reviewer/src/review.rs b/codex-rs/ext/guardian-reviewer/src/review.rs index 7d7787cd6c..aed7b0e65e 100644 --- a/codex-rs/ext/guardian-reviewer/src/review.rs +++ b/codex-rs/ext/guardian-reviewer/src/review.rs @@ -5,6 +5,7 @@ use crate::GuardianReviewError; use crate::GuardianReviewOutcome; use crate::GuardianReviewSessionLimits; +use crate::ReviewDeadline; use crate::ReviewDenials; use crate::ReviewReport; use crate::ReviewRequest; @@ -24,7 +25,6 @@ use std::future::Future; use std::sync::Arc; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; -use tokio::time::Instant; use tokio_util::sync::CancellationToken; /// Operations bound to one immutable action and its issuing context. @@ -46,14 +46,14 @@ pub trait ReviewHost: Clone + Send + Sync + 'static { &self, approval_id: &str, reason: GuardianReviewReason, - deadline: Instant, + deadline: ReviewDeadline, cancellation: &CancellationToken, ) -> impl Future> + Send; /// Captures fresh authorization evidence and rejects stale approvals for each attempt. fn attempt( &self, prepared: &Self::Prepared, - deadline: Instant, + deadline: ReviewDeadline, cancellation: &CancellationToken, ) -> impl Future< Output = ( @@ -106,7 +106,7 @@ impl SynchronousApprovalReviewer for ReviewRequest<'_, H> { &[("route", "required_sync"), ("reason", reason.into())], ); } - let deadline = Instant::now() + crate::REVIEW_TIMEOUT; + let deadline = ReviewDeadline::from_timeout(self.timeout); let (prepared, report) = match self .host .prepare(self.approval_id, reason, deadline, &self.cancellation) @@ -166,7 +166,7 @@ impl SynchronousApprovalReviewer for ReviewRequest<'_, H> { let outcome = if !self.host.cached_approval_is_current(&prepared).await || self.cancellation.is_cancelled() { GuardianReviewOutcome::Error(GuardianReviewError::Cancelled) - } else if Instant::now() >= deadline { + } else if deadline.is_elapsed() { GuardianReviewOutcome::Error(GuardianReviewError::Timeout) } else { GuardianReviewOutcome::CachedApproval diff --git a/codex-rs/ext/guardian-reviewer/src/review_tests.rs b/codex-rs/ext/guardian-reviewer/src/review_tests.rs index f4660f048e..ae8cbd1fb3 100644 --- a/codex-rs/ext/guardian-reviewer/src/review_tests.rs +++ b/codex-rs/ext/guardian-reviewer/src/review_tests.rs @@ -50,7 +50,7 @@ impl ReviewHost for GatedHost { &self, approval_id: &str, reason: GuardianReviewReason, - _: Instant, + _: ReviewDeadline, _: &CancellationToken, ) -> Result<((), ReviewReport), ReviewDecision> { Ok(( @@ -69,6 +69,7 @@ impl ReviewHost for GatedHost { permissions: Default::default(), }, review_reason: reason, + timeout: Some(crate::REVIEW_TIMEOUT), model_context: ModelInvocationContext { model_slug: "test".into(), reasoning_effort: None, @@ -79,7 +80,7 @@ impl ReviewHost for GatedHost { async fn attempt( &self, _: &(), - _: Instant, + _: ReviewDeadline, cancellation: &CancellationToken, ) -> ( GuardianReviewOutcome, @@ -176,6 +177,7 @@ async fn async_approval_revalidates_without_waiting_for_tracked_cleanup(state: A escalated_exec: false, full_access: false, cancellation: CancellationToken::new(), + timeout: Some(crate::REVIEW_TIMEOUT), model: &model, telemetry: &telemetry, log_assessments: false, diff --git a/codex-rs/ext/guardian-reviewer/src/routing.rs b/codex-rs/ext/guardian-reviewer/src/routing.rs index 28d97774ba..557ea1ba7e 100644 --- a/codex-rs/ext/guardian-reviewer/src/routing.rs +++ b/codex-rs/ext/guardian-reviewer/src/routing.rs @@ -43,6 +43,8 @@ pub struct ReviewRequest<'a, H> { pub escalated_exec: bool, pub full_access: bool, pub cancellation: CancellationToken, + /// Overall review budget; None disables only the review deadline. + pub timeout: Option, pub model: &'a ModelInfo, pub telemetry: &'a SessionTelemetry, /// The host verified the assessment opt-in and an explicit OTLP log destination. diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index af50bd5f44..eff25f5926 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -274,6 +274,7 @@ async fn new_config( guardian_conversation_history_prompt: None, guardian_conversation_history_max_output_tokens: None, guardian_circuit_break_action: Default::default(), + guardian_review_timeout: Some(std::time::Duration::from_secs(90)), include_permissions_instructions: false, include_apps_instructions: false, include_collaboration_mode_instructions: false,