mirror of
https://github.com/openai/codex.git
synced 2026-10-31 19:28:49 +00:00
Serialize manual network approval prompts per endpoint (#58351)
## Why Concurrent executions accessing the same endpoint can produce duplicate network approval prompts before a session decision is recorded. ## What changed - Queue manual prompts by environment, host, protocol, and port while keeping execution decisions separate. - After a one-time approval or denial, let the next execution prompt. Reuse committed session grants and deny rules without another prompt. - Keep tool results paused while requests are queued and attribute the wait to the triggering tool call. ## Testing Extend coverage for concurrent requests, one-time and session decisions, independent ports, queued cancellation and timing, and isolation between local and remote environments. GitOrigin-RevId: ebf580dddc39a9a9dc3b0d1acfa21438fbc6e100
This commit is contained in:
@@ -3070,6 +3070,22 @@ impl Session {
|
||||
timing
|
||||
}
|
||||
|
||||
/// Keep tool results paused and attribute a queued network prompt to its triggering call.
|
||||
pub(crate) fn hold_queued_network_approval(
|
||||
&self,
|
||||
turn: &TurnContext,
|
||||
call_id: Option<&str>,
|
||||
) -> (
|
||||
crate::elicitation::ElicitationRegistration,
|
||||
Option<codex_direct_tool_timing::ToolBlockingTimingGuard>,
|
||||
) {
|
||||
let owner = self.user_wait_timing_owner(turn, call_id);
|
||||
(
|
||||
self.services.elicitations.register(),
|
||||
self.start_user_wait_timing(turn, owner.as_deref()),
|
||||
)
|
||||
}
|
||||
|
||||
/// Emit an exec approval request event and await the user's decision.
|
||||
///
|
||||
/// The request is keyed by `call_id` + `approval_id` so matching responses
|
||||
|
||||
@@ -474,6 +474,14 @@ pub(crate) struct ApprovalResolution {
|
||||
}
|
||||
|
||||
impl ApprovalResolution {
|
||||
pub(super) fn user(decision: ReviewDecision) -> Self {
|
||||
Self {
|
||||
decision,
|
||||
source: ApprovalResolutionSource::User,
|
||||
cancellation_reason: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn into_tool_result(self, model_info: &ModelInfo) -> Result<ReviewDecision, ToolError> {
|
||||
let source = self.source;
|
||||
match self.decision {
|
||||
@@ -538,7 +546,21 @@ impl Session {
|
||||
action: ApprovalAction,
|
||||
ctx: ApprovalContext,
|
||||
) -> ApprovalResult {
|
||||
if let Err(err) = self.check_approval_requirement(&ctx).await {
|
||||
self.request_approval_with_user_prompt(
|
||||
&action,
|
||||
&ctx,
|
||||
self.request_user_approval(&action, &ctx),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(super) async fn request_approval_with_user_prompt(
|
||||
self: &Arc<Self>,
|
||||
action: &ApprovalAction,
|
||||
ctx: &ApprovalContext,
|
||||
user_prompt: impl Future<Output = ApprovalResolution>,
|
||||
) -> ApprovalResult {
|
||||
if let Err(err) = self.check_approval_requirement(ctx).await {
|
||||
return ApprovalResult {
|
||||
result: Err(err),
|
||||
source: None,
|
||||
@@ -548,7 +570,7 @@ impl Session {
|
||||
// Retained authority outside current grants needs a new sandbox approval.
|
||||
// Current grants permit review of the input; they never approve the input itself.
|
||||
let policy = ctx.review_context.turn().approval_policy();
|
||||
if matches!(&action, ApprovalAction::WriteStdin { sandbox_permissions, additional_permissions_preapproved, .. }
|
||||
if matches!(action, ApprovalAction::WriteStdin { sandbox_permissions, additional_permissions_preapproved, .. }
|
||||
if sandbox_permissions.requests_sandbox_override()
|
||||
&& (!additional_permissions_preapproved || matches!(policy, AskForApproval::Never)))
|
||||
&& !(ctx.strict_auto_review && matches!(policy, AskForApproval::Never))
|
||||
@@ -561,9 +583,9 @@ impl Session {
|
||||
cancellation_reason: None,
|
||||
};
|
||||
}
|
||||
let is_mcp_tool_call = matches!(&action, ApprovalAction::McpToolCall { .. });
|
||||
let is_network_approval = matches!(&action, ApprovalAction::NetworkAccess { .. });
|
||||
let permission_request_run_id = match &action {
|
||||
let is_mcp_tool_call = matches!(action, ApprovalAction::McpToolCall { .. });
|
||||
let is_network_approval = matches!(action, ApprovalAction::NetworkAccess { .. });
|
||||
let permission_request_run_id = match action {
|
||||
ApprovalAction::NetworkAccess { hook_run_id, .. } => hook_run_id.clone(),
|
||||
_ if ctx.retry_reason.is_some() => format!("{}:retry", ctx.call_id),
|
||||
_ => ctx.call_id.clone(),
|
||||
@@ -593,11 +615,12 @@ impl Session {
|
||||
cancellation_reason: None,
|
||||
},
|
||||
Some(PermissionRequestDecision::Allow) | None => {
|
||||
self.request_reviewer_approval(action, &ctx).await
|
||||
self.request_reviewer_approval(action, ctx, user_prompt)
|
||||
.await
|
||||
}
|
||||
};
|
||||
if resolution.decision == ReviewDecision::Approved {
|
||||
if let Err(ToolError::Rejected(reason)) = self.check_approval_requirement(&ctx).await {
|
||||
if let Err(ToolError::Rejected(reason)) = self.check_approval_requirement(ctx).await {
|
||||
resolution.decision = ReviewDecision::denied(reason);
|
||||
}
|
||||
if ctx.requires_human_approval()
|
||||
@@ -611,7 +634,7 @@ impl Session {
|
||||
}
|
||||
// Network approvals record their final telemetry after validation and persistence.
|
||||
if !is_network_approval {
|
||||
record_resolution(&ctx, &resolution);
|
||||
record_resolution(ctx, &resolution);
|
||||
}
|
||||
let source = Some(resolution.source);
|
||||
let cancellation_reason = resolution.cancellation_reason;
|
||||
@@ -657,11 +680,12 @@ impl Session {
|
||||
|
||||
async fn request_reviewer_approval(
|
||||
self: &Arc<Self>,
|
||||
action: ApprovalAction,
|
||||
action: &ApprovalAction,
|
||||
ctx: &ApprovalContext,
|
||||
user_prompt: impl Future<Output = ApprovalResolution>,
|
||||
) -> ApprovalResolution {
|
||||
if ctx.requires_human_approval() {
|
||||
self.request_user_approval(&action, ctx).await
|
||||
user_prompt.await
|
||||
} else if let Some(decision) = self.request_guardian_approval(action.clone(), ctx).await {
|
||||
ApprovalResolution {
|
||||
cancellation_reason: (decision == ReviewDecision::Abort)
|
||||
@@ -670,7 +694,7 @@ impl Session {
|
||||
source: ApprovalResolutionSource::Guardian,
|
||||
}
|
||||
} else {
|
||||
self.request_user_approval(&action, ctx).await
|
||||
user_prompt.await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -754,7 +778,7 @@ impl Session {
|
||||
}
|
||||
}
|
||||
|
||||
async fn request_user_approval(
|
||||
pub(super) async fn request_user_approval(
|
||||
&self,
|
||||
action: &ApprovalAction,
|
||||
ctx: &ApprovalContext,
|
||||
@@ -948,11 +972,7 @@ impl Session {
|
||||
unreachable!("permission requests are routed directly to Guardian")
|
||||
}
|
||||
};
|
||||
ApprovalResolution {
|
||||
decision,
|
||||
source: ApprovalResolutionSource::User,
|
||||
cancellation_reason: None,
|
||||
}
|
||||
ApprovalResolution::user(decision)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -154,7 +154,8 @@ async fn non_utf8_cwd_preserves_approval_routing(
|
||||
tty: false,
|
||||
proposed_execpolicy_amendment: None,
|
||||
};
|
||||
let approval = session.request_reviewer_approval(action, &context);
|
||||
let prompt = session.request_user_approval(&action, &context);
|
||||
let approval = session.request_reviewer_approval(&action, &context, prompt);
|
||||
tokio::pin!(approval);
|
||||
let expected = if reviewer == ApprovalsReviewer::User {
|
||||
tokio::select! {
|
||||
@@ -225,8 +226,9 @@ async fn explicit_mcp_reviewer_override_takes_precedence_over_action_context() {
|
||||
network_approval_context: None,
|
||||
};
|
||||
|
||||
let prompt = session.request_user_approval(&action, &context);
|
||||
tokio::select! {
|
||||
resolution = session.request_reviewer_approval(action, &context) => {
|
||||
resolution = session.request_reviewer_approval(&action, &context, prompt) => {
|
||||
panic!("expected a user approval request, got {resolution:?}");
|
||||
}
|
||||
event = events.recv() => {
|
||||
|
||||
@@ -8,6 +8,7 @@ use crate::session::turn_context::TurnContext;
|
||||
use crate::session::turn_context::TurnEnvironment;
|
||||
use crate::tools::approvals::ApprovalAction;
|
||||
use crate::tools::approvals::ApprovalContext;
|
||||
use crate::tools::approvals::ApprovalResolution;
|
||||
use crate::tools::events::truncate_rejection_message;
|
||||
use crate::tools::sandboxing::ToolError;
|
||||
use codex_features::Feature;
|
||||
@@ -47,12 +48,15 @@ use std::sync::OnceLock as SyncOnceLock;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::Notify;
|
||||
use tokio::sync::OnceCell;
|
||||
use tokio::sync::OwnedSemaphorePermit;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::error;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
const REASON_NOT_ALLOWED: &str = "not_allowed";
|
||||
const ABANDONED_NETWORK_APPROVAL_MESSAGE: &str =
|
||||
"network approval was cancelled before a decision was returned";
|
||||
|
||||
@@ -204,7 +208,7 @@ impl PendingApprovalDecision {
|
||||
fn to_network_decision(self) -> NetworkDecision {
|
||||
match self {
|
||||
Self::AllowOnce | Self::AllowForSession => NetworkDecision::Allow,
|
||||
Self::Deny => NetworkDecision::deny("not_allowed"),
|
||||
Self::Deny => NetworkDecision::deny(REASON_NOT_ALLOWED),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -268,11 +272,17 @@ struct NetworkApprovalCallState {
|
||||
call_outcomes: HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct PendingNetworkApprovals {
|
||||
executions: HashMap<PendingHostApprovalKey, Arc<PendingHostApproval>>,
|
||||
manual_prompts: HashMap<HostApprovalKey, Arc<Semaphore>>,
|
||||
}
|
||||
|
||||
pub(crate) struct NetworkApprovalService {
|
||||
// Pending-request cleanup must publish a tool outcome synchronously before cancelling it.
|
||||
calls: SyncMutex<NetworkApprovalCallState>,
|
||||
// Owner cleanup runs from Drop, so this lock cannot require an async task.
|
||||
pending_host_approvals: SyncMutex<HashMap<PendingHostApprovalKey, Arc<PendingHostApproval>>>,
|
||||
pending_host_approvals: SyncMutex<PendingNetworkApprovals>,
|
||||
// Keep persisted session policy and the in-memory approval caches in the same order.
|
||||
session_policy_commit_lock: Mutex<()>,
|
||||
session_approved_hosts: Mutex<HashSet<HostApprovalKey>>,
|
||||
@@ -291,6 +301,7 @@ struct PendingHostApprovalOwner<'a> {
|
||||
disconnect: Option<NetworkRequestDisconnect>,
|
||||
cancellation: Option<NetworkRequestCancellation>,
|
||||
decision_on_drop: PendingApprovalDecision,
|
||||
prompt_permit: Option<OwnedSemaphorePermit>,
|
||||
completed: bool,
|
||||
}
|
||||
|
||||
@@ -309,10 +320,62 @@ impl<'a> PendingHostApprovalOwner<'a> {
|
||||
disconnect: None,
|
||||
cancellation: None,
|
||||
decision_on_drop: PendingApprovalDecision::Deny,
|
||||
prompt_permit: None,
|
||||
completed: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Keep the endpoint's prompt until this execution commits or abandons its decision.
|
||||
async fn review_user_prompt(
|
||||
&mut self,
|
||||
session: &Session,
|
||||
turn: &TurnContext,
|
||||
timing_call_id: Option<&str>,
|
||||
policy_denial: &str,
|
||||
prompt: impl Future<Output = ApprovalResolution>,
|
||||
) -> ApprovalResolution {
|
||||
let queued = {
|
||||
let mut approvals = self
|
||||
.service
|
||||
.pending_host_approvals
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let queue = approvals
|
||||
.manual_prompts
|
||||
.entry(self.key.host.clone())
|
||||
.or_insert_with(|| Arc::new(Semaphore::new(/*permits*/ 1)));
|
||||
if let Ok(permit) = Arc::clone(queue).try_acquire_owned() {
|
||||
self.prompt_permit = Some(permit);
|
||||
None
|
||||
} else {
|
||||
Some((
|
||||
Arc::clone(queue),
|
||||
session.hold_queued_network_approval(turn, timing_call_id),
|
||||
))
|
||||
}
|
||||
};
|
||||
// Keep the queue's elicitation and timing guards through any handoff to the visible prompt.
|
||||
if let Some((queue, _hold)) = &queued {
|
||||
tracing::debug!(
|
||||
name: "codex.network_approval.queued",
|
||||
execution_id = self.key.execution_id.as_deref(),
|
||||
"waiting for the current endpoint approval"
|
||||
);
|
||||
let Ok(permit) = Arc::clone(queue).acquire_owned().await else {
|
||||
// The prompt owner closes the queue when it commits a session grant.
|
||||
return ApprovalResolution::user(ReviewDecision::Approved);
|
||||
};
|
||||
self.prompt_permit = Some(permit);
|
||||
}
|
||||
match self.service.session_host_decision(&self.key.host).await {
|
||||
Some(NetworkDecision::Allow) => ApprovalResolution::user(ReviewDecision::Approved),
|
||||
Some(NetworkDecision::Deny { .. }) => {
|
||||
ApprovalResolution::user(ReviewDecision::denied(policy_denial))
|
||||
}
|
||||
None => prompt.await,
|
||||
}
|
||||
}
|
||||
|
||||
fn set_decision_on_drop(&mut self, decision: PendingApprovalDecision) {
|
||||
self.decision_on_drop = decision;
|
||||
}
|
||||
@@ -331,7 +394,7 @@ impl<'a> PendingHostApprovalOwner<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
fn publish_and_remove(&self, decision: PendingApprovalDecision) {
|
||||
fn publish_and_remove(&mut self, decision: PendingApprovalDecision) {
|
||||
// Remove this generation before waking its waiters so a concurrent retry
|
||||
// cannot attach to an already-completed approval.
|
||||
let mut approvals = self
|
||||
@@ -339,11 +402,24 @@ impl<'a> PendingHostApprovalOwner<'a> {
|
||||
.pending_host_approvals
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(queue) = approvals.manual_prompts.get(&self.key.host) {
|
||||
let permit = self.prompt_permit.take();
|
||||
if permit.is_some() && decision == PendingApprovalDecision::AllowForSession {
|
||||
queue.close();
|
||||
}
|
||||
// Returning the permit lets one waiter take over the prompt.
|
||||
drop(permit);
|
||||
// Each queued or active execution holds a reference in addition to this map.
|
||||
if queue.is_closed() || Arc::strong_count(queue) == 1 {
|
||||
approvals.manual_prompts.remove(&self.key.host);
|
||||
}
|
||||
}
|
||||
if approvals
|
||||
.executions
|
||||
.get(&self.key)
|
||||
.is_some_and(|current| Arc::ptr_eq(current, &self.pending))
|
||||
{
|
||||
approvals.remove(&self.key);
|
||||
approvals.executions.remove(&self.key);
|
||||
}
|
||||
drop(approvals);
|
||||
self.pending.set_decision(decision);
|
||||
@@ -386,7 +462,7 @@ impl Default for NetworkApprovalService {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
calls: SyncMutex::new(NetworkApprovalCallState::default()),
|
||||
pending_host_approvals: SyncMutex::new(HashMap::new()),
|
||||
pending_host_approvals: SyncMutex::default(),
|
||||
session_policy_commit_lock: Mutex::new(()),
|
||||
session_approved_hosts: Mutex::new(HashSet::new()),
|
||||
session_denied_hosts: Mutex::new(HashSet::new()),
|
||||
@@ -520,15 +596,31 @@ impl NetworkApprovalService {
|
||||
.pending_host_approvals
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(existing) = pending.get(&key).cloned() {
|
||||
if let Some(existing) = pending.executions.get(&key).cloned() {
|
||||
return (existing, false);
|
||||
}
|
||||
|
||||
let created = Arc::new(PendingHostApproval::new());
|
||||
pending.insert(key, Arc::clone(&created));
|
||||
pending.executions.insert(key, Arc::clone(&created));
|
||||
(created, true)
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "read both host caches under the session policy commit lock"
|
||||
)]
|
||||
async fn session_host_decision(&self, key: &HostApprovalKey) -> Option<NetworkDecision> {
|
||||
let _commit_guard = self.session_policy_commit_lock.lock().await;
|
||||
if self.session_denied_hosts.lock().await.contains(key) {
|
||||
return Some(NetworkDecision::deny(REASON_NOT_ALLOWED));
|
||||
}
|
||||
self.session_approved_hosts
|
||||
.lock()
|
||||
.await
|
||||
.contains(key)
|
||||
.then_some(NetworkDecision::Allow)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
async fn take_call_outcome(&self, registration_id: &str) -> Option<String> {
|
||||
let mut calls = self
|
||||
@@ -622,17 +714,11 @@ impl NetworkApprovalService {
|
||||
)
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "persisted session policy and its in-memory caches must commit atomically"
|
||||
)]
|
||||
pub(crate) async fn handle_inline_policy_request(
|
||||
&self,
|
||||
session: Arc<Session>,
|
||||
request: NetworkPolicyRequest,
|
||||
) -> NetworkDecision {
|
||||
const REASON_NOT_ALLOWED: &str = "not_allowed";
|
||||
|
||||
let protocol = match request.protocol {
|
||||
NetworkProtocol::Http => NetworkApprovalProtocol::Http,
|
||||
NetworkProtocol::HttpsConnect => NetworkApprovalProtocol::Https,
|
||||
@@ -663,19 +749,8 @@ impl NetworkApprovalService {
|
||||
};
|
||||
let key = HostApprovalKey::from_request(&request, protocol, environment_id.clone());
|
||||
|
||||
{
|
||||
let _commit_guard = self.session_policy_commit_lock.lock().await;
|
||||
{
|
||||
let denied_hosts = self.session_denied_hosts.lock().await;
|
||||
if denied_hosts.contains(&key) {
|
||||
return NetworkDecision::deny(REASON_NOT_ALLOWED);
|
||||
}
|
||||
}
|
||||
|
||||
let approved_hosts = self.session_approved_hosts.lock().await;
|
||||
if approved_hosts.contains(&key) {
|
||||
return NetworkDecision::Allow;
|
||||
}
|
||||
if let Some(decision) = self.session_host_decision(&key).await {
|
||||
return decision;
|
||||
}
|
||||
|
||||
let target = Self::format_network_target(key.protocol, request.host.as_str(), key.port);
|
||||
@@ -809,7 +884,17 @@ impl NetworkApprovalService {
|
||||
retry_reason: Some(policy_denial_message.clone()),
|
||||
network_approval_context: Some(network_approval_context.clone()),
|
||||
};
|
||||
let approval = session.request_approval(action, approval_context);
|
||||
let user_prompt = pending_owner.review_user_prompt(
|
||||
&session,
|
||||
&turn_context,
|
||||
owner_call
|
||||
.as_ref()
|
||||
.map(|call| call.trigger.call_id.as_str()),
|
||||
&policy_denial_message,
|
||||
session.request_user_approval(&action, &approval_context),
|
||||
);
|
||||
let approval =
|
||||
session.request_approval_with_user_prompt(&action, &approval_context, user_prompt);
|
||||
let approval_result = if let Some(owner_call) = &owner_call {
|
||||
// A ready review result takes precedence over concurrent cancellation.
|
||||
let Some(result) = owner_call
|
||||
|
||||
@@ -1,15 +1,18 @@
|
||||
use super::*;
|
||||
use crate::sandboxing::SandboxPermissions;
|
||||
use codex_direct_tool_timing::ToolBatchTiming;
|
||||
use codex_network_proxy::BlockedRequestArgs;
|
||||
use codex_protocol::models::PermissionProfile;
|
||||
use codex_protocol::permissions::NetworkSandboxPolicy;
|
||||
use codex_protocol::protocol::AskForApproval;
|
||||
use codex_utils_critical_path::CriticalPath;
|
||||
use codex_utils_path_uri::PathUri;
|
||||
use core_test_support::PathBufExt;
|
||||
use core_test_support::test_path_buf;
|
||||
use futures::poll;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
use test_case::test_case;
|
||||
use tokio::time::timeout;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
@@ -60,7 +63,7 @@ async fn execution_cancellation_respects_network_approval_boundary(
|
||||
}
|
||||
|
||||
let (session, _, events) = make_session_and_context_with_rx().await;
|
||||
let (turn, _) = session
|
||||
let (mut turn, _) = session
|
||||
.new_turn_with_sub_id(
|
||||
"active-turn".to_string(),
|
||||
crate::session::SessionSettingsUpdate {
|
||||
@@ -75,9 +78,17 @@ async fn execution_cancellation_respects_network_approval_boundary(
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
Arc::make_mut(&mut Arc::get_mut(&mut turn).expect("unique test turn").config)
|
||||
.features
|
||||
.enable(Feature::CriticalPathTiming)
|
||||
.unwrap();
|
||||
let finish_turn = CancellationToken::new();
|
||||
session
|
||||
.start_task(turn, Vec::new(), PendingTask(finish_turn.clone()))
|
||||
.start_task(
|
||||
Arc::clone(&turn),
|
||||
Vec::new(),
|
||||
PendingTask(finish_turn.clone()),
|
||||
)
|
||||
.await;
|
||||
let service = &session.services.network_approval;
|
||||
register_call_with_default_shell_trigger(service, "execution-1").await;
|
||||
@@ -110,6 +121,7 @@ async fn execution_cancellation_respects_network_approval_boundary(
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let mut queued = None;
|
||||
let expected_decision = match timing {
|
||||
ExecutionCancellationTiming::TurnFinished => {
|
||||
finish_turn.cancel();
|
||||
@@ -127,6 +139,33 @@ async fn execution_cancellation_respects_network_approval_boundary(
|
||||
NetworkDecision::deny("not_allowed")
|
||||
}
|
||||
ExecutionCancellationTiming::AcceptedReview => {
|
||||
let mut queue_events = core_test_support::tracing::NamedTraceEvents::subscribe(
|
||||
"codex.network_approval.queued",
|
||||
);
|
||||
let started = Instant::now();
|
||||
let timing = turn
|
||||
.turn_timing_state
|
||||
.direct_tool_timings
|
||||
.start_turn_direct_call(
|
||||
"execution-2",
|
||||
&turn.sub_id,
|
||||
ToolBatchTiming::default(),
|
||||
started,
|
||||
Some(CriticalPath::default()),
|
||||
);
|
||||
register_call_with_default_shell_trigger(service, "execution-2").await;
|
||||
let mut waiter = Box::pin(service.handle_inline_policy_request(
|
||||
Arc::clone(&session),
|
||||
NetworkPolicyRequest {
|
||||
execution_id: Some("execution-2".to_string()),
|
||||
..request.clone()
|
||||
},
|
||||
));
|
||||
tokio::select! {
|
||||
result = waiter.as_mut() => panic!("expected queued request, got {result:?}"),
|
||||
event = queue_events.next() => event.expect("execution joins the queue"),
|
||||
}
|
||||
queued = Some((waiter, timing, started));
|
||||
let commit_guard = service.session_policy_commit_lock.lock().await;
|
||||
session
|
||||
.notify_approval(
|
||||
@@ -136,6 +175,8 @@ async fn execution_cancellation_respects_network_approval_boundary(
|
||||
.await;
|
||||
// Consume the approval and stop at the held policy commit lock.
|
||||
assert!(poll!(decision.as_mut()).is_pending());
|
||||
// The visible prompt closed, but the queued execution still holds tool results.
|
||||
assert!(*session.subscribe_elicitation_pause_state().borrow());
|
||||
service.record_call_outcome("execution-1", "original execution denial".to_string());
|
||||
assert!(poll!(decision.as_mut()).is_pending());
|
||||
drop(commit_guard);
|
||||
@@ -148,6 +189,21 @@ async fn execution_cancellation_respects_network_approval_boundary(
|
||||
.expect("finish cancelled or accepted review"),
|
||||
expected_decision,
|
||||
);
|
||||
if let Some((queued, timing, started)) = queued {
|
||||
assert_eq!(
|
||||
timeout(Duration::from_secs(5), queued).await.unwrap(),
|
||||
NetworkDecision::Allow
|
||||
);
|
||||
assert!(!*session.subscribe_elicitation_pause_state().borrow());
|
||||
let work = timing.finish().unwrap();
|
||||
assert!(
|
||||
work.path
|
||||
.unwrap()
|
||||
.durations(started, work.ended_at)
|
||||
.user_wait
|
||||
> Duration::ZERO
|
||||
);
|
||||
}
|
||||
let key =
|
||||
HostApprovalKey::from_request(&request, NetworkApprovalProtocol::Http, "local".to_string());
|
||||
assert_eq!(
|
||||
@@ -166,7 +222,11 @@ async fn execution_cancellation_respects_network_approval_boundary(
|
||||
service.take_call_outcome("execution-1").await,
|
||||
Some("original execution denial".to_string())
|
||||
);
|
||||
assert!(service.pending_host_approvals.lock().unwrap().is_empty());
|
||||
{
|
||||
let pending = service.pending_host_approvals.lock().unwrap();
|
||||
assert!(pending.executions.is_empty());
|
||||
assert!(pending.manual_prompts.is_empty());
|
||||
}
|
||||
let mut resolved = Vec::new();
|
||||
while let Ok(event) = events.try_recv() {
|
||||
match event.msg {
|
||||
@@ -438,19 +498,67 @@ async fn sync_session_approved_hosts_to_replaces_existing_target_hosts() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test_case(PendingApprovalDecision::AllowOnce, 1; "approve once wakes one waiter")]
|
||||
#[test_case(PendingApprovalDecision::Deny, 1; "denial wakes one waiter")]
|
||||
#[test_case(PendingApprovalDecision::AllowForSession, 3; "session grant wakes all waiters")]
|
||||
#[tokio::test]
|
||||
async fn pending_waiters_receive_owner_decision() {
|
||||
let pending = Arc::new(PendingHostApproval::new());
|
||||
async fn manual_prompt_wakeups(decision: PendingApprovalDecision, expected: usize) {
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::task::Context;
|
||||
use std::task::Wake;
|
||||
use std::task::Waker;
|
||||
|
||||
let waiter = {
|
||||
let pending = Arc::clone(&pending);
|
||||
tokio::spawn(async move { pending.wait_for_decision().await })
|
||||
#[derive(Default)]
|
||||
struct WakeCount(AtomicUsize);
|
||||
impl Wake for WakeCount {
|
||||
fn wake(self: Arc<Self>) {
|
||||
self.0.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
let (session, turn, _events) = crate::session::tests::make_session_and_context_with_rx().await;
|
||||
let service = &session.services.network_approval;
|
||||
let owner = |id: &str| {
|
||||
let host = HostApprovalKey {
|
||||
environment_id: "local".to_string(),
|
||||
host: "example.com".to_string(),
|
||||
protocol: "https",
|
||||
port: 443,
|
||||
};
|
||||
let key = pending_key(host, &turn.sub_id, id);
|
||||
let (pending, _) = service.get_or_create_pending_approval(key.clone());
|
||||
PendingHostApprovalOwner::new(service, key, pending, /*execution_cancellation*/ None)
|
||||
};
|
||||
let mut first = owner("first");
|
||||
first
|
||||
.review_user_prompt(
|
||||
&session,
|
||||
&turn,
|
||||
/*timing_call_id*/ None,
|
||||
"denied",
|
||||
std::future::ready(ApprovalResolution::user(ReviewDecision::Approved)),
|
||||
)
|
||||
.await;
|
||||
let mut owners = ["second", "third", "fourth"].map(owner);
|
||||
let mut waiters = owners.each_mut().map(|owner| {
|
||||
Box::pin(owner.review_user_prompt(
|
||||
&session,
|
||||
&turn,
|
||||
/*timing_call_id*/ None,
|
||||
"denied",
|
||||
std::future::pending(),
|
||||
))
|
||||
});
|
||||
let wakes = Arc::new(WakeCount::default());
|
||||
let waker = Waker::from(Arc::clone(&wakes));
|
||||
let mut cx = Context::from_waker(&waker);
|
||||
for waiter in &mut waiters {
|
||||
assert!(waiter.as_mut().poll(&mut cx).is_pending());
|
||||
}
|
||||
|
||||
pending.set_decision(PendingApprovalDecision::AllowOnce);
|
||||
|
||||
let decision = waiter.await.expect("waiter should complete");
|
||||
assert_eq!(decision, PendingApprovalDecision::AllowOnce);
|
||||
first.complete(decision);
|
||||
assert_eq!(wakes.0.load(Ordering::Relaxed), expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -659,7 +767,7 @@ async fn register_call_with_default_shell_trigger(
|
||||
turn_id: "turn-1".to_string(),
|
||||
tool_name: ToolName::plain("exec_command"),
|
||||
trigger: GuardianNetworkAccessTrigger {
|
||||
call_id: "call-1".to_string(),
|
||||
call_id: registration_id.to_string(),
|
||||
tool_name: "exec_command".to_string(),
|
||||
command: vec!["curl".to_string(), "https://example.com".to_string()],
|
||||
cwd: PathUri::from_abs_path(&test_path_buf("/tmp").abs()),
|
||||
|
||||
Reference in New Issue
Block a user