Fix browser wait during daemon startup

Co-authored-by: Cursor <cursoragent@cursor.com>

crates/bsk-cli/src/cli/browser_wait.rs

@@ -0,0 +1,56 @@
+//! Shared browser-connect wait durations for status/doctor/browsers commands.
+
+use std::time::Duration;
+
+use crate::daemon::browsers::EXTENSION_CONNECT_WAIT;
+
+/// Default wait passed to daemon status/list RPCs when the registry is empty.
+pub const DEFAULT_BROWSER_CONNECT_WAIT: Duration = EXTENSION_CONNECT_WAIT;
+/// Keep CLI-side waits aligned with the daemon's IPC clamp so a bad env value
+/// cannot inflate read timeouts after the daemon has already returned.
+const MAX_BROWSER_CONNECT_WAIT: Duration = Duration::from_secs(60);
+
+/// Browser wait for `bsk status` / `bsk browsers`. Override in tests via
+/// `BSK_BROWSER_WAIT_MS`.
+pub fn browser_connect_wait() -> Duration {
+    wait_from_env("BSK_BROWSER_WAIT_MS").unwrap_or(DEFAULT_BROWSER_CONNECT_WAIT)
+}
+
+/// Browser wait for `bsk doctor`. Falls back to [`browser_connect_wait`] when
+/// `BSK_DOCTOR_BROWSER_WAIT_MS` is unset so tests can shorten only doctor.
+pub fn doctor_browser_connect_wait() -> Duration {
+    wait_from_env("BSK_DOCTOR_BROWSER_WAIT_MS").unwrap_or_else(browser_connect_wait)
+}
+
+/// Encodes a wait duration for protocol params. A zero duration preserves the
+/// legacy immediate-snapshot semantics.
+pub fn wait_for_browser_ms(wait: Duration) -> Option<u64> {
+    let wait = clamp_browser_wait(wait);
+    if wait.is_zero() {
+        None
+    } else {
+        Some(wait.as_millis() as u64)
+    }
+}
+
+/// IPC call budget: the command's normal timeout, or browser wait plus
+/// scheduling slack when a wait is requested.
+pub fn browser_query_ipc_timeout(wait: Duration, default_timeout: Duration) -> Duration {
+    let wait = clamp_browser_wait(wait);
+    if wait.is_zero() {
+        default_timeout
+    } else {
+        wait.saturating_add(Duration::from_secs(2))
+    }
+}
+
+fn wait_from_env(key: &str) -> Option<Duration> {
+    std::env::var(key)
+        .ok()
+        .and_then(|v| v.parse::<u64>().ok())
+        .map(|ms| clamp_browser_wait(Duration::from_millis(ms)))
+}
+
+fn clamp_browser_wait(wait: Duration) -> Duration {
+    wait.min(MAX_BROWSER_CONNECT_WAIT)
+}

crates/bsk-cli/src/cli/browsers.rs

@@ -5,9 +5,13 @@ use std::time::Duration;
 
 use anyhow::Context;
 use bsk_protocol::Method;
+use bsk_protocol::system::BrowserListParams;
 use bsk_protocol::system::BrowserStatusEntry;
 use serde::Deserialize;
 
+use crate::cli::browser_wait::{
+    browser_connect_wait, browser_query_ipc_timeout, wait_for_browser_ms,
+};
 use crate::cli::ensure_daemon::ensure_daemon;
 use crate::cli::error::{CliError, Format};
 
@@ -22,7 +26,12 @@ pub fn dispatch(format: Format) -> Result<(), CliError> {
 }
 
 fn run_list(sock: PathBuf, format: Format) -> Result<(), CliError> {
-    let reply: ListReply = call(sock, Duration::from_secs(5))?;
+    let wait = browser_connect_wait();
+    let params = BrowserListParams {
+        wait_for_browser_ms: wait_for_browser_ms(wait),
+    };
+    let timeout = browser_query_ipc_timeout(wait, Duration::from_secs(5));
+    let reply: ListReply = call(sock, params, timeout)?;
     match format {
         Format::Json => {
             println!(
@@ -92,7 +101,11 @@ fn run_list(sock: PathBuf, format: Format) -> Result<(), CliError> {
     Ok(())
 }
 
-fn call(sock: PathBuf, timeout: Duration) -> Result<ListReply, CliError> {
+fn call(
+    sock: PathBuf,
+    params: BrowserListParams,
+    timeout: Duration,
+) -> Result<ListReply, CliError> {
     let rt = tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()
@@ -101,7 +114,7 @@ fn call(sock: PathBuf, timeout: Duration) -> Result<ListReply, CliError> {
     rt.block_on(async move {
         let mut client = crate::ipc_client::IpcClient::connect(sock).await?;
         let outcome = client
-            .call::<(), ListReply>("browser-list-1", Method::BrowserList, None, timeout)
+            .call("browser-list-1", Method::BrowserList, Some(params), timeout)
             .await?;
         outcome.map_err(CliError::from_rpc)
     })

crates/bsk-cli/src/cli/doctor.rs

@@ -1,23 +1,19 @@
 //! `bsk doctor` — guided diagnostics + repair hints.
 
-use std::time::{Duration, Instant};
+use std::time::Duration;
 
 use anyhow::Result;
 use bsk_protocol::StatusResult;
 use console::style;
 use serde::Serialize;
 
+use crate::cli::browser_wait::doctor_browser_connect_wait;
 use crate::cli::ensure_daemon::ensure_daemon;
 use crate::cli::status::{self, Output};
-use crate::daemon::browsers::EXTENSION_CONNECT_WAIT;
 use crate::daemon::info::{self, DaemonInfo};
 use crate::daemon::paths;
 use crate::daemon::state::PROTOCOL_VERSION;
 
-/// How often `warm_up` polls `system.status` while waiting for the
-/// browser extension to complete its WS handshake.
-const BROWSER_CONNECT_POLL: Duration = Duration::from_millis(200);
-
 /// Chrome Web Store listing for the browser-skill extension.
 const EXTENSION_STORE_URL: &str =
     "https://chromewebstore.google.com/detail/hhcmgoofomhgciiibhipgmgkgnoenaoi";
@@ -99,8 +95,8 @@ impl CheckResult {
 }
 
 pub fn run(output: Output) -> Result<Vec<CheckResult>> {
-    warm_up(output);
-    let checks = collect_checks();
+    let state = resolve_daemon_state(output);
+    let checks = collect_checks(state);
     match output {
         Output::Human => render_human(&checks),
         Output::Json => render_json(&checks)?,
@@ -108,74 +104,47 @@ pub fn run(output: Output) -> Result<Vec<CheckResult>> {
     Ok(checks)
 }
 
-/// Ensure the daemon is up and give the browser extension time to connect.
-///
-/// All errors are silently swallowed — failures are reported by the
-/// individual checks in `collect_checks()` with actionable hints.
-///
-/// Flow:
-/// 1. Read the current daemon state. **Only** proceed for `Missing` and
-///    `StaleDead`: broken states (`ReadError`, `IpcUnreachable`,
-///    `PidMismatch`) are real diagnostic findings that `collect_checks()`
-///    should surface — fixing them here would mask the real problem.
-/// 2. Call `ensure_daemon()` to auto-start the daemon (spawns a detached
-///    child and waits up to 3 s for it to write a valid `daemon.json`).
-/// 3. Poll `system.status` for up to [`EXTENSION_CONNECT_WAIT`] so the MV3
-///    service worker has time to wake up, discover the WS port, and finish
-///    its handshake before the checks are evaluated.
-///    The wait duration can be shortened for tests by setting the env var
-///    `BSK_DOCTOR_BROWSER_WAIT_MS` to a millisecond count.
-fn warm_up(output: Output) {
-    // Only spawn for states where the daemon is genuinely absent or dead.
-    match current_state() {
-        DaemonState::Missing | DaemonState::StaleDead(_) => {}
-        _ => return,
+/// Ensure the daemon is reachable and give the browser extension time to
+/// connect before checks run. Returns a single [`DaemonState`] snapshot
+/// for check evaluation.
+fn resolve_daemon_state(output: Output) -> DaemonState {
+    let mut state = current_state(Duration::ZERO);
+
+    if matches!(state, DaemonState::Missing | DaemonState::StaleDead(_)) && ensure_daemon().is_err()
+    {
+        return current_state(Duration::ZERO);
     }
 
-    let info = match ensure_daemon() {
-        Ok(info) => info,
-        Err(_) => return,
+    if matches!(state, DaemonState::Missing | DaemonState::StaleDead(_)) {
+        state = current_state(Duration::ZERO);
+    }
+
+    let wait = if needs_browser_wait(&state) {
+        doctor_browser_connect_wait()
+    } else {
+        Duration::ZERO
     };
 
-    // Daemon was freshly spawned. Give the MV3 service worker time to
-    // wake up, discover the WS port, and complete its handshake.
-    if output == Output::Human {
+    if wait > Duration::ZERO && output == Output::Human {
         eprintln!("waiting for browser extension to connect…");
     }
 
-    let wait = browser_connect_wait();
-    let deadline = Instant::now() + wait;
-    loop {
-        // Check deadline before sleeping so a zero-duration wait exits
-        // immediately and we never start a query_sock call past the deadline.
-        if Instant::now() >= deadline {
-            break;
-        }
-        std::thread::sleep(BROWSER_CONNECT_POLL);
-        if Instant::now() >= deadline {
-            break;
-        }
-        if let Ok(s) = status::query_sock(info.sock_path.clone()) {
-            if !s.browsers.is_empty() {
-                break;
-            }
-        }
+    if wait.is_zero() {
+        state
+    } else {
+        current_state(wait)
     }
 }
 
-/// The browser connect wait duration for `warm_up`. Defaults to
-/// [`EXTENSION_CONNECT_WAIT`]; can be overridden in tests via the
-/// `BSK_DOCTOR_BROWSER_WAIT_MS` environment variable.
-fn browser_connect_wait() -> Duration {
-    std::env::var("BSK_DOCTOR_BROWSER_WAIT_MS")
-        .ok()
-        .and_then(|v| v.parse::<u64>().ok())
-        .map(Duration::from_millis)
-        .unwrap_or(EXTENSION_CONNECT_WAIT)
+fn needs_browser_wait(state: &DaemonState) -> bool {
+    match state {
+        DaemonState::Verified { status } => status.browsers.is_empty(),
+        _ => false,
+    }
 }
 
 /// What the disk + IPC says about a possibly-running daemon. Threaded
-/// through every check so they share one IPC round-trip.
+/// through every check so they share one snapshot.
 enum DaemonState {
     /// No `daemon.json` at all.
     Missing,
@@ -208,8 +177,7 @@ impl DaemonState {
     }
 }
 
-fn collect_checks() -> Vec<CheckResult> {
-    let state = current_state();
+fn collect_checks(state: DaemonState) -> Vec<CheckResult> {
     vec![
         check_home_writable(),
         check_skill_up_to_date(),
@@ -220,7 +188,7 @@ fn collect_checks() -> Vec<CheckResult> {
     ]
 }
 
-fn current_state() -> DaemonState {
+fn current_state(browser_wait: Duration) -> DaemonState {
     let info = match info::read() {
         Ok(Some(info)) => info,
         Ok(None) => return DaemonState::Missing,
@@ -229,7 +197,7 @@ fn current_state() -> DaemonState {
     if !crate::daemon::lockfile::pid_alive(info.pid) {
         return DaemonState::StaleDead(info);
     }
-    match status::query_sock(info.sock_path.clone()) {
+    match status::query_sock_with_wait(info.sock_path.clone(), browser_wait) {
         Ok(status) => {
             if status.pid == info.pid {
                 DaemonState::Verified { status }

crates/bsk-cli/src/cli/mod.rs

@@ -2,6 +2,7 @@
 
 use std::time::Duration;
 
+pub mod browser_wait;
 pub mod browsers;
 pub mod business_rpc;
 pub mod daemon;

crates/bsk-cli/src/cli/status.rs

@@ -5,6 +5,9 @@ use std::time::Duration;
 use anyhow::Context;
 use bsk_protocol::{Method, StatusParams, StatusResult};
 
+use crate::cli::browser_wait::{
+    browser_connect_wait, browser_query_ipc_timeout, wait_for_browser_ms,
+};
 use crate::cli::ensure_daemon::ensure_daemon;
 use crate::cli::error::CliError;
 
@@ -19,7 +22,8 @@ pub enum Output {
 /// reuse it.
 pub fn run(output: Output) -> Result<StatusResult, CliError> {
     let info = ensure_daemon().context("ensure daemon is running")?;
-    let result = query_sock(info.sock_path)?;
+    let wait = browser_connect_wait();
+    let result = query_sock_with_wait(info.sock_path, wait)?;
     match output {
         Output::Human => render_human(&result),
         Output::Json => render_json(&result).map_err(CliError::Local)?,
@@ -27,7 +31,14 @@ pub fn run(output: Output) -> Result<StatusResult, CliError> {
     Ok(result)
 }
 
-pub(crate) fn query_sock(sock: std::path::PathBuf) -> Result<StatusResult, CliError> {
+pub(crate) fn query_sock_with_wait(
+    sock: std::path::PathBuf,
+    wait: Duration,
+) -> Result<StatusResult, CliError> {
+    let params = StatusParams {
+        wait_for_browser_ms: wait_for_browser_ms(wait),
+    };
+    let timeout = browser_query_ipc_timeout(wait, Duration::from_secs(2));
     let rt = tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()
@@ -36,11 +47,7 @@ pub(crate) fn query_sock(sock: std::path::PathBuf) -> Result<StatusResult, CliEr
     rt.block_on(async move {
         let mut client = crate::ipc_client::Client::connect_path(sock).await?;
         let outcome = client
-            .call::<_, StatusResult>(
-                Method::SystemStatus,
-                &StatusParams::default(),
-                Duration::from_secs(2),
-            )
+            .call::<_, StatusResult>(Method::SystemStatus, &params, timeout)
             .await?;
         outcome.map_err(CliError::from_rpc)
     })