Autosuggest env vars as first words

src/completions/context.rs

@@ -2015,4 +2015,79 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn test_quoted_env_var_command_completion() {
+        let ctx_unquoted_space = run_inline("FOO=1 █");
+        assert_eq!(ctx_unquoted_space.word_under_cursor.as_ref(), "");
+        assert_eq!(
+            ctx_unquoted_space.comp_types(),
+            vec![
+                CompType::FirstWord,
+                CompType::FuzzyFirstWord,
+                CompType::FilenameExpansion,
+                CompType::FuzzyFilenameExpansion,
+            ]
+        );
+
+        let ctx_quoted_space = run_inline(r#"FOO="1" █"#);
+        assert_eq!(ctx_quoted_space.word_under_cursor.as_ref(), "");
+        assert_eq!(
+            ctx_quoted_space.comp_types(),
+            vec![
+                CompType::FirstWord,
+                CompType::FuzzyFirstWord,
+                CompType::FilenameExpansion,
+                CompType::FuzzyFilenameExpansion,
+            ]
+        );
+
+        let ctx_quoted = run_inline(r#"FOO="1" gre█"#);
+        assert_eq!(ctx_quoted.word_under_cursor.as_ref(), "gre");
+        assert_eq!(
+            ctx_quoted.comp_types(),
+            vec![
+                CompType::FirstWord,
+                CompType::FuzzyFirstWord,
+                CompType::FilenameExpansion,
+                CompType::FuzzyFilenameExpansion,
+            ]
+        );
+
+        let ctx_subshell = run_inline(r#"echo $( FOO="1" gr█"#);
+        assert_eq!(ctx_subshell.word_under_cursor.as_ref(), "gr");
+        assert_eq!(
+            ctx_subshell.comp_types(),
+            vec![
+                CompType::FirstWord,
+                CompType::FuzzyFirstWord,
+                CompType::FilenameExpansion,
+                CompType::FuzzyFilenameExpansion,
+            ]
+        );
+
+        let ctx_subshell_unquoted = run_inline(r#"echo $( FOO=1 gr█"#);
+        assert_eq!(ctx_subshell_unquoted.word_under_cursor.as_ref(), "gr");
+        assert_eq!(
+            ctx_subshell_unquoted.comp_types(),
+            vec![
+                CompType::FirstWord,
+                CompType::FuzzyFirstWord,
+                CompType::FilenameExpansion,
+                CompType::FuzzyFilenameExpansion,
+            ]
+        );
+
+        let ctx_subshell_space = run_inline(r#"echo $( FOO="1" █"#);
+        assert_eq!(ctx_subshell_space.word_under_cursor.as_ref(), "");
+        assert_eq!(
+            ctx_subshell_space.comp_types(),
+            vec![
+                CompType::FirstWord,
+                CompType::FuzzyFirstWord,
+                CompType::FilenameExpansion,
+                CompType::FuzzyFilenameExpansion,
+            ]
+        );
+    }
 }

src/completions/tab_completion.rs

@@ -608,10 +608,11 @@ fn tab_complete_first_word(command: &str, word_under_cursor: &str) -> ActiveSugg
     }
 
     let mut res = vec![];
-    let mut seen: HashSet<String> = HashSet::new();
+    let mut seen: HashSet<(String, bool)> = HashSet::new();
     for poss_info in shell::backend().possible_command_words() {
         let cmd_name = poss_info.command();
-        if cmd_name.starts_with(command) && seen.insert(cmd_name.to_string()) {
+        let is_env_var = matches!(poss_info, shell::CommandWordInfo::EnvVar { .. });
+        if cmd_name.starts_with(command) && seen.insert((cmd_name.to_string(), is_env_var)) {
             res.push(poss_info);
         }
     }
@@ -635,7 +636,9 @@ fn processed_suggestions_from_command_info(
         .into_iter()
         .map(|info| {
             let s = info.command().to_string();
-            let new_suffix = if s.ends_with(' ') {
+            let new_suffix = if matches!(info, shell::CommandWordInfo::EnvVar { .. }) {
+                "=".to_string()
+            } else if s.ends_with(' ') {
                 "".to_string()
             } else {
                 " ".to_string()
@@ -666,10 +669,11 @@ fn tab_complete_fuzzy_first_word(command: &str) -> ActiveSuggestionsBuilder {
     let matcher = ArinaeMatcher::new(skim::CaseMatching::Smart, true);
     let mut scored = vec![];
 
-    let mut seen: HashSet<String> = HashSet::new();
+    let mut seen: HashSet<(String, bool)> = HashSet::new();
     for poss_info in shell::backend().possible_command_words() {
         let cmd_name = poss_info.command();
-        if seen.insert(cmd_name.to_string())
+        let is_env_var = matches!(poss_info, shell::CommandWordInfo::EnvVar { .. });
+        if seen.insert((cmd_name.to_string(), is_env_var))
             && let Some(score) =
                 fuzzy_match_with_threshold(&matcher, cmd_name, command, FuzzyMatchThreshold::High)
         {
@@ -2193,5 +2197,52 @@ mod tab_completion_tests {
             let item = builder.processed.first().unwrap();
             assert_eq!(item.suffix, " ");
         }
+
+        #[test]
+        fn test_first_word_env_var_completion_has_equal_suffix() {
+            crate::shell::backend()
+                .export_env_var("MY_CUSTOM_ENV_VAR", "value123")
+                .unwrap();
+            let builder = tab_complete_first_word("MY_CUSTOM", "MY_CUSTOM");
+            let item = builder
+                .processed
+                .iter()
+                .find(|s| s.s == "MY_CUSTOM_ENV_VAR")
+                .expect("Should find MY_CUSTOM_ENV_VAR");
+            assert_eq!(item.suffix, "=");
+            assert_eq!(
+                item.description,
+                SuggestionDescription::Static(vec![ratatui::text::Span::raw("env var")])
+            );
+        }
+
+        #[test]
+        fn test_fuzzy_first_word_env_var_completion_has_equal_suffix() {
+            crate::shell::backend()
+                .export_env_var("MY_LONG_VARIABLE_NAME", "hello")
+                .unwrap();
+            let builder = tab_complete_fuzzy_first_word("MYLGVAR");
+            let item = builder
+                .processed
+                .iter()
+                .find(|s| s.s == "MY_LONG_VARIABLE_NAME")
+                .expect("Should fuzzy match MY_LONG_VARIABLE_NAME");
+            assert_eq!(item.suffix, "=");
+            assert_eq!(
+                item.description,
+                SuggestionDescription::Static(vec![ratatui::text::Span::raw("env var")])
+            );
+        }
+
+        #[test]
+        fn test_first_word_executable_completion_has_space_suffix() {
+            let builder = tab_complete_first_word("ech", "ech");
+            let item = builder
+                .processed
+                .iter()
+                .find(|s| s.s == "echo")
+                .expect("Should find echo");
+            assert_eq!(item.suffix, " ");
+        }
     }
 }

src/grammar/dparser.rs

@@ -200,13 +200,8 @@ impl DParser {
         new_tokens
     }
 
-    fn nested_opening_satisfied(
-        token: &Token,
-        current_nesting: Option<&TokenKind>,
-        is_command_extraction: bool,
-    ) -> bool {
+    fn nested_opening_satisfied(token: &Token, current_nesting: Option<&TokenKind>) -> bool {
         match token.kind {
-            TokenKind::Quote | TokenKind::SingleQuote if is_command_extraction => false,
             TokenKind::Backtick | TokenKind::Quote | TokenKind::SingleQuote => {
                 if Some(&token.kind) == current_nesting {
                     // backtick or quote is acting as closer
@@ -402,11 +397,6 @@ impl DParser {
 
         let mut previous_token: Option<AnnotatedToken> = None;
 
-        // Set to true when a closing nesting restores a command range whose first token
-        // is an env-var name (e.g. closing `"` in `FOO="bar"`).  The next non-whitespace
-        // word token will then reset current_command_range to None so that it can be
-        // recognised as a fresh command word.
-        let mut assignment_value_just_closed = false;
         let mut cursor_token_idx = None;
 
         let mut idx = 0;
@@ -426,21 +416,6 @@ impl DParser {
                 self.tokens[idx].token.kind = TokenKind::DoubleRParen;
             }
 
-            // If the previous env-var value nesting just closed, reset the command range
-            // now (before the arg-merging check below) so that the next word token is
-            // treated as the start of a fresh command rather than as another argument to
-            // the assignment statement.  The reset is deferred until here so that it skips
-            // over any intervening whitespace tokens.  For non-word, non-whitespace tokens
-            // (e.g. redirects) the flag is cleared without resetting the range.
-            if assignment_value_just_closed {
-                if self.tokens[idx].token.kind.is_word() {
-                    self.current_command_range = None;
-                }
-                if !matches!(self.tokens[idx].token.kind, TokenKind::Whitespace(_)) {
-                    assignment_value_just_closed = false;
-                }
-            }
-
             // Something like `echo foo=bar` is not an assignment.
             if self.current_command_range.is_some()
                 && self.tokens[idx].token.kind.is_word()
@@ -554,11 +529,7 @@ impl DParser {
                 | TokenKind::For
                 | TokenKind::While
                 | TokenKind::Until
-                    if Self::nested_opening_satisfied(
-                        &token,
-                        nestings.last().map(|(_, k)| k),
-                        cursor_byte_pos.is_some(),
-                    ) =>
+                    if Self::nested_opening_satisfied(&token, nestings.last().map(|(_, k)| k)) =>
                 {
                     let depth = nestings.len();
                     self.tokens[idx].annotations.opening = Some(OpeningState::Unmatched);
@@ -675,19 +646,6 @@ impl DParser {
                     } else if let Some(range) = &mut self.current_command_range {
                         *range = *range.start()..=idx;
                     }
-
-                    // If the restored range begins with an env-var token (e.g. the `FOO` in
-                    // `FOO="bar"`), the nesting we just closed was the value of an env-var
-                    // assignment.  The next word token should start a fresh command, so defer
-                    // the reset of current_command_range until then.
-                    if self
-                        .current_command_range
-                        .as_ref()
-                        .and_then(|r| self.tokens.get(*r.start()))
-                        .is_some_and(|t| t.annotations.is_env_var)
-                    {
-                        assignment_value_just_closed = true;
-                    }
                 }
                 TokenKind::Assignment => {
                     // When an assignment operator immediately follows a word (e.g. `FOO=1`),
@@ -797,7 +755,20 @@ impl DParser {
                     self.current_command_range = None;
                 }
                 TokenKind::Whitespace(_) => {
-                    if token_inclusively_contains_cursor
+                    let env_var_assignment_closed =
+                        self.current_command_range.as_ref().is_some_and(|r| {
+                            self.tokens
+                                .get(*r.start())
+                                .is_some_and(|t| t.annotations.is_env_var)
+                                && nestings
+                                    .last()
+                                    .map(|(open_idx, _)| *open_idx < *r.start())
+                                    .unwrap_or(true)
+                        });
+
+                    if env_var_assignment_closed {
+                        self.current_command_range = None;
+                    } else if token_inclusively_contains_cursor
                         && let Some(range) = &mut self.current_command_range
                     {
                         *range = *range.start()..=idx;

src/shell/bash/funcs.rs

@@ -978,12 +978,49 @@ fn get_cached_builtins() -> Vec<CommandWordInfo> {
         .clone()
 }
 
-/// Get all potential first word completions (aliases, reserved words, functions, builtins, executables)
+pub fn get_cached_environment_variables() -> Vec<CommandWordInfo> {
+    let _guard = super::symbols::BASH_LOCK.lock();
+    let mut variables = Vec::new();
+    let prefix_c_str = std::ffi::CString::new("").unwrap();
+
+    unsafe {
+        let var_ptr = bash_symbols::all_variables_matching_prefix(prefix_c_str.as_ptr());
+        if var_ptr.is_null() {
+            return variables;
+        }
+
+        let mut offset = 0;
+        let mut ptrs_to_free = Vec::new();
+        loop {
+            let ptr = *var_ptr.add(offset);
+            if ptr.is_null() {
+                break;
+            }
+            let c_str = std::ffi::CStr::from_ptr(ptr);
+            if let Ok(name) = c_str.to_str() {
+                variables.push(CommandWordInfo::EnvVar {
+                    name: name.to_string(),
+                });
+            }
+            ptrs_to_free.push(ptr);
+            offset += 1;
+        }
+        for str_ptr in ptrs_to_free {
+            bash_symbols::locked_xfree(str_ptr as *mut libc::c_void);
+        }
+        bash_symbols::locked_xfree(var_ptr as *mut libc::c_void);
+    }
+
+    variables
+}
+
+/// Get all potential first word completions (aliases, reserved words, functions, builtins, executables, env vars)
 pub fn get_possible_command_words() -> impl Iterator<Item = CommandWordInfo> {
     let aliases = get_cached_aliases();
     let reserved_words = get_cached_reserved_words();
     let shell_functions = get_cached_shell_functions();
     let builtins = get_cached_builtins();
+    let env_vars = get_cached_environment_variables();
     // This should be pre warmed by warm_completion_caches
     // We don't update the executables cache here to avoid hitting the filesystem
     // when we are just tab completing
@@ -999,6 +1036,7 @@ pub fn get_possible_command_words() -> impl Iterator<Item = CommandWordInfo> {
         .chain(shell_functions)
         .chain(builtins)
         .chain(executables)
+        .chain(env_vars)
 }
 
 pub fn warm_bash_caches() {
@@ -1007,6 +1045,7 @@ pub fn warm_bash_caches() {
     let _ = get_cached_reserved_words();
     let _ = get_cached_shell_functions();
     let _ = get_cached_builtins();
+    let _ = get_cached_environment_variables();
 }
 
 pub fn read_terminating_signal() -> c_int {

src/shell/mod.rs

@@ -50,6 +50,9 @@ pub enum CommandWordInfo {
         command: String,
         path: String,
     },
+    EnvVar {
+        name: String,
+    },
 }
 
 impl CommandWordInfo {
@@ -65,6 +68,7 @@ impl CommandWordInfo {
             CommandWordInfo::Function { command, .. } => command,
             CommandWordInfo::Builtin { command, .. } => command,
             CommandWordInfo::File { command, .. } => command,
+            CommandWordInfo::EnvVar { name } => name,
         }
     }
 
@@ -87,6 +91,7 @@ impl CommandWordInfo {
                 }
             }
             CommandWordInfo::File { path, .. } => path.clone(),
+            CommandWordInfo::EnvVar { .. } => "env var".to_string(),
             CommandWordInfo::Function {
                 source_file, line, ..
             } => match (source_file, line) {

src/shell/test_backend.rs

@@ -538,6 +538,11 @@ impl ShellBackend for TestBackend {
                 expansion: expansion.clone(),
             });
         }
+        for (env_k, _) in self.env_vars.read().iter() {
+            words.push(CommandWordInfo::EnvVar {
+                name: env_k.clone(),
+            });
+        }
         words
     }