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)
     })

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

@@ -214,6 +214,24 @@ impl BrowserRegistry {
         self.len() == 0
     }
 
+    /// Poll until at least one browser is registered or `connect_wait` elapses.
+    /// Returns immediately when the registry is already non-empty.
+    pub async fn wait_for_any_connected(&self, connect_wait: Duration) {
+        if connect_wait.is_zero() || !self.is_empty() {
+            return;
+        }
+        let deadline = Instant::now() + connect_wait;
+        loop {
+            if !self.is_empty() {
+                return;
+            }
+            if Instant::now() >= deadline {
+                return;
+            }
+            tokio::time::sleep(EXTENSION_CONNECT_POLL).await;
+        }
+    }
+
     /// Pick the only connected browser when exactly one is online, or
     /// match a specific id / label when supplied. Returns an error code
     /// aligned with §4.5.
@@ -511,21 +529,29 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn select_with_connect_wait_does_not_wait_on_ambiguous_label() {
+    async fn wait_for_any_connected_returns_when_browser_registers_late() {
+        let reg = std::sync::Arc::new(BrowserRegistry::new());
+        let reg_bg = reg.clone();
+        let waiter = tokio::spawn(async move {
+            reg_bg
+                .wait_for_any_connected(Duration::from_millis(500))
+                .await;
+        });
+        tokio::time::sleep(Duration::from_millis(80)).await;
+        reg.insert(fake_client("late", ""));
+        waiter.await.expect("join");
+        assert_eq!(reg.len(), 1);
+    }
+
+    #[tokio::test]
+    async fn wait_for_any_connected_skips_when_already_populated() {
         let reg = BrowserRegistry::new();
-        reg.insert(fake_client("alpha", "Shared"));
-        reg.insert(fake_client("beta", "Shared"));
+        reg.insert(fake_client("a", ""));
         let started = Instant::now();
-        match reg
-            .select_with_connect_wait(Some("Shared"), Duration::from_secs(1))
-            .await
-        {
-            Err(SelectError::AmbiguousLabel { label, .. }) => assert_eq!(label, "Shared"),
-            other => panic!("expected AmbiguousLabel, got {other:?}"),
-        }
+        reg.wait_for_any_connected(Duration::from_secs(1)).await;
         assert!(
             started.elapsed() < Duration::from_millis(50),
-            "ambiguous label should fail immediately"
+            "should return immediately when browsers already connected"
         );
     }
 }

crates/bsk-cli/src/daemon/ipc.rs

@@ -23,12 +23,14 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use bsk_protocol::system::{
-    BrowserStatusEntry, SessionStatusEntry, StatusParams, StatusResult, VersionSkewEntry,
+    BrowserListParams, BrowserStatusEntry, SessionStatusEntry, StatusParams, StatusResult,
+    VersionSkewEntry,
 };
 use bsk_protocol::tools::{ReturnFailure, WaitMsParams, WaitMsResult};
 use bsk_protocol::{
     CancelParams, CancelResult, ErrorCode, Method, PingResult, ResponseBody, RpcError, RpcId,
 };
+use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use tokio::sync::Notify;
@@ -57,6 +59,8 @@ pub type RpcHandler = Arc<
 >;
 
 const DEFAULT_RPC_TIMEOUT: Duration = Duration::from_secs(15);
+/// Upper bound on `wait_for_browser_ms` accepted over IPC.
+const MAX_BROWSER_WAIT: Duration = Duration::from_secs(60);
 // `session.stop` fast-fails while another tool is active; this budget
 // only needs to cover the stop RPC itself and IPC scheduling grace.
 const DEFAULT_SESSION_STOP_TIMEOUT: Duration =
@@ -213,7 +217,10 @@ pub fn full_handler(status: DaemonStatus, state: Arc<DaemonState>) -> RpcHandler
                     Err(e) => ResponseBody::Err(e),
                 },
                 Method::SessionList => handle_session_list(&state),
-                Method::BrowserList => handle_browser_list(&state),
+                Method::BrowserList => match handle_browser_list(&state, params).await {
+                    Ok(v) => ResponseBody::Ok(v),
+                    Err(e) => ResponseBody::Err(e),
+                },
                 Method::ToolTabList
                 | Method::ToolTabCreate
                 | Method::ToolTabClose
@@ -554,14 +561,40 @@ async fn handle_status(
     state: &Arc<DaemonState>,
     params: Value,
 ) -> Result<Value, RpcError> {
-    if !params.is_null() {
-        let _: StatusParams = serde_json::from_value(params).map_err(|err| RpcError {
+    let params: StatusParams = parse_params_or_default(params)?;
+    maybe_wait_for_browser(state, params.wait_for_browser_ms).await;
+    Ok(serde_json::to_value(status.snapshot_with(state)).unwrap_or(Value::Null))
+}
+
+fn parse_params_or_default<T>(params: Value) -> Result<T, RpcError>
+where
+    T: DeserializeOwned + Default,
+{
+    if params.is_null() {
+        Ok(T::default())
+    } else {
+        serde_json::from_value(params).map_err(|err| RpcError {
             code: ErrorCode::InvalidParams,
             message: err.to_string(),
             data: None,
-        })?;
+        })
     }
-    Ok(serde_json::to_value(status.snapshot_with(state)).unwrap_or(Value::Null))
+}
+
+async fn maybe_wait_for_browser(state: &Arc<DaemonState>, wait_ms: Option<u64>) {
+    if let Some(wait) = clamp_browser_wait(wait_ms) {
+        state.browsers.wait_for_any_connected(wait).await;
+    }
+}
+
+fn clamp_browser_wait(wait_ms: Option<u64>) -> Option<Duration> {
+    let ms = wait_ms?;
+    if ms == 0 {
+        return None;
+    }
+    Some(Duration::from_millis(
+        ms.min(MAX_BROWSER_WAIT.as_millis() as u64),
+    ))
 }
 
 async fn handle_session_start(state: &Arc<DaemonState>, params: Value) -> Result<Value, RpcError> {
@@ -788,14 +821,16 @@ fn handle_session_list(state: &Arc<DaemonState>) -> ResponseBody {
     ResponseBody::Ok(serde_json::to_value(SessionListResult { sessions }).unwrap_or(Value::Null))
 }
 
-fn handle_browser_list(state: &Arc<DaemonState>) -> ResponseBody {
+async fn handle_browser_list(state: &Arc<DaemonState>, params: Value) -> Result<Value, RpcError> {
+    let params: BrowserListParams = parse_params_or_default(params)?;
+    maybe_wait_for_browser(state, params.wait_for_browser_ms).await;
     // Reuse the same helper that produces the
     // `multiple_browsers_online.error.data.browsers` payload so the
     // two surfaces always agree on order — sorted by `connected_at_ms`
     // ascending, with `instance_id` as a deterministic tiebreaker
     // (review I1).
     let browsers = snapshot_status_entries(&state.browsers, &state.sessions);
-    ResponseBody::Ok(serde_json::to_value(BrowserListResult { browsers }).unwrap_or(Value::Null))
+    Ok(serde_json::to_value(BrowserListResult { browsers }).unwrap_or(Value::Null))
 }
 
 // ----- Test-helper IpcServer wrapper around the transport layer -----