diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 7f52da219b..c27820525e 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -1,4 +1,5 @@ use super::*; +use codex_extension_api::ApprovalReviewSource; use codex_mcp::ElicitationReviewRequest; use codex_mcp::ElicitationReviewer; use codex_mcp::ElicitationReviewerHandle; @@ -450,19 +451,20 @@ async fn review_guardian_mcp_elicitation( GuardianElicitationReview::ApprovalRequest(guardian_request) => *guardian_request, }; - let review_id = crate::guardian::new_guardian_review_id(); - let decision = crate::guardian::review_approval_request( + let automated = crate::tools::approval_dispatch::request_automated_approval( &session, &turn_context, - review_id.clone(), + crate::guardian::new_guardian_review_id(), guardian_request, + approvals_reviewer, /*retry_reason*/ None, + ApprovalReviewSource::MainTurn, ) .await; - Ok(Some( - mcp_elicitation_response_from_guardian_decision(session.as_ref(), &review_id, decision) - .await, - )) + Ok(Some(match automated { + Ok(automated) => mcp_elicitation_response_from_automated(&automated), + Err(message) => mcp_elicitation_decline_with_message(message), + })) } fn guardian_elicitation_review_request( @@ -604,18 +606,17 @@ fn mcp_elicitation_request_id(id: &RequestId) -> String { } } -async fn mcp_elicitation_response_from_guardian_decision( - session: &Session, - review_id: &str, - decision: ReviewDecision, +fn mcp_elicitation_response_from_automated( + automated: &crate::tools::approval_dispatch::AutomatedApprovalDecision, ) -> ElicitationResponse { - let denial_message = match decision { - ReviewDecision::Denied => { - Some(crate::guardian::guardian_rejection_message(session, review_id).await) - } + let denial_message = match &automated.decision { + ReviewDecision::Denied | ReviewDecision::TimedOut => Some(automated.denial_message()), _ => None, }; - mcp_elicitation_response_from_guardian_decision_parts(decision, denial_message) + mcp_elicitation_response_from_guardian_decision_parts( + automated.decision.clone(), + denial_message, + ) } fn mcp_elicitation_response_from_guardian_decision_parts( diff --git a/codex-rs/core/src/session/mcp_tests.rs b/codex-rs/core/src/session/mcp_tests.rs index 0e3664019b..698a5997c4 100644 --- a/codex-rs/core/src/session/mcp_tests.rs +++ b/codex-rs/core/src/session/mcp_tests.rs @@ -1,4 +1,6 @@ use super::*; +use crate::tools::approval_dispatch::AutomatedApprovalDecision; +use crate::tools::approval_dispatch::AutomatedApprovalSource; use rmcp::model::BooleanSchema; use rmcp::model::ElicitationSchema; use rmcp::model::PrimitiveSchema; @@ -267,3 +269,25 @@ fn guardian_decisions_map_to_elicitation_responses_without_session_state() { } ); } + +#[test] +fn extension_denial_maps_to_mcp_elicitation_response() { + let rationale = "the browser origin is outside the user's request"; + let response = mcp_elicitation_response_from_automated(&AutomatedApprovalDecision { + decision: ReviewDecision::Denied, + denial_message: Some(rationale.to_string()), + source: AutomatedApprovalSource::Extension, + }); + + assert_eq!( + response, + ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: Some(json!({ + "approvals_reviewer": ApprovalsReviewer::AutoReview, + "message": rationale, + })), + } + ); +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 08c691b9c1..12bc8d7f40 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2207,29 +2207,27 @@ impl Session { let active = self.active_turn.lock().await; active.as_ref().map(|active| Arc::clone(&active.turn_state)) }; - let review_id = crate::guardian::new_guardian_review_id(); - let session = Arc::clone(self); - let turn = Arc::clone(turn_context); let request = crate::guardian::GuardianApprovalRequest::RequestPermissions { id: call_id, turn_id: turn_context.sub_id.clone(), reason: args.reason, permissions: requested_permissions.clone(), }; - let review_rx = crate::guardian::spawn_approval_request_review( - session, - turn, - review_id, - request, - /*retry_reason*/ None, - codex_analytics::GuardianApprovalRequestSource::MainTurn, - cancellation_token.clone(), - ); - let decision = tokio::select! { - biased; - _ = cancellation_token.cancelled() => return None, - decision = review_rx => decision.unwrap_or(ReviewDecision::Denied), - }; + let automated = + crate::tools::approval_dispatch::request_automated_approval_with_cancel( + self, + turn_context, + crate::guardian::new_guardian_review_id(), + request, + turn_context.config.approvals_reviewer, + /*retry_reason*/ None, + codex_extension_api::ApprovalReviewSource::MainTurn, + cancellation_token, + ) + .await?; + let decision = automated + .map(|automated| automated.decision) + .unwrap_or(ReviewDecision::Denied); let response = match decision { ReviewDecision::Approved | ReviewDecision::ApprovedExecpolicyAmendment { .. } => { RequestPermissionsResponse { diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 0d3c9c033b..53371542ab 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -17,6 +17,12 @@ use codex_exec_server::EnvironmentManager; use codex_execpolicy::Decision; use codex_execpolicy::Evaluation; use codex_execpolicy::RuleMatch; +use codex_extension_api::ApprovalReviewContributor; +use codex_extension_api::ApprovalReviewError; +use codex_extension_api::ApprovalReviewInput; +use codex_extension_api::ApprovalReviewOutcome; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistryBuilder; use codex_features::Feature; use codex_model_provider::create_model_provider; use codex_protocol::config_types::ApprovalsReviewer; @@ -49,6 +55,22 @@ use tempfile::tempdir; use tokio::time::timeout; use tokio_util::sync::CancellationToken; +struct ApprovingPermissionsContributor; + +impl ApprovalReviewContributor for ApprovingPermissionsContributor { + fn review<'a>( + &'a self, + input: ApprovalReviewInput<'a>, + ) -> ExtensionFuture<'a, Result> { + Box::pin(async move { + assert_eq!(input.reviewer, ApprovalsReviewer::AutoReview); + Ok(ApprovalReviewOutcome::decision( + ReviewDecision::ApprovedForSession, + )) + }) + } +} + fn expect_text_output(output: &T) -> String where T: ToolOutput + ?Sized, @@ -167,6 +189,89 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { assert!(guardian_request.body_contains_text("need network")); } +#[tokio::test] +async fn request_permissions_uses_extension_decision_and_records_session_grant() { + let (mut session, mut turn_context) = make_session_and_context().await; + *session.active_turn.lock().await = Some(ActiveTurn::default()); + turn_context + .approval_policy + .set(AskForApproval::OnRequest) + .expect("test setup should allow updating approval policy"); + turn_context + .features + .enable(Feature::GuardianApproval) + .expect("test setup should allow enabling guardian approvals"); + let mut config = (*turn_context.config).clone(); + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + turn_context.config = Arc::new(config); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.approval_review_contributor(Arc::new(ApprovingPermissionsContributor)); + session.services.extensions = Arc::new(extensions.build()); + + let session = Arc::new(session); + let turn_context = Arc::new(turn_context); + let requested_permissions = RequestPermissionProfile { + network: Some(NetworkPermissions { + enabled: Some(true), + }), + ..RequestPermissionProfile::default() + }; + let environment = turn_context + .environments + .primary() + .expect("primary environment") + .selection(); + let response = session + .request_permissions_for_environment( + &turn_context, + "perm-call-extension".to_string(), + RequestPermissionsArgs { + environment_id: None, + reason: Some("need network".to_string()), + permissions: requested_permissions.clone(), + }, + environment, + CancellationToken::new(), + ) + .await; + + assert_eq!( + response, + Some(RequestPermissionsResponse { + permissions: requested_permissions.clone(), + scope: PermissionGrantScope::Session, + strict_auto_review: false, + }) + ); + assert_eq!( + session + .granted_session_permissions(codex_exec_server::LOCAL_ENVIRONMENT_ID) + .await, + Some(requested_permissions.into()) + ); + + let cancellation_token = CancellationToken::new(); + cancellation_token.cancel(); + let cancelled_response = session + .request_permissions_for_environment( + &turn_context, + "perm-call-extension-cancelled".to_string(), + RequestPermissionsArgs { + environment_id: None, + reason: Some("need more network".to_string()), + permissions: RequestPermissionProfile::default(), + }, + turn_context + .environments + .primary() + .expect("primary environment") + .selection(), + cancellation_token, + ) + .await; + assert_eq!(cancelled_response, None); +} + #[tokio::test] async fn request_permissions_guardian_review_stops_when_cancelled() { let server = start_mock_server().await; diff --git a/codex-rs/core/src/tools/approval_dispatch.rs b/codex-rs/core/src/tools/approval_dispatch.rs index 428cd7c861..fe27f10210 100644 --- a/codex-rs/core/src/tools/approval_dispatch.rs +++ b/codex-rs/core/src/tools/approval_dispatch.rs @@ -20,6 +20,7 @@ use codex_extension_api::ApprovalReviewSource; use codex_hooks::PermissionRequestDecision; use codex_otel::ToolDecisionSource; use codex_protocol::protocol::ReviewDecision; +use tokio_util::sync::CancellationToken; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum AutomatedApprovalSource { @@ -83,7 +84,6 @@ pub(crate) async fn request_automated_approval( retry_reason: retry_reason.as_deref(), source, }; - match session .services .extensions @@ -119,6 +119,93 @@ pub(crate) async fn request_automated_approval( } } +#[allow(clippy::too_many_arguments)] +pub(crate) async fn request_automated_approval_with_cancel( + session: &std::sync::Arc, + turn: &std::sync::Arc, + review_id: String, + request: GuardianApprovalRequest, + reviewer: codex_protocol::config_types::ApprovalsReviewer, + retry_reason: Option, + source: ApprovalReviewSource, + cancellation_token: CancellationToken, +) -> Option> { + let action = guardian_assessment_action(&request); + let prompt = match format_guardian_action_pretty(&request) { + Ok(prompt) => prompt.text, + Err(error) => { + return Some(Err(format!("failed to render review prompt: {error}"))); + } + }; + let approval_policy = turn.approval_policy.value(); + let review_input = ApprovalReviewInput { + session_store: &session.services.session_extension_data, + thread_store: &session.services.thread_extension_data, + turn_store: turn.extension_data.as_ref(), + review_id: &review_id, + turn_id: guardian_request_turn_id(&request, &turn.sub_id), + target_item_id: guardian_request_target_item_id(&request), + prompt: &prompt, + action: &action, + reviewer, + approval_policy: &approval_policy, + retry_reason: retry_reason.as_deref(), + source, + }; + + let extension_outcome = tokio::select! { + biased; + _ = cancellation_token.cancelled() => return None, + outcome = session.services.extensions.approval_review(review_input) => outcome, + }; + match extension_outcome { + Ok(ApprovalReviewOutcome::Decision { + decision, + denial_message, + }) => Some(Ok(AutomatedApprovalDecision { + decision, + denial_message, + source: AutomatedApprovalSource::Extension, + })), + Ok(ApprovalReviewOutcome::Abstain) => { + let review_rx = crate::guardian::spawn_approval_request_review( + std::sync::Arc::clone(session), + std::sync::Arc::clone(turn), + review_id.clone(), + request, + retry_reason, + match source { + ApprovalReviewSource::MainTurn => { + codex_analytics::GuardianApprovalRequestSource::MainTurn + } + ApprovalReviewSource::DelegatedSubagent => { + codex_analytics::GuardianApprovalRequestSource::DelegatedSubagent + } + }, + cancellation_token.clone(), + ); + let decision = tokio::select! { + biased; + _ = cancellation_token.cancelled() => return None, + decision = review_rx => decision.unwrap_or(ReviewDecision::Denied), + }; + let denial_message = match decision { + ReviewDecision::Denied | ReviewDecision::Abort => { + Some(guardian_rejection_message(session, &review_id).await) + } + ReviewDecision::TimedOut => Some(guardian_timeout_message()), + _ => None, + }; + Some(Ok(AutomatedApprovalDecision { + decision, + denial_message, + source: AutomatedApprovalSource::Guardian, + })) + } + Err(error) => Some(Err(format!("automatic approval review failed: {error}"))), + } +} + pub(crate) async fn request_approval( tool: &mut T, req: &Rq,