mirror of
https://github.com/Tencent/BrowserSkill.git
synced 2026-09-10 13:01:23 +08:00
fix(record): recover failed Trace exports and hint --url when the default start page fails
Reject legacy JSON --output before recording starts, persist the completed Trace before writing the bundle so `bsk record stop --output <dir>` can retry after an export failure, and tell users to pass --url when https://example.com/ does not load.
This commit is contained in:
“haonan”
@@ -0,0 +1,327 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use bsk::daemon::info::DaemonInfo;
|
||||
use bsk::daemon::paths::BSK_HOME_ENV;
|
||||
use bsk::daemon::{self, DaemonConfig};
|
||||
use bsk::ipc_client::IpcClient;
|
||||
use bsk_protocol::system::{HandshakeParams, HandshakeResult};
|
||||
use bsk_protocol::tools::{
|
||||
RecordStopResult, RecordedTrace, RecorderInfo, SessionStartParams, SessionStartResult,
|
||||
StopReason, TraceEntry, TraceStateV3, TraceV3, VOM_FORMAT_VERSION,
|
||||
};
|
||||
use bsk_protocol::{BrowserPeerInfo, Frame, Method, RequestFrame, ResponseBody, ResponseFrame};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use serde_json::json;
|
||||
use tokio::process::Command;
|
||||
use tokio_tungstenite::tungstenite::handshake::client::generate_key;
|
||||
use tokio_tungstenite::tungstenite::http::Request;
|
||||
use tokio_tungstenite::tungstenite::protocol::Message;
|
||||
|
||||
const TEST_EXT_ID: &str = "abcdefghijklmnopabcdefghijklmnop";
|
||||
|
||||
type Ws =
|
||||
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
|
||||
|
||||
fn bsk_bin() -> PathBuf {
|
||||
PathBuf::from(env!("CARGO_BIN_EXE_bsk"))
|
||||
}
|
||||
|
||||
async fn connect_extension(addr: std::net::SocketAddr) -> Ws {
|
||||
let origin = format!("chrome-extension://{TEST_EXT_ID}");
|
||||
let request = Request::builder()
|
||||
.method("GET")
|
||||
.uri(format!("ws://{addr}/"))
|
||||
.header("Host", addr.to_string())
|
||||
.header("Upgrade", "websocket")
|
||||
.header("Connection", "Upgrade")
|
||||
.header("Sec-WebSocket-Version", "13")
|
||||
.header("Sec-WebSocket-Key", generate_key())
|
||||
.header("Origin", origin)
|
||||
.body(())
|
||||
.unwrap();
|
||||
tokio_tungstenite::connect_async(request).await.unwrap().0
|
||||
}
|
||||
|
||||
async fn handshake(ws: &mut Ws) {
|
||||
let request = RequestFrame {
|
||||
id: "handshake".into(),
|
||||
method: Method::SystemHandshake,
|
||||
params: Some(
|
||||
serde_json::to_value(HandshakeParams {
|
||||
client: "browser-skill-extension".into(),
|
||||
version: "0.1.10".parse().unwrap(),
|
||||
protocol_version: "1.1".into(),
|
||||
instance_id: TEST_EXT_ID.into(),
|
||||
browser: BrowserPeerInfo {
|
||||
name: "chrome".into(),
|
||||
version: "131.0".into(),
|
||||
},
|
||||
min_compatible_peer: Some("0.1.10".parse().unwrap()),
|
||||
min_compatible_protocol: Some("1.1".into()),
|
||||
label: "Record export recovery test".into(),
|
||||
})
|
||||
.unwrap(),
|
||||
),
|
||||
};
|
||||
ws.send(Message::Text(serde_json::to_string(&request).unwrap()))
|
||||
.await
|
||||
.unwrap();
|
||||
let Message::Text(response) = ws.next().await.unwrap().unwrap() else {
|
||||
panic!("expected text handshake response");
|
||||
};
|
||||
let response: ResponseFrame = serde_json::from_str(&response).unwrap();
|
||||
let ResponseBody::Ok(value) = response.body else {
|
||||
panic!("handshake failed");
|
||||
};
|
||||
let _: HandshakeResult = serde_json::from_value(value).unwrap();
|
||||
}
|
||||
|
||||
fn sample_trace() -> TraceV3 {
|
||||
TraceV3 {
|
||||
version: 3,
|
||||
purpose: Some("export recovery".into()),
|
||||
started_at: Some("2026-08-12T10:00:00Z".into()),
|
||||
recorded_at: "2026-08-12T10:01:00Z".into(),
|
||||
stopped_by: StopReason::CliStop,
|
||||
entry: TraceEntry {
|
||||
start_url: "https://example.com/".into(),
|
||||
},
|
||||
recorder: RecorderInfo {
|
||||
bsk: "0.1.10".into(),
|
||||
vom: VOM_FORMAT_VERSION,
|
||||
},
|
||||
states: vec![TraceStateV3 {
|
||||
id: "s1".into(),
|
||||
url: "https://example.com/".into(),
|
||||
title: Some("Example".into()),
|
||||
body: "@vom 1\nL1 page".into(),
|
||||
truncated: false,
|
||||
}],
|
||||
steps: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
fn run_extension(
|
||||
mut ws: Ws,
|
||||
record_stop_calls: Arc<AtomicUsize>,
|
||||
session_stop_calls: Arc<AtomicUsize>,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
while let Some(Ok(Message::Text(text))) = ws.next().await {
|
||||
let Ok(Frame::Request(request)) = serde_json::from_str::<Frame>(&text) else {
|
||||
continue;
|
||||
};
|
||||
let body = match request.method {
|
||||
Method::ToolSessionStart => {
|
||||
let _: SessionStartParams =
|
||||
serde_json::from_value(request.params.clone().unwrap()).unwrap();
|
||||
ResponseBody::Ok(
|
||||
serde_json::to_value(SessionStartResult {
|
||||
agent_window_id: Some(100),
|
||||
})
|
||||
.unwrap(),
|
||||
)
|
||||
}
|
||||
Method::ToolRecordStop => {
|
||||
record_stop_calls.fetch_add(1, Ordering::SeqCst);
|
||||
ResponseBody::Ok(
|
||||
serde_json::to_value(RecordStopResult {
|
||||
trace: RecordedTrace::V3(sample_trace()),
|
||||
})
|
||||
.unwrap(),
|
||||
)
|
||||
}
|
||||
Method::ToolSessionStop => {
|
||||
session_stop_calls.fetch_add(1, Ordering::SeqCst);
|
||||
ResponseBody::Ok(json!({}))
|
||||
}
|
||||
_ => ResponseBody::Err(bsk_protocol::RpcError {
|
||||
code: bsk_protocol::ErrorCode::ProtocolError,
|
||||
message: format!("unexpected method {:?}", request.method),
|
||||
data: None,
|
||||
}),
|
||||
};
|
||||
let response = ResponseFrame {
|
||||
id: request.id,
|
||||
body,
|
||||
};
|
||||
ws.send(Message::Text(serde_json::to_string(&response).unwrap()))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async fn start_session(sock: &Path) -> String {
|
||||
#[derive(serde::Deserialize)]
|
||||
struct StartResult {
|
||||
session_id: String,
|
||||
}
|
||||
|
||||
let mut client = IpcClient::connect(sock).await.unwrap();
|
||||
client
|
||||
.call::<(), StartResult>(
|
||||
"session-start",
|
||||
Method::SessionStart,
|
||||
None,
|
||||
Duration::from_secs(5),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.session_id
|
||||
}
|
||||
|
||||
fn write_record_session(home: &Path, session_id: &str) {
|
||||
std::fs::create_dir_all(home).unwrap();
|
||||
std::fs::write(
|
||||
home.join("record-session.json"),
|
||||
format!("{{\"session_id\":\"{session_id}\",\"started_at\":\"1\"}}\n"),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn run_record_stop(home: &Path, output: &Path) -> std::process::Output {
|
||||
Command::new(bsk_bin())
|
||||
.arg("record")
|
||||
.arg("stop")
|
||||
.arg("--output")
|
||||
.arg(output)
|
||||
.env(BSK_HOME_ENV, home)
|
||||
.output()
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn setup_daemon(temp: &Path) -> (bsk::daemon::DaemonHandle, PathBuf, PathBuf) {
|
||||
let home = temp.join("bsk-home");
|
||||
let sock = temp.join("daemon.sock");
|
||||
let daemon = daemon::run(DaemonConfig::new(0), Some(sock.clone()))
|
||||
.await
|
||||
.unwrap();
|
||||
bsk::daemon::info::write_to_path(
|
||||
&DaemonInfo::now(
|
||||
std::process::id(),
|
||||
sock.clone(),
|
||||
daemon.ws_addr().port(),
|
||||
env!("CARGO_PKG_VERSION"),
|
||||
),
|
||||
&home.join("daemon.json"),
|
||||
)
|
||||
.unwrap();
|
||||
(daemon, home, sock)
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn existing_trace_json_file_is_rejected_before_record_stop() {
|
||||
let temp = tempfile::Builder::new()
|
||||
.prefix("bsk-record-json-")
|
||||
.tempdir()
|
||||
.unwrap();
|
||||
let (daemon, home, sock) = setup_daemon(temp.path()).await;
|
||||
let output = temp.path().join("trace.json");
|
||||
std::fs::write(&output, "{}\n").unwrap();
|
||||
|
||||
let mut extension = connect_extension(daemon.ws_addr()).await;
|
||||
handshake(&mut extension).await;
|
||||
let record_stop_calls = Arc::new(AtomicUsize::new(0));
|
||||
let session_stop_calls = Arc::new(AtomicUsize::new(0));
|
||||
run_extension(
|
||||
extension,
|
||||
Arc::clone(&record_stop_calls),
|
||||
Arc::clone(&session_stop_calls),
|
||||
);
|
||||
|
||||
let session_id = start_session(&sock).await;
|
||||
write_record_session(&home, &session_id);
|
||||
|
||||
let first = run_record_stop(&home, &output).await;
|
||||
let stderr = String::from_utf8_lossy(&first.stderr);
|
||||
assert!(
|
||||
!first.status.success(),
|
||||
"legacy JSON --output should fail: {stderr}"
|
||||
);
|
||||
assert!(
|
||||
stderr.contains("JSON file") || stderr.contains("bundle directory"),
|
||||
"{stderr}"
|
||||
);
|
||||
let state: serde_json::Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(home.join("record-session.json")).unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(state["session_id"], session_id);
|
||||
assert_eq!(
|
||||
record_stop_calls.load(Ordering::SeqCst),
|
||||
0,
|
||||
"existing trace.json must be rejected before the extension returns the Trace"
|
||||
);
|
||||
assert_eq!(session_stop_calls.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(std::fs::read_to_string(&output).unwrap(), "{}\n");
|
||||
|
||||
daemon.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn failed_bundle_export_can_be_recovered() {
|
||||
let temp = tempfile::Builder::new()
|
||||
.prefix("bsk-record-recov-")
|
||||
.tempdir()
|
||||
.unwrap();
|
||||
let (daemon, home, sock) = setup_daemon(temp.path()).await;
|
||||
let output = temp.path().join("trace");
|
||||
std::fs::create_dir_all(&output).unwrap();
|
||||
std::fs::write(output.join("states"), "not a directory\n").unwrap();
|
||||
|
||||
let mut extension = connect_extension(daemon.ws_addr()).await;
|
||||
handshake(&mut extension).await;
|
||||
let record_stop_calls = Arc::new(AtomicUsize::new(0));
|
||||
let session_stop_calls = Arc::new(AtomicUsize::new(0));
|
||||
run_extension(
|
||||
extension,
|
||||
Arc::clone(&record_stop_calls),
|
||||
Arc::clone(&session_stop_calls),
|
||||
);
|
||||
|
||||
let session_id = start_session(&sock).await;
|
||||
write_record_session(&home, &session_id);
|
||||
|
||||
let first = run_record_stop(&home, &output).await;
|
||||
let stderr = String::from_utf8_lossy(&first.stderr);
|
||||
assert!(
|
||||
!first.status.success(),
|
||||
"blocked states path should fail export: {stderr}"
|
||||
);
|
||||
assert!(
|
||||
home.join("record-recovery.json").exists(),
|
||||
"completed Trace must be saved before the failed export"
|
||||
);
|
||||
assert!(
|
||||
!home.join("record-session.json").exists(),
|
||||
"session state is cleared after the extension returns the Trace"
|
||||
);
|
||||
assert_eq!(record_stop_calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(session_stop_calls.load(Ordering::SeqCst), 1);
|
||||
|
||||
std::fs::remove_file(output.join("states")).unwrap();
|
||||
let second = run_record_stop(&home, &output).await;
|
||||
assert!(
|
||||
second.status.success(),
|
||||
"recovery export failed: {}",
|
||||
String::from_utf8_lossy(&second.stderr)
|
||||
);
|
||||
assert_eq!(
|
||||
record_stop_calls.load(Ordering::SeqCst),
|
||||
1,
|
||||
"recovery must not call the extension again"
|
||||
);
|
||||
assert!(output.join("trace.json").exists());
|
||||
assert!(output.join("states/s1.txt").exists());
|
||||
assert!(
|
||||
!home.join("record-recovery.json").exists(),
|
||||
"successful export must delete recovery data"
|
||||
);
|
||||
|
||||
daemon.shutdown().await;
|
||||
}
|
||||
Reference in New Issue
Block a user