fix(desktop): launch Databricks OAuth from passive model discovery

desktop/src-tauri/src/commands/agent_models_databricks.rs

@@ -1,7 +1,8 @@
 //! Databricks v1/v2 model discovery and interactive reauthentication.
 
-use std::collections::BTreeMap;
-use std::sync::LazyLock;
+use std::collections::{BTreeMap, HashMap};
+use std::sync::{LazyLock, Mutex, MutexGuard};
+use std::time::{Duration, Instant};
 
 use crate::commands::agent_models_env::{
     env_or_process_value, redaction_env_with_value, DiscoveryProvider,
@@ -13,6 +14,77 @@ use crate::managed_agents::AgentModelsResponse;
 // callback listener/browser flow for the process-wide OAuth cache.
 static AUTH_GATE: LazyLock<tokio::sync::Mutex<()>> = LazyLock::new(|| tokio::sync::Mutex::new(()));
 
+// Hard cap on the interactive browser flow launched from a discovery surface.
+// An abandoned SSO tab must fail discovery cleanly rather than wedge the
+// dropdown forever. (`authenticate_databricks` has its own 60s callback wait;
+// this outer bound also covers endpoint discovery and token exchange.)
+const AUTH_FLOW_TIMEOUT: Duration = Duration::from_secs(150);
+
+// How long a failed/cancelled interactive sign-in suppresses re-launching the
+// browser from passive surfaces.
+pub(super) const AUTH_COOLDOWN: Duration = Duration::from_secs(5 * 60);
+
+/// Per-host record of a recently failed, cancelled, or timed-out interactive
+/// sign-in.
+///
+/// Passive discovery surfaces fire on every form-state change, so without this
+/// a cancelled SSO page would re-pop the browser on the very next keystroke.
+/// Entries expire so a genuine later retry still launches; the saved-model
+/// picker bypasses the cooldown and a success clears it.
+#[derive(Default)]
+pub(super) struct AuthCooldown {
+    until: Mutex<HashMap<String, Instant>>,
+}
+
+impl AuthCooldown {
+    fn map(&self) -> MutexGuard<'_, HashMap<String, Instant>> {
+        // The critical sections below are panic-free map ops, so recover from a
+        // poisoned lock rather than wedge every future sign-in on one panic.
+        self.until
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    pub(super) fn is_active(&self, host: &str, now: Instant) -> bool {
+        let mut map = self.map();
+        match map.get(host) {
+            Some(&expiry) if now < expiry => true,
+            Some(_) => {
+                map.remove(host);
+                false
+            }
+            None => false,
+        }
+    }
+
+    pub(super) fn record(&self, host: &str, now: Instant) {
+        self.map().insert(host.to_string(), now + AUTH_COOLDOWN);
+    }
+
+    pub(super) fn clear(&self, host: &str) {
+        self.map().remove(host);
+    }
+
+    /// Whether the interactive browser flow may launch now under `auth_intent`.
+    /// Passive surfaces are suppressed while a per-host cooldown is active; the
+    /// explicit picker path always launches and clears any stale suppression.
+    pub(super) fn permits_launch(
+        &self,
+        auth_intent: DatabricksAuthIntent,
+        host: &str,
+        now: Instant,
+    ) -> bool {
+        if auth_intent.respects_cooldown() {
+            !self.is_active(host, now)
+        } else {
+            self.clear(host);
+            true
+        }
+    }
+}
+
+static AUTH_COOLDOWNS: LazyLock<AuthCooldown> = LazyLock::new(AuthCooldown::default);
+
 pub(super) fn is_databricks_provider(provider: Option<&str>) -> bool {
     matches!(
         provider
@@ -50,8 +122,14 @@ pub(super) enum DatabricksAuthIntent {
 }
 
 impl DatabricksAuthIntent {
-    fn allows_interactive_auth(self) -> bool {
-        matches!(self, Self::InteractiveModelPicker)
+    /// Passive draft discovery honors (and, on failure, writes) the per-host
+    /// cooldown so a cancelled SSO page does not re-pop on the next form
+    /// keystroke. The saved-model picker is an explicit user action, so it
+    /// bypasses the cooldown and clears it before launching. Both surfaces
+    /// launch the browser flow (Phase 2 goose-parity); this predicate is the
+    /// only behavioral difference between them.
+    fn respects_cooldown(self) -> bool {
+        matches!(self, Self::PassiveDraftDiscovery)
     }
 }
 
@@ -60,11 +138,16 @@ pub(super) fn databricks_sign_in_required_error() -> String {
         .to_string()
 }
 
-pub(super) fn should_start_interactive_auth(
-    api_key: &str,
-    auth_intent: DatabricksAuthIntent,
-) -> bool {
-    api_key.is_empty() && auth_intent.allows_interactive_auth()
+pub(super) fn databricks_sign_in_timed_out_error() -> String {
+    "Databricks sign-in timed out; open the model picker to retry, or run `buzz-agent auth databricks`"
+        .to_string()
+}
+
+pub(super) fn should_start_interactive_auth(api_key: &str) -> bool {
+    // Phase 2: both discovery surfaces launch the browser flow when no static
+    // token is configured. Which surface is allowed to actually pop the browser
+    // (vs. respect a cooldown) is decided via `AuthCooldown::permits_launch`.
+    api_key.is_empty()
 }
 
 pub(super) async fn discover_databricks_models(
@@ -93,22 +176,26 @@ pub(super) async fn discover_databricks_models(
 
     let entries = match buzz_agent_pkg::discover_databricks_models(&config).await {
         Ok(entries) => entries,
-        Err(buzz_agent_pkg::AgentError::LlmAuth(_))
-            if should_start_interactive_auth(&api_key, auth_intent) =>
-        {
+        Err(buzz_agent_pkg::AgentError::LlmAuth(_)) if should_start_interactive_auth(&api_key) => {
             let _auth = AUTH_GATE.lock().await;
             match buzz_agent_pkg::discover_databricks_models(&config).await {
+                // A peer sign-in under the gate already succeeded.
                 Ok(entries) => entries,
                 Err(buzz_agent_pkg::AgentError::LlmAuth(_)) => {
-                    buzz_agent_pkg::authenticate_databricks(&host)
-                        .await
-                        .map_err(|error| {
-                            format_redacted_error(
-                                "Databricks sign-in failed",
-                                &error,
-                                &redaction_env,
-                            )
-                        })?;
+                    // Passive surfaces suppress the browser while a recent
+                    // failure/cancel is cooling down; the explicit picker path
+                    // always launches (and clears any stale cooldown).
+                    if !AUTH_COOLDOWNS.permits_launch(auth_intent, &host, Instant::now()) {
+                        return Err(databricks_sign_in_required_error());
+                    }
+                    run_interactive_databricks_auth(
+                        buzz_agent_pkg::authenticate_databricks(&host),
+                        AUTH_FLOW_TIMEOUT,
+                        &AUTH_COOLDOWNS,
+                        &host,
+                        &redaction_env,
+                    )
+                    .await?;
                     buzz_agent_pkg::discover_databricks_models(&config)
                         .await
                         .map_err(|error| {
@@ -172,3 +259,43 @@ fn format_redacted_error(
     let message = crate::managed_agents::redact_env_values_in(&error.to_string(), redaction_env);
     format!("{context}: {message}")
 }
+
+/// Run the interactive browser OAuth flow under a hard timeout and maintain the
+/// per-host cooldown. Success clears the cooldown; a failure, cancel, or
+/// timeout records it so passive surfaces stop re-launching the browser on the
+/// next form keystroke. `timeout` is injected (production passes
+/// [`AUTH_FLOW_TIMEOUT`]) so the timeout/cooldown policy is unit-testable
+/// without a live browser.
+pub(super) async fn run_interactive_databricks_auth<Fut>(
+    auth: Fut,
+    timeout: Duration,
+    cooldowns: &AuthCooldown,
+    host: &str,
+    redaction_env: &BTreeMap<String, String>,
+) -> Result<(), String>
+where
+    Fut: std::future::Future<Output = Result<(), buzz_agent_pkg::AgentError>>,
+{
+    match tokio::time::timeout(timeout, auth).await {
+        Ok(Ok(())) => {
+            cooldowns.clear(host);
+            Ok(())
+        }
+        Ok(Err(error)) => {
+            cooldowns.record(host, Instant::now());
+            Err(format_redacted_error(
+                "Databricks sign-in failed",
+                &error,
+                redaction_env,
+            ))
+        }
+        Err(_elapsed) => {
+            cooldowns.record(host, Instant::now());
+            Err(databricks_sign_in_timed_out_error())
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "agent_models_databricks_tests.rs"]
+mod tests;

desktop/src-tauri/src/commands/agent_models_databricks_tests.rs

@@ -0,0 +1,109 @@
+//! Cooldown and interactive-auth policy tests for Databricks discovery.
+//!
+//! Housed as a child of `agent_models_databricks` (not the shared
+//! `agent_models_tests`) so the async timeout/cooldown cases sit next to the
+//! code they exercise and reach its `pub(super)` items directly via
+//! `use super::*` — and so the shared test file stays under its size ratchet.
+
+use super::*;
+
+#[test]
+fn databricks_cooldown_suppresses_passive_relaunch_but_never_the_picker() {
+    let cooldowns = AuthCooldown::default();
+    let host = "https://example.cloud.databricks.com";
+    let now = Instant::now();
+
+    // A fresh host permits either surface to launch.
+    assert!(cooldowns.permits_launch(DatabricksAuthIntent::PassiveDraftDiscovery, host, now));
+    assert!(cooldowns.permits_launch(DatabricksAuthIntent::InteractiveModelPicker, host, now));
+
+    // After a failed/cancelled attempt, passive discovery must NOT re-pop the
+    // browser while the window is active...
+    cooldowns.record(host, now);
+    assert!(!cooldowns.permits_launch(DatabricksAuthIntent::PassiveDraftDiscovery, host, now));
+
+    // ...but an explicit picker click always launches, and clears the window so
+    // a later passive read is unblocked too.
+    assert!(cooldowns.permits_launch(DatabricksAuthIntent::InteractiveModelPicker, host, now));
+    assert!(cooldowns.permits_launch(DatabricksAuthIntent::PassiveDraftDiscovery, host, now));
+}
+
+#[test]
+fn databricks_cooldown_expires_after_its_window_and_is_host_scoped() {
+    let cooldowns = AuthCooldown::default();
+    let host = "https://a.cloud.databricks.com";
+    let other = "https://b.cloud.databricks.com";
+    let now = Instant::now();
+
+    cooldowns.record(host, now);
+    // A cooldown on one host never suppresses another.
+    assert!(!cooldowns.is_active(other, now));
+    assert!(cooldowns.is_active(host, now));
+
+    // The window is closed the instant it elapses, so a genuine later retry
+    // launches again.
+    let after = now + AUTH_COOLDOWN;
+    assert!(!cooldowns.is_active(host, after));
+}
+
+#[tokio::test]
+async fn databricks_interactive_auth_success_clears_a_prior_cooldown() {
+    let cooldowns = AuthCooldown::default();
+    let host = "https://example.cloud.databricks.com";
+    let redaction = BTreeMap::new();
+    cooldowns.record(host, Instant::now());
+
+    let result = run_interactive_databricks_auth(
+        async { Ok(()) },
+        Duration::from_secs(150),
+        &cooldowns,
+        host,
+        &redaction,
+    )
+    .await;
+
+    assert!(result.is_ok());
+    assert!(!cooldowns.is_active(host, Instant::now()));
+}
+
+#[tokio::test]
+async fn databricks_interactive_auth_failure_records_a_cooldown() {
+    let cooldowns = AuthCooldown::default();
+    let host = "https://example.cloud.databricks.com";
+    let redaction = BTreeMap::new();
+
+    let result = run_interactive_databricks_auth(
+        async { Err(buzz_agent_pkg::AgentError::LlmAuth("closed the tab".into())) },
+        Duration::from_secs(150),
+        &cooldowns,
+        host,
+        &redaction,
+    )
+    .await;
+
+    let error = result.expect_err("a failed sign-in must surface an error");
+    assert!(error.contains("Databricks sign-in failed"));
+    assert!(cooldowns.is_active(host, Instant::now()));
+}
+
+#[tokio::test(start_paused = true)]
+async fn databricks_interactive_auth_timeout_records_cooldown_and_returns_timeout_copy() {
+    let cooldowns = AuthCooldown::default();
+    let host = "https://example.cloud.databricks.com";
+    let redaction = BTreeMap::new();
+
+    // An abandoned SSO tab: the flow never resolves. Under the paused clock the
+    // injected timeout fires deterministically without real waiting.
+    let result = run_interactive_databricks_auth(
+        std::future::pending::<Result<(), buzz_agent_pkg::AgentError>>(),
+        Duration::from_secs(150),
+        &cooldowns,
+        host,
+        &redaction,
+    )
+    .await;
+
+    let error = result.expect_err("a timed-out sign-in must surface an error");
+    assert_eq!(error, databricks_sign_in_timed_out_error());
+    assert!(cooldowns.is_active(host, Instant::now()));
+}

desktop/src-tauri/src/commands/agent_models_tests.rs

@@ -577,19 +577,12 @@ fn is_databricks_provider_matches_both_variants() {
 }
 
 #[test]
-fn databricks_interactive_auth_requires_explicit_intent_and_no_static_token() {
-    assert!(should_start_interactive_auth(
-        "",
-        DatabricksAuthIntent::InteractiveModelPicker
-    ));
-    assert!(!should_start_interactive_auth(
-        "",
-        DatabricksAuthIntent::PassiveDraftDiscovery
-    ));
-    assert!(!should_start_interactive_auth(
-        "static-token",
-        DatabricksAuthIntent::InteractiveModelPicker
-    ));
+fn databricks_interactive_auth_launches_only_without_a_static_token() {
+    // Phase 2: both surfaces launch the browser flow when the token is empty;
+    // the surface distinction is now cooldown-only (asserted separately). A
+    // configured static token still short-circuits interactive auth entirely.
+    assert!(should_start_interactive_auth(""));
+    assert!(!should_start_interactive_auth("static-token"));
 }
 
 #[test]