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
2026-06-22 07:20:36 +08:00
parent f0a76fd82c
commit 5501558354
3 changed files with 314 additions and 63 deletions
+28 -12
View File
@@ -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
}
}
})