Add a custom host picker for orchestration (#11080)

Author: advait-m

app/src/ai/blocklist/inline_action/host_picker.rs

@@ -0,0 +1,118 @@
+//! Unit tests for the pure helpers in `host_picker.rs`. View-driven
+//! behaviors (custom-mode commit, blur, etc.) are covered by manual smoke
+//! testing.
+
+use super::{
+    build_menu_items, menu_label_for, normalize_slug, DropdownAction, InternalAction, MenuItem,
+    ORCHESTRATION_WARP_WORKER_HOST,
+};
+
+/// Extracts the visible label text out of a `MenuItem::Item`, panicking
+/// on the unreachable `Header` / `Separator` cases that our builder
+/// doesn't emit.
+fn item_label(item: &MenuItem<DropdownAction<InternalAction>>) -> &str {
+    match item {
+        MenuItem::Item(fields) => fields.label(),
+        other => panic!("expected MenuItem::Item, got {other:?}"),
+    }
+}
+
+/// Extracts the on-select action from a `MenuItem::Item`.
+fn item_action(item: &MenuItem<DropdownAction<InternalAction>>) -> &DropdownAction<InternalAction> {
+    match item {
+        MenuItem::Item(fields) => fields
+            .on_select_action()
+            .expect("test items always have a select action"),
+        other => panic!("expected MenuItem::Item, got {other:?}"),
+    }
+}
+
+#[test]
+fn build_menu_items_with_no_defaults_shows_warp_and_custom() {
+    let items = build_menu_items(None, None);
+    assert_eq!(items.len(), 2, "expected warp + custom-host entries");
+    assert_eq!(item_label(&items[0]), ORCHESTRATION_WARP_WORKER_HOST);
+    assert_eq!(item_label(&items[1]), "Custom host\u{2026}");
+}
+
+#[test]
+fn build_menu_items_promotes_default_to_top() {
+    // Workspace default sits above warp and gets the "Default" badge,
+    // matching the Oz webapp's HostSelector layout.
+    let items = build_menu_items(Some("my-corp"), None);
+    assert_eq!(items.len(), 3);
+    assert_eq!(item_label(&items[0]), "my-corp  (Default)");
+    assert_eq!(item_label(&items[1]), ORCHESTRATION_WARP_WORKER_HOST);
+    assert_eq!(item_label(&items[2]), "Custom host\u{2026}");
+}
+
+#[test]
+fn build_menu_items_adds_recent_after_warp() {
+    let items = build_menu_items(None, Some("other-host"));
+    assert_eq!(items.len(), 3);
+    assert_eq!(item_label(&items[0]), ORCHESTRATION_WARP_WORKER_HOST);
+    // Recent hosts render as plain slugs (no "(Recent)" suffix).
+    assert_eq!(item_label(&items[1]), "other-host");
+    assert_eq!(item_label(&items[2]), "Custom host\u{2026}");
+}
+
+#[test]
+fn build_menu_items_dedups_recent_when_it_matches_default_or_warp() {
+    // Same as the workspace default → no duplicate "Recent" row.
+    let items = build_menu_items(Some("my-corp"), Some("my-corp"));
+    assert_eq!(items.len(), 3);
+    assert_eq!(item_label(&items[0]), "my-corp  (Default)");
+    assert_eq!(item_label(&items[1]), ORCHESTRATION_WARP_WORKER_HOST);
+    assert_eq!(item_label(&items[2]), "Custom host\u{2026}");
+
+    // Recent == "warp" is also skipped (warp is already a row).
+    let items = build_menu_items(Some("my-corp"), Some("warp"));
+    assert_eq!(items.len(), 3, "warp recent should not double-add");
+}
+
+#[test]
+fn build_menu_items_warp_entry_dispatches_select_known_warp() {
+    let items = build_menu_items(None, None);
+    match item_action(&items[0]) {
+        DropdownAction::SelectActionAndClose(InternalAction::SelectKnown(slug)) => {
+            assert_eq!(slug, ORCHESTRATION_WARP_WORKER_HOST);
+        }
+        other => panic!("expected SelectActionAndClose(SelectKnown), got {other:?}"),
+    }
+}
+
+#[test]
+fn build_menu_items_custom_entry_dispatches_enter_custom_mode() {
+    let items = build_menu_items(None, None);
+    let custom = items.last().expect("custom entry is always last");
+    match item_action(custom) {
+        DropdownAction::SelectActionAndClose(InternalAction::EnterCustomMode) => {}
+        other => panic!("expected EnterCustomMode, got {other:?}"),
+    }
+}
+
+#[test]
+fn menu_label_for_picks_default_badge_when_slug_matches_default() {
+    let label = menu_label_for("my-corp", Some("my-corp"));
+    assert_eq!(label, "my-corp  (Default)");
+}
+
+#[test]
+fn menu_label_for_returns_plain_slug_for_warp() {
+    let label = menu_label_for(ORCHESTRATION_WARP_WORKER_HOST, Some("my-corp"));
+    assert_eq!(label, ORCHESTRATION_WARP_WORKER_HOST);
+}
+
+#[test]
+fn menu_label_for_returns_plain_slug_for_unknown_value() {
+    // A slug typed via custom mode that we haven't promoted to "recent" yet.
+    let label = menu_label_for("typed-once", Some("my-corp"));
+    assert_eq!(label, "typed-once");
+}
+
+#[test]
+fn normalize_slug_trims_whitespace_and_falls_back_to_warp_when_empty() {
+    assert_eq!(normalize_slug("  my-corp  "), "my-corp");
+    assert_eq!(normalize_slug(""), ORCHESTRATION_WARP_WORKER_HOST);
+    assert_eq!(normalize_slug("   "), ORCHESTRATION_WARP_WORKER_HOST);
+}

app/src/ai/blocklist/inline_action/host_picker_tests.rs

@@ -2,6 +2,7 @@ pub(crate) mod ask_user_question_view;
 pub(super) mod aws_bedrock_credentials_error;
 pub(crate) mod code_diff_view;
 pub(crate) mod create_or_edit_document;
+pub(crate) mod host_picker;
 pub(crate) mod inline_action_header;
 pub(crate) mod inline_action_icons;
 mod malformed_line_heuristics;

app/src/ai/blocklist/inline_action/mod.rs

@@ -28,11 +28,11 @@ use warpui::{
 
 use settings::Setting;
 use warp_cli::agent::Harness;
-use warp_core::channel::{Channel, ChannelState};
 use warp_core::features::FeatureFlag;
 use warp_core::ui::theme::Fill;
 
 use crate::ai::auth_secret_types::auth_secret_types_for_harness;
+use crate::ai::blocklist::inline_action::host_picker::HostPicker;
 use crate::ai::cloud_agent_settings::CloudAgentSettings;
 use crate::ai::cloud_environments::CloudAmbientAgentEnvironment;
 use crate::ai::execution_profiles::model_menu_items::available_model_menu_items;
@@ -47,8 +47,13 @@ use crate::report_if_error;
 use crate::ui_components::blended_colors;
 use crate::view_components::dropdown::{Dropdown, DropdownAction, DropdownStyle};
 use crate::view_components::FilterableDropdown;
+use crate::workspaces::user_workspaces::UserWorkspaces;
 use crate::LLMPreferences;
 
+/// Env var override for the workspace default host (developer testing).
+/// Mirrors the single-agent ambient flow.
+const DEFAULT_HOST_ENV_VAR: &str = "WARP_CLOUD_MODE_DEFAULT_HOST";
+
 // ── Shared constants ────────────────────────────────────────────────
 
 pub const ORCHESTRATION_WARP_WORKER_HOST: &str = "warp";
@@ -78,7 +83,6 @@ pub trait OrchestrationControlAction: Clone + Debug + Send + Sync + 'static {
     fn model_changed(model_id: String) -> Self;
     fn harness_changed(harness_type: String) -> Self;
     fn environment_changed(environment_id: String) -> Self;
-    fn worker_host_changed(worker_host: String) -> Self;
     /// Fires when the auth secret picker selects a managed secret.
     /// `None` means "clear the selection / inherit from environment".
     fn auth_secret_changed(name: Option<String>) -> Self;
@@ -291,7 +295,7 @@ pub struct OrchestrationPickerHandles<A: OrchestrationControlAction> {
     pub model_picker: Option<ViewHandle<Dropdown<A>>>,
     pub harness_picker: Option<ViewHandle<Dropdown<A>>>,
     pub environment_picker: Option<ViewHandle<FilterableDropdown<A>>>,
-    pub host_picker: Option<ViewHandle<Dropdown<A>>>,
+    pub host_picker: Option<ViewHandle<HostPicker>>,
     /// Picker for the managed auth secret used by non-Oz cloud children.
     /// `None` when the picker hasn't been built yet (e.g. harness is Oz or
     /// execution mode is Local), or when the harness has no supported
@@ -688,37 +692,72 @@ pub fn create_environment_picker<A: OrchestrationControlAction, V: View>(
     dropdown_handle
 }
 
-pub fn populate_host_picker<A: OrchestrationControlAction, V: View>(
-    dropdown: &ViewHandle<Dropdown<A>>,
+/// Repopulates the host picker with the workspace default (if any) and
+/// the user's last-selected custom host (if any), then sets the current
+/// selection to `initial_host`.
+pub fn populate_host_picker<V: View>(
+    picker: &ViewHandle<HostPicker>,
     initial_host: &str,
     ctx: &mut ViewContext<V>,
 ) {
-    let initial_host = if initial_host.is_empty() {
+    let default_host = resolve_default_host_slug(ctx);
+    let recent_host = resolve_recent_host_slug(ctx);
+    let initial = if initial_host.trim().is_empty() {
         ORCHESTRATION_WARP_WORKER_HOST.to_string()
     } else {
         initial_host.to_string()
     };
-    dropdown.update(ctx, |dropdown, ctx_dropdown| {
-        let hosts: &[&str] = if matches!(ChannelState::channel(), Channel::Local) {
-            &["warp", "local-dev"]
-        } else {
-            &["warp"]
-        };
-        let mut items: Vec<MenuItem<DropdownAction<A>>> = Vec::new();
-        let mut selected_idx = None;
-        for (idx, &host) in hosts.iter().enumerate() {
-            let fields = MenuItemFields::new(host).with_on_select_action(
-                DropdownAction::SelectActionAndClose(A::worker_host_changed(host.to_string())),
-            );
-            if host.eq_ignore_ascii_case(&initial_host) {
-                selected_idx = Some(idx);
-            }
-            items.push(MenuItem::Item(fields));
-        }
-        dropdown.set_rich_items(items, ctx_dropdown);
-        if let Some(idx) = selected_idx {
-            dropdown.set_selected_by_index(idx, ctx_dropdown);
+    picker.update(ctx, |picker, picker_ctx| {
+        picker.set_options(default_host, recent_host, picker_ctx);
+        picker.set_selected(&initial, picker_ctx);
+    });
+}
+
+/// Resolves the workspace-configured default host slug, honoring the
+/// `WARP_CLOUD_MODE_DEFAULT_HOST` env var override for developer
+/// testing. Mirrors the single-agent ambient flow.
+pub fn resolve_default_host_slug(ctx: &AppContext) -> Option<String> {
+    if let Ok(slug) = std::env::var(DEFAULT_HOST_ENV_VAR) {
+        let trimmed = slug.trim();
+        if !trimmed.is_empty() {
+            return Some(trimmed.to_string());
         }
+    }
+    UserWorkspaces::as_ref(ctx)
+        .default_host_slug()
+        .map(str::to_string)
+        .filter(|s| !s.trim().is_empty())
+}
+
+/// Returns the user's last-selected custom host slug from
+/// `CloudAgentSettings.last_selected_host`, excluding `"warp"` and the
+/// workspace default (those are surfaced as separate menu rows).
+pub fn resolve_recent_host_slug(ctx: &AppContext) -> Option<String> {
+    let last = CloudAgentSettings::as_ref(ctx)
+        .last_selected_host
+        .value()
+        .clone()
+        .filter(|s| !s.trim().is_empty())?;
+    if last.eq_ignore_ascii_case(ORCHESTRATION_WARP_WORKER_HOST) {
+        return None;
+    }
+    if resolve_default_host_slug(ctx).as_deref() == Some(last.as_str()) {
+        return None;
+    }
+    Some(last)
+}
+
+/// Persists the user's most-recent host selection to
+/// `CloudAgentSettings.last_selected_host`. Skipped for `"warp"` and
+/// empty values (those don't represent a custom slug worth remembering).
+pub fn persist_host_selection<V: View>(worker_host: &str, ctx: &mut ViewContext<V>) {
+    let trimmed = worker_host.trim();
+    if trimmed.is_empty() || trimmed.eq_ignore_ascii_case(ORCHESTRATION_WARP_WORKER_HOST) {
+        return;
+    }
+    let value = trimmed.to_string();
+    CloudAgentSettings::handle(ctx).update(ctx, |settings, ctx| {
+        report_if_error!(settings.last_selected_host.set_value(Some(value), ctx));
     });
 }
 
@@ -1188,8 +1227,8 @@ pub fn sync_picker_selections<A: OrchestrationControlAction, V: View>(
             RunAgentsExecutionMode::Remote { worker_host, .. } => worker_host.clone(),
             RunAgentsExecutionMode::Local => ORCHESTRATION_WARP_WORKER_HOST.to_string(),
         };
-        host_picker.update(ctx, |dropdown, ctx_dropdown| {
-            dropdown.set_selected_by_name(&worker_host, ctx_dropdown);
+        host_picker.update(ctx, |picker, picker_ctx| {
+            picker.set_selected(&worker_host, picker_ctx);
         });
     }
     if let Some(auth_secret_picker) = handles.auth_secret_picker.clone() {

app/src/ai/blocklist/inline_action/orchestration_controls.rs

@@ -34,6 +34,7 @@ use crate::ai::blocklist::agent_view::orchestration_pill_bar::render_static_agen
 use crate::ai::blocklist::block::model::AIBlockModel;
 use crate::ai::blocklist::block::view_impl::WithContentItemSpacing;
 use crate::ai::blocklist::block::AIBlock;
+use crate::ai::blocklist::inline_action::host_picker::{HostPicker, HostPickerEvent};
 use crate::ai::blocklist::inline_action::inline_action_header::{HeaderConfig, InteractionMode};
 use crate::ai::blocklist::inline_action::inline_action_icons;
 use crate::ai::blocklist::inline_action::orchestration_controls::{
@@ -146,9 +147,6 @@ impl OrchestrationControlAction for RunAgentsCardViewAction {
     fn environment_changed(environment_id: String) -> Self {
         Self::EnvironmentChanged { environment_id }
     }
-    fn worker_host_changed(worker_host: String) -> Self {
-        Self::WorkerHostChanged { worker_host }
-    }
     fn auth_secret_changed(auth_secret_name: Option<String>) -> Self {
         Self::AuthSecretChanged { auth_secret_name }
     }
@@ -254,9 +252,13 @@ fn resolve_interactive_defaults(
         let needs_host = worker_host.is_empty();
         let needs_env = environment_id.is_empty();
         if needs_host {
-            state
-                .orch
-                .set_worker_host(oc::ORCHESTRATION_WARP_WORKER_HOST.to_string());
+            // Prefer the workspace default (or the dev env-var override)
+            // over the bare "warp" fallback so self-hosted teams see
+            // their default pre-selected. Mirrors the Oz webapp's
+            // `HostSelector` initial-selection behavior.
+            let default_host = oc::resolve_default_host_slug(ctx)
+                .unwrap_or_else(|| oc::ORCHESTRATION_WARP_WORKER_HOST.to_string());
+            state.orch.set_worker_host(default_host);
         }
         if needs_env {
             if let Some(default_env) = oc::resolve_default_environment_id(ctx) {
@@ -668,10 +670,27 @@ impl RunAgentsCardView {
                 RunAgentsExecutionMode::Remote { worker_host, .. } => worker_host.as_str(),
                 RunAgentsExecutionMode::Local => oc::ORCHESTRATION_WARP_WORKER_HOST,
             };
-            let handle = oc::new_standard_picker_dropdown(&colors, ctx);
-            Self::set_upward_menu_position(&handle, ctx);
+            let handle = ctx.add_typed_action_view(HostPicker::new);
+            // Open upward so the menu doesn't overlap pickers below it,
+            // matching the other dropdowns in this card.
+            handle.update(ctx, |picker, picker_ctx| {
+                picker.set_menu_position(
+                    warpui::elements::PositionedElementAnchor::TopLeft,
+                    warpui::elements::ChildAnchor::BottomLeft,
+                    picker_ctx,
+                );
+            });
             oc::populate_host_picker(&handle, initial_host, ctx);
-            Self::subscribe_picker_close(&handle, ctx);
+            ctx.subscribe_to_view(&handle, |me, _, event, ctx| match event {
+                HostPickerEvent::HostChanged { slug } => {
+                    ctx.dispatch_typed_action(&RunAgentsCardViewAction::WorkerHostChanged {
+                        worker_host: slug.clone(),
+                    });
+                }
+                HostPickerEvent::Closed => {
+                    me.refocus_after_picker_close(ctx);
+                }
+            });
             self.handles.pickers.host_picker = Some(handle);
         }
 
@@ -920,6 +939,7 @@ impl TypedActionView for RunAgentsCardView {
             }
             RunAgentsCardViewAction::WorkerHostChanged { worker_host } => {
                 self.state.orch.set_worker_host(worker_host.clone());
+                oc::persist_host_selection(worker_host, ctx);
                 ctx.notify();
             }
             RunAgentsCardViewAction::AuthSecretChanged { auth_secret_name } => {