TUI: Align shell completions with GUI

Co-Authored-By: Oz <oz-agent@warp.dev>

Author: moirahuang

app/src/completer/completion_source.rs

@@ -0,0 +1,134 @@
+use std::collections::HashMap;
+use std::future::Future;
+
+use warp_completer::completer::{
+    self, CompleterOptions, CompletionContext, CompletionsFallbackStrategy, MatchStrategy,
+    SuggestionResults,
+};
+use warp_core::features::FeatureFlag;
+use warp_core::user_preferences::GetUserPreferences;
+use warpui::AppContext;
+
+use crate::terminal::model::completions::ShellCompletion;
+use crate::terminal::model::session::Session;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct CompletionSourcePolicy {
+    use_native_shell_completions: bool,
+    force_native_shell_completions: bool,
+}
+
+impl CompletionSourcePolicy {
+    pub fn for_session(session: &Session, buffer_text: &str, ctx: &AppContext) -> Self {
+        let force_native_shell_completions = ctx
+            .private_user_preferences()
+            .read_value("ForceNativeShellCompletions")
+            .ok()
+            .flatten()
+            .and_then(|value| value.parse().ok())
+            .unwrap_or(false);
+        Self::from_inputs(
+            FeatureFlag::NativeShellCompletions.is_enabled(),
+            force_native_shell_completions,
+            session.shell().supports_native_shell_completions(),
+            buffer_text.contains('\n'),
+        )
+    }
+
+    fn from_inputs(
+        native_shell_completions_enabled: bool,
+        force_native_shell_completions: bool,
+        shell_supports_native_completions: bool,
+        is_multiline: bool,
+    ) -> Self {
+        Self {
+            use_native_shell_completions: (native_shell_completions_enabled
+                || force_native_shell_completions)
+                && shell_supports_native_completions
+                && !is_multiline,
+            force_native_shell_completions,
+        }
+    }
+
+    pub fn should_request_native_shell_completions(self) -> bool {
+        self.use_native_shell_completions
+    }
+
+    pub fn fallback_strategy(
+        self,
+        fallback_when_native_is_unavailable: CompletionsFallbackStrategy,
+    ) -> CompletionsFallbackStrategy {
+        if self.use_native_shell_completions {
+            CompletionsFallbackStrategy::None
+        } else {
+            fallback_when_native_is_unavailable
+        }
+    }
+}
+
+pub async fn completion_suggestions_with_native_fallback<T, F>(
+    buffer_text: &str,
+    cursor_position: usize,
+    session_env_vars: Option<&HashMap<String, String>>,
+    mut options: CompleterOptions,
+    policy: CompletionSourcePolicy,
+    native_results: F,
+    ctx: &T,
+) -> Option<SuggestionResults>
+where
+    T: CompletionContext,
+    F: Future<Output = Option<Vec<ShellCompletion>>>,
+{
+    let before_cursor_text = buffer_text.get(..cursor_position)?;
+    options.fallback_strategy = policy.fallback_strategy(options.fallback_strategy);
+    let warp_results = completer::suggestions(
+        before_cursor_text,
+        before_cursor_text.len(),
+        session_env_vars,
+        options,
+        ctx,
+    )
+    .await;
+
+    resolve_completion_results(
+        warp_results,
+        native_results,
+        before_cursor_text,
+        policy.force_native_shell_completions,
+    )
+    .await
+}
+
+async fn resolve_completion_results<F>(
+    warp_results: Option<SuggestionResults>,
+    native_results: F,
+    before_cursor_text: &str,
+    force_native_shell_completions: bool,
+) -> Option<SuggestionResults>
+where
+    F: Future<Output = Option<Vec<ShellCompletion>>>,
+{
+    if let Some(warp_results) = warp_results
+        && !warp_results.suggestions.is_empty()
+        && !force_native_shell_completions
+    {
+        return Some(warp_results);
+    }
+
+    native_results.await.map(|results| {
+        let token_end = before_cursor_text.len();
+        let token_start = before_cursor_text
+            .rfind(char::is_whitespace)
+            .map(|position| position + 1)
+            .unwrap_or_default();
+        SuggestionResults {
+            replacement_span: (token_start, token_end).into(),
+            suggestions: results.into_iter().map(Into::into).collect(),
+            match_strategy: MatchStrategy::Fuzzy,
+        }
+    })
+}
+
+#[cfg(test)]
+#[path = "completion_source_tests.rs"]
+mod tests;

app/src/completer/completion_source_tests.rs

@@ -0,0 +1,93 @@
+use std::future;
+
+use warp_completer::completer::{
+    CompletionsFallbackStrategy, Match, MatchStrategy, MatchedSuggestion, Priority, Suggestion,
+    SuggestionResults, SuggestionType,
+};
+use warp_completer::meta::Span;
+use warpui::r#async::block_on;
+
+use super::{CompletionSourcePolicy, resolve_completion_results};
+use crate::terminal::model::completions::ShellCompletion;
+
+fn suggestion_results(name: &str) -> SuggestionResults {
+    SuggestionResults {
+        replacement_span: Span::new(0, 1),
+        suggestions: vec![MatchedSuggestion::new(
+            Suggestion::with_same_display_and_replacement(
+                name,
+                None,
+                SuggestionType::Argument,
+                Priority::default(),
+            ),
+            Match::Prefix {
+                is_case_sensitive: true,
+            },
+        )],
+        match_strategy: MatchStrategy::Fuzzy,
+    }
+}
+
+#[test]
+fn policy_requires_enablement_shell_support_and_single_line_input() {
+    let disabled = CompletionSourcePolicy::from_inputs(false, false, true, false);
+    assert!(!disabled.should_request_native_shell_completions());
+    assert!(matches!(
+        disabled.fallback_strategy(CompletionsFallbackStrategy::FilePaths),
+        CompletionsFallbackStrategy::FilePaths
+    ));
+
+    let forced = CompletionSourcePolicy::from_inputs(false, true, true, false);
+    assert!(forced.should_request_native_shell_completions());
+    assert!(matches!(
+        forced.fallback_strategy(CompletionsFallbackStrategy::FilePaths),
+        CompletionsFallbackStrategy::None
+    ));
+
+    let unsupported = CompletionSourcePolicy::from_inputs(true, false, false, false);
+    assert!(!unsupported.should_request_native_shell_completions());
+
+    let multiline = CompletionSourcePolicy::from_inputs(true, false, true, true);
+    assert!(!multiline.should_request_native_shell_completions());
+}
+
+#[test]
+fn nonempty_warp_results_win_without_force() {
+    let results = block_on(resolve_completion_results(
+        Some(suggestion_results("warp")),
+        future::ready(Some(vec![ShellCompletion::new("native".to_owned())])),
+        "w",
+        false,
+    ))
+    .expect("Warp results should be retained");
+
+    assert_eq!(results.suggestions[0].replacement(), "warp");
+}
+
+#[test]
+fn native_results_replace_empty_or_forced_warp_results() {
+    let empty_warp_results = SuggestionResults {
+        replacement_span: Span::new(0, 0),
+        suggestions: vec![],
+        match_strategy: MatchStrategy::Fuzzy,
+    };
+    let native_after_empty = block_on(resolve_completion_results(
+        Some(empty_warp_results),
+        future::ready(Some(vec![ShellCompletion::new("native".to_owned())])),
+        "command na",
+        false,
+    ))
+    .expect("native results should replace empty Warp results");
+    assert_eq!(native_after_empty.suggestions[0].replacement(), "native");
+    assert_eq!(native_after_empty.replacement_span, Span::new(8, 10));
+
+    let native_after_force = block_on(resolve_completion_results(
+        Some(suggestion_results("warp")),
+        future::ready(Some(vec![ShellCompletion::new("native".to_owned())])),
+        "λ native",
+        true,
+    ))
+    .expect("forced native results should replace Warp results");
+    assert_eq!(native_after_force.suggestions[0].replacement(), "native");
+    assert_eq!(native_after_force.replacement_span, Span::new(3, 9));
+}

app/src/completer/mod.rs

@@ -1,3 +1,4 @@
+mod completion_source;
 #[cfg(feature = "completions_v2")]
 mod js;
 
@@ -8,6 +9,7 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use async_trait::async_trait;
+pub use completion_source::{CompletionSourcePolicy, completion_suggestions_with_native_fallback};
 use lazy_static::lazy_static;
 use smol_str::SmolStr;
 use typed_path::{TypedPath, TypedPathBuf};

app/src/input_suggestions.rs

@@ -10,7 +10,7 @@ use itertools::Itertools;
 use pathfinder_geometry::vector::vec2f;
 use warp_command_signatures::IconType;
 use warp_completer::completer::{
-    MatchType, PathSeparators, Suggestion, SuggestionResults, SuggestionType,
+    MatchType, PathSeparators, PreparedSuggestion, Suggestion, SuggestionResults, SuggestionType,
 };
 use warp_core::features::FeatureFlag;
 use warp_core::ui::theme::AnsiColorIdentifier;
@@ -273,9 +273,15 @@ fn filter_tab_suggestions(
     suggestions: &SuggestionResults,
     query: &str,
     path_separators: &[char],
+) -> Vec<Item> {
+    items_from_prepared_suggestions(suggestions.prepare_for_query(query, path_separators))
+}
+
+fn items_from_prepared_suggestions(
+    suggestions: impl IntoIterator<Item = PreparedSuggestion>,
 ) -> Vec<Item> {
     suggestions
-        .filter_by_query(query, path_separators)
+        .into_iter()
         .map(|suggestion| Item {
             // TODO(vorporeal): Consider changing the type of `text` and `display` here to be `SmolStr`.
             text: suggestion.suggestion.replacement.to_string(),
@@ -286,7 +292,7 @@ fn filter_tab_suggestions(
                 .as_ref()
                 .map(|desc| desc.clone().into()),
             matches: Some(suggestion.matching_indices),
-            icon_type: Some(icon_type(suggestion.suggestion)),
+            icon_type: Some(icon_type(&suggestion.suggestion)),
             match_type: suggestion.match_type,
             is_ai_query: false,
             is_history_item: false,
@@ -339,6 +345,27 @@ impl InputSuggestions {
         ctx.notify();
     }
 
+    pub fn set_prepared_tab_completions(
+        &mut self,
+        suggestions: Vec<PreparedSuggestion>,
+        preselect_option: TabCompletionsPreselectOption,
+        ctx: &mut ViewContext<Self>,
+    ) {
+        self.set_items(items_from_prepared_suggestions(suggestions));
+        if self.items.is_empty() {
+            return;
+        }
+
+        match preselect_option {
+            TabCompletionsPreselectOption::First => self.select_first_item(ctx),
+            TabCompletionsPreselectOption::Unselected => self.selected_index = None,
+            TabCompletionsPreselectOption::Unchanged => {}
+        }
+
+        self.cycle = true;
+        ctx.notify();
+    }
+
     pub fn position_id_at_index(index: usize) -> String {
         format!("input_suggestions:index_{index}")
     }