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:
+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())
|
||||
// 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")
|
||||
}
|
||||
// Sender dropped (session ended) or timed out.
|
||||
_ => {
|
||||
approval::cancel(&request_id);
|
||||
PermissionDecision::Deny("Approval timed out".to_string())
|
||||
}
|
||||
}
|
||||
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
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user