fix(agent): make tool approval race-free with an atomic state machine
The previous approval registry used a oneshot channel, so two independent signals decided an approval: the HTTP handler's result (from Sender::send's bool) and the agent loop's allow/deny (from which timeout arm fired). At the 300s boundary these could disagree — a late /approve could see send() return Ok (receiver not yet dropped) and report ok:true to the user, while the timeout's final poll simultaneously fired and denied the tool. The frontend was told "approved" while nothing ran. The mirror case (ok:false but allowed) was impossible, so it never executed an unapproved tool, but the user-facing signal could lie. Replace the channel with a single source of truth: each entry is an AtomicU8 state machine (Pending -> Approved/Denied/TimedOut) plus a Notify. The decision is made by exactly one compare_exchange from Pending — whoever wins (approver or timeout) settles it, everyone else observes a losing CAS and no-ops. resolve() now returns whether *its* CAS won, so /approve tells the frontend the truth; the loop decides purely from the settled atomic. They read the same atomic and cannot disagree. A PendingGuard removes the registry entry on drop (normal completion or aborted/disconnected session), so there is no leak on any path. The callback also emits an approval_resolved SSE event with the final decision so a late click — already told ok:false — can dismiss its modal with the real outcome. Verified with an integration test that races resolve() against a near-zero timeout 5000x and asserts ok==true iff the executed decision is Approved (never a mixed state). Run via tests/ because the in-crate lib test binary is currently blocked by unrelated AppConfig test fixtures in other modules.
This commit is contained in:
boos
+28
-12
@@ -23,7 +23,7 @@
|
||||
|
||||
use crate::auth::jwt::validate_jwt;
|
||||
use crate::services::agent_workbench::{
|
||||
approval::{self, ApprovalOutcome},
|
||||
approval::{self, ApprovalOutcome, SettledState},
|
||||
permission::{classify_tool_risk, decide_permission, PermissionAction, PermissionMode},
|
||||
sandbox,
|
||||
session_store::{self, UsageDelta},
|
||||
@@ -189,7 +189,9 @@ fn build_can_use_tool(
|
||||
PermissionAction::AutoAllow => PermissionDecision::Allow,
|
||||
PermissionAction::RequireApproval => {
|
||||
let request_id = Uuid::new_v4().to_string();
|
||||
let rx = approval::register(&request_id);
|
||||
// `_guard` removes the registry entry when this future ends
|
||||
// (normal completion or drop on session abort) — no leak.
|
||||
let (entry, _guard) = approval::register(&request_id);
|
||||
send_sse(
|
||||
&sse_tx,
|
||||
json!({
|
||||
@@ -203,17 +205,31 @@ fn build_can_use_tool(
|
||||
}),
|
||||
);
|
||||
|
||||
match tokio::time::timeout(APPROVAL_TIMEOUT, rx).await {
|
||||
Ok(Ok(ApprovalOutcome::Approve)) => PermissionDecision::Allow,
|
||||
Ok(Ok(ApprovalOutcome::Deny)) => {
|
||||
PermissionDecision::Deny("Denied by user".to_string())
|
||||
}
|
||||
// Sender dropped (session ended) or timed out.
|
||||
_ => {
|
||||
approval::cancel(&request_id);
|
||||
PermissionDecision::Deny("Approval timed out".to_string())
|
||||
}
|
||||
// Single source of truth: the settled atomic state decides
|
||||
// allow vs deny. `resolve` reports back to the frontend
|
||||
// whether its CAS won, so the two can never disagree.
|
||||
let settled = approval::wait_for_decision(&entry, APPROVAL_TIMEOUT).await;
|
||||
let (decision, wire) = match settled {
|
||||
SettledState::Approved => (PermissionDecision::Allow, "approved"),
|
||||
SettledState::Denied => {
|
||||
(PermissionDecision::Deny("Denied by user".to_string()), "denied")
|
||||
}
|
||||
SettledState::TimedOut => (
|
||||
PermissionDecision::Deny("Approval timed out".to_string()),
|
||||
"timed_out",
|
||||
),
|
||||
};
|
||||
// Tell the client how this prompt actually resolved, so a
|
||||
// late click (which the /approve call already reported as
|
||||
// ok:false) can dismiss the modal with the real outcome.
|
||||
send_sse(
|
||||
&sse_tx,
|
||||
json!({
|
||||
"type": "approval_resolved",
|
||||
"data": { "request_id": request_id, "decision": wire }
|
||||
}),
|
||||
);
|
||||
decision
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
@@ -1,19 +1,36 @@
|
||||
//! Pending-approval registry for the Agent Workbench.
|
||||
//!
|
||||
//! When a session runs in a mode that requires approval for a risky tool, the
|
||||
//! permission callback emits an `approval_request` SSE event and then blocks on
|
||||
//! a [`tokio::sync::oneshot`] channel keyed by a freshly minted request id. The
|
||||
//! `POST /api/agent/approve` endpoint looks the id up here and delivers the
|
||||
//! user's decision, unblocking the agent loop.
|
||||
//! permission callback registers a pending approval, emits an `approval_request`
|
||||
//! SSE event, and waits for either the user's decision or a timeout. The
|
||||
//! `POST /api/agent/approve` endpoint resolves the pending entry.
|
||||
//!
|
||||
//! The registry is process-global and self-cleaning: a request is removed the
|
||||
//! moment its decision is delivered, when it times out, or when the waiting side
|
||||
//! is dropped.
|
||||
//! ## Single source of truth
|
||||
//!
|
||||
//! Each entry carries an [`AtomicU8`] state machine (`Pending → Approved /
|
||||
//! Denied / TimedOut`) plus a [`Notify`]. The decision is made by exactly one
|
||||
//! `compare_exchange` from `Pending`: whichever side wins (the approver or the
|
||||
//! timeout) settles the state; everyone else observes a losing CAS and no-ops.
|
||||
//!
|
||||
//! This closes the race that a bare `oneshot` channel had, where the HTTP
|
||||
//! response (derived from `Sender::send`'s bool) and the loop's allow/deny
|
||||
//! decision (derived from which `select`/`timeout` arm fired) were two
|
||||
//! independent signals that could disagree. Now both read the same atomic:
|
||||
//! `resolve` reports back whether *its* CAS won, and the callback decides purely
|
||||
//! from the settled state — so the frontend's `ok` and the executed decision can
|
||||
//! never contradict each other.
|
||||
|
||||
use once_cell::sync::Lazy;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex;
|
||||
use tokio::sync::oneshot;
|
||||
use std::sync::atomic::{AtomicU8, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::sync::Notify;
|
||||
|
||||
/// Settled decision states. `Pending` is the only non-terminal value.
|
||||
const PENDING: u8 = 0;
|
||||
const APPROVED: u8 = 1;
|
||||
const DENIED: u8 = 2;
|
||||
const TIMED_OUT: u8 = 3;
|
||||
|
||||
/// User's decision on a pending tool approval.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -22,68 +39,203 @@ pub enum ApprovalOutcome {
|
||||
Deny,
|
||||
}
|
||||
|
||||
static REGISTRY: Lazy<Mutex<HashMap<String, oneshot::Sender<ApprovalOutcome>>>> =
|
||||
Lazy::new(|| Mutex::new(HashMap::new()));
|
||||
|
||||
/// Register a new pending approval, returning the receiver the caller should
|
||||
/// await. The returned id is what the frontend echoes back to `/approve`.
|
||||
pub fn register(request_id: &str) -> oneshot::Receiver<ApprovalOutcome> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let mut map = REGISTRY.lock().unwrap();
|
||||
map.insert(request_id.to_string(), tx);
|
||||
rx
|
||||
/// Final, settled state of an approval — the single source of truth read by the
|
||||
/// permission callback to decide allow vs deny.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum SettledState {
|
||||
Approved,
|
||||
Denied,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
/// Resolve a pending approval. Returns `true` if a waiter was found and
|
||||
/// notified, `false` if the id was unknown (already resolved, timed out, or
|
||||
/// belonged to another process instance).
|
||||
pub fn resolve(request_id: &str, outcome: ApprovalOutcome) -> bool {
|
||||
let sender = {
|
||||
let mut map = REGISTRY.lock().unwrap();
|
||||
map.remove(request_id)
|
||||
};
|
||||
match sender {
|
||||
Some(tx) => tx.send(outcome).is_ok(),
|
||||
None => false,
|
||||
/// Shared approval entry: an atomic state machine plus a wakeup primitive.
|
||||
#[derive(Debug)]
|
||||
pub struct ApprovalEntry {
|
||||
state: AtomicU8,
|
||||
notify: Notify,
|
||||
}
|
||||
|
||||
impl ApprovalEntry {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
state: AtomicU8::new(PENDING),
|
||||
notify: Notify::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to settle this entry to a terminal state from `Pending`.
|
||||
/// Returns `true` iff this caller's CAS won the race.
|
||||
fn try_settle(&self, target: u8) -> bool {
|
||||
self.state
|
||||
.compare_exchange(PENDING, target, Ordering::SeqCst, Ordering::SeqCst)
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
fn settled(&self) -> Option<SettledState> {
|
||||
match self.state.load(Ordering::SeqCst) {
|
||||
APPROVED => Some(SettledState::Approved),
|
||||
DENIED => Some(SettledState::Denied),
|
||||
TIMED_OUT => Some(SettledState::TimedOut),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop a pending approval without delivering a decision (e.g. on timeout or
|
||||
/// when the session ends). No-op if already resolved.
|
||||
pub fn cancel(request_id: &str) {
|
||||
let mut map = REGISTRY.lock().unwrap();
|
||||
map.remove(request_id);
|
||||
static REGISTRY: Lazy<Mutex<HashMap<String, Arc<ApprovalEntry>>>> =
|
||||
Lazy::new(|| Mutex::new(HashMap::new()));
|
||||
|
||||
/// RAII guard that removes its entry from the registry when dropped, so a
|
||||
/// pending approval is cleaned up whether the callback finishes normally or its
|
||||
/// future is dropped (session aborted/disconnected). No leak on any path.
|
||||
pub struct PendingGuard {
|
||||
request_id: String,
|
||||
}
|
||||
|
||||
impl Drop for PendingGuard {
|
||||
fn drop(&mut self) {
|
||||
REGISTRY.lock().unwrap().remove(&self.request_id);
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a new pending approval. Returns the shared entry (await its
|
||||
/// [`wait_settled`](ApprovalEntry)/notify and read its settled state) plus a
|
||||
/// [`PendingGuard`] that the caller must hold for the lifetime of the wait; when
|
||||
/// the guard drops, the registry entry is removed.
|
||||
pub fn register(request_id: &str) -> (Arc<ApprovalEntry>, PendingGuard) {
|
||||
let entry = Arc::new(ApprovalEntry::new());
|
||||
REGISTRY
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(request_id.to_string(), entry.clone());
|
||||
(
|
||||
entry,
|
||||
PendingGuard {
|
||||
request_id: request_id.to_string(),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Wait until the entry is settled or `timeout` elapses, then return the settled
|
||||
/// state. On timeout this claims `TimedOut` (if still pending) so the result is
|
||||
/// always a terminal state read from the single atomic.
|
||||
pub async fn wait_for_decision(
|
||||
entry: &ApprovalEntry,
|
||||
timeout: std::time::Duration,
|
||||
) -> SettledState {
|
||||
// Create the notified future *before* checking state to avoid a lost wakeup.
|
||||
let notified = entry.notify.notified();
|
||||
|
||||
// Fast path: already settled by a resolve that raced ahead of us.
|
||||
if let Some(s) = entry.settled() {
|
||||
return s;
|
||||
}
|
||||
|
||||
if tokio::time::timeout(timeout, notified).await.is_err() {
|
||||
// Timer fired first — claim TimedOut if nobody settled in the meantime.
|
||||
entry.try_settle(TIMED_OUT);
|
||||
}
|
||||
|
||||
// Single source of truth: whatever the atomic settled to. If a wakeup was
|
||||
// lost but the state is set, we still read the correct decision; if somehow
|
||||
// still pending (shouldn't happen — we just CAS'd), treat as timed out.
|
||||
entry.settled().unwrap_or(SettledState::TimedOut)
|
||||
}
|
||||
|
||||
/// Resolve a pending approval with the user's decision.
|
||||
///
|
||||
/// Returns `true` iff this call's `compare_exchange` won — i.e. it is the call
|
||||
/// that actually settled the decision. Returns `false` truthfully when the id is
|
||||
/// unknown (already cleaned up) or when the timeout (or another resolve) already
|
||||
/// claimed the entry, so the frontend is never told its vote took effect when it
|
||||
/// did not.
|
||||
pub fn resolve(request_id: &str, outcome: ApprovalOutcome) -> bool {
|
||||
let entry = REGISTRY.lock().unwrap().get(request_id).cloned();
|
||||
let Some(entry) = entry else {
|
||||
return false;
|
||||
};
|
||||
let target = match outcome {
|
||||
ApprovalOutcome::Approve => APPROVED,
|
||||
ApprovalOutcome::Deny => DENIED,
|
||||
};
|
||||
if entry.try_settle(target) {
|
||||
entry.notify.notify_one();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_delivers_outcome_to_waiter() {
|
||||
let rx = register("req-1");
|
||||
assert!(resolve("req-1", ApprovalOutcome::Approve));
|
||||
assert_eq!(rx.await.unwrap(), ApprovalOutcome::Approve);
|
||||
async fn approve_before_wait_is_observed() {
|
||||
let (entry, _guard) = register("r1");
|
||||
assert!(resolve("r1", ApprovalOutcome::Approve));
|
||||
let s = wait_for_decision(&entry, Duration::from_secs(5)).await;
|
||||
assert_eq!(s, SettledState::Approved);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_unknown_id_returns_false() {
|
||||
assert!(!resolve("does-not-exist", ApprovalOutcome::Deny));
|
||||
async fn deny_is_observed() {
|
||||
let (entry, _guard) = register("r2");
|
||||
assert!(resolve("r2", ApprovalOutcome::Deny));
|
||||
assert_eq!(
|
||||
wait_for_decision(&entry, Duration::from_secs(5)).await,
|
||||
SettledState::Denied
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancel_removes_pending() {
|
||||
let rx = register("req-2");
|
||||
cancel("req-2");
|
||||
// Sender dropped → receiver errors out.
|
||||
assert!(rx.await.is_err());
|
||||
assert!(!resolve("req-2", ApprovalOutcome::Approve));
|
||||
async fn timeout_settles_timed_out_and_later_resolve_loses() {
|
||||
let (entry, _guard) = register("r3");
|
||||
// Very short timeout → times out with no resolve.
|
||||
let s = wait_for_decision(&entry, Duration::from_millis(20)).await;
|
||||
assert_eq!(s, SettledState::TimedOut);
|
||||
// A late approve must report false: the CAS already went to TimedOut.
|
||||
assert!(!resolve("r3", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn double_resolve_only_first_wins() {
|
||||
let _rx = register("req-3");
|
||||
assert!(resolve("req-3", ApprovalOutcome::Approve));
|
||||
assert!(!resolve("req-3", ApprovalOutcome::Deny));
|
||||
async fn first_cas_wins_second_resolve_loses() {
|
||||
let (_entry, _guard) = register("r4");
|
||||
assert!(resolve("r4", ApprovalOutcome::Approve));
|
||||
assert!(!resolve("r4", ApprovalOutcome::Deny));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_unknown_id_is_false() {
|
||||
assert!(!resolve("nope", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn guard_drop_removes_entry() {
|
||||
{
|
||||
let (_entry, _guard) = register("r5");
|
||||
} // guard dropped here
|
||||
// Entry gone → resolve can't find it.
|
||||
assert!(!resolve("r5", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_resolve_and_timeout_never_both_win() {
|
||||
// Stress the boundary: race a resolve against a near-zero timeout many
|
||||
// times and assert the frontend-ok and the settled decision always agree.
|
||||
for i in 0..2000 {
|
||||
let id = format!("race-{i}");
|
||||
let (entry, _guard) = register(&id);
|
||||
let id2 = id.clone();
|
||||
let resolver = tokio::spawn(async move { resolve(&id2, ApprovalOutcome::Approve) });
|
||||
let decided = wait_for_decision(&entry, Duration::from_nanos(1)).await;
|
||||
let won = resolver.await.unwrap();
|
||||
// The single-atomic invariant: ok==true <=> decision==Approved.
|
||||
match (won, decided) {
|
||||
(true, SettledState::Approved) => {}
|
||||
(false, SettledState::TimedOut) => {}
|
||||
other => panic!("inconsistent outcome: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
//! Integration tests for the Agent Workbench approval registry.
|
||||
//!
|
||||
//! These run against the library's public API (not the in-crate `#[cfg(test)]`
|
||||
//! modules), so they compile and execute independently of unrelated test
|
||||
//! fixtures elsewhere in the crate. They pin the core safety invariant of the
|
||||
//! atomic approval state machine:
|
||||
//!
|
||||
//! the value `/api/agent/approve` reports back (`resolve` -> bool) and the
|
||||
//! decision the agent loop executes (`wait_for_decision` -> SettledState)
|
||||
//! are derived from the *same* atomic, so they can never disagree.
|
||||
|
||||
use boosapi::services::agent_workbench::approval::{
|
||||
register, resolve, wait_for_decision, ApprovalOutcome, SettledState,
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
#[tokio::test]
|
||||
async fn approve_then_wait_observes_approved() {
|
||||
let (entry, _guard) = register("it-approve");
|
||||
assert!(resolve("it-approve", ApprovalOutcome::Approve));
|
||||
assert_eq!(
|
||||
wait_for_decision(&entry, Duration::from_secs(5)).await,
|
||||
SettledState::Approved
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deny_then_wait_observes_denied() {
|
||||
let (entry, _guard) = register("it-deny");
|
||||
assert!(resolve("it-deny", ApprovalOutcome::Deny));
|
||||
assert_eq!(
|
||||
wait_for_decision(&entry, Duration::from_secs(5)).await,
|
||||
SettledState::Denied
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn timeout_denies_and_late_click_is_truthfully_rejected() {
|
||||
let (entry, _guard) = register("it-timeout");
|
||||
// No resolve arrives before the (tiny) timeout.
|
||||
assert_eq!(
|
||||
wait_for_decision(&entry, Duration::from_millis(20)).await,
|
||||
SettledState::TimedOut
|
||||
);
|
||||
// A late click must be told ok:false — the timeout already claimed it.
|
||||
assert!(!resolve("it-timeout", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_id_resolves_false() {
|
||||
assert!(!resolve("it-unknown", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn guard_drop_cleans_up_registry() {
|
||||
{
|
||||
let (_entry, _guard) = register("it-guard");
|
||||
}
|
||||
assert!(!resolve("it-guard", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
/// The invariant under contention: race a `resolve(Approve)` against a
|
||||
/// near-zero timeout thousands of times. The frontend-reported success
|
||||
/// (`resolve` -> true) must hold if and only if the executed decision is
|
||||
/// `Approved`; otherwise the decision must be `TimedOut`. The mixed state
|
||||
/// (frontend told "approved" while the loop denied, or vice versa) must never
|
||||
/// occur — that was the bug the atomic state machine fixes.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn resolve_and_timeout_never_disagree() {
|
||||
for i in 0..5000 {
|
||||
let id = format!("it-race-{i}");
|
||||
let (entry, _guard) = register(&id);
|
||||
let id2 = id.clone();
|
||||
let resolver = tokio::spawn(async move { resolve(&id2, ApprovalOutcome::Approve) });
|
||||
let decided = wait_for_decision(&entry, Duration::from_nanos(1)).await;
|
||||
let frontend_ok = resolver.await.unwrap();
|
||||
match (frontend_ok, decided) {
|
||||
(true, SettledState::Approved) => {}
|
||||
(false, SettledState::TimedOut) => {}
|
||||
other => panic!("inconsistent outcome at i={i}: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user