figmar/BrowserSkill
1
0
Fork 0
mirror of https://github.com/Tencent/BrowserSkill.git synced 2026-09-10 13:01:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(extension): make tool cancellation cooperative (#89)

Browse Source 
Tool cancellation is now cooperative end-to-end: the queue stays busy until compensation completes, cancellation is checked cooperatively across the pipeline, the fallback window is compensated too, and cleanup_failed is surfaced explicitly. Includes a gated delayed-compensation E2E test.
This commit is contained in:
NianJiu
2026-08-13 17:31:23 +08:00
committed by GitHub
parent 37fbaacc94
commit 4549ed0892
20 changed files with 904 additions and 214 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/bsk-cli/src/daemon/ipc.rs
+4 -4
View File
@@ -443,10 +443,10 @@ async fn handle_wait_ms(
/// * **Queued** — the worker's pre-flight observes the cancelled
/// token and short-circuits with `cancelled` before any WS
/// frame leaves the daemon (review C2 fix).
/// * **Forwarded** — the worker's `tokio::select!` returns
/// `cancelled` immediately, AND we additionally push a WS-side
/// `cancel { rpc_id: ws_rpc_id }` frame so the extension's
/// dispatcher can trip its `AbortController`.
/// * **Forwarded** — we push a WS-side `cancel { rpc_id: ws_rpc_id }`
/// frame so the extension's dispatcher can trip its
/// `AbortController`; the worker keeps the session busy until the
/// extension returns its final result or the cleanup timeout expires.
///
/// Returns `{ cancelled }` reflecting whether either surface
/// matched. The RPC itself never errors — a cancelled tool surfaces
crates/bsk-cli/src/daemon/queue.rs
+122 -48
View File
@@ -58,6 +58,11 @@ pub const QUEUE_CAPACITY: usize = 64;
/// so tests can default it.
pub const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(30);
/// Once a forwarded request is cancelled, keep the session queue busy while
/// the extension finishes handler-side compensation. This is deliberately
/// bounded so a wedged extension cannot pin the session forever.
pub const CANCEL_CLEANUP_TIMEOUT: Duration = Duration::from_secs(2);
/// Job submitted into a session queue. Carries everything the worker
/// needs to forward one RPC and one oneshot to deliver the answer.
pub struct ToolJob {
@@ -70,9 +75,9 @@ pub struct ToolJob {
/// the queued/forwarded state machine; the worker checks it on
/// pre-flight (so a cancel arriving while the job was still
/// queued short-circuits without ever touching the extension)
/// and selects on `cancel.cancelled()` while awaiting the WS
/// response (so a cancel arriving mid-flight unblocks the worker
/// the same way the WS-side cancel frame already did).
/// and observes `cancel.cancelled()` while awaiting the WS response.
/// A forwarded cancel keeps the worker busy until the extension's
/// final response arrives or the bounded cleanup timeout expires.
///
/// `None` for daemon-internal callers that do not flow through an
/// IPC request id (e.g. `session.stop`'s queued teardown call —
@@ -571,20 +576,41 @@ async fn forward_one(
None
}
};
let waited = await_with_optional_cancel(job.timeout, waiter, cancel_token.as_ref()).await;
let waited = await_with_optional_cancel(
job.timeout,
CANCEL_CLEANUP_TIMEOUT,
waiter,
cancel_token.as_ref(),
)
.await;
let response = match waited {
WaitOutcome::Response(resp) => resp,
WaitOutcome::Cancelled => {
// Drop the WS waiter so a late extension reply is dropped
// cleanly (otherwise pending.resolve would log a stale
// entry). The CLI caller will get a synthesised
// `cancelled` here even if the extension never answers.
client.pending.lock().unwrap().cancel(&rpc_id);
WaitOutcome::CancelledAfterResponse(resp) => {
// Cancellation wins the external verdict, but only after the
// extension's original RPC has settled. Preserve any non-cancel
// error so compensation failures remain explicit instead of
// being hidden behind a generic cancelled result.
if let ResponseBody::Err(err) = resp.body
&& !matches!(err.code, ErrorCode::Cancelled | ErrorCode::UserAborted)
{
return Err(err);
}
return Err(cancelled_error(
job.inflight.as_deref(),
"tool dispatch cancelled mid-flight",
"tool dispatch cancelled after extension cleanup",
));
}
WaitOutcome::CleanupTimeout => {
client.pending.lock().unwrap().cancel(&rpc_id);
return Err(RpcError {
code: ErrorCode::Timeout,
message: format!(
"cancelled tool did not finish cleanup within {:?}",
CANCEL_CLEANUP_TIMEOUT
),
data: Some(serde_json::json!({ "reason": "cancel_cleanup_timeout" })),
});
}
WaitOutcome::WaiterClosed => {
client.pending.lock().unwrap().cancel(&rpc_id);
return Err(RpcError {
@@ -611,47 +637,42 @@ async fn forward_one(
#[derive(Debug)]
enum WaitOutcome {
Response(bsk_protocol::ResponseFrame),
Cancelled,
CancelledAfterResponse(bsk_protocol::ResponseFrame),
CleanupTimeout,
WaiterClosed,
Timeout,
}
async fn await_with_optional_cancel(
timeout: Duration,
waiter: oneshot::Receiver<bsk_protocol::ResponseFrame>,
cleanup_timeout: Duration,
mut waiter: oneshot::Receiver<bsk_protocol::ResponseFrame>,
cancel: Option<&super::abort::AbortToken>,
) -> WaitOutcome {
match cancel {
// Cancel wins if both ready: when a cancel notification and
// the extension's response are both observable on the same
// tokio tick, `biased;` polls `token.cancelled()` first and
// resolves the whole `select!` to `WaitOutcome::Cancelled`.
// The already-arrived tool result is intentionally dropped.
//
// This extends design §4.6 — which only specifies "CLI sends
// `cancel` on SIGINT and waits for the extension's
// `cancelled` reply (≤ 2s before forced exit)" — by pinning
// the same-tick race to a single, observable verdict.
// Without `biased;`, an extension racing a fast successful
// reply against a cancel notification could occasionally
// resolve as `ok` even though the caller had already moved on
// to compensation logic, leaving the agent trusting state the
// daemon had just been asked to roll back. Picking
// cancel-wins keeps the external observation rule simple
// ("once a cancel is in flight, the in-flight RPC's verdict
// is `cancelled`") and matches what the CLI / agent already
// assumes after firing SIGINT. Round 3 M1 / round 4 M2
// nail-down.
Some(token) => tokio::select! {
biased;
_ = token.cancelled() => WaitOutcome::Cancelled,
outcome = tokio::time::timeout(timeout, waiter) => match outcome {
Ok(Ok(resp)) => WaitOutcome::Response(resp),
Ok(Err(_)) => WaitOutcome::WaiterClosed,
Err(_) => WaitOutcome::Timeout,
},
},
None => match tokio::time::timeout(timeout, waiter).await {
Some(token) => {
let deadline = tokio::time::sleep(timeout);
tokio::pin!(deadline);
tokio::select! {
// Cancel keeps same-tick priority, but it no longer drops the
// waiter. The worker remains busy until the extension replies
// after compensation or the cleanup deadline expires.
biased;
_ = token.cancelled() => {
match tokio::time::timeout(cleanup_timeout, &mut waiter).await {
Ok(Ok(resp)) => WaitOutcome::CancelledAfterResponse(resp),
Ok(Err(_)) => WaitOutcome::WaiterClosed,
Err(_) => WaitOutcome::CleanupTimeout,
}
},
outcome = &mut waiter => match outcome {
Ok(resp) => WaitOutcome::Response(resp),
Err(_) => WaitOutcome::WaiterClosed,
},
_ = &mut deadline => WaitOutcome::Timeout,
}
}
None => match tokio::time::timeout(timeout, &mut waiter).await {
Ok(Ok(resp)) => WaitOutcome::Response(resp),
Ok(Err(_)) => WaitOutcome::WaiterClosed,
Err(_) => WaitOutcome::Timeout,
@@ -733,10 +754,16 @@ mod await_with_optional_cancel_tests {
let (tx, rx) = oneshot::channel();
tx.send(dummy_response()).unwrap();
let outcome = await_with_optional_cancel(Duration::from_secs(10), rx, Some(&token)).await;
let outcome = await_with_optional_cancel(
Duration::from_secs(10),
Duration::from_secs(1),
rx,
Some(&token),
)
.await;
assert!(
matches!(outcome, WaitOutcome::Cancelled),
"expected Cancelled when both arms are ready under biased; (got something else)"
matches!(outcome, WaitOutcome::CancelledAfterResponse(_)),
"expected cancellation to win while retaining the ready response"
);
}
@@ -749,7 +776,13 @@ mod await_with_optional_cancel_tests {
let (tx, rx) = oneshot::channel();
tx.send(dummy_response()).unwrap();
let outcome = await_with_optional_cancel(Duration::from_secs(10), rx, Some(&token)).await;
let outcome = await_with_optional_cancel(
Duration::from_secs(10),
Duration::from_secs(1),
rx,
Some(&token),
)
.await;
match outcome {
WaitOutcome::Response(frame) => {
assert!(matches!(frame.body, ResponseBody::Ok(_)));
@@ -762,9 +795,50 @@ mod await_with_optional_cancel_tests {
async fn no_cancel_token_still_returns_response() {
let (tx, rx) = oneshot::channel();
tx.send(dummy_response()).unwrap();
let outcome = await_with_optional_cancel(Duration::from_secs(10), rx, None).await;
let outcome =
await_with_optional_cancel(Duration::from_secs(10), Duration::from_secs(1), rx, None)
.await;
assert!(matches!(outcome, WaitOutcome::Response(_)));
}
#[tokio::test]
async fn cancelled_waiter_stays_pending_until_cleanup_response_arrives() {
let token = AbortToken::new();
let (tx, rx) = oneshot::channel();
token.cancel();
let pending = tokio::spawn(async move {
await_with_optional_cancel(
Duration::from_secs(10),
Duration::from_secs(1),
rx,
Some(&token),
)
.await
});
tokio::task::yield_now().await;
assert!(!pending.is_finished());
tx.send(dummy_response()).unwrap();
assert!(matches!(
pending.await.unwrap(),
WaitOutcome::CancelledAfterResponse(_)
));
}
#[tokio::test]
async fn cancelled_waiter_releases_after_cleanup_timeout() {
let token = AbortToken::new();
let (_tx, rx) = oneshot::channel();
token.cancel();
let outcome = await_with_optional_cancel(
Duration::from_secs(10),
Duration::from_millis(10),
rx,
Some(&token),
)
.await;
assert!(matches!(outcome, WaitOutcome::CleanupTimeout));
}
}
#[cfg(test)]
crates/bsk-cli/tests/cancel_forwarding.rs
+60 -5
View File
@@ -583,11 +583,11 @@ async fn cancel_for_unknown_rpc_id_returns_false() {
}
#[tokio::test]
async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() {
// With session busy fast-fail, a second tool RPC submitted while
// the worker is occupied returns `session_busy` immediately
// instead of queuing behind the slow call. Cancel still unblocks
// the in-flight RPC through the existing forwarded-cancel path.
async fn cancel_keeps_session_busy_until_delayed_extension_cleanup_finishes() {
// A forwarded cancel is acknowledged immediately, but the extension
// deliberately delays the original RPC's final response to model
// handler-side compensation. The daemon must keep the session busy
// throughout that delay and release it only after cleanup settles.
use std::sync::atomic::{AtomicUsize, Ordering as AOrdering};
let (handle, sock) = spawn_daemon().await;
let mut ws = connect_ext(handle.ws_addr()).await;
@@ -596,6 +596,10 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() {
let ws = Arc::new(tokio::sync::Mutex::new(ws));
let snapshots_seen = Arc::new(AtomicUsize::new(0));
let snapshots_seen_clone = Arc::clone(&snapshots_seen);
let cleanup_started = Arc::new(tokio::sync::Notify::new());
let cleanup_started_clone = Arc::clone(&cleanup_started);
let cleanup_release = Arc::new(tokio::sync::Notify::new());
let cleanup_release_clone = Arc::clone(&cleanup_release);
let ws_clone = Arc::clone(&ws);
let responder = tokio::spawn(async move {
let mut pending_snapshot: Option<String> = None;
@@ -636,6 +640,16 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() {
snapshots_seen_clone.fetch_add(1, AOrdering::SeqCst);
pending_snapshot = Some(req.id);
}
Method::ToolTabList => {
let reply = ResponseFrame {
id: req.id,
body: ResponseBody::Ok(json!({ "tabs": [] })),
};
let mut g = ws_clone.lock().await;
g.send(Message::Text(serde_json::to_string(&reply).unwrap()))
.await
.unwrap();
}
Method::Cancel => {
let target = req
.params
@@ -647,6 +661,8 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() {
if let (Some(target), Some(snap)) = (target, snapshot_id)
&& target == snap
{
cleanup_started_clone.notify_one();
cleanup_release_clone.notified().await;
let reply = ResponseFrame {
id: snap,
body: ResponseBody::Err(RpcError {
@@ -753,6 +769,33 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() {
.expect("cancel slow rpc succeeds");
assert!(cancel_slow.cancelled);
cleanup_started.notified().await;
let busy_after_cancel = tokio::time::timeout(
Duration::from_millis(200),
busy_ipc.call_with_id::<_, serde_json::Value>(
"snap-busy-after-cancel".into(),
Method::ToolSnapshot,
Some(json!({"session_id": session_id.clone()})),
Duration::from_secs(10),
),
)
.await
.expect("post-cancel RPC should fast-fail while cleanup is pending")
.unwrap();
let busy_after_cancel_err =
busy_after_cancel.expect_err("session must remain busy during delayed compensation");
assert_eq!(busy_after_cancel_err.code, ErrorCode::Timeout);
assert_eq!(
busy_after_cancel_err
.data
.as_ref()
.and_then(|d| d.get("reason"))
.and_then(|v| v.as_str()),
Some(bsk::rpc_reason::SESSION_BUSY)
);
cleanup_release.notify_one();
let slow_outcome = tokio::time::timeout(Duration::from_secs(5), slow_handle)
.await
.expect("slow snapshot did not resolve")
@@ -761,6 +804,18 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() {
let slow_err = slow_outcome.expect_err("slow snapshot must surface cancelled");
assert_eq!(slow_err.code, ErrorCode::Cancelled);
let after_cleanup: serde_json::Value = busy_ipc
.call(
"tabs-after-cleanup",
Method::ToolTabList,
Some(json!({"session_id": session_id})),
Duration::from_secs(3),
)
.await
.unwrap()
.expect("session queue should reopen after extension cleanup");
assert_eq!(after_cleanup, json!({ "tabs": [] }));
assert_eq!(
snapshots_seen.load(AOrdering::SeqCst),
1,