use anyhow::Result;
use crate::registry::Registry;
use crate::session::backend::TabBackend;
pub async fn with_scratch_recovery<F, T, Fut>(
backend: &TabBackend,
registry: &Registry,
browser_name: &str,
mut op: F,
) -> Result<T>
where
F: FnMut(TabBackend, String) -> Fut,
Fut: std::future::Future<Output = Result<T>>,
{
let attempt_one_target = match registry.scratch_get(browser_name)? {
Some(row) => row.target_id,
None => {
let new_target = backend.create_tab("about:blank").await?;
registry.scratch_upsert(browser_name, &new_target)?;
new_target
}
};
match op(backend.clone(), attempt_one_target.clone()).await {
Ok(value) => {
let _ = registry.scratch_touch(browser_name);
return Ok(value);
}
Err(e) if is_scratch_failure(&e) => {
let _ = backend.close_tab(&attempt_one_target).await;
registry.scratch_delete(browser_name)?;
}
Err(e) => return Err(e),
}
let new_target = backend.create_tab("about:blank").await?;
registry.scratch_upsert(browser_name, &new_target)?;
let result = op(backend.clone(), new_target).await;
if result.is_ok() {
let _ = registry.scratch_touch(browser_name);
}
result
}
fn is_scratch_failure(err: &anyhow::Error) -> bool {
crate::errors::is_recoverable_tab_failure(err)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cdp::CdpClient;
use crate::errors::SessionError;
use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::sync::oneshot;
use tokio_tungstenite::tungstenite::Message;
async fn spawn_mock(
eval_behaviour: EvalBehaviour,
) -> (String, Arc<AtomicU32>, oneshot::Sender<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let create_count = Arc::new(AtomicU32::new(0));
let cc = create_count.clone();
let (stop_tx, mut stop_rx) = oneshot::channel::<()>();
tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap();
let mut next_target = 0u32;
let mut next_session = 0u32;
let mut evals_seen = 0u32;
loop {
tokio::select! {
_ = &mut stop_rx => break,
msg = ws.next() => {
let msg = match msg {
Some(Ok(m)) => m,
_ => break,
};
if let Message::Text(t) = msg {
let req: Value = serde_json::from_str(&t).unwrap();
let id = req["id"].as_u64().unwrap();
let method = req["method"].as_str().unwrap_or("");
if method == "Runtime.evaluate" {
evals_seen += 1;
if eval_behaviour.should_drop(evals_seen) {
continue;
}
}
let result = match method {
"Target.createTarget" => {
next_target += 1;
cc.fetch_add(1, Ordering::SeqCst);
serde_json::json!({"targetId": format!("T{next_target}")})
}
"Target.closeTarget" => serde_json::json!({"success": true}),
"Target.attachToTarget" => {
next_session += 1;
serde_json::json!({"sessionId": format!("S{next_session}")})
}
"Target.detachFromTarget" => serde_json::json!({}),
"Runtime.evaluate" => {
serde_json::json!({"result": {"value": 42}})
}
_ => serde_json::json!({}),
};
let resp = serde_json::json!({"id": id, "result": result});
ws.send(Message::Text(resp.to_string())).await.unwrap();
}
}
}
}
});
(format!("ws://{addr}"), create_count, stop_tx)
}
#[derive(Copy, Clone)]
enum EvalBehaviour {
Always,
Wedge(u32),
}
impl EvalBehaviour {
fn should_drop(self, eval_index: u32) -> bool {
match self {
EvalBehaviour::Always => false,
EvalBehaviour::Wedge(n) => eval_index <= n,
}
}
}
async fn eval_op(backend: TabBackend, target_id: String) -> Result<Value> {
backend
.evaluate(
&target_id,
"1",
false,
std::time::Duration::from_millis(200),
)
.await
}
#[tokio::test]
async fn first_call_creates_scratch_and_returns_value() {
let (url, create_count, _stop) = spawn_mock(EvalBehaviour::Always).await;
let client = Arc::new(CdpClient::connect(&url).await.unwrap());
let backend = TabBackend::Cdp(client);
let reg = Registry::open_in_memory().unwrap();
let v = with_scratch_recovery(&backend, ®, "brave-twilight", eval_op)
.await
.unwrap();
assert_eq!(v, json!(42));
assert_eq!(create_count.load(Ordering::SeqCst), 1, "one scratch tab");
let row = reg.scratch_get("brave-twilight").unwrap().unwrap();
assert_eq!(row.target_id, "T1");
}
#[tokio::test]
async fn second_call_reuses_scratch_row() {
let (url, create_count, _stop) = spawn_mock(EvalBehaviour::Always).await;
let client = Arc::new(CdpClient::connect(&url).await.unwrap());
let backend = TabBackend::Cdp(client);
let reg = Registry::open_in_memory().unwrap();
with_scratch_recovery(&backend, ®, "b", eval_op)
.await
.unwrap();
with_scratch_recovery(&backend, ®, "b", eval_op)
.await
.unwrap();
assert_eq!(
create_count.load(Ordering::SeqCst),
1,
"second call reused the row, no new target"
);
}
#[tokio::test]
async fn recovers_after_one_wedge() {
let (url, create_count, _stop) = spawn_mock(EvalBehaviour::Wedge(1)).await;
let client = Arc::new(CdpClient::connect(&url).await.unwrap());
let backend = TabBackend::Cdp(client);
let reg = Registry::open_in_memory().unwrap();
let v = with_scratch_recovery(&backend, ®, "b", eval_op)
.await
.unwrap();
assert_eq!(v, json!(42));
assert_eq!(
create_count.load(Ordering::SeqCst),
2,
"second target created after the wedge"
);
let row = reg.scratch_get("b").unwrap().unwrap();
assert_eq!(row.target_id, "T2");
}
#[test]
fn is_scratch_failure_recognizes_typed_target_gone() {
use crate::errors::TargetKind;
let typed: anyhow::Error = SessionError::TargetGone {
kind: TargetKind::Cdp,
details: "CDP error -32000: No target with given id found: T1".into(),
}
.into();
assert!(is_scratch_failure(&typed));
let typed_bidi: anyhow::Error = SessionError::TargetGone {
kind: TargetKind::Bidi,
details: "BiDi error no such context: C1".into(),
}
.into();
assert!(is_scratch_failure(&typed_bidi));
let hung: anyhow::Error = SessionError::TabHung {
target_id: None,
url: None,
timeout_ms: 100,
hint: "test",
}
.into();
assert!(is_scratch_failure(&hung));
let raw: anyhow::Error = anyhow::anyhow!("No target with given id found: T1");
assert!(is_scratch_failure(&raw));
let unrelated: anyhow::Error = anyhow::anyhow!("network unreachable");
assert!(!is_scratch_failure(&unrelated));
}
#[tokio::test]
async fn escalates_after_second_wedge() {
let (url, create_count, _stop) = spawn_mock(EvalBehaviour::Wedge(99)).await;
let client = Arc::new(CdpClient::connect(&url).await.unwrap());
let backend = TabBackend::Cdp(client);
let reg = Registry::open_in_memory().unwrap();
let err = with_scratch_recovery(&backend, ®, "b", eval_op)
.await
.expect_err("must escalate");
let typed = err
.downcast_ref::<SessionError>()
.expect("typed SessionError");
assert!(
matches!(typed, SessionError::TabHung { .. }),
"expected TabHung, got {typed:?}"
);
assert_eq!(create_count.load(Ordering::SeqCst), 2);
}
}