diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 5424355fbe..93c66e97c6 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3558,6 +3558,7 @@ version = "0.0.0" dependencies = [ "anyhow", "codex-analytics", + "codex-extension-api", "codex-protocol", "pretty_assertions", "rand 0.9.3", @@ -3578,6 +3579,7 @@ dependencies = [ "codex-extension-api", "codex-features", "codex-guardian-context", + "codex-guardian-reviewer", "codex-history", "codex-http-client", "codex-login", diff --git a/codex-rs/app-server/tests/suite/v2/guardian_sync_session_tests.rs b/codex-rs/app-server/tests/suite/v2/guardian_sync_session_tests.rs index a3dfa3b6ab..a226cd900b 100644 --- a/codex-rs/app-server/tests/suite/v2/guardian_sync_session_tests.rs +++ b/codex-rs/app-server/tests/suite/v2/guardian_sync_session_tests.rs @@ -16,10 +16,13 @@ use codex_app_server_protocol::ThreadReadResponse; use codex_app_server_protocol::ThreadSourceKind; use codex_app_server_protocol::ThreadUnsubscribeParams; use codex_app_server_protocol::ThreadUnsubscribeResponse; +use codex_app_server_protocol::TurnInterruptParams; +use codex_app_server_protocol::TurnInterruptResponse; use codex_protocol::protocol::SubAgentSource; use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; use pretty_assertions::assert_eq; +use test_case::test_case; use tokio::sync::oneshot; const ALLOW: &str = r#"{"outcome":"allow","rationale":"completed seed assessment"}"#; @@ -48,31 +51,44 @@ fn tool_response(server: &str, tool: &str, calls: &[&str]) -> String { responses::sse(events) } +#[derive(Clone, Copy)] +enum ReviewEnding { + Complete, + Interrupt, +} + +#[test_case(ReviewEnding::Complete; "completed reviews")] +#[test_case(ReviewEnding::Interrupt; "cancelled concurrent reviews")] #[tokio::test] -async fn managed_reviewers_reuse_fork_and_resume_after_parent_shutdown() -> Result<()> { +async fn managed_reviewers_reuse_fork_and_resume_after_parent_shutdown( + ending: ReviewEnding, +) -> Result<()> { skip_if_no_network!(Ok(())); let (continue_parent, parent_gate) = oneshot::channel(); let (finish_first, first_gate) = oneshot::channel(); let (finish_second, second_gate) = oneshot::channel(); + let mut replies = vec![ + ( + tool_response(TEST_SERVER_NAME, TEST_TOOL_NAME, &["seed"]), + None, + ), + (assistant_response(ALLOW)?, None), + ( + tool_response(TEST_SERVER_NAME, TEST_TOOL_NAME, &["first", "second"]), + Some(parent_gate), + ), + (assistant_response(ALLOW)?, Some(first_gate)), + (assistant_response(ALLOW)?, Some(second_gate)), + ]; + if matches!(ending, ReviewEnding::Complete) { + replies.push((assistant_response("Done.")?, None)); + } + replies.push((assistant_response("The user authorized the checks.")?, None)); let (server, _completions) = start_streaming_sse_server( - vec![ - ( - tool_response(TEST_SERVER_NAME, TEST_TOOL_NAME, &["seed"]), - None, - ), - (assistant_response(ALLOW)?, None), - ( - tool_response(TEST_SERVER_NAME, TEST_TOOL_NAME, &["first", "second"]), - Some(parent_gate), - ), - (assistant_response(ALLOW)?, Some(first_gate)), - (assistant_response(ALLOW)?, Some(second_gate)), - (assistant_response("Done.")?, None), - (assistant_response("The user authorized the checks.")?, None), - ] - .into_iter() - .map(|(body, gate)| vec![StreamingSseChunk { gate, body }]) - .collect(), + replies + .into_iter() + .map(|(body, gate)| vec![StreamingSseChunk { gate, body }]) + .collect(), ) .await; let (mcp_url, mcp_server) = start_mcp_server(/*sensitive_action*/ None).await?; @@ -91,7 +107,7 @@ async fn managed_reviewers_reuse_fork_and_resume_after_parent_shutdown() -> Resu .build_initialized_with_timeout(TIMEOUT) .await?; let parent = app.start_thread(ThreadStartParams::default()).await?.thread; - let _: TurnStartResponse = app + let started: TurnStartResponse = app .request(|request_id| ClientRequest::TurnStart { request_id, params: TurnStartParams { @@ -184,11 +200,27 @@ async fn managed_reviewers_reuse_fork_and_resume_after_parent_shutdown() -> Resu .contains("completed seed assessment") ); } - finish_first.send(()).expect("finish first review"); - finish_second.send(()).expect("finish second review"); + if matches!(ending, ReviewEnding::Interrupt) { + let _: TurnInterruptResponse = app + .request(|request_id| ClientRequest::TurnInterrupt { + request_id, + params: TurnInterruptParams { + thread_id: parent.id.clone(), + turn_id: started.turn.id.clone(), + }, + }) + .await?; + } else { + finish_first.send(()).expect("finish first review"); + finish_second.send(()).expect("finish second review"); + } let completed: TurnCompletedNotification = timeout(TIMEOUT, app.read_notification("turn/completed")).await??; - assert_eq!(completed.turn.status, TurnStatus::Completed); + let expected_status = match ending { + ReviewEnding::Complete => TurnStatus::Completed, + ReviewEnding::Interrupt => TurnStatus::Interrupted, + }; + assert_eq!(completed.turn.status, expected_status); let read: ThreadReadResponse = app .request(|request_id| ClientRequest::ThreadRead { @@ -246,7 +278,13 @@ async fn managed_reviewers_reuse_fork_and_resume_after_parent_shutdown() -> Resu ) .await??; assert_eq!(completed.turn.status, TurnStatus::Completed); - assert_eq!(server.requests().await.len(), 7); + assert_eq!( + server.requests().await.len(), + match ending { + ReviewEnding::Complete => 7, + ReviewEnding::Interrupt => 6, + } + ); // Once the owner releases it, the resumed reviewer follows normal client removal. for method in ["thread/archive", "thread/delete"] { let id = app diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 37dc635aa0..3491aa30d2 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -85,9 +85,9 @@ pub(crate) async fn run_codex_thread_interactive( Arc::clone(&parent_session.services.extensions) }; // Inline delegates never register with ThreadManager or receive on_thread_ready. - // Seed their standalone Guardian manager before inherited extensions run. + // Bind their standalone spawn path before inherited extensions run. let mut thread_extension_init = codex_extension_api::ExtensionDataInit::default(); - thread_extension_init.insert(crate::guardian::GuardianReviewSessionManager::default()); + thread_extension_init.insert(crate::guardian::GuardianReviewSessionHost::default()); let (session, io) = Session::spawn(SessionSpawnArgs { config, allow_provider_model_fallback: false, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 99185e498a..a630749a4b 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -684,7 +684,9 @@ impl CodexThread { pub async fn guardian_trunk_rollout_path(&self) -> Option { self.session .guardian_review_session() - .trunk_rollout_path() + .trunk() + .await? + .rollout_path() .await } diff --git a/codex-rs/core/src/guardian/decision.rs b/codex-rs/core/src/guardian/decision.rs index aa3d039cde..fe5329d5ae 100644 --- a/codex-rs/core/src/guardian/decision.rs +++ b/codex-rs/core/src/guardian/decision.rs @@ -109,14 +109,14 @@ pub(crate) async fn decide_approval( )); } }; - let runtime = ReviewRuntime { + let runtime = codex_guardian_reviewer::SynchronousReview::new(ReviewRuntime { session: Arc::clone(&session), context: context.clone(), review_id: review_id.clone(), request: request.clone(), reasons, options, - }; + }); let input = ApprovalDecisionInput { approval_id: &review_id, tool_call_id: request diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 99d0713ae3..32a3f2f5e3 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -35,20 +35,23 @@ pub(crate) use approval_request::GuardianMcpAnnotations; pub(crate) use approval_request::GuardianNetworkAccessTrigger; #[cfg(test)] pub(crate) use approval_request::guardian_approval_request_to_json; +pub(crate) use codex_guardian_reviewer::guardian_timeout_message; pub(crate) use decision::decide_approval; pub(crate) use decision::spawn_approval_decision; pub(crate) use prompt::BUNDLED_GUARDIAN_POLICY; pub(crate) use prompt::BUNDLED_GUARDIAN_POLICY_TEMPLATE; pub(crate) use prompt::guardian_truncate_text; pub(crate) use review::GuardianReviewOptions; -pub(crate) use review::guardian_timeout_message; pub(crate) use review::is_basic_session_source; 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::GuardianReviewSessionManager; +pub use review_session::GuardianReviewSession; +pub use review_session::GuardianReviewSessionHost; +pub(crate) use review_session::GuardianReviewSessionManager; +pub(crate) use review_session::prewarm_guardian_review_session; pub(crate) use review_session::prompt_cache_key_override_for_review_session; pub(crate) use runtime::ReviewAction; @@ -160,10 +163,8 @@ impl From<&Arc> for GuardianReviewContext { } } -pub use codex_guardian_reviewer::GuardianAssessment; -pub use codex_guardian_reviewer::guardian_output_schema; -pub use codex_guardian_reviewer::parse_guardian_assessment; -pub use reviewer_config::build_guardian_review_session_config; +#[cfg(test)] +use codex_guardian_reviewer::guardian_output_schema; pub(crate) use codex_guardian_reviewer::GuardianRejectionCircuitBreaker; pub(crate) use codex_guardian_reviewer::GuardianRejectionCircuitBreakerAction; diff --git a/codex-rs/core/src/guardian/prompt.rs b/codex-rs/core/src/guardian/prompt.rs index f22a3bcafd..ca89fb4b3a 100644 --- a/codex-rs/core/src/guardian/prompt.rs +++ b/codex-rs/core/src/guardian/prompt.rs @@ -363,7 +363,7 @@ pub(crate) fn guardian_truncate_text(content: &str, token_cap: usize) -> (String } use codex_guardian_reviewer::guardian_output_contract_prompt; -pub use codex_guardian_reviewer::parse_guardian_assessment; +pub(super) use codex_guardian_reviewer::parse_guardian_assessment; pub(crate) const BUNDLED_GUARDIAN_POLICY: &str = include_str!("../../assets/guardian/policy.md"); pub(crate) const BUNDLED_GUARDIAN_POLICY_TEMPLATE: &str = diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 1520d30c42..840f031508 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -1,3 +1,9 @@ +//! Supplies host review preparation, telemetry and runtime configuration. +//! Guardian's extension owns the synchronous review loop and pool. + +#[path = "review_request.rs"] +mod request; + use crate::context::GuardianContextMode; use codex_analytics::GuardianApprovalRequestSource; use codex_analytics::GuardianReviewAnalyticsResult; @@ -11,17 +17,14 @@ use codex_extension_api::ThreadIdleCause; use codex_features::Feature; use codex_guardian_reviewer::GuardianReviewError; use codex_guardian_reviewer::GuardianReviewOutcome; +#[cfg(test)] use codex_guardian_reviewer::GuardianReviewSessionLimits; use codex_protocol::config_types::ApprovalsReviewer; -use codex_protocol::openai_models::MODEL_SPECIALTY_CYBER; -use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GuardianAssessmentDecisionSource; use codex_protocol::protocol::GuardianAssessmentEvent; use codex_protocol::protocol::GuardianAssessmentStatus; -use codex_protocol::protocol::GuardianRiskLevel; -use codex_protocol::protocol::GuardianUserAuthorization; use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::SessionSource; @@ -44,7 +47,6 @@ use super::ApprovalRequestReasons; use super::GUARDIAN_REVIEW_TIMEOUT; use super::GUARDIAN_REVIEWER_NAME; use super::GuardianApprovalRequest; -use super::GuardianAssessment; use super::GuardianAssessmentOutcome; use super::GuardianRejectionCircuitBreakerAction; use super::GuardianRejectionCircuitBreakerPolicy; @@ -59,20 +61,6 @@ use super::review_session::GuardianReviewSessionParams; use super::review_session::build_guardian_review_session_config; use codex_guardian_reviewer::guardian_output_schema; -const GUARDIAN_REJECTION_INSTRUCTIONS: &str = concat!( - "The agent must not attempt to achieve the same outcome via workaround, ", - "indirect execution, or policy circumvention. ", - "Proceed only with a materially safer alternative, ", - "or if the user explicitly approves the action after being informed of the risk. ", - "Otherwise, stop and request user input.", -); - -const GUARDIAN_TIMEOUT_INSTRUCTIONS: &str = concat!( - "The automatic permission approval review did not finish before its deadline. ", - "Do not assume the action is unsafe based on the timeout alone. ", - "You may retry once, or ask the user for guidance or explicit approval.", -); - const GUARDIAN_PLUGIN_ATTRIBUTION_TIMEOUT: Duration = Duration::from_secs(5); async fn plugin_attribution_for_guardian_request( @@ -125,25 +113,6 @@ pub(crate) fn new_guardian_review_id() -> String { uuid::Uuid::new_v4().to_string() } -pub(crate) fn guardian_timeout_message(model_info: &ModelInfo) -> String { - model_info - .model_messages - .as_ref() - .and_then(|messages| messages.auto_review.as_ref()) - .and_then(|messages| messages.timeout_instructions.as_deref()) - .unwrap_or(GUARDIAN_TIMEOUT_INSTRUCTIONS) - .to_string() -} - -fn guardian_risk_level_str(level: GuardianRiskLevel) -> &'static str { - match level { - GuardianRiskLevel::Low => "low", - GuardianRiskLevel::Medium => "medium", - GuardianRiskLevel::High => "high", - GuardianRiskLevel::Critical => "critical", - } -} - /// 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. @@ -209,11 +178,7 @@ pub(super) async fn record_guardian_non_denial(session: &Arc, turn_id: } async fn record_guardian_denial(session: &Arc, turn: &Arc, turn_id: &str) { - let policy = if turn.model_info().model_specialty.as_deref() == Some(MODEL_SPECIALTY_CYBER) { - GuardianRejectionCircuitBreakerPolicy::CyberModel - } else { - GuardianRejectionCircuitBreakerPolicy::Standard - }; + let policy = GuardianRejectionCircuitBreakerPolicy::from(turn.model_info().as_ref()); let action = session .services .guardian_rejection_circuit_breaker @@ -269,457 +234,6 @@ pub(crate) async fn record_guardian_denial_for_test( record_guardian_denial(session, turn, turn_id).await; } -/// Runs the existing synchronous reviewer without choosing policy or reusing a score. -pub(super) async fn run_synchronous_review( - runtime: super::runtime::ReviewRuntime, - review_reason: codex_protocol::approvals::GuardianReviewReason, -) -> ReviewDecision { - let super::runtime::ReviewRuntime { - session, - context, - review_id, - request, - reasons, - options, - } = runtime; - let request = match request.validate(&context) { - Ok(request) => request.clone(), - Err(decision) => return decision, - }; - let turn = Arc::clone(context.turn()); - let GuardianReviewOptions { - plugin_attribution_override, - approval_request_source, - external_cancel, - require_synchronous_review: _, - require_guardian: _, - } = options; - // Bound executor attribution and the Guardian session with one review deadline. - let deadline = Instant::now() + GUARDIAN_REVIEW_TIMEOUT; - let target_item_id = guardian_request_target_item_id(&request).map(str::to_string); - let assessment_turn_id = guardian_request_turn_id(&request, &turn.sub_id).to_string(); - let plugin_attribution = match plugin_attribution_override { - Some(attribution) => Some(attribution), - None if matches!(&request, GuardianApprovalRequest::ExecCommand { .. }) => { - let cancellation = external_cancel - .clone() - .unwrap_or_else(CancellationToken::new); - let attribution_deadline = std::cmp::min( - deadline, - Instant::now() + GUARDIAN_PLUGIN_ATTRIBUTION_TIMEOUT, - ); - let attribution = tokio::select! { - biased; - _ = cancellation.cancelled() => return ReviewDecision::Abort, - attribution = tokio::time::timeout_at( - attribution_deadline, - plugin_attribution_for_guardian_request(&context, &request), - ) => attribution, - }; - match attribution { - Ok(attribution) => attribution, - Err(_) => { - tracing::warn!( - timeout_ms = GUARDIAN_PLUGIN_ATTRIBUTION_TIMEOUT.as_millis(), - "Guardian plugin attribution timed out" - ); - None - } - } - } - None => plugin_attribution_for_guardian_request(&context, &request).await, - }; - let (plugin_id, script_path) = plugin_attribution - .as_ref() - .map(PluginCommandAttribution::serialized_fields) - .unzip(); - let action_summary = guardian_assessment_action(&request); - let reviewed_action = guardian_reviewed_action(&request); - let review_tracking = GuardianReviewTrackContext::new( - session.thread_id.to_string(), - assessment_turn_id.clone(), - review_id.clone(), - target_item_id.clone(), - approval_request_source, - reviewed_action.clone(), - GUARDIAN_REVIEW_TIMEOUT.as_millis() as u64, - ); - let started_at_ms = review_tracking.started_at_ms.try_into().unwrap_or_default(); - session - .send_event( - turn.as_ref(), - EventMsg::GuardianAssessment(GuardianAssessmentEvent { - review_reason: Some(review_reason), - id: review_id.clone(), - target_item_id: target_item_id.clone(), - plugin_id: plugin_id.clone(), - script_path: script_path.clone(), - turn_id: assessment_turn_id.clone(), - started_at_ms, - completed_at_ms: None, - status: GuardianAssessmentStatus::InProgress, - risk_level: None, - user_authorization: None, - rationale: None, - decision_source: None, - action: action_summary.clone(), - }), - ) - .await; - - if external_cancel - .as_ref() - .is_some_and(CancellationToken::is_cancelled) - { - let completed_at_ms = now_unix_timestamp_ms(); - track_guardian_review( - session.as_ref(), - &review_tracking, - approval_request_source, - &reviewed_action, - GuardianReviewAnalyticsResult { - decision: GuardianReviewDecision::Aborted, - terminal_status: GuardianReviewTerminalStatus::Aborted, - failure_reason: Some(GuardianReviewFailureReason::Cancelled), - ..GuardianReviewAnalyticsResult::without_session() - }, - completed_at_ms.try_into().unwrap_or_default(), - ); - session - .send_event( - turn.as_ref(), - EventMsg::GuardianAssessment(GuardianAssessmentEvent { - review_reason: Some(review_reason), - id: review_id, - target_item_id, - plugin_id: plugin_id.clone(), - script_path: script_path.clone(), - turn_id: assessment_turn_id.clone(), - started_at_ms, - completed_at_ms: Some(completed_at_ms), - status: GuardianAssessmentStatus::Aborted, - risk_level: None, - user_authorization: None, - rationale: None, - decision_source: Some(GuardianAssessmentDecisionSource::Agent), - action: action_summary, - }), - ) - .await; - record_guardian_non_denial(&session, &assessment_turn_id).await; - return ReviewDecision::Abort; - } - - let schema = guardian_output_schema(); - let terminal_action = action_summary.clone(); - let root_authorization_version = session - .services - .agent_control - .root_user_authorization(session.thread_id) - .await - .map(|snapshot| snapshot.authorization_version); - // Keep the authorization revision even when no cacheable review evidence exists. - let history = session.conversation_history_snapshot().await; - let user_message_revision = history.user_message_revision(); - let review_evidence = if let Some(evidence) = session - .services - .thread_extension_data - .get::() - { - let authorization_version = evidence.authorization_version(history.as_ref()); - format_guardian_action_pretty(&request).ok().map(|action| { - ( - evidence, - action.text, - authorization_version, - root_authorization_version, - ) - }) - } else { - None - }; - drop(history); - let (mut outcome, analytics_result) = - Box::pin(run_guardian_review_session_with_retry_before_deadline( - session.clone(), - context, - request, - reasons, - schema, - external_cancel, - GuardianReviewSessionLimits { - max_attempts: codex_guardian_reviewer::MAX_REVIEW_ATTEMPTS, - deadline, - }, - )) - .await; - 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 - != session - .conversation_history_snapshot() - .await - .user_message_revision()) - { - // A completed approval cannot outlive the owning-session or root evidence - // it evaluated, including when either changed before prompt construction. - outcome = GuardianReviewOutcome::Error(GuardianReviewError::Cancelled); - } - - let completed_at_ms = now_unix_timestamp_ms(); - let completed_review = matches!(&outcome, GuardianReviewOutcome::Completed(_)); - let (assessment, count_denial_for_circuit_breaker) = match outcome { - GuardianReviewOutcome::Completed(assessment) => { - let approved = matches!(assessment.outcome, GuardianAssessmentOutcome::Allow); - track_guardian_review( - session.as_ref(), - &review_tracking, - approval_request_source, - &reviewed_action, - GuardianReviewAnalyticsResult { - decision: if approved { - GuardianReviewDecision::Approved - } else { - GuardianReviewDecision::Denied - }, - terminal_status: if approved { - GuardianReviewTerminalStatus::Approved - } else { - GuardianReviewTerminalStatus::Denied - }, - failure_reason: None, - risk_level: Some(assessment.risk_level), - user_authorization: Some(assessment.user_authorization), - outcome: Some(assessment.outcome), - ..analytics_result - }, - completed_at_ms.try_into().unwrap_or_default(), - ); - let count_denial_for_circuit_breaker = - matches!(assessment.outcome, GuardianAssessmentOutcome::Deny); - (assessment, count_denial_for_circuit_breaker) - } - GuardianReviewOutcome::Error(error) => match error { - GuardianReviewError::Timeout => { - let rationale = - "Automatic approval review timed out while evaluating the requested approval." - .to_string(); - track_guardian_review( - session.as_ref(), - &review_tracking, - approval_request_source, - &reviewed_action, - GuardianReviewAnalyticsResult { - decision: GuardianReviewDecision::Denied, - terminal_status: GuardianReviewTerminalStatus::TimedOut, - failure_reason: Some(error.failure_reason()), - ..analytics_result - }, - completed_at_ms.try_into().unwrap_or_default(), - ); - session - .send_event( - turn.as_ref(), - EventMsg::GuardianWarning(WarningEvent { - message: rationale.clone(), - }), - ) - .await; - session - .send_event( - turn.as_ref(), - EventMsg::GuardianAssessment(GuardianAssessmentEvent { - review_reason: Some(review_reason), - id: review_id, - target_item_id, - plugin_id: plugin_id.clone(), - script_path: script_path.clone(), - turn_id: assessment_turn_id.clone(), - started_at_ms, - completed_at_ms: Some(completed_at_ms), - status: GuardianAssessmentStatus::TimedOut, - risk_level: None, - user_authorization: None, - rationale: Some(rationale), - decision_source: Some(GuardianAssessmentDecisionSource::Agent), - action: terminal_action, - }), - ) - .await; - record_guardian_non_denial(&session, &assessment_turn_id).await; - return ReviewDecision::TimedOut; - } - GuardianReviewError::Cancelled => { - track_guardian_review( - session.as_ref(), - &review_tracking, - approval_request_source, - &reviewed_action, - GuardianReviewAnalyticsResult { - decision: GuardianReviewDecision::Aborted, - terminal_status: GuardianReviewTerminalStatus::Aborted, - failure_reason: Some(error.failure_reason()), - ..analytics_result - }, - completed_at_ms.try_into().unwrap_or_default(), - ); - session - .send_event( - turn.as_ref(), - EventMsg::GuardianAssessment(GuardianAssessmentEvent { - review_reason: Some(review_reason), - id: review_id, - target_item_id, - plugin_id: plugin_id.clone(), - script_path: script_path.clone(), - turn_id: assessment_turn_id.clone(), - started_at_ms, - completed_at_ms: Some(completed_at_ms), - status: GuardianAssessmentStatus::Aborted, - risk_level: None, - user_authorization: None, - rationale: None, - decision_source: Some(GuardianAssessmentDecisionSource::Agent), - action: action_summary, - }), - ) - .await; - record_guardian_non_denial(&session, &assessment_turn_id).await; - return ReviewDecision::Abort; - } - GuardianReviewError::PromptBuild { .. } - | GuardianReviewError::Session { .. } - | GuardianReviewError::Parse { .. } => { - let message = match &error { - GuardianReviewError::PromptBuild { message } - | GuardianReviewError::Session { message, .. } - | GuardianReviewError::Parse { message } => message, - GuardianReviewError::Timeout | GuardianReviewError::Cancelled => { - "guardian review failed" - } - }; - let rationale = format!("Automatic approval review failed: {message}"); - track_guardian_review( - session.as_ref(), - &review_tracking, - approval_request_source, - &reviewed_action, - GuardianReviewAnalyticsResult { - decision: GuardianReviewDecision::Denied, - terminal_status: GuardianReviewTerminalStatus::FailedClosed, - failure_reason: Some(error.failure_reason()), - ..analytics_result - }, - completed_at_ms.try_into().unwrap_or_default(), - ); - ( - GuardianAssessment { - risk_level: GuardianRiskLevel::High, - user_authorization: GuardianUserAuthorization::Unknown, - outcome: GuardianAssessmentOutcome::Deny, - rationale, - }, - false, - ) - } - }, - }; - - let approved = match assessment.outcome { - GuardianAssessmentOutcome::Allow => true, - GuardianAssessmentOutcome::Deny => false, - }; - let verdict = if approved { "approved" } else { "denied" }; - let user_authorization = match assessment.user_authorization { - GuardianUserAuthorization::Unknown => "unknown", - GuardianUserAuthorization::Low => "low", - GuardianUserAuthorization::Medium => "medium", - GuardianUserAuthorization::High => "high", - }; - let warning = format!( - "Automatic approval review {verdict} (risk: {}, authorization: {user_authorization}): {}", - guardian_risk_level_str(assessment.risk_level), - assessment.rationale - ); - session - .send_event( - turn.as_ref(), - EventMsg::GuardianWarning(WarningEvent { message: warning }), - ) - .await; - let status = if approved { - GuardianAssessmentStatus::Approved - } else { - GuardianAssessmentStatus::Denied - }; - let assessment_event = GuardianAssessmentEvent { - review_reason: Some(review_reason), - id: review_id, - target_item_id, - plugin_id: plugin_id.clone(), - script_path: script_path.clone(), - turn_id: assessment_turn_id.clone(), - started_at_ms, - completed_at_ms: Some(completed_at_ms), - status, - risk_level: Some(assessment.risk_level), - user_authorization: Some(assessment.user_authorization), - rationale: Some(assessment.rationale.clone()), - decision_source: Some(GuardianAssessmentDecisionSource::Agent), - action: terminal_action, - }; - if completed_review - && let Some((evidence, action, authorization_version, root_authorization_version)) = - review_evidence - { - evidence.record( - &assessment_event, - &action, - authorization_version, - root_authorization_version, - ); - } - session - .send_event( - turn.as_ref(), - EventMsg::GuardianAssessment(assessment_event), - ) - .await; - - if count_denial_for_circuit_breaker { - record_guardian_denial(&session, &turn, &assessment_turn_id).await; - } else { - record_guardian_non_denial(&session, &assessment_turn_id).await; - } - - if approved { - ReviewDecision::Approved - } else { - let rationale = if assessment.rationale.trim().is_empty() { - "Auto-reviewer denied the action without a specific rationale." - } else { - assessment.rationale.trim() - }; - let rejection_instructions = turn - .model_info() - .model_messages - .as_ref() - .and_then(|messages| messages.auto_review.as_ref()) - .and_then(|messages| messages.rejection_instructions.as_deref()) - .unwrap_or(GUARDIAN_REJECTION_INSTRUCTIONS); - ReviewDecision::denied(format!( - "This action was rejected due to unacceptable risk.\nReason: {rationale}\n{rejection_instructions}" - )) - } -} - #[derive(Clone)] pub(crate) struct GuardianReviewOptions { /// Requires Guardian rather than a manual approval; cached evidence may still satisfy it. @@ -853,10 +367,10 @@ async fn run_guardian_review_session_before_deadline( ); } }; - let (session_outcome, session_analytics_result) = Box::pin( - session - .guardian_review_session() - .run_review(GuardianReviewSessionParams { + let (session_outcome, session_analytics_result) = + Box::pin(super::review_session::run_guardian_review_session( + session.guardian_review_session(), + GuardianReviewSessionParams { parent_session: Arc::clone(&session), parent_context: context.clone(), parent_history: session.clone_history().await, @@ -876,9 +390,9 @@ async fn run_guardian_review_session_before_deadline( personality: turn.personality(), external_cancel, deadline, - }), - ) - .await; + }, + )) + .await; (session_outcome.into(), session_analytics_result) } @@ -908,6 +422,7 @@ pub(super) async fn run_guardian_review_session_with_retry( .await } +#[cfg(test)] async fn run_guardian_review_session_with_retry_before_deadline( session: Arc, context: impl Into, diff --git a/codex-rs/core/src/guardian/review_request.rs b/codex-rs/core/src/guardian/review_request.rs new file mode 100644 index 0000000000..478dbc9e08 --- /dev/null +++ b/codex-rs/core/src/guardian/review_request.rs @@ -0,0 +1,345 @@ +//! Captures and reports one review on the host's original action and authorization state. +//! Guardian's extension owns orchestration; this adapter preserves event and evidence behavior. + +use super::*; +use crate::codex_thread::GuardianAuthorizationVersion; +use codex_guardian_reviewer::ReviewHost; +use codex_protocol::approvals::GuardianAssessmentAction; +use codex_protocol::approvals::GuardianReviewReason; + +pub(in crate::guardian) struct PreparedApproval { + request: GuardianApprovalRequest, + turn: Arc, + review_reason: GuardianReviewReason, + assessment_turn_id: String, + target_item_id: Option, + plugin_id: Option, + script_path: Option, + action_summary: GuardianAssessmentAction, + reviewed_action: GuardianReviewedAction, + review_tracking: GuardianReviewTrackContext, + started_at_ms: i64, + root_authorization_version: Option, + user_message_revision: u64, + review_evidence: Option<( + Arc, + String, + GuardianAuthorizationVersion, + Option, + )>, +} + +impl ReviewHost for super::super::runtime::ReviewRuntime { + type Prepared = PreparedApproval; + + fn cancellation(&self) -> Option<&CancellationToken> { + self.options.external_cancel.as_ref() + } + + async fn prepare( + &self, + review_reason: GuardianReviewReason, + deadline: Instant, + ) -> Result { + let super::super::runtime::ReviewRuntime { + session, + context, + review_id, + request, + reasons: _, + options, + } = self.clone(); + let request = match request.validate(&context) { + Ok(request) => request.clone(), + Err(decision) => return Err(decision), + }; + let turn = Arc::clone(context.turn()); + let GuardianReviewOptions { + plugin_attribution_override, + approval_request_source, + external_cancel, + require_synchronous_review: _, + require_guardian: _, + } = options; + let target_item_id = guardian_request_target_item_id(&request).map(str::to_string); + let assessment_turn_id = guardian_request_turn_id(&request, &turn.sub_id).to_string(); + let plugin_attribution = match plugin_attribution_override { + Some(attribution) => Some(attribution), + None if matches!(&request, GuardianApprovalRequest::ExecCommand { .. }) => { + let cancellation = external_cancel + .clone() + .unwrap_or_else(CancellationToken::new); + let attribution_deadline = std::cmp::min( + deadline, + Instant::now() + GUARDIAN_PLUGIN_ATTRIBUTION_TIMEOUT, + ); + let attribution = tokio::select! { + biased; + _ = cancellation.cancelled() => return Err(ReviewDecision::Abort), + attribution = tokio::time::timeout_at( + attribution_deadline, + plugin_attribution_for_guardian_request(&context, &request), + ) => attribution, + }; + match attribution { + Ok(attribution) => attribution, + Err(_) => { + tracing::warn!( + timeout_ms = GUARDIAN_PLUGIN_ATTRIBUTION_TIMEOUT.as_millis(), + "Guardian plugin attribution timed out" + ); + None + } + } + } + None => plugin_attribution_for_guardian_request(&context, &request).await, + }; + let (plugin_id, script_path) = plugin_attribution + .as_ref() + .map(PluginCommandAttribution::serialized_fields) + .unzip(); + let action_summary = guardian_assessment_action(&request); + let reviewed_action = guardian_reviewed_action(&request); + let review_tracking = GuardianReviewTrackContext::new( + session.thread_id.to_string(), + assessment_turn_id.clone(), + review_id.clone(), + target_item_id.clone(), + approval_request_source, + reviewed_action.clone(), + GUARDIAN_REVIEW_TIMEOUT.as_millis() as u64, + ); + let started_at_ms = review_tracking.started_at_ms.try_into().unwrap_or_default(); + session + .send_event( + turn.as_ref(), + EventMsg::GuardianAssessment(GuardianAssessmentEvent { + review_reason: Some(review_reason), + id: review_id.clone(), + target_item_id: target_item_id.clone(), + plugin_id: plugin_id.clone(), + script_path: script_path.clone(), + turn_id: assessment_turn_id.clone(), + started_at_ms, + completed_at_ms: None, + status: GuardianAssessmentStatus::InProgress, + risk_level: None, + user_authorization: None, + rationale: None, + decision_source: None, + action: action_summary.clone(), + }), + ) + .await; + + if external_cancel + .as_ref() + .is_some_and(CancellationToken::is_cancelled) + { + let completed_at_ms = now_unix_timestamp_ms(); + track_guardian_review( + session.as_ref(), + &review_tracking, + approval_request_source, + &reviewed_action, + GuardianReviewAnalyticsResult { + decision: GuardianReviewDecision::Aborted, + terminal_status: GuardianReviewTerminalStatus::Aborted, + failure_reason: Some(GuardianReviewFailureReason::Cancelled), + ..GuardianReviewAnalyticsResult::without_session() + }, + completed_at_ms.try_into().unwrap_or_default(), + ); + session + .send_event( + turn.as_ref(), + EventMsg::GuardianAssessment(GuardianAssessmentEvent { + review_reason: Some(review_reason), + id: review_id, + target_item_id, + plugin_id: plugin_id.clone(), + script_path: script_path.clone(), + turn_id: assessment_turn_id.clone(), + started_at_ms, + completed_at_ms: Some(completed_at_ms), + status: GuardianAssessmentStatus::Aborted, + risk_level: None, + user_authorization: None, + rationale: None, + decision_source: Some(GuardianAssessmentDecisionSource::Agent), + action: action_summary, + }), + ) + .await; + record_guardian_non_denial(&session, &assessment_turn_id).await; + return Err(ReviewDecision::Abort); + } + + let root_authorization_version = session + .services + .agent_control + .root_user_authorization(session.thread_id) + .await + .map(|snapshot| snapshot.authorization_version); + // Keep the authorization revision even when no cacheable review evidence exists. + let history = session.conversation_history_snapshot().await; + let user_message_revision = history.user_message_revision(); + let review_evidence = if let Some(evidence) = session + .services + .thread_extension_data + .get::() + { + let authorization_version = evidence.authorization_version(history.as_ref()); + format_guardian_action_pretty(&request).ok().map(|action| { + ( + evidence, + action.text, + authorization_version, + root_authorization_version, + ) + }) + } else { + None + }; + drop(history); + Ok(PreparedApproval { + request, + turn, + review_reason, + assessment_turn_id, + target_item_id, + plugin_id, + script_path, + action_summary, + reviewed_action, + review_tracking, + started_at_ms, + root_authorization_version, + user_message_revision, + review_evidence, + }) + } + + async fn attempt( + &self, + prepared: &PreparedApproval, + deadline: Instant, + ) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) { + run_guardian_review_session_before_deadline( + Arc::clone(&self.session), + self.context.clone(), + prepared.request.clone(), + self.reasons.clone(), + guardian_output_schema(), + self.options.external_cancel.clone(), + deadline, + ) + .await + } + + async fn complete( + &self, + prepared: PreparedApproval, + mut outcome: GuardianReviewOutcome, + analytics_result: GuardianReviewAnalyticsResult, + ) -> ReviewDecision { + let PreparedApproval { + request: _, + turn, + review_reason, + assessment_turn_id, + target_item_id, + plugin_id, + script_path, + action_summary, + reviewed_action, + review_tracking, + started_at_ms, + root_authorization_version, + user_message_revision, + review_evidence, + } = prepared; + let session = Arc::clone(&self.session); + let review_id = self.review_id.clone(); + let approval_request_source = self.options.approval_request_source; + let terminal_action = action_summary.clone(); + 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 + != session + .conversation_history_snapshot() + .await + .user_message_revision()) + { + // A completed approval cannot outlive the owning-session or root evidence + // it evaluated, including when either changed before prompt construction. + outcome = GuardianReviewOutcome::Error(GuardianReviewError::Cancelled); + } + + let completed_at_ms = now_unix_timestamp_ms(); + let completed = codex_guardian_reviewer::complete_review( + outcome, + turn.model_info(), + GuardianAssessmentEvent { + review_reason: Some(review_reason), + id: review_id, + target_item_id, + plugin_id, + script_path, + turn_id: assessment_turn_id.clone(), + started_at_ms, + completed_at_ms: Some(completed_at_ms), + status: GuardianAssessmentStatus::InProgress, + risk_level: None, + user_authorization: None, + rationale: None, + decision_source: Some(GuardianAssessmentDecisionSource::Agent), + action: terminal_action, + }, + analytics_result, + ); + track_guardian_review( + session.as_ref(), + &review_tracking, + approval_request_source, + &reviewed_action, + completed.analytics, + completed_at_ms.try_into().unwrap_or_default(), + ); + if let Some(message) = completed.warning { + session + .send_event( + turn.as_ref(), + EventMsg::GuardianWarning(WarningEvent { message }), + ) + .await; + } + if completed.assessment_outcome.is_some() + && let Some((evidence, action, authorization_version, root_authorization_version)) = + review_evidence + { + evidence.record( + &completed.event, + &action, + authorization_version, + root_authorization_version, + ); + } + session + .send_event(turn.as_ref(), EventMsg::GuardianAssessment(completed.event)) + .await; + if completed.assessment_outcome == Some(GuardianAssessmentOutcome::Deny) { + record_guardian_denial(&session, &turn, &assessment_turn_id).await; + } else { + record_guardian_non_denial(&session, &assessment_turn_id).await; + } + completed.decision + } +} diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index d792e2870b..c753599571 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -1,5 +1,11 @@ -//! Shared synchronous reviewer lifecycle. The extension owns its instance; standalone -//! callers use the same reuse, fork, timeout and cancellation rules. +//! Host adapter for synchronous Guardian sessions and the existing context builder. +//! The extension owns review policy and pooling; this module binds runtime operations +//! to the captured parent action, environments, authorization and context snapshots. + +#[path = "review_session_factory.rs"] +mod factory; +pub(crate) use factory::prewarm_guardian_review_session; +pub(crate) use factory::run_guardian_review_session; #[path = "review_session_threads.rs"] mod managed_threads; @@ -10,7 +16,6 @@ use context_policy::ReviewContextPolicy; use std::collections::HashMap; use std::collections::HashSet; -use std::future::Future; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -49,7 +54,6 @@ use codex_protocol::protocol::TokenUsage; use futures::future::BoxFuture; use serde_json::Value; use tokio::sync::Mutex; -use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; use tracing::warn; @@ -98,6 +102,9 @@ use super::prompt::build_guardian_prompt_items_with_parent_turn; use super::review::guardian_review_session_config; pub(crate) use super::reviewer_config::build_guardian_review_session_config; use super::reviewer_config::read_only_guardian_permission_profile; +use codex_guardian_reviewer::run_before_review_deadline; +#[cfg(test)] +use codex_guardian_reviewer::run_before_review_deadline_with_cancel; const GUARDIAN_INTERRUPT_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); const GUARDIAN_MAX_IMAGE_ITEM_TOKENS: i64 = 10_000; @@ -126,27 +133,36 @@ pub(crate) struct GuardianReviewSessionParams { pub(crate) deadline: tokio::time::Instant, } -/// One reviewer pool per parent thread, shared by extension and standalone hosts. -/// The extension supplies its thread manager without replacing the review algorithm. +/// 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 GuardianReviewSessionManager { +pub struct GuardianReviewSessionHost { managed_threads: Option, - state: Arc>, - cancellation_token: CancellationToken, } -#[derive(Default)] -struct GuardianReviewSessionState { - trunk: Option>, - ephemeral_reviews: Vec>, +impl GuardianReviewSessionHost { + pub fn with_thread_manager(manager: std::sync::Weak) -> Self { + Self { + managed_threads: Some(managed_threads::ManagedReviewerThreads::new(manager)), + } + } + + pub fn mark_ready(&self) { + if let Some(threads) = &self.managed_threads { + threads.mark_ready(); + } + } } -struct GuardianReviewSession { +pub(crate) type GuardianReviewSessionManager = + codex_guardian_reviewer::ReviewerPool; + +/// Opaque host session handle. Its state belongs to the existing context builder. +pub struct GuardianReviewSession { session: Arc, io: SessionIo, cancel_token: CancellationToken, reuse_key: GuardianReviewSessionReuseKey, - review_lock: Semaphore, state: Mutex, } @@ -181,21 +197,18 @@ fn token_usage_delta(start: &TokenUsage, end: &TokenUsage) -> TokenUsage { } } -struct EphemeralReviewCleanup { - state: Arc>, - review_session: Option>, -} - +/// Committed context used to seed a private reviewer fork. #[derive(Clone)] -struct GuardianReviewForkSnapshot { +pub struct GuardianReviewForkSnapshot { initial_history: InitialHistory, prior_review_count: usize, last_reviewed_transcript_cursor: Option, last_admitted_node_repl_response_sequence: u64, } +/// Opaque compatibility key derived by the existing context builder. #[derive(Debug, Clone, PartialEq)] -struct GuardianReviewSessionReuseKey { +pub struct GuardianReviewSessionReuseKey { // Only include settings that affect spawned-session behavior and parent // history rewrites that invalidate existing reviewer context. parent_history_version: u64, @@ -301,45 +314,6 @@ pub(crate) fn prompt_cache_key_override_for_review_session( } impl GuardianReviewSession { - async fn shutdown(&self) { - self.cancel_token.cancel(); - let _ = self.io.shutdown_and_wait().await; - } - - fn shutdown_in_background(self: &Arc) { - let review_session = Arc::clone(self); - drop(tokio::spawn(async move { - review_session.shutdown().await; - })); - } - - async fn fork_snapshot(&self) -> Option { - self.state.lock().await.last_committed_fork_snapshot.clone() - } - - async fn refresh_last_committed_fork_snapshot(&self) { - match load_rollout_items_for_fork(&self.session).await { - Ok(Some(items)) if !items.is_empty() => { - let mut state = self.state.lock().await; - let prior_review_count = state.prior_review_count; - let last_reviewed_transcript_cursor = state.last_reviewed_transcript_cursor; - let last_admitted_node_repl_response_sequence = - state.last_admitted_node_repl_response_sequence; - state.last_committed_fork_snapshot = Some(GuardianReviewForkSnapshot { - initial_history: InitialHistory::Forked(items), - prior_review_count, - last_reviewed_transcript_cursor, - last_admitted_node_repl_response_sequence, - }); - } - Ok(Some(_)) => {} - Ok(None) => {} - Err(err) => { - warn!("failed to refresh guardian trunk rollout snapshot: {err}"); - } - } - } - async fn admit_node_repl_evidence(&self, event: &Event) { let EventMsg::ItemCompleted(completed) = &event.msg else { return; @@ -364,580 +338,6 @@ impl GuardianReviewSession { } } -impl EphemeralReviewCleanup { - fn new( - state: Arc>, - review_session: Arc, - ) -> Self { - Self { - state, - review_session: Some(review_session), - } - } - - fn disarm(&mut self) { - self.review_session = None; - } -} - -impl Drop for EphemeralReviewCleanup { - fn drop(&mut self) { - let Some(review_session) = self.review_session.take() else { - return; - }; - let state = Arc::clone(&self.state); - drop(tokio::spawn(async move { - let review_session = { - let mut state = state.lock().await; - state - .ephemeral_reviews - .iter() - .position(|active_review| Arc::ptr_eq(active_review, &review_session)) - .map(|index| state.ephemeral_reviews.swap_remove(index)) - }; - if let Some(review_session) = review_session { - review_session.shutdown().await; - } - })); - } -} - -impl GuardianReviewSessionManager { - /// Creates the extension-owned pool before its parent has been registered. - pub fn with_thread_manager(manager: std::sync::Weak) -> Self { - Self { - managed_threads: Some(managed_threads::ManagedReviewerThreads::new(manager)), - ..Self::default() - } - } - - /// Allows startup prewarming to spawn after the parent is registered. - pub fn mark_ready(&self) { - if let Some(threads) = &self.managed_threads { - threads.mark_ready(); - } - } - - pub(crate) fn initialize( - &self, - parent_session: Arc, - parent_turn: Arc, - ) -> BoxFuture<'_, anyhow::Result<()>> { - // Boxing breaks the Session::new -> Guardian -> Session::new future recursion. - Box::pin(async move { - let session_config = - guardian_review_session_config(&parent_session, &parent_turn).await?; - let spawn_config = session_config.spawn_config; - let parent_history = parent_session.clone_history().await; - let context_policy = ReviewContextPolicy::for_context( - parent_session.guardian_context_mode, - &spawn_config.features, - ); - let root_authorization_version = context_policy - .root_authorization_version(&parent_session) - .await; - let parent_compaction = context_policy.parent_compaction( - &parent_history, - session_config.compaction_model_hash.as_deref(), - )?; - let parent_context = GuardianReviewContext::from(parent_turn); - let mut reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( - &spawn_config, - parent_session.user_instructions().await, - parent_history.history_version(), - parent_session.guardian_context_mode, - ) - .with_environments(parent_context.environments()) - .with_node_repl_policy_eligibility( - parent_context - .turn() - .model_info() - .computer_use_review_required(), - ) - .with_node_repl_policy(&session_config.node_repl_policy); - reuse_key.root_authorization_version = root_authorization_version; - let spawn_cancel_token = self.cancellation_token.child_token(); - let spawn_cancel_guard = spawn_cancel_token.clone().drop_guard(); - let review_session = spawn_guardian_review_session( - self, - &parent_session, - &parent_context, - spawn_config, - reuse_key, - spawn_cancel_token.clone(), - parent_compaction, - /*fork_snapshot*/ None, - ) - .await?; - // A first review or shutdown may win while eager initialization is in flight; - // install only if neither has happened. - let mut state = self.state.lock().await; - if !spawn_cancel_token.is_cancelled() && state.trunk.is_none() { - state.trunk = Some(Arc::new(review_session)); - drop(spawn_cancel_guard.disarm()); - } - Ok(()) - }) - } - - pub(crate) async fn trunk_rollout_path(&self) -> Option { - let trunk = self.state.lock().await.trunk.clone()?; - trunk - .session - .ensure_rollout_materialized(PersistContext::Standard) - .await; - match trunk.session.current_rollout_path().await { - Ok(path) => path, - Err(err) => { - warn!("failed to resolve guardian trunk rollout path: {err}"); - None - } - } - } - - pub(crate) async fn shutdown(&self) { - self.cancellation_token.cancel(); - self.invalidate().await; - } - - pub(crate) async fn invalidate(&self) { - let (review_session, ephemeral_reviews) = { - let mut state = self.state.lock().await; - ( - state.trunk.take(), - std::mem::take(&mut state.ephemeral_reviews), - ) - }; - for review_session in review_session.into_iter().chain(ephemeral_reviews) { - if self.cancellation_token.is_cancelled() { - review_session.shutdown().await; - } else { - review_session.cancel_token.cancel(); - review_session.shutdown_in_background(); - } - } - } - - #[expect( - clippy::await_holding_invalid_type, - reason = "review session selection and trunk spawning must stay serialized" - )] - pub(super) async fn run_review( - &self, - params: GuardianReviewSessionParams, - ) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult) { - let deadline = params.deadline; - let parent_history = ¶ms.parent_history; - let context_policy = ReviewContextPolicy::for_context( - params.parent_session.guardian_context_mode, - ¶ms.spawn_config.features, - ); - let root_authorization_version = context_policy - .root_authorization_version(¶ms.parent_session) - .await; - let parent_compaction = match context_policy - .parent_compaction(parent_history, params.compaction_model_hash.as_deref()) - { - Ok(checkpoint) => checkpoint, - Err(error) => { - return ( - GuardianReviewSessionOutcome::PromptBuildFailed(error), - GuardianReviewAnalyticsResult::without_session(), - ); - } - }; - let mut next_reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( - ¶ms.spawn_config, - params.parent_session.user_instructions().await, - parent_history.history_version(), - params.parent_session.guardian_context_mode, - ) - .with_environments(params.parent_context.environments()) - .with_node_repl_policy_eligibility( - params - .parent_context - .turn() - .model_info() - .computer_use_review_required(), - ) - .with_node_repl_policy(¶ms.node_repl_policy); - next_reuse_key.root_authorization_version = root_authorization_version; - let mut spawned_trunk = false; - let trunk_candidate = match run_before_review_deadline( - deadline, - params.external_cancel.as_ref(), - self.state.lock(), - ) - .await - { - Ok(mut state) => { - if context_policy != ReviewContextPolicy::ThreadOwned - && parent_compaction.is_none() - && let Some(trunk) = state.trunk.as_ref() - { - // Without a decryptable summary, the existing reviewer may - // hold the only remaining authorization or restriction. - next_reuse_key.parent_history_version = trunk.reuse_key.parent_history_version; - } - if let Some(trunk) = state.trunk.as_ref() - && trunk.reuse_key != next_reuse_key - && trunk.review_lock.try_acquire().is_ok() - && let Some(stale_trunk) = state.trunk.take() - { - stale_trunk.shutdown_in_background(); - } - - if state.trunk.is_none() { - let spawn_cancel_token = self.cancellation_token.child_token(); - let review_session = match run_before_review_deadline_with_cancel( - deadline, - params.external_cancel.as_ref(), - &spawn_cancel_token, - Box::pin(spawn_guardian_review_session( - self, - ¶ms.parent_session, - ¶ms.parent_context, - params.spawn_config.clone(), - next_reuse_key.clone(), - spawn_cancel_token.clone(), - parent_compaction.clone(), - /*fork_snapshot*/ None, - )), - ) - .await - { - Ok(Ok(review_session)) => Arc::new(review_session), - Ok(Err(err)) => { - return ( - GuardianReviewSessionOutcome::PromptBuildFailed(err), - GuardianReviewAnalyticsResult::without_session(), - ); - } - Err(outcome) => { - return (outcome, GuardianReviewAnalyticsResult::without_session()); - } - }; - state.trunk = Some(Arc::clone(&review_session)); - spawned_trunk = true; - } - - state.trunk.as_ref().cloned() - } - Err(outcome) => { - return (outcome, GuardianReviewAnalyticsResult::without_session()); - } - }; - - let Some(trunk) = trunk_candidate else { - return ( - GuardianReviewSessionOutcome::Completed(Err(anyhow!( - "guardian review session was not available after spawn" - ))), - GuardianReviewAnalyticsResult::without_session(), - ); - }; - - if trunk.reuse_key != next_reuse_key { - return Box::pin(self.run_ephemeral_review( - params, - next_reuse_key, - deadline, - parent_compaction, - /*fork_snapshot*/ None, - )) - .await; - } - - let trunk_guard = match trunk.review_lock.try_acquire() { - Ok(trunk_guard) => trunk_guard, - Err(_) => { - return Box::pin(self.run_ephemeral_review( - params, - next_reuse_key, - deadline, - parent_compaction, - trunk.fork_snapshot().await, - )) - .await; - } - }; - - let guardian_session_kind = if spawned_trunk { - GuardianReviewSessionKind::TrunkNew - } else { - GuardianReviewSessionKind::TrunkReused - }; - let (outcome, keep_review_session, analytics_result) = Box::pin(run_review_on_session( - trunk.as_ref(), - ¶ms, - guardian_session_kind, - deadline, - )) - .await; - record_failed_review(&trunk.session, ¶ms, &outcome).await; - if keep_review_session && matches!(outcome, GuardianReviewSessionOutcome::Completed(_)) { - trunk.refresh_last_committed_fork_snapshot().await; - } - drop(trunk_guard); - - if keep_review_session { - (outcome, analytics_result) - } else { - if let Some(review_session) = self.remove_trunk_if_current(&trunk).await { - review_session.shutdown_in_background(); - } - (outcome, analytics_result) - } - } - - #[cfg(test)] - pub(crate) async fn cache_for_test(&self, session: Arc, io: SessionIo) { - let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( - session.get_config().await.as_ref(), - session.user_instructions().await, - session.clone_history().await.history_version(), - session.guardian_context_mode, - ); - self.state.lock().await.trunk = Some(Arc::new(GuardianReviewSession { - reuse_key, - session, - io, - cancel_token: CancellationToken::new(), - review_lock: Semaphore::new(/*permits*/ 1), - state: Mutex::new(GuardianReviewState { - prior_review_count: 0, - last_reviewed_transcript_cursor: None, - last_admitted_node_repl_response_sequence: 0, - pending_node_repl_evidence_admission: None, - last_committed_fork_snapshot: None, - }), - })); - } - - #[cfg(test)] - pub(crate) async fn register_ephemeral_for_test(&self, session: Arc, io: SessionIo) { - let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( - session.get_config().await.as_ref(), - session.user_instructions().await, - session.clone_history().await.history_version(), - session.guardian_context_mode, - ); - self.state - .lock() - .await - .ephemeral_reviews - .push(Arc::new(GuardianReviewSession { - reuse_key, - session, - io, - cancel_token: CancellationToken::new(), - review_lock: Semaphore::new(/*permits*/ 1), - state: Mutex::new(GuardianReviewState { - prior_review_count: 0, - last_reviewed_transcript_cursor: None, - last_admitted_node_repl_response_sequence: 0, - pending_node_repl_evidence_admission: None, - last_committed_fork_snapshot: None, - }), - })); - } - - #[cfg(test)] - pub(crate) async fn committed_fork_rollout_items_for_test(&self) -> Option> { - let trunk = self.state.lock().await.trunk.clone()?; - let state = trunk.state.lock().await; - let snapshot = state.last_committed_fork_snapshot.as_ref()?; - match &snapshot.initial_history { - InitialHistory::Forked(items) => Some(items.clone()), - InitialHistory::New | InitialHistory::Cleared | InitialHistory::Resumed(_) => None, - } - } - - #[cfg(test)] - pub(crate) async fn send_trunk_event_raw_for_test(&self, event: Event) { - let trunk = self - .state - .lock() - .await - .trunk - .clone() - .expect("guardian trunk should exist"); - trunk.session.send_event_raw(event).await; - } - - async fn remove_trunk_if_current( - &self, - trunk: &Arc, - ) -> Option> { - let mut state = self.state.lock().await; - if state - .trunk - .as_ref() - .is_some_and(|current| Arc::ptr_eq(current, trunk)) - { - state.trunk.take() - } else { - None - } - } - - async fn register_active_ephemeral(&self, review_session: Arc) { - self.state - .lock() - .await - .ephemeral_reviews - .push(review_session); - } - - async fn take_active_ephemeral( - &self, - review_session: &Arc, - ) -> Option> { - let mut state = self.state.lock().await; - let ephemeral_review_index = state - .ephemeral_reviews - .iter() - .position(|active_review| Arc::ptr_eq(active_review, review_session))?; - Some(state.ephemeral_reviews.swap_remove(ephemeral_review_index)) - } - - async fn run_ephemeral_review( - &self, - params: GuardianReviewSessionParams, - reuse_key: GuardianReviewSessionReuseKey, - deadline: tokio::time::Instant, - parent_compaction: Option, - fork_snapshot: Option, - ) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult) { - let spawn_cancel_token = self.cancellation_token.child_token(); - let mut fork_config = params.spawn_config.clone(); - fork_config.ephemeral = true; - let review_session = match run_before_review_deadline_with_cancel( - deadline, - params.external_cancel.as_ref(), - &spawn_cancel_token, - Box::pin(spawn_guardian_review_session( - self, - ¶ms.parent_session, - ¶ms.parent_context, - fork_config, - reuse_key, - spawn_cancel_token.clone(), - parent_compaction, - fork_snapshot, - )), - ) - .await - { - Ok(Ok(review_session)) => Arc::new(review_session), - Ok(Err(err)) => { - return ( - GuardianReviewSessionOutcome::PromptBuildFailed(err), - GuardianReviewAnalyticsResult::without_session(), - ); - } - Err(outcome) => { - return (outcome, GuardianReviewAnalyticsResult::without_session()); - } - }; - self.register_active_ephemeral(Arc::clone(&review_session)) - .await; - let mut cleanup = - EphemeralReviewCleanup::new(Arc::clone(&self.state), Arc::clone(&review_session)); - - let (outcome, _, analytics_result) = Box::pin(run_review_on_session( - review_session.as_ref(), - ¶ms, - GuardianReviewSessionKind::EphemeralForked, - deadline, - )) - .await; - record_failed_review(&review_session.session, ¶ms, &outcome).await; - if let Some(review_session) = self.take_active_ephemeral(&review_session).await { - cleanup.disarm(); - review_session.shutdown_in_background(); - } - (outcome, analytics_result) - } -} - -#[allow(clippy::too_many_arguments)] -async fn spawn_guardian_review_session( - manager: &GuardianReviewSessionManager, - parent_session: &Arc, - parent_context: &GuardianReviewContext, - spawn_config: Config, - reuse_key: GuardianReviewSessionReuseKey, - cancel_token: CancellationToken, - parent_compaction: Option, - fork_snapshot: Option, -) -> anyhow::Result { - let ( - initial_history, - prior_review_count, - initial_transcript_cursor, - last_admitted_node_repl_response_sequence, - ) = match fork_snapshot { - Some(fork_snapshot) => ( - Some(fork_snapshot.initial_history), - fork_snapshot.prior_review_count, - fork_snapshot.last_reviewed_transcript_cursor, - fork_snapshot.last_admitted_node_repl_response_sequence, - ), - None => ( - parent_compaction - .map(|item| InitialHistory::Forked(vec![RolloutItem::ResponseItem(item.into())])), - 0, - None, - 0, - ), - }; - let (session, io) = match &manager.managed_threads { - Some(threads) => { - threads - .spawn( - parent_session, - parent_context, - spawn_config, - cancel_token.clone(), - initial_history, - ) - .await? - } - None => { - Box::pin(run_codex_thread_interactive( - spawn_config, - parent_session.services.auth_manager.clone(), - parent_session.services.models_manager.clone(), - Arc::clone(parent_session), - Arc::clone(parent_context.turn()), - parent_context.environments().clone(), - cancel_token.clone(), - SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_string()), - initial_history, - GitEnrichmentPolicy::Skip, - codex_sandboxing::WindowsSandboxProxySettingsMode::Preserve, - )) - .await? - } - }; - - Ok(GuardianReviewSession { - session, - io, - cancel_token, - reuse_key, - review_lock: Semaphore::new(/*permits*/ 1), - state: Mutex::new(GuardianReviewState { - prior_review_count, - last_reviewed_transcript_cursor: initial_transcript_cursor, - last_admitted_node_repl_response_sequence, - pending_node_repl_evidence_admission: None, - last_committed_fork_snapshot: None, - }), - }) -} - async fn run_review_on_session( review_session: &GuardianReviewSession, params: &GuardianReviewSessionParams, @@ -1561,37 +961,6 @@ fn event_matches_turn(event: &Event, expected_turn_id: &str) -> bool { } } -async fn run_before_review_deadline( - deadline: tokio::time::Instant, - external_cancel: Option<&CancellationToken>, - future: impl Future, -) -> Result { - tokio::select! { - _ = tokio::time::sleep_until(deadline) => Err(GuardianReviewSessionOutcome::TimedOut), - result = future => Ok(result), - _ = async { - if let Some(cancel_token) = external_cancel { - cancel_token.cancelled().await; - } else { - std::future::pending::<()>().await; - } - } => Err(GuardianReviewSessionOutcome::Aborted), - } -} - -async fn run_before_review_deadline_with_cancel( - deadline: tokio::time::Instant, - external_cancel: Option<&CancellationToken>, - cancel_token: &CancellationToken, - future: impl Future, -) -> Result { - let result = run_before_review_deadline(deadline, external_cancel, future).await; - if result.is_err() { - cancel_token.cancel(); - } - result -} - async fn interrupt_and_drain_turn( review_session: &GuardianReviewSession, expected_turn_id: &str, @@ -1622,3 +991,78 @@ async fn interrupt_and_drain_turn( #[cfg(test)] #[path = "review_session_tests.rs"] mod tests; + +impl codex_guardian_reviewer::ReviewerSession for GuardianReviewSession { + type Context = GuardianReviewSessionReuseKey; + type Snapshot = GuardianReviewForkSnapshot; + + fn context(&self) -> &Self::Context { + &self.reuse_key + } + fn cancel(&self) { + self.cancel_token.cancel(); + } + + async fn shutdown(&self) { + self.cancel_token.cancel(); + let _ = self.io.shutdown_and_wait().await; + } + + async fn snapshot(&self) -> Option { + self.state.lock().await.last_committed_fork_snapshot.clone() + } + + async fn commit_snapshot(&self) { + match load_rollout_items_for_fork(&self.session).await { + Ok(Some(items)) if !items.is_empty() => { + let mut state = self.state.lock().await; + let prior_review_count = state.prior_review_count; + let last_reviewed_transcript_cursor = state.last_reviewed_transcript_cursor; + let last_admitted_node_repl_response_sequence = + state.last_admitted_node_repl_response_sequence; + state.last_committed_fork_snapshot = Some(GuardianReviewForkSnapshot { + initial_history: InitialHistory::Forked(items), + prior_review_count, + last_reviewed_transcript_cursor, + last_admitted_node_repl_response_sequence, + }); + } + Ok(Some(_)) => {} + Ok(None) => {} + Err(err) => { + warn!("failed to refresh guardian trunk rollout snapshot: {err}"); + } + } + } +} + +impl GuardianReviewSession { + pub(crate) async fn rollout_path(&self) -> Option { + self.session + .ensure_rollout_materialized(PersistContext::Standard) + .await; + match self.session.current_rollout_path().await { + Ok(path) => path, + Err(error) => { + warn!("failed to resolve guardian trunk rollout path: {error}"); + None + } + } + } +} + +#[cfg(test)] +impl GuardianReviewSession { + pub(crate) async fn committed_fork_rollout_items_for_test(&self) -> Option> { + let state = self.state.lock().await; + let snapshot = state.last_committed_fork_snapshot.as_ref()?; + match &snapshot.initial_history { + InitialHistory::Forked(items) => Some(items.clone()), + InitialHistory::New | InitialHistory::Cleared | InitialHistory::Resumed(_) => None, + } + } + + pub(crate) async fn send_trunk_event_raw_for_test(&self, event: Event) { + self.session.send_event_raw(event).await; + } +} diff --git a/codex-rs/core/src/guardian/review_session_factory.rs b/codex-rs/core/src/guardian/review_session_factory.rs new file mode 100644 index 0000000000..b58db97bc9 --- /dev/null +++ b/codex-rs/core/src/guardian/review_session_factory.rs @@ -0,0 +1,245 @@ +//! Prepares opaque session inputs for the extension-owned reviewer pool. +//! Context selection and assembly stay on the existing path pending its replacement. + +use super::*; +use codex_guardian_reviewer::ReviewerPool; +use codex_guardian_reviewer::ReviewerRequest; +use codex_guardian_reviewer::ReviewerSessionFactory; +use codex_guardian_reviewer::SessionDisposition; + +pub(super) struct PreparedSession { + parent: Arc, + context: GuardianReviewContext, + config: Config, + context_policy: ReviewContextPolicy, + key: GuardianReviewSessionReuseKey, + parent_compaction: Option, + host: Arc, +} + +impl PreparedSession { + async fn prepare( + parent: Arc, + context: GuardianReviewContext, + config: Config, + history: &ContextManager, + node_repl_policy: &GuardianNodeReplPolicy, + compaction_model_hash: Option<&str>, + ) -> anyhow::Result { + let context_policy = + ReviewContextPolicy::for_context(parent.guardian_context_mode, &config.features); + let root_authorization_version = context_policy.root_authorization_version(&parent).await; + let parent_compaction = context_policy.parent_compaction(history, compaction_model_hash)?; + let mut key = GuardianReviewSessionReuseKey::from_spawn_config( + &config, + parent.user_instructions().await, + history.history_version(), + parent.guardian_context_mode, + ) + .with_environments(context.environments()) + .with_node_repl_policy_eligibility( + context.turn().model_info().computer_use_review_required(), + ) + .with_node_repl_policy(node_repl_policy); + key.root_authorization_version = root_authorization_version; + let host = parent + .services + .thread_extension_data + .get_or_init(GuardianReviewSessionHost::default); + Ok(Self { + parent, + context, + config, + context_policy, + key, + parent_compaction, + host, + }) + } +} + +impl ReviewerSessionFactory for PreparedSession { + type Session = GuardianReviewSession; + + fn context(&self, previous: Option<&GuardianReviewSession>) -> GuardianReviewSessionReuseKey { + let mut key = self.key.clone(); + if self.context_policy != ReviewContextPolicy::ThreadOwned + && self.parent_compaction.is_none() + && let Some(previous) = previous + { + // Without a decryptable summary, the existing reviewer may hold the + // only remaining authorization or restriction from parent history. + key.parent_history_version = previous.reuse_key.parent_history_version; + } + key + } + + async fn spawn( + &self, + context: GuardianReviewSessionReuseKey, + kind: GuardianReviewSessionKind, + snapshot: Option, + cancellation: CancellationToken, + ) -> anyhow::Result { + let mut config = self.config.clone(); + if matches!(kind, GuardianReviewSessionKind::EphemeralForked) { + config.ephemeral = true; + } + let ( + initial_history, + prior_review_count, + initial_transcript_cursor, + last_admitted_node_repl_response_sequence, + ) = match snapshot { + Some(snapshot) => ( + Some(snapshot.initial_history), + snapshot.prior_review_count, + snapshot.last_reviewed_transcript_cursor, + snapshot.last_admitted_node_repl_response_sequence, + ), + None => ( + self.parent_compaction.clone().map(|item| { + InitialHistory::Forked(vec![RolloutItem::ResponseItem(item.into())]) + }), + 0, + None, + 0, + ), + }; + let (session, io) = match &self.host.managed_threads { + Some(threads) => { + threads + .spawn( + &self.parent, + &self.context, + config, + cancellation.clone(), + initial_history, + ) + .await? + } + None => { + Box::pin(run_codex_thread_interactive( + config, + Arc::clone(&self.parent.services.auth_manager), + self.parent.services.models_manager.clone(), + Arc::clone(&self.parent), + Arc::clone(self.context.turn()), + self.context.environments().clone(), + cancellation.clone(), + SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_owned()), + initial_history, + GitEnrichmentPolicy::Skip, + codex_sandboxing::WindowsSandboxProxySettingsMode::Preserve, + )) + .await? + } + }; + Ok(GuardianReviewSession { + session, + io, + cancel_token: cancellation, + reuse_key: context, + state: Mutex::new(GuardianReviewState { + prior_review_count, + last_reviewed_transcript_cursor: initial_transcript_cursor, + last_admitted_node_repl_response_sequence, + pending_node_repl_evidence_admission: None, + last_committed_fork_snapshot: None, + }), + }) + } +} + +pub(super) struct PreparedReview { + factory: PreparedSession, + params: GuardianReviewSessionParams, +} + +impl ReviewerRequest for PreparedReview { + type Factory = PreparedSession; + + fn factory(&self) -> &PreparedSession { + &self.factory + } + fn deadline(&self) -> tokio::time::Instant { + self.params.deadline + } + fn cancellation(&self) -> Option<&CancellationToken> { + self.params.external_cancel.as_ref() + } + + async fn run( + &self, + session: &GuardianReviewSession, + kind: GuardianReviewSessionKind, + ) -> ( + GuardianReviewSessionOutcome, + SessionDisposition, + GuardianReviewAnalyticsResult, + ) { + let (outcome, keep_session, analytics) = Box::pin(run_review_on_session( + session, + &self.params, + kind, + self.params.deadline, + )) + .await; + record_failed_review(&session.session, &self.params, &outcome).await; + let disposition = if keep_session { + SessionDisposition::Reusable + } else { + SessionDisposition::Discard + }; + (outcome, disposition, analytics) + } +} + +pub(crate) async fn run_guardian_review_session( + pool: Arc>, + params: GuardianReviewSessionParams, +) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult) { + match prepare_review(params).await { + Ok(prepared) => pool.review(prepared).await, + Err(error) => ( + GuardianReviewSessionOutcome::PromptBuildFailed(error), + GuardianReviewAnalyticsResult::without_session(), + ), + } +} + +pub(super) async fn prepare_review( + params: GuardianReviewSessionParams, +) -> anyhow::Result { + let factory = PreparedSession::prepare( + Arc::clone(¶ms.parent_session), + params.parent_context.clone(), + params.spawn_config.clone(), + ¶ms.parent_history, + ¶ms.node_repl_policy, + params.compaction_model_hash.as_deref(), + ) + .await?; + Ok(PreparedReview { factory, params }) +} + +pub(crate) fn prewarm_guardian_review_session( + parent: Arc, + turn: Arc, +) -> BoxFuture<'static, anyhow::Result<()>> { + // Keep the Session -> Guardian -> Session startup future on the heap. + Box::pin(async move { + let config = guardian_review_session_config(&parent, &turn).await?; + let history = parent.clone_history().await; + let factory = PreparedSession::prepare( + Arc::clone(&parent), + GuardianReviewContext::from(turn), + config.spawn_config, + &history, + &config.node_repl_policy, + config.compaction_model_hash.as_deref(), + ) + .await?; + parent.guardian_review_session().prewarm(&factory).await + }) +} diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index e028d32028..6490c838aa 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -14,10 +14,6 @@ use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; #[tokio::test] -#[allow( - clippy::await_holding_invalid_type, - reason = "hold session selection while the parent compacts to reproduce the race" -)] async fn run_review_preserves_evidence_during_parent_compaction() { const EVIDENCE: &str = "The inspected repository is public."; let (parent, turn, _events) = @@ -71,19 +67,10 @@ async fn run_review_preserves_evidence_during_parent_compaction() { .node_repl_auto_review_required, ) .with_node_repl_policy(¶ms.node_repl_policy); - let manager = GuardianReviewSessionManager { - state: Arc::new(Mutex::new(GuardianReviewSessionState { - trunk: Some(Arc::new(reviewer)), - ephemeral_reviews: Vec::new(), - })), - ..Default::default() - }; - let gate = manager.state.lock().await; - let review = manager.run_review(params); - tokio::pin!(review); - // Earlier awaits use unlocked in-memory state; this polls through checkpoint - // selection and deterministically stops at the held manager mutex. - assert!(futures::poll!(&mut review).is_pending()); + let manager = GuardianReviewSessionManager::default(); + prewarm_test_session(&manager, reviewer).await; + // Capture the review context, then compact the parent before the reviewer builds its prompt. + let prepared = factory::prepare_review(params).await.unwrap(); let checkpoint: ResponseItem = serde_json::from_value(serde_json::json!({ "type": "compaction", "id": "cmp_new", "encrypted_content": "new-checkpoint" })) @@ -103,9 +90,7 @@ async fn run_review_preserves_evidence_during_parent_compaction() { }, ) .await; - drop(gate); - - let ((outcome, _), submitted_text) = tokio::join!(review, async { + let ((outcome, _), submitted_text) = tokio::join!(manager.review(prepared), async { let submission = rx_sub.recv().await.unwrap(); let id = submission.id; let Op::TurnInput { request, reply, .. } = submission.op else { @@ -167,7 +152,6 @@ async fn test_review_session() -> ( }, cancel_token: CancellationToken::new(), reuse_key, - review_lock: Semaphore::new(/*permits*/ 1), state: Mutex::new(GuardianReviewState { prior_review_count: 0, last_reviewed_transcript_cursor: None, @@ -266,20 +250,16 @@ 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 = GuardianReviewSessionManager::default(); - manager - .initialize( - params.parent_session, - Arc::clone(params.parent_context.turn()), - ) - .await - .expect("initialize Guardian session"); + 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 - .state - .lock() + .trunk() .await - .trunk - .as_ref() .expect("Guardian session") .session .windows_sandbox_proxy_settings_mode; @@ -920,15 +900,11 @@ async fn run_review_removes_trunk_when_event_stream_is_broken() { ) .with_environments(params.parent_context.environments()) .with_node_repl_policy(¶ms.node_repl_policy); - let manager = Arc::new(GuardianReviewSessionManager { - state: Arc::new(Mutex::new(GuardianReviewSessionState { - trunk: Some(Arc::new(review_session)), - ephemeral_reviews: Vec::new(), - })), - ..Default::default() - }); + let manager = Arc::new(GuardianReviewSessionManager::default()); + prewarm_test_session(&manager, review_session).await; let manager_for_review = Arc::clone(&manager); - let review = tokio::spawn(async move { manager_for_review.run_review(params).await }); + let review = + tokio::spawn(async move { run_guardian_review_session(manager_for_review, params).await }); let submission = rx_sub.recv().await.expect("guardian submission"); let id = submission.id; let Op::TurnInput { reply, .. } = submission.op else { @@ -945,7 +921,7 @@ async fn run_review_removes_trunk_when_event_stream_is_broken() { outcome, GuardianReviewSessionOutcome::Completed(Err(_)) )); - assert!(manager.state.lock().await.trunk.is_none()); + assert!(manager.trunk().await.is_none()); } #[tokio::test] @@ -1184,3 +1160,34 @@ async fn interrupt_and_drain_turn_ignores_prior_turn_completion() { assert!(review_session.io.rx_event.try_recv().is_err()); } + +// Reuse the existing in-memory reviewer fixture through the production prewarm path. +async fn prewarm_test_session(pool: &GuardianReviewSessionManager, session: GuardianReviewSession) { + struct ReadySession { + context: GuardianReviewSessionReuseKey, + session: Mutex>, + } + impl codex_guardian_reviewer::ReviewerSessionFactory for ReadySession { + type Session = GuardianReviewSession; + fn context( + &self, + _previous: Option<&GuardianReviewSession>, + ) -> GuardianReviewSessionReuseKey { + self.context.clone() + } + async fn spawn( + &self, + _context: GuardianReviewSessionReuseKey, + _kind: GuardianReviewSessionKind, + _snapshot: Option, + _cancellation: CancellationToken, + ) -> anyhow::Result { + Ok(self.session.lock().await.take().expect("one fixture spawn")) + } + } + let factory = ReadySession { + context: session.reuse_key.clone(), + session: Mutex::new(Some(session)), + }; + pool.prewarm(&factory).await.unwrap(); +} diff --git a/codex-rs/core/src/guardian/runtime.rs b/codex-rs/core/src/guardian/runtime.rs index 29772ec5e2..9abbb59d3d 100644 --- a/codex-rs/core/src/guardian/runtime.rs +++ b/codex-rs/core/src/guardian/runtime.rs @@ -1,7 +1,5 @@ -//! Binds the existing synchronous reviewer to one approval action. +//! Captures one approval action for the extension-owned synchronous reviewer. -use codex_extension_api::ExtensionFuture; -use codex_extension_api::SynchronousApprovalReviewer; use codex_protocol::protocol::ReviewDecision; use std::sync::Arc; @@ -10,7 +8,6 @@ use super::GuardianApprovalRequest; use super::GuardianReviewContext; use super::GuardianReviewOptions; use super::approval_request::guardian_approval_request_to_json; -use super::review::run_synchronous_review; use crate::session::session::Session; /// Carries the original action even when the legacy synchronous renderer cannot handle its paths. @@ -83,12 +80,3 @@ pub(super) struct ReviewRuntime { pub(super) reasons: ApprovalRequestReasons, pub(super) options: GuardianReviewOptions, } - -impl SynchronousApprovalReviewer for ReviewRuntime { - fn review( - &self, - reason: codex_protocol::approvals::GuardianReviewReason, - ) -> ExtensionFuture<'_, ReviewDecision> { - Box::pin(run_synchronous_review(self.clone(), reason)) - } -} diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index c4d85a5e88..4489eecbae 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -2491,6 +2491,9 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: .await; let committed_rollout_items = session .guardian_review_session() + .trunk() + .await + .expect("reviewer") .committed_fork_rollout_items_for_test() .await .expect("committed guardian fork snapshot"); @@ -2815,6 +2818,9 @@ async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow:: session .guardian_review_session() + .trunk() + .await + .expect("reviewer") .send_trunk_event_raw_for_test(Event { id: "stale-turn".to_string(), msg: EventMsg::TurnComplete(TurnCompleteEvent { diff --git a/codex-rs/core/src/guardian_review.rs b/codex-rs/core/src/guardian_review.rs index 0a6dab4f62..4778ece8e5 100644 --- a/codex-rs/core/src/guardian_review.rs +++ b/codex-rs/core/src/guardian_review.rs @@ -1,11 +1,5 @@ -//! Production synchronous reviewer helpers shared with the Guardian extension. -//! These use the current policy, output contract and reviewer configuration; -//! making them available does not select a new transcript mode or start a review. +//! Opaque host capabilities for the synchronous Guardian extension. +//! Core supplies context and runtime operations; the extension owns review policy and pooling. -pub use crate::guardian::GuardianAssessment; -pub use crate::guardian::build_guardian_review_session_config; -pub use crate::guardian::guardian_output_schema; -pub use crate::guardian::parse_guardian_assessment; - -// The extension owns this instance; the existing review implementation stays shared. -pub use crate::guardian::GuardianReviewSessionManager; +pub use crate::guardian::GuardianReviewSession; +pub use crate::guardian::GuardianReviewSessionHost; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0968a54b03..f50da45137 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -8609,145 +8609,6 @@ async fn shutdown_and_wait_waits_when_shutdown_is_already_in_progress() { .expect("shutdown waiter"); } -#[tokio::test] -async fn shutdown_and_wait_shuts_down_cached_guardian_subagent() { - let (parent_session, parent_turn_context) = make_session_and_context().await; - let parent_session = Arc::new(parent_session); - let parent_config = Arc::clone(&parent_turn_context.config); - let (parent_tx_sub, parent_rx_sub) = async_channel::bounded(4); - let (_parent_tx_event, parent_rx_event) = async_channel::unbounded(); - let parent_session_for_loop = Arc::clone(&parent_session); - let parent_session_loop_handle = tokio::spawn(async move { - submission_loop(parent_session_for_loop, parent_config, parent_rx_sub).await; - }); - let parent_io = SessionIo { - tx_sub: parent_tx_sub, - rx_event: parent_rx_event, - agent_status: watch::channel(AgentStatus::PendingInit).1, - session_loop_termination: session_loop_termination_from_handle(parent_session_loop_handle), - }; - - let (child_session, _child_turn_context) = make_session_and_context().await; - let (child_tx_sub, child_rx_sub) = async_channel::bounded::(4); - let (_child_tx_event, child_rx_event) = async_channel::unbounded(); - let (child_shutdown_tx, child_shutdown_rx) = tokio::sync::oneshot::channel(); - let child_session_loop_handle = tokio::spawn(async move { - let shutdown = child_rx_sub - .recv() - .await - .expect("child shutdown submission"); - assert!(matches!(shutdown.op, Op::Shutdown)); - child_shutdown_tx - .send(()) - .expect("child shutdown signal should be delivered"); - }); - let child_session = Arc::new(child_session); - let child_io = SessionIo { - tx_sub: child_tx_sub, - rx_event: child_rx_event, - agent_status: watch::channel(AgentStatus::PendingInit).1, - session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), - }; - parent_session - .guardian_review_session() - .cache_for_test(child_session, child_io) - .await; - - parent_io - .shutdown_and_wait() - .await - .expect("parent shutdown should succeed"); - - child_shutdown_rx - .await - .expect("guardian subagent should receive a shutdown op"); -} - -#[tokio::test] -async fn cached_guardian_subagent_exposes_its_rollout_path() { - let (parent_session, _parent_turn_context) = make_session_and_context().await; - let parent_session = Arc::new(parent_session); - - let (mut child_session, _child_turn_context) = make_session_and_context().await; - let child_rollout_path = attach_thread_persistence(&mut child_session).await; - let (child_tx_sub, _child_rx_sub) = async_channel::bounded(4); - let (_child_tx_event, child_rx_event) = async_channel::unbounded(); - let child_session_loop_handle = tokio::spawn(async {}); - let child_session = Arc::new(child_session); - let child_io = SessionIo { - tx_sub: child_tx_sub, - rx_event: child_rx_event, - agent_status: watch::channel(AgentStatus::PendingInit).1, - session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), - }; - parent_session - .guardian_review_session() - .cache_for_test(child_session, child_io) - .await; - - assert_eq!( - parent_session - .guardian_review_session() - .trunk_rollout_path() - .await, - Some(child_rollout_path) - ); -} - -#[tokio::test] -async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() { - let (parent_session, parent_turn_context) = make_session_and_context().await; - let parent_session = Arc::new(parent_session); - let parent_config = Arc::clone(&parent_turn_context.config); - let (parent_tx_sub, parent_rx_sub) = async_channel::bounded(4); - let (_parent_tx_event, parent_rx_event) = async_channel::unbounded(); - let parent_session_for_loop = Arc::clone(&parent_session); - let parent_session_loop_handle = tokio::spawn(async move { - submission_loop(parent_session_for_loop, parent_config, parent_rx_sub).await; - }); - let parent_io = SessionIo { - tx_sub: parent_tx_sub, - rx_event: parent_rx_event, - agent_status: watch::channel(AgentStatus::PendingInit).1, - session_loop_termination: session_loop_termination_from_handle(parent_session_loop_handle), - }; - - let (child_session, _child_turn_context) = make_session_and_context().await; - let (child_tx_sub, child_rx_sub) = async_channel::bounded::(4); - let (_child_tx_event, child_rx_event) = async_channel::unbounded(); - let (child_shutdown_tx, child_shutdown_rx) = tokio::sync::oneshot::channel(); - let child_session_loop_handle = tokio::spawn(async move { - let shutdown = child_rx_sub - .recv() - .await - .expect("child shutdown submission"); - assert!(matches!(shutdown.op, Op::Shutdown)); - child_shutdown_tx - .send(()) - .expect("child shutdown signal should be delivered"); - }); - let child_session = Arc::new(child_session); - let child_io = SessionIo { - tx_sub: child_tx_sub, - rx_event: child_rx_event, - agent_status: watch::channel(AgentStatus::PendingInit).1, - session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), - }; - parent_session - .guardian_review_session() - .register_ephemeral_for_test(child_session, child_io) - .await; - - parent_io - .shutdown_and_wait() - .await - .expect("parent shutdown should succeed"); - - child_shutdown_rx - .await - .expect("ephemeral guardian review should receive a shutdown op"); -} - pub(crate) async fn make_session_and_context_with_auth_and_config_and_rx( auth: CodexAuth, dynamic_tools: Vec, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index e0dff00a6f..2c04f88a8f 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -277,10 +277,11 @@ async fn schedule_startup_prewarm_inner( let guardian_session = Arc::clone(&session); let guardian_parent_turn = Arc::clone(&startup_turn_context); drop(tokio::spawn(async move { - if let Err(err) = guardian_session - .guardian_review_session() - .initialize(Arc::clone(&guardian_session), guardian_parent_turn) - .await + if let Err(err) = crate::guardian::prewarm_guardian_review_session( + guardian_session, + guardian_parent_turn, + ) + .await { warn!("failed to initialize guardian review session: {err:#}"); } diff --git a/codex-rs/ext/extension-api/src/contributors/approval_review.rs b/codex-rs/ext/extension-api/src/contributors/approval_review.rs index 218998c61e..6b3f836d88 100644 --- a/codex-rs/ext/extension-api/src/contributors/approval_review.rs +++ b/codex-rs/ext/extension-api/src/contributors/approval_review.rs @@ -62,8 +62,8 @@ pub enum ApprovalDecision { AskUser, } -/// Runs the existing synchronous review for the bound action and cancellation token. -/// Implementations must not resolve policy or reuse an async score. +/// Runs an extension-owned synchronous review already bound to an action by the host. +/// Implementations must not resolve approval routing or reuse an async score. pub trait SynchronousApprovalReviewer: Send + Sync { fn review( &self, diff --git a/codex-rs/ext/guardian-reviewer/Cargo.toml b/codex-rs/ext/guardian-reviewer/Cargo.toml index 381e341f31..a19fe9bbb4 100644 --- a/codex-rs/ext/guardian-reviewer/Cargo.toml +++ b/codex-rs/ext/guardian-reviewer/Cargo.toml @@ -15,6 +15,7 @@ workspace = true [dependencies] anyhow = { workspace = true } codex-analytics = { workspace = true } +codex-extension-api = { workspace = true } codex-protocol = { workspace = true } rand = { workspace = true } serde = { workspace = true } diff --git a/codex-rs/ext/guardian-reviewer/src/circuit_breaker.rs b/codex-rs/ext/guardian-reviewer/src/circuit_breaker.rs index 7f1e5ee4eb..4034d50c01 100644 --- a/codex-rs/ext/guardian-reviewer/src/circuit_breaker.rs +++ b/codex-rs/ext/guardian-reviewer/src/circuit_breaker.rs @@ -24,6 +24,18 @@ pub enum GuardianRejectionCircuitBreakerPolicy { CyberModel, } +impl From<&codex_protocol::openai_models::ModelInfo> for GuardianRejectionCircuitBreakerPolicy { + fn from(model: &codex_protocol::openai_models::ModelInfo) -> Self { + if model.model_specialty.as_deref() + == Some(codex_protocol::openai_models::MODEL_SPECIALTY_CYBER) + { + Self::CyberModel + } else { + Self::Standard + } + } +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum GuardianRejectionCircuitBreakerAction { Continue, diff --git a/codex-rs/ext/guardian-reviewer/src/completion.rs b/codex-rs/ext/guardian-reviewer/src/completion.rs new file mode 100644 index 0000000000..2d19b17ad1 --- /dev/null +++ b/codex-rs/ext/guardian-reviewer/src/completion.rs @@ -0,0 +1,175 @@ +//! Converts reviewer outcomes into decisions and reports. The host publishes these +//! reports and applies them only to the action whose evidence was reviewed. + +use codex_analytics::GuardianReviewAnalyticsResult; +use codex_analytics::GuardianReviewDecision; +use codex_analytics::GuardianReviewTerminalStatus; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::protocol::GuardianAssessmentEvent; +use codex_protocol::protocol::GuardianAssessmentOutcome; +use codex_protocol::protocol::GuardianAssessmentStatus; +use codex_protocol::protocol::GuardianRiskLevel; +use codex_protocol::protocol::GuardianUserAuthorization; +use codex_protocol::protocol::ReviewDecision; + +use crate::GuardianAssessment; +use crate::GuardianReviewError; +use crate::GuardianReviewOutcome; + +const REJECTION_INSTRUCTIONS: &str = concat!( + "The agent must not attempt to achieve the same outcome via workaround, ", + "indirect execution, or policy circumvention. ", + "Proceed only with a materially safer alternative, ", + "or if the user explicitly approves the action after being informed of the risk. ", + "Otherwise, stop and request user input.", +); +const TIMEOUT_INSTRUCTIONS: &str = concat!( + "The automatic permission approval review did not finish before its deadline. ", + "Do not assume the action is unsafe based on the timeout alone. ", + "You may retry once, or ask the user for guidance or explicit approval.", +); + +pub fn guardian_timeout_message(model_info: &ModelInfo) -> String { + model_info + .model_messages + .as_ref() + .and_then(|messages| messages.auto_review.as_ref()) + .and_then(|messages| messages.timeout_instructions.as_deref()) + .unwrap_or(TIMEOUT_INSTRUCTIONS) + .to_string() +} + +pub struct ReviewCompletion { + pub decision: ReviewDecision, + pub event: GuardianAssessmentEvent, + pub warning: Option, + pub analytics: GuardianReviewAnalyticsResult, + /// Only completed assessments may enter the evidence cache or count as policy denials. + pub assessment_outcome: Option, +} + +pub fn complete_review( + outcome: GuardianReviewOutcome, + model: &ModelInfo, + mut event: GuardianAssessmentEvent, + mut analytics: GuardianReviewAnalyticsResult, +) -> ReviewCompletion { + let completed_assessment = match &outcome { + GuardianReviewOutcome::Completed(assessment) => Some(assessment.outcome), + GuardianReviewOutcome::Error(_) => None, + }; + let assessment = match outcome { + GuardianReviewOutcome::Completed(assessment) => { + let approved = matches!(assessment.outcome, GuardianAssessmentOutcome::Allow); + analytics.decision = if approved { + GuardianReviewDecision::Approved + } else { + GuardianReviewDecision::Denied + }; + analytics.terminal_status = if approved { + GuardianReviewTerminalStatus::Approved + } else { + GuardianReviewTerminalStatus::Denied + }; + analytics.failure_reason = None; + analytics.risk_level = Some(assessment.risk_level); + analytics.user_authorization = Some(assessment.user_authorization); + analytics.outcome = Some(assessment.outcome); + assessment + } + GuardianReviewOutcome::Error(error) => { + analytics.failure_reason = Some(error.failure_reason()); + match error { + GuardianReviewError::Timeout => { + let rationale = "Automatic approval review timed out while evaluating the requested approval.".to_string(); + analytics.decision = GuardianReviewDecision::Denied; + analytics.terminal_status = GuardianReviewTerminalStatus::TimedOut; + event.status = GuardianAssessmentStatus::TimedOut; + event.rationale = Some(rationale.clone()); + return ReviewCompletion { + decision: ReviewDecision::TimedOut, + event, + warning: Some(rationale), + analytics, + assessment_outcome: None, + }; + } + GuardianReviewError::Cancelled => { + analytics.decision = GuardianReviewDecision::Aborted; + analytics.terminal_status = GuardianReviewTerminalStatus::Aborted; + event.status = GuardianAssessmentStatus::Aborted; + return ReviewCompletion { + decision: ReviewDecision::Abort, + event, + warning: None, + analytics, + assessment_outcome: None, + }; + } + GuardianReviewError::PromptBuild { message } + | GuardianReviewError::Session { message, .. } + | GuardianReviewError::Parse { message } => { + analytics.decision = GuardianReviewDecision::Denied; + analytics.terminal_status = GuardianReviewTerminalStatus::FailedClosed; + GuardianAssessment { + risk_level: GuardianRiskLevel::High, + user_authorization: GuardianUserAuthorization::Unknown, + outcome: GuardianAssessmentOutcome::Deny, + rationale: format!("Automatic approval review failed: {message}"), + } + } + } + } + }; + let approved = matches!(assessment.outcome, GuardianAssessmentOutcome::Allow); + let verdict = if approved { "approved" } else { "denied" }; + let authorization = match assessment.user_authorization { + GuardianUserAuthorization::Unknown => "unknown", + GuardianUserAuthorization::Low => "low", + GuardianUserAuthorization::Medium => "medium", + GuardianUserAuthorization::High => "high", + }; + let risk = match assessment.risk_level { + GuardianRiskLevel::Low => "low", + GuardianRiskLevel::Medium => "medium", + GuardianRiskLevel::High => "high", + GuardianRiskLevel::Critical => "critical", + }; + let warning = format!( + "Automatic approval review {verdict} (risk: {risk}, authorization: {authorization}): {}", + assessment.rationale + ); + event.status = if approved { + GuardianAssessmentStatus::Approved + } else { + GuardianAssessmentStatus::Denied + }; + event.risk_level = Some(assessment.risk_level); + event.user_authorization = Some(assessment.user_authorization); + event.rationale = Some(assessment.rationale.clone()); + let decision = if approved { + ReviewDecision::Approved + } else { + let rationale = if assessment.rationale.trim().is_empty() { + "Auto-reviewer denied the action without a specific rationale." + } else { + assessment.rationale.trim() + }; + let instructions = model + .model_messages + .as_ref() + .and_then(|messages| messages.auto_review.as_ref()) + .and_then(|messages| messages.rejection_instructions.as_deref()) + .unwrap_or(REJECTION_INSTRUCTIONS); + ReviewDecision::denied(format!( + "This action was rejected due to unacceptable risk.\nReason: {rationale}\n{instructions}" + )) + }; + ReviewCompletion { + decision, + event, + warning: Some(warning), + analytics, + assessment_outcome: completed_assessment, + } +} diff --git a/codex-rs/ext/guardian-reviewer/src/deadline.rs b/codex-rs/ext/guardian-reviewer/src/deadline.rs new file mode 100644 index 0000000000..6a91bfac5d --- /dev/null +++ b/codex-rs/ext/guardian-reviewer/src/deadline.rs @@ -0,0 +1,36 @@ +//! Bounds session startup and context work by the review's existing deadline and cancellation. + +use crate::GuardianReviewSessionOutcome; +use std::future::Future; +use tokio_util::sync::CancellationToken; + +pub async fn run_before_review_deadline( + deadline: tokio::time::Instant, + external_cancel: Option<&CancellationToken>, + future: impl Future, +) -> Result { + tokio::select! { + _ = tokio::time::sleep_until(deadline) => Err(GuardianReviewSessionOutcome::TimedOut), + result = future => Ok(result), + _ = async { + if let Some(cancel_token) = external_cancel { + cancel_token.cancelled().await; + } else { + std::future::pending::<()>().await; + } + } => Err(GuardianReviewSessionOutcome::Aborted), + } +} + +pub async fn run_before_review_deadline_with_cancel( + deadline: tokio::time::Instant, + external_cancel: Option<&CancellationToken>, + cancel_token: &CancellationToken, + future: impl Future, +) -> Result { + let result = run_before_review_deadline(deadline, external_cancel, future).await; + if result.is_err() { + cancel_token.cancel(); + } + result +} diff --git a/codex-rs/ext/guardian-reviewer/src/lib.rs b/codex-rs/ext/guardian-reviewer/src/lib.rs index 4bb0ae8bcb..b697541485 100644 --- a/codex-rs/ext/guardian-reviewer/src/lib.rs +++ b/codex-rs/ext/guardian-reviewer/src/lib.rs @@ -3,9 +3,13 @@ mod assessment; mod circuit_breaker; +mod completion; +mod deadline; mod model; mod outcome; +mod pool; mod retry; +mod review; pub use assessment::GuardianAssessment; pub use assessment::guardian_output_contract_prompt; @@ -25,3 +29,18 @@ pub use retry::run_with_retry; pub const MAX_REVIEW_ATTEMPTS: i64 = 3; pub const REVIEW_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90); + +pub use deadline::run_before_review_deadline; +pub use deadline::run_before_review_deadline_with_cancel; +pub use pool::ReviewerPool; +pub use pool::ReviewerRequest; +pub use pool::ReviewerSession; +pub use pool::ReviewerSessionFactory; +pub use pool::SessionDisposition; + +pub use review::ReviewHost; +pub use review::SynchronousReview; + +pub use completion::ReviewCompletion; +pub use completion::complete_review; +pub use completion::guardian_timeout_message; diff --git a/codex-rs/ext/guardian-reviewer/src/pool.rs b/codex-rs/ext/guardian-reviewer/src/pool.rs new file mode 100644 index 0000000000..e3f1a1894d --- /dev/null +++ b/codex-rs/ext/guardian-reviewer/src/pool.rs @@ -0,0 +1,382 @@ +//! Owns the reusable reviewer and temporary forks for one parent thread. +//! The host owns session execution and context construction. Selection stays serialized; +//! concurrent reviews fork committed context and shutdown joins every tracked session. +//! Startup and fork futures stay boxed to bound the orchestration stack frames. + +use std::future::Future; +use std::sync::Arc; + +use codex_analytics::GuardianReviewAnalyticsResult; +use codex_analytics::GuardianReviewSessionKind; +use tokio::sync::Mutex; +use tokio::sync::Semaphore; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; + +use crate::GuardianReviewSessionOutcome; +use crate::run_before_review_deadline; +use crate::run_before_review_deadline_with_cancel; + +/// A host-owned reviewer session. Context and snapshots remain opaque to the pool. +/// Shutdown must cancel the runtime and await its termination. +pub trait ReviewerSession: Send + Sync + 'static { + type Context: Clone + PartialEq + Send + Sync; + type Snapshot: Send + Sync; + + fn context(&self) -> &Self::Context; + fn cancel(&self); + fn snapshot(&self) -> impl Future> + Send; + fn commit_snapshot(&self) -> impl Future + Send; + fn shutdown(&self) -> impl Future + Send; +} + +/// Builds a session from one captured parent context. The host must preserve that +/// context's authority and honor cancellation during spawning, including partial startup. +pub trait ReviewerSessionFactory: Send + Sync { + type Session: ReviewerSession; + + fn context( + &self, + previous: Option<&Self::Session>, + ) -> ::Context; + + fn spawn( + &self, + context: ::Context, + kind: GuardianReviewSessionKind, + snapshot: Option<::Snapshot>, + cancellation: CancellationToken, + ) -> impl Future> + Send; +} + +/// Executes one approval on a selected session. The host must drain the submitted +/// turn before returning Reusable, and must keep the issuing action and permissions bound. +pub trait ReviewerRequest: Send + Sync { + type Factory: ReviewerSessionFactory; + + fn factory(&self) -> &Self::Factory; + fn deadline(&self) -> Instant; + fn cancellation(&self) -> Option<&CancellationToken>; + fn run( + &self, + session: &::Session, + kind: GuardianReviewSessionKind, + ) -> impl Future< + Output = ( + GuardianReviewSessionOutcome, + SessionDisposition, + GuardianReviewAnalyticsResult, + ), + > + Send; +} + +/// Whether the host drained the session sufficiently for another review to use it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SessionDisposition { + Reusable, + Discard, +} + +/// Per-parent reviewer state. The same pool serves prewarm, review, invalidation and shutdown. +pub struct ReviewerPool { + state: Arc>>, + cancellation: CancellationToken, +} + +struct PoolState { + trunk: Option>>, + ephemeral_reviews: Vec>, +} + +struct Trunk { + session: Arc, + review_lock: Semaphore, +} + +impl Default for ReviewerPool { + fn default() -> Self { + Self { + state: Arc::new(Mutex::new(PoolState { + trunk: None, + ephemeral_reviews: Vec::new(), + })), + cancellation: CancellationToken::new(), + } + } +} + +impl ReviewerPool { + /// Returns the current reviewer handle for host inspection and feedback collection. + pub async fn trunk(&self) -> Option> { + self.state + .lock() + .await + .trunk + .as_ref() + .map(|trunk| Arc::clone(&trunk.session)) + } + + /// Prepares the first reviewer without replacing a review that won the startup race. + pub async fn prewarm( + &self, + factory: &impl ReviewerSessionFactory, + ) -> anyhow::Result<()> { + let cancellation = self.cancellation.child_token(); + let guard = cancellation.clone().drop_guard(); + let session = factory + .spawn( + factory.context(/*previous*/ None), + GuardianReviewSessionKind::TrunkNew, + /*snapshot*/ None, + cancellation.clone(), + ) + .await?; + let mut state = self.state.lock().await; + if !cancellation.is_cancelled() && state.trunk.is_none() { + state.trunk = Some(Arc::new(Trunk { + session: Arc::new(session), + review_lock: Semaphore::new(/*permits*/ 1), + })); + drop(guard.disarm()); + } + Ok(()) + } + + /// Permanently stops this parent's reviewer pool and waits for tracked runtimes. + pub async fn shutdown(&self) { + self.cancellation.cancel(); + self.invalidate().await; + } + + /// Drops reusable context after parent history rollback or another host invalidation. + pub async fn invalidate(&self) { + let (trunk, ephemeral) = { + let mut state = self.state.lock().await; + ( + state.trunk.take(), + std::mem::take(&mut state.ephemeral_reviews), + ) + }; + for session in trunk + .into_iter() + .map(|trunk| Arc::clone(&trunk.session)) + .chain(ephemeral) + { + if self.cancellation.is_cancelled() { + session.shutdown().await; + } else { + session.cancel(); + shutdown_in_background(session); + } + } + } + + /// Selects one reviewer; busy or incompatible trunks use an isolated temporary session. + #[expect( + clippy::await_holding_invalid_type, + reason = "reviewer selection and spawning stay serialized" + )] + pub async fn review( + &self, + request: R, + ) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult) + where + R: ReviewerRequest, + R::Factory: ReviewerSessionFactory, + { + let mut spawned_trunk = false; + let (trunk, context) = match run_before_review_deadline( + request.deadline(), + request.cancellation(), + self.state.lock(), + ) + .await + { + Ok(mut state) => { + let context = request + .factory() + .context(state.trunk.as_ref().map(|trunk| trunk.session.as_ref())); + if let Some(trunk) = state.trunk.as_ref() + && trunk.session.context() != &context + && trunk.review_lock.try_acquire().is_ok() + && let Some(stale) = state.trunk.take() + { + shutdown_in_background(Arc::clone(&stale.session)); + } + if state.trunk.is_none() { + let cancellation = self.cancellation.child_token(); + let session = match run_before_review_deadline_with_cancel( + request.deadline(), + request.cancellation(), + &cancellation, + Box::pin(request.factory().spawn( + context.clone(), + GuardianReviewSessionKind::TrunkNew, + /*snapshot*/ None, + cancellation.clone(), + )), + ) + .await + { + Ok(Ok(session)) => Arc::new(session), + Ok(Err(error)) => { + return ( + GuardianReviewSessionOutcome::PromptBuildFailed(error), + GuardianReviewAnalyticsResult::without_session(), + ); + } + Err(outcome) => { + return (outcome, GuardianReviewAnalyticsResult::without_session()); + } + }; + state.trunk = Some(Arc::new(Trunk { + session, + review_lock: Semaphore::new(/*permits*/ 1), + })); + spawned_trunk = true; + } + (state.trunk.as_ref().cloned(), context) + } + Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()), + }; + let Some(trunk) = trunk else { + return ( + GuardianReviewSessionOutcome::Completed(Err(anyhow::anyhow!( + "guardian review session was not available after spawn" + ))), + GuardianReviewAnalyticsResult::without_session(), + ); + }; + if trunk.session.context() != &context { + return Box::pin(self.review_ephemeral(&request, context, /*snapshot*/ None)).await; + } + let guard = match trunk.review_lock.try_acquire() { + Ok(guard) => guard, + Err(_) => { + return Box::pin(self.review_ephemeral( + &request, + context, + trunk.session.snapshot().await, + )) + .await; + } + }; + let kind = if spawned_trunk { + GuardianReviewSessionKind::TrunkNew + } else { + GuardianReviewSessionKind::TrunkReused + }; + let (outcome, disposition, analytics) = request.run(&trunk.session, kind).await; + if disposition == SessionDisposition::Reusable + && matches!(outcome, GuardianReviewSessionOutcome::Completed(_)) + { + trunk.session.commit_snapshot().await; + } + drop(guard); + if disposition == SessionDisposition::Discard { + let mut state = self.state.lock().await; + if state + .trunk + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &trunk)) + && let Some(removed) = state.trunk.take() + { + shutdown_in_background(Arc::clone(&removed.session)); + } + } + (outcome, analytics) + } + + async fn review_ephemeral( + &self, + request: &R, + context: S::Context, + snapshot: Option, + ) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult) + where + R: ReviewerRequest, + R::Factory: ReviewerSessionFactory, + { + let cancellation = self.cancellation.child_token(); + let session = match run_before_review_deadline_with_cancel( + request.deadline(), + request.cancellation(), + &cancellation, + Box::pin(request.factory().spawn( + context, + GuardianReviewSessionKind::EphemeralForked, + snapshot, + cancellation.clone(), + )), + ) + .await + { + Ok(Ok(session)) => Arc::new(session), + Ok(Err(error)) => { + return ( + GuardianReviewSessionOutcome::PromptBuildFailed(error), + GuardianReviewAnalyticsResult::without_session(), + ); + } + Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()), + }; + self.state + .lock() + .await + .ephemeral_reviews + .push(Arc::clone(&session)); + let mut cleanup = EphemeralCleanup { + state: Arc::clone(&self.state), + session: Some(Arc::clone(&session)), + }; + let (outcome, _, analytics) = request + .run(&session, GuardianReviewSessionKind::EphemeralForked) + .await; + let removed = { + let mut state = self.state.lock().await; + state + .ephemeral_reviews + .iter() + .position(|active| Arc::ptr_eq(active, &session)) + .map(|index| state.ephemeral_reviews.swap_remove(index)) + }; + if let Some(removed) = removed { + cleanup.session = None; + shutdown_in_background(removed); + } + (outcome, analytics) + } +} + +fn shutdown_in_background(session: Arc) { + drop(tokio::spawn(async move { + session.shutdown().await; + })); +} + +struct EphemeralCleanup { + state: Arc>>, + session: Option>, +} + +impl Drop for EphemeralCleanup { + fn drop(&mut self) { + let Some(session) = self.session.take() else { + return; + }; + let state = Arc::clone(&self.state); + drop(tokio::spawn(async move { + let removed = { + let mut state = state.lock().await; + state + .ephemeral_reviews + .iter() + .position(|active| Arc::ptr_eq(active, &session)) + .map(|index| state.ephemeral_reviews.swap_remove(index)) + }; + if let Some(removed) = removed { + removed.shutdown().await; + } + })); + } +} diff --git a/codex-rs/ext/guardian-reviewer/src/review.rs b/codex-rs/ext/guardian-reviewer/src/review.rs new file mode 100644 index 0000000000..cdb49634fa --- /dev/null +++ b/codex-rs/ext/guardian-reviewer/src/review.rs @@ -0,0 +1,73 @@ +//! Orchestrates one synchronous approval review over a host-bound action. +//! All callers share this policy; the host captures evidence and enforces the result. + +use std::future::Future; + +use codex_analytics::GuardianReviewAnalyticsResult; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::SynchronousApprovalReviewer; +use codex_protocol::approvals::GuardianReviewReason; +use codex_protocol::protocol::ReviewDecision; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; + +use crate::GuardianReviewOutcome; +use crate::GuardianReviewSessionLimits; + +/// Runtime operations bound to one immutable approval action and its issuing context. +/// Preparation captures trusted evidence. Completion must invalidate stale approvals +/// and may only satisfy the review gate, never expand the action's execution authority. +pub trait ReviewHost: Send + Sync { + type Prepared: Send + Sync; + + fn cancellation(&self) -> Option<&CancellationToken>; + fn prepare( + &self, + reason: GuardianReviewReason, + deadline: Instant, + ) -> impl Future> + Send; + fn attempt( + &self, + prepared: &Self::Prepared, + deadline: Instant, + ) -> impl Future + Send; + fn complete( + &self, + prepared: Self::Prepared, + outcome: GuardianReviewOutcome, + analytics: GuardianReviewAnalyticsResult, + ) -> impl Future + Send; +} + +/// One review bound by the host before Guardian's approval policy chooses to run it. +pub struct SynchronousReview { + host: H, +} + +impl SynchronousReview { + pub fn new(host: H) -> Self { + Self { host } + } +} + +impl SynchronousApprovalReviewer for SynchronousReview { + fn review(&self, reason: GuardianReviewReason) -> ExtensionFuture<'_, ReviewDecision> { + Box::pin(async move { + let deadline = Instant::now() + crate::REVIEW_TIMEOUT; + let prepared = match self.host.prepare(reason, deadline).await { + Ok(prepared) => prepared, + Err(decision) => return decision, + }; + let (outcome, analytics) = Box::pin(crate::run_with_retry( + GuardianReviewSessionLimits { + max_attempts: crate::MAX_REVIEW_ATTEMPTS, + deadline, + }, + self.host.cancellation(), + |deadline| self.host.attempt(&prepared, deadline), + )) + .await; + self.host.complete(prepared, outcome, analytics).await + }) + } +} diff --git a/codex-rs/ext/guardian-v2/Cargo.toml b/codex-rs/ext/guardian-v2/Cargo.toml index b2e57173b9..661f78fb6a 100644 --- a/codex-rs/ext/guardian-v2/Cargo.toml +++ b/codex-rs/ext/guardian-v2/Cargo.toml @@ -19,6 +19,7 @@ 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 } diff --git a/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs b/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs index 266314d157..3176a311e6 100644 --- a/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs +++ b/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs @@ -1,17 +1,19 @@ //! Installs the synchronous reviewer independently of async scorer startup. -//! Reviewer state lives in the existing thread store and uses the shared ThreadManager. +//! Reviewer policy and pool state live in the extension; core supplies context and runtime handles. use std::sync::Arc; use std::sync::Weak; use codex_core::ThreadManager; use codex_core::config::Config; -use codex_core::guardian_review::GuardianReviewSessionManager; +use codex_core::guardian_review::GuardianReviewSession; +use codex_core::guardian_review::GuardianReviewSessionHost; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadReadyInput; use codex_extension_api::ThreadStartInput; +use codex_guardian_reviewer::ReviewerPool; /// Owns reviewer state through the same thread manager as the parent conversation. #[derive(Clone, Debug)] @@ -34,8 +36,11 @@ impl ThreadLifecycleContributor for GuardianExtension { if input.session_source.is_internal() { return; } + input + .thread_store + .get_or_init(ReviewerPool::::default); input.thread_store.get_or_init(|| { - GuardianReviewSessionManager::with_thread_manager(self.thread_manager.clone()) + GuardianReviewSessionHost::with_thread_manager(self.thread_manager.clone()) }); }) } @@ -45,7 +50,7 @@ impl ThreadLifecycleContributor for GuardianExtension { input: ThreadReadyInput<'a, Config>, ) -> ExtensionFuture<'a, ()> { Box::pin(async move { - if let Some(sessions) = input.thread_store.get::() { + if let Some(sessions) = input.thread_store.get::() { sessions.mark_ready(); } })