diff --git a/crates/bsk-cli/src/cli/business_rpc.rs b/crates/bsk-cli/src/cli/business_rpc.rs index 46c7efe..ff93c04 100644 --- a/crates/bsk-cli/src/cli/business_rpc.rs +++ b/crates/bsk-cli/src/cli/business_rpc.rs @@ -78,7 +78,7 @@ where return Err(CliError::from_rpc(RpcError { code: ErrorCode::Cancelled, message: "parent closed the cancellation pipe".into(), - data: None, + data: Some(serde_json::json!({"reason": "cancelled_before_dispatch"})), })); } let rpc_id: RpcId = format!("{}-{}", rpc_id_prefix, random_short_id()); @@ -120,6 +120,46 @@ where outcome.map_err(CliError::from_rpc) } +/// Release caller-owned staging even after cancellation. This deliberately +/// exposes only transfer.release, not a general bypass for business RPCs. +/// The entire batch, including connection setup, shares one bounded budget. +pub fn release_transfers<'a>( + sock: &Path, + transfer_ids: impl IntoIterator, +) -> Result<(), CliError> { + use bsk_protocol::tools::{TransferIdParams, TransferReleaseResult}; + const CLEANUP_TIMEOUT: Duration = Duration::from_secs(5); + let ids: Vec<_> = transfer_ids.into_iter().collect(); + if ids.is_empty() { + return Ok(()); + } + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .context("build runtime for transfer cleanup")?; + rt.block_on(async { + tokio::time::timeout(CLEANUP_TIMEOUT, async { + let mut client = IpcClient::connect(sock).await?; + for id in ids { + client + .call::<_, TransferReleaseResult>( + "transfer-release", + Method::TransferRelease, + Some(TransferIdParams { + transfer_id: id.to_string(), + }), + CLEANUP_TIMEOUT, + ) + .await? + .map_err(CliError::from_rpc)?; + } + Ok::<_, CliError>(()) + }) + .await + .context("transfer cleanup exceeded its total time budget")? + }) +} + /// Send a `cancel { rpc_id }` frame over a fresh connection so it /// lands on the daemon while the original call is still parked. /// diff --git a/crates/bsk-cli/src/cli/download.rs b/crates/bsk-cli/src/cli/download.rs index 5a97b28..b1ef44d 100644 --- a/crates/bsk-cli/src/cli/download.rs +++ b/crates/bsk-cli/src/cli/download.rs @@ -9,8 +9,7 @@ use anyhow::Context; use base64::Engine; use bsk_protocol::Method; use bsk_protocol::tools::{ - DownloadParams, DownloadResult, TransferChunkParams, TransferChunkResult, TransferIdParams, - TransferReleaseResult, + DownloadParams, DownloadResult, TransferChunkParams, TransferChunkResult, }; use clap::Args; use uuid::Uuid; @@ -70,13 +69,7 @@ pub fn dispatch(args: DownloadArgs, format: Format) -> Result<(), CliError> { CliError::Local(anyhow::anyhow!("daemon returned no download transfer id")) })?; let write_result = write_transfer(&info.sock_path, &transfer_id, &args.out, args.overwrite); - let _: Result = crate::cli::business_rpc::call( - info.sock_path, - "transfer-release", - Method::TransferRelease, - Some(TransferIdParams { transfer_id }), - Duration::from_secs(5), - ); + let _ = crate::cli::business_rpc::release_transfers(&info.sock_path, [transfer_id.as_str()]); write_result?; match format { Format::Json => { diff --git a/crates/bsk-cli/src/cli/upload.rs b/crates/bsk-cli/src/cli/upload.rs index a602732..7249a8d 100644 --- a/crates/bsk-cli/src/cli/upload.rs +++ b/crates/bsk-cli/src/cli/upload.rs @@ -10,8 +10,7 @@ use base64::Engine; use bsk_protocol::Method; use bsk_protocol::tools::{ TransferBeginParams, TransferBeginResult, TransferChunkParams, TransferChunkResult, - TransferIdParams, TransferReadyResult, TransferReleaseResult, UploadFile, UploadMode, - UploadParams, UploadResult, + TransferIdParams, TransferReadyResult, UploadFile, UploadMode, UploadParams, UploadResult, }; use clap::{Args, ValueEnum}; @@ -69,14 +68,12 @@ pub fn dispatch(args: UploadArgs, format: Format) -> Result<(), CliError> { let (ref_, selector) = split_target(args.target, args.ref_, args.selector)?; let mut staged = Vec::new(); for path in &args.files { - match stage_file(&info.sock_path, &args.session, path) { - Ok(file) => staged.push(file), - Err(err) => { - for (id, _) in &staged { - let _ = release(&info.sock_path, id); - } - return Err(err); - } + if let Err(err) = stage_file(&info.sock_path, &args.session, path, &mut staged) { + let _ = crate::cli::business_rpc::release_transfers( + &info.sock_path, + staged.iter().map(|(id, _)| id.as_str()), + ); + return Err(err); } } let params = UploadParams { @@ -102,6 +99,17 @@ pub fn dispatch(args: UploadArgs, format: Format) -> Result<(), CliError> { Some(params), ipc_timeout(args.timeout), ); + // The helper can reject cancellation before handing any request to the + // daemon. In that case staging still belongs to this CLI. + if result.as_ref().is_err_and(|err| { + err.data() + .is_some_and(|data| data["reason"] == "cancelled_before_dispatch") + }) { + let _ = crate::cli::business_rpc::release_transfers( + &info.sock_path, + staged.iter().map(|(id, _)| id.as_str()), + ); + } // Once tool.upload is dispatched, staging ownership belongs to the // session. A transport timeout cannot prove that Chrome did not attach // the file, so releasing here could invalidate a late successful attach. @@ -118,7 +126,12 @@ pub fn dispatch(args: UploadArgs, format: Format) -> Result<(), CliError> { Ok(()) } -fn stage_file(sock: &Path, session: &str, path: &PathBuf) -> Result<(String, String), CliError> { +fn stage_file( + sock: &Path, + session: &str, + path: &PathBuf, + staged: &mut Vec<(String, String)>, +) -> Result<(), CliError> { let mut file = File::open(path) .with_context(|| format!("open upload file {}", path.display())) .map_err(CliError::Local)?; @@ -149,55 +162,39 @@ fn stage_file(sock: &Path, session: &str, path: &PathBuf) -> Result<(String, Str }), Duration::from_secs(10), )?; - let staged = (|| { - let mut offset = 0u64; - let mut buf = vec![0u8; begin.chunk_size as usize]; - loop { - let n = file.read(&mut buf).map_err(|e| CliError::Local(e.into()))?; - if n == 0 { - break; - } - let reply: TransferChunkResult = crate::cli::business_rpc::call( - sock.to_path_buf(), - "transfer-chunk", - Method::TransferChunk, - Some(TransferChunkParams { - transfer_id: begin.transfer_id.clone(), - offset, - data_base64: base64::engine::general_purpose::STANDARD.encode(&buf[..n]), - }), - Duration::from_secs(30), - )?; - offset = reply.next_offset; + // Record ownership immediately so one error path releases both completed + // files and the current partially staged file. + staged.push((begin.transfer_id.clone(), name)); + let mut offset = 0u64; + let mut buf = vec![0u8; begin.chunk_size as usize]; + loop { + let n = file.read(&mut buf).map_err(|e| CliError::Local(e.into()))?; + if n == 0 { + break; } - let _: TransferReadyResult = crate::cli::business_rpc::call( + let reply: TransferChunkResult = crate::cli::business_rpc::call( sock.to_path_buf(), - "transfer-finish", - Method::TransferFinish, - Some(TransferIdParams { + "transfer-chunk", + Method::TransferChunk, + Some(TransferChunkParams { transfer_id: begin.transfer_id.clone(), + offset, + data_base64: base64::engine::general_purpose::STANDARD.encode(&buf[..n]), }), - Duration::from_secs(10), + Duration::from_secs(30), )?; - Ok::<_, CliError>(()) - })(); - if let Err(err) = staged { - let _ = release(sock, &begin.transfer_id); - return Err(err); + offset = reply.next_offset; } - Ok((begin.transfer_id, name)) -} - -fn release(sock: &Path, id: &str) -> Result { - crate::cli::business_rpc::call( + let _: TransferReadyResult = crate::cli::business_rpc::call( sock.to_path_buf(), - "transfer-release", - Method::TransferRelease, + "transfer-finish", + Method::TransferFinish, Some(TransferIdParams { - transfer_id: id.to_string(), + transfer_id: begin.transfer_id.clone(), }), - Duration::from_secs(5), - ) + Duration::from_secs(10), + )?; + Ok(()) } fn ipc_timeout(timeout_ms: u32) -> Duration { diff --git a/crates/bsk-cli/tests/windows_parent_cancel.rs b/crates/bsk-cli/tests/windows_parent_cancel.rs index 44b4452..77d2738 100644 --- a/crates/bsk-cli/tests/windows_parent_cancel.rs +++ b/crates/bsk-cli/tests/windows_parent_cancel.rs @@ -1,10 +1,12 @@ -//! Real Windows CLI processes cancel over IPC when their parent closes stdin. +//! Real Windows CLI cancellation, including caller-owned transfer cleanup. #![cfg(windows)] +use std::path::Path; use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::time::{Duration, Instant}; -use bsk::daemon::info::DaemonInfo; +use bsk::daemon::{file_transfer::TransferRegistry, info::DaemonInfo}; +use bsk_protocol::tools::{TransferBeginParams, TransferIdParams}; use bsk_protocol::{ErrorCode, Frame, Method, ResponseBody, ResponseFrame, RpcError}; use serde_json::json; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; @@ -12,9 +14,49 @@ use tokio::net::windows::named_pipe::ServerOptions; use tokio::process::Command; use tokio::sync::{Notify, watch}; -#[tokio::test] -async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { - for mode in ["wait", "upload", "ordinary"] { +fn response(result: Result) -> ResponseBody { + match result { + Ok(value) => ResponseBody::Ok(serde_json::to_value(value).unwrap()), + Err(err) => ResponseBody::Err(err), + } +} + +#[test] +fn stdin_close_cancels_business_but_releases_resources_with_a_bounded_budget() { + // Resolve the production registry's root before starting any runtime threads. + let registry_home = tempfile::tempdir().unwrap(); + let original_home = std::env::var_os("BSK_HOME"); + unsafe { + std::env::set_var("BSK_HOME", registry_home.path()); + } + let transfers = Arc::new(TransferRegistry::new().unwrap()); + unsafe { + match original_home { + Some(value) => std::env::set_var("BSK_HOME", value), + None => std::env::remove_var("BSK_HOME"), + } + } + transfers.initialize().unwrap(); + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(exercise( + transfers, + ®istry_home.path().join("run/transfers"), + )); +} + +async fn exercise(transfers: Arc, staging_root: &Path) { + for mode in [ + "wait", + "upload", + "upload-multi", + "upload-before-dispatch", + "download", + "upload-multi-stall", + "ordinary", + ] { let home = tempfile::tempdir().unwrap(); let pipe_name = format!(r"\\.\pipe\bsk-parent-cancel-{}", uuid::Uuid::new_v4()); let info = DaemonInfo::now(std::process::id(), (&pipe_name).into(), 0, "0.2.0"); @@ -29,9 +71,30 @@ async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { .unwrap(); let registered = Arc::new(Notify::new()); let pending = Arc::new(Mutex::new(None::)); + let ids = Arc::new(Mutex::new(Vec::::new())); + let released = Arc::new(Mutex::new(Vec::::new())); + let received = Arc::new(Mutex::new(Vec::::new())); + let download_id = if mode == "download" { + let staging = transfers.begin_download("fixture").unwrap(); + let browser_dir = home.path().join("BrowserSkill").join(&staging.transfer_id); + std::fs::create_dir_all(&browser_dir).unwrap(); + let file = browser_dir.join("download.txt"); + std::fs::write(&file, b"download fixture").unwrap(); + transfers + .import_download(&staging.transfer_id, &file) + .unwrap(); + ids.lock().unwrap().push(staging.transfer_id.clone()); + Some(staging.transfer_id) + } else { + None + }; let (cancel, cancelled) = watch::channel(false); let ready = Arc::clone(®istered); let name = pipe_name.clone(); + let registry = Arc::clone(&transfers); + let server_ids = Arc::clone(&ids); + let server_released = Arc::clone(&released); + let server_received = Arc::clone(&received); let server = tokio::spawn(async move { let mut connections = tokio::task::JoinSet::new(); loop { @@ -43,6 +106,11 @@ async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { let cancel = cancel.clone(); let mut cancelled = cancelled.clone(); let name = name.clone(); + let registry = Arc::clone(®istry); + let ids = Arc::clone(&server_ids); + let released = Arc::clone(&server_released); + let received = Arc::clone(&server_received); + let download_id = download_id.clone(); connections.spawn(async move { let (read, mut write) = tokio::io::split(connection); let mut reader = BufReader::new(read); @@ -53,6 +121,8 @@ async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { panic!("expected request"); }; line.clear(); + received.lock().unwrap().push(req.method.clone()); + let params = req.params.clone().unwrap_or_default(); let body = match req.method { Method::SystemStatus => ResponseBody::Ok(json!({ "daemon_version": "0.2.0", "protocol_version": "1.1", @@ -60,29 +130,69 @@ async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { "ws_port": 0, "sock_path": name, "browsers": [], "sessions": [], "version_skew_browsers": [] })), - Method::TransferBegin => ResponseBody::Ok(json!({ - "transfer_id": "fixture", "chunk_size": 512 + Method::TransferBegin => { + let begin = registry + .begin_upload(serde_json::from_value(params).unwrap()) + .unwrap(); + ids.lock().unwrap().push(begin.transfer_id.clone()); + response(Ok(begin)) + } + Method::ToolDownload => ResponseBody::Ok(json!({ + "tab_id": 1, "suggested_filename": "download.txt", + "byte_size": 16, "transfer_id": download_id, })), Method::ToolWaitMs if mode == "ordinary" => { ResponseBody::Ok(json!({"waited_ms": 1})) } - Method::ToolWaitMs | Method::TransferChunk => { + Method::TransferChunk + if mode == "upload-before-dispatch" + || (mode.contains("multi") + && ids.lock().unwrap().len() == 1) => + { + response( + registry.write_chunk(serde_json::from_value(params).unwrap()), + ) + } + Method::TransferFinish if mode != "upload-before-dispatch" => response( + registry.finish_upload(serde_json::from_value(params).unwrap()), + ), + Method::ToolWaitMs + | Method::TransferChunk + | Method::TransferRead + | Method::TransferFinish => { *pending.lock().unwrap() = Some(req.id.clone()); ready.notify_one(); cancelled.wait_for(|value| *value).await.unwrap(); - ResponseBody::Err(RpcError { - code: ErrorCode::Cancelled, - message: "parent cancellation reached daemon".into(), - data: None, - }) + if mode == "upload-before-dispatch" { + // Finish won the race with cancellation, but tool.upload must + // not be dispatched and the CLI still owns the staging. + response( + registry + .finish_upload(serde_json::from_value(params).unwrap()), + ) + } else { + ResponseBody::Err(RpcError { + code: ErrorCode::Cancelled, + message: "parent cancellation reached daemon".into(), + data: None, + }) + } } Method::Cancel => { - let target = - req.params.as_ref().unwrap()["rpc_id"].as_str().unwrap(); + let target = params["rpc_id"].as_str().unwrap(); assert_eq!(pending.lock().unwrap().as_deref(), Some(target)); cancel.send(true).unwrap(); ResponseBody::Ok(json!({"cancelled": true})) } + Method::TransferRelease => { + let params: TransferIdParams = + serde_json::from_value(params).unwrap(); + released.lock().unwrap().push(params.transfer_id.clone()); + if mode.contains("stall") { + std::future::pending::<()>().await; + } + response(Ok(registry.release(params))) + } other => panic!("unexpected request after cancellation: {other:?}"), }; let response = Frame::Response(ResponseFrame { id: req.id, body }); @@ -104,18 +214,37 @@ async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { .stderr(std::process::Stdio::piped()) .creation_flags(0x0800_0000) .kill_on_drop(true); - if mode == "upload" { + if mode.starts_with("upload") { let source = home.path().join("upload.txt"); - std::fs::write(&source, b"upload fixture").unwrap(); + // Declare 300 MiB without allocating or transferring that much data. + std::fs::File::create(&source) + .unwrap() + .set_len(if mode == "upload-before-dispatch" { + 16 + } else { + 300 * 1024 * 1024 + }) + .unwrap(); + cmd.args(["upload", "--session", "fixture", "--selector", "input"]); + if mode.contains("multi") { + let first = home.path().join("first.txt"); + std::fs::write(&first, b"first file").unwrap(); + cmd.arg("--file").arg(first); + } + cmd.arg("--file").arg(source); + } else if mode == "download" { + let out = home.path().join("out.txt"); + std::fs::write(&out, b"existing output").unwrap(); cmd.args([ - "upload", + "download", "--session", "fixture", "--selector", - "input", - "--file", + "a", + "--overwrite", + "--out", ]) - .arg(source); + .arg(out); } else { cmd.args(["wait-ms", if mode == "ordinary" { "1ms" } else { "60s" }]); } @@ -131,7 +260,6 @@ async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { tokio::time::timeout(Duration::from_secs(10), registered.notified()) .await .expect("CLI registered business RPC"); - // Normal stdin bytes are not a cancellation; only closing the pipe is. child .stdin .as_mut() @@ -141,9 +269,10 @@ async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { .unwrap(); drop(child.stdin.take()); } - let output = tokio::time::timeout(Duration::from_secs(5), child.wait_with_output()) + let start = Instant::now(); + let output = tokio::time::timeout(Duration::from_secs(8), child.wait_with_output()) .await - .expect("CLI must settle promptly") + .expect("CLI must settle promptly, including cleanup") .unwrap(); server.abort(); let _ = server.await; @@ -157,7 +286,57 @@ async fn stdin_close_cancels_single_and_multi_rpc_commands_but_is_opt_in() { if mode != "ordinary" { let body: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); assert_eq!(body["code"], "cancelled"); - assert_eq!(body["message"], "parent cancellation reached daemon"); + if mode != "upload-before-dispatch" { + assert_eq!(body["message"], "parent cancellation reached daemon"); + } + let received = received.lock().unwrap(); + let cancel = received.iter().position(|m| *m == Method::Cancel).unwrap(); + assert!( + received[cancel + 1..] + .iter() + .all(|m| *m == Method::TransferRelease), + "only cleanup may follow cancellation: {received:?}" + ); + } + let ids = ids.lock().unwrap(); + if mode.contains("stall") { + assert_eq!(released.lock().unwrap().len(), 1); + assert!(start.elapsed() >= Duration::from_secs(4)); + // The CLI cannot force an unresponsive daemon to clean up. + transfers.release_session("fixture"); + } else { + assert_eq!( + *released.lock().unwrap(), + *ids, + "{mode}: every allocation must be released" + ); + for id in ids.iter() { + assert!(!staging_root.join(id).exists(), "staging leaked: {id}"); + } + if mode.starts_with("upload") { + let retry = transfers + .begin_upload(TransferBeginParams { + session_id: "fixture".into(), + name: "retry.txt".into(), + byte_size: 300 * 1024 * 1024, + }) + .expect("same session must have room for another 300 MiB upload"); + transfers.release(TransferIdParams { + transfer_id: retry.transfer_id, + }); + } + } + if mode == "download" { + assert_eq!( + std::fs::read(home.path().join("out.txt")).unwrap(), + b"existing output" + ); + assert!( + !std::fs::read_dir(home.path()) + .unwrap() + .flatten() + .any(|entry| entry.file_name().to_string_lossy().ends_with(".part")) + ); } } } diff --git a/packages/dsh-plugin-browserskill/src/runner.ts b/packages/dsh-plugin-browserskill/src/runner.ts index c80080e..9ec4adf 100644 --- a/packages/dsh-plugin-browserskill/src/runner.ts +++ b/packages/dsh-plugin-browserskill/src/runner.ts @@ -71,8 +71,9 @@ export interface BskRunner { // Business RPCs translate Ctrl-C / opt-in stdin EOF into cancel(rpc_id). // Allow reconciliation before hard-killing an old or unresponsive CLI. const KILL_GRACE_MS = 3000; -// Windows IPC may spend 5s connecting, then 2s cancelling and 2s settling. -const WINDOWS_KILL_GRACE_MS = 10_000; +// Windows IPC may spend 5s connecting, 2s cancelling, 2s settling, +// and up to 5s releasing the entire batch of caller-owned transfers. +const WINDOWS_KILL_GRACE_MS = 15_000; const SESSION_BUSY_RETRY_DELAY_MS = 100; export function createBskRunner(bskPath: string, spawnImpl: SpawnImpl = spawn): BskRunner { diff --git a/packages/dsh-plugin-browserskill/tests/runner.test.ts b/packages/dsh-plugin-browserskill/tests/runner.test.ts index 8ea7af4..219f290 100644 --- a/packages/dsh-plugin-browserskill/tests/runner.test.ts +++ b/packages/dsh-plugin-browserskill/tests/runner.test.ts @@ -133,6 +133,8 @@ describe("Windows parent cancellation", () => { expect(child.stdin.writableEnded).toBe(true); expect(child.killedWith).toEqual([]); await vi.advanceTimersByTimeAsync(10_000); + expect(child.killedWith).toEqual([]); + await vi.advanceTimersByTimeAsync(5_000); expect(child.killedWith).toEqual(["SIGKILL"]); expect(await result).toMatchObject({ timedOut: true }); });