completion: fix completion of arguments for aliases/default-command in bash and zsh

This also adds a test case for the completion of arguments following
multi-argument aliases, to cover the bug reported in issue #5377.

The default command is like a special kind of alias, one which is
expanded from zero tokens. Consequently, it also triggers the bug
#5377 on bash/zsh, even if the `default-command` is just a single token.

The fix is along the lines sketched out by the "TODO" comment. Bash and
Zsh don't behave identical, so the padding ([""]) still needs to be
applied (and removed) conditionally in a disciplined manner.

Author: jgreitemann

CHANGELOG.md

@@ -72,6 +72,9 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 * The builtin diff editor now tries to preserve unresolved conflicts.
   [#4963](https://github.com/jj-vcs/jj/issues/4963)
 
+* Fixed bash and zsh shell completion when completing aliases of multiple arguments.
+  [#5377](https://github.com/jj-vcs/jj/issues/5377)
+
 ### Packaging changes
 
 * Jujutsu now uses

cli/src/cli_util.rs

@@ -3551,29 +3551,55 @@ fn handle_shell_completion(
     config: &StackedConfig,
     cwd: &Path,
 ) -> Result<(), CommandError> {
+    let mut orig_args = env::args_os();
+
     let mut args = vec![];
     // Take the first two arguments as is, they must be passed to clap_complete
     // without any changes. They are usually "jj --".
-    args.extend(env::args_os().take(2));
+    args.extend(orig_args.by_ref().take(2));
 
     // Make sure aliases are expanded before passing them to clap_complete. We
     // skip the first two args ("jj" and "--") for alias resolution, then we
     // stitch the args back together, like clap_complete expects them.
-    let orig_args = env::args_os().skip(2);
     if orig_args.len() > 0 {
-        let arg_index: Option<usize> = env::var("_CLAP_COMPLETE_INDEX")
+        let complete_index: Option<usize> = env::var("_CLAP_COMPLETE_INDEX")
             .ok()
             .and_then(|s| s.parse().ok());
-        let resolved_aliases = if let Some(index) = arg_index {
+        let resolved_aliases = if let Some(index) = complete_index {
             // As of clap_complete 4.5.38, zsh completion script doesn't pad an
             // empty arg at the complete position. If the args doesn't include a
             // command name, the default command would be expanded at that
             // position. Therefore, no other command names would be suggested.
-            // TODO: Maybe we should instead expand args[..index] + [""], adjust
-            // the index accordingly, strip the last "", and append remainder?
             let pad_len = usize::saturating_sub(index + 1, orig_args.len());
-            let padded_args = orig_args.chain(std::iter::repeat_n(OsString::new(), pad_len));
-            expand_args(ui, app, padded_args, config)?
+            let padded_args = orig_args
+                .by_ref()
+                .chain(std::iter::repeat_n(OsString::new(), pad_len));
+
+            // Expand aliases left of the completion index.
+            let mut expanded_args = expand_args(ui, app, padded_args.take(index + 1), config)?;
+
+            // Adjust env var to compensate for shift of the completion point in the
+            // expanded command line.
+            // SAFETY: Program is running single-threaded at this point.
+            unsafe {
+                env::set_var(
+                    "_CLAP_COMPLETE_INDEX",
+                    (expanded_args.len() - 1).to_string(),
+                );
+            }
+
+            // Remove extra padding again to align with clap_complete's expectations for
+            // zsh.
+            let split_off_padding = expanded_args.split_off(expanded_args.len() - pad_len);
+            assert!(
+                split_off_padding.iter().all(|s| s.is_empty()),
+                "split-off padding should only consist of empty strings but was \
+                 {split_off_padding:?}",
+            );
+
+            // Append the remaining arguments to the right of the completion point.
+            expanded_args.extend(to_string_args(orig_args)?);
+            expanded_args
         } else {
             expand_args(ui, app, orig_args, config)?
         };
@@ -3598,19 +3624,24 @@ pub fn expand_args(
     args_os: impl IntoIterator<Item = OsString>,
     config: &StackedConfig,
 ) -> Result<Vec<String>, CommandError> {
-    let mut string_args: Vec<String> = vec![];
-    for arg_os in args_os {
-        if let Some(string_arg) = arg_os.to_str() {
-            string_args.push(string_arg.to_owned());
-        } else {
-            return Err(cli_error("Non-utf8 argument"));
-        }
-    }
-
+    let string_args = to_string_args(args_os)?;
     let string_args = resolve_default_command(ui, config, app, string_args)?;
     resolve_aliases(ui, config, app, string_args)
 }
 
+fn to_string_args(
+    args_os: impl IntoIterator<Item = OsString>,
+) -> Result<Vec<String>, CommandError> {
+    args_os
+        .into_iter()
+        .map(|arg_os| {
+            arg_os
+                .into_string()
+                .map_err(|_| cli_error("Non-UTF-8 argument"))
+        })
+        .collect()
+}
+
 fn parse_args(app: &Command, string_args: &[String]) -> Result<(ArgMatches, Args), clap::Error> {
     let matches = app
         .clone()

cli/tests/common/command_output.rs

@@ -57,6 +57,16 @@ impl CommandOutput {
         }
     }
 
+    /// Removes all but the first `n` lines from normalized stdout text.
+    #[must_use]
+    pub fn take_stdout_n_lines(self, n: usize) -> Self {
+        CommandOutput {
+            stdout: self.stdout.take_n_lines(n),
+            stderr: self.stderr,
+            status: self.status,
+        }
+    }
+
     #[must_use]
     pub fn normalize_stdout_with(self, f: impl FnOnce(String) -> String) -> Self {
         CommandOutput {
@@ -141,6 +151,12 @@ impl CommandOutputString {
         })
     }
 
+    /// Removes all but the first `n` lines from the normalized text.
+    #[must_use]
+    pub fn take_n_lines(self, n: usize) -> Self {
+        self.normalize_with(|s| s.split_inclusive("\n").take(n).collect())
+    }
+
     #[must_use]
     pub fn normalize_with(mut self, f: impl FnOnce(String) -> String) -> Self {
         self.normalized = f(self.normalized);

cli/tests/test_completion.rs

@@ -321,6 +321,7 @@ fn test_aliases_are_resolved(shell: Shell) {
 
     // user config alias
     test_env.add_config(r#"aliases.b = ["bookmark"]"#);
+    test_env.add_config(r#"aliases.rlog = ["log", "--reversed"]"#);
     // repo config alias
     work_dir
         .run_jj(["config", "set", "--repo", "aliases.b2", "['bookmark']"])
@@ -359,6 +360,30 @@ fn test_aliases_are_resolved(shell: Shell) {
         }
         _ => unimplemented!("unexpected shell '{shell}'"),
     }
+
+    let output = work_dir.complete_at(shell, 2, ["rlog", "--rev"]);
+    match shell {
+        Shell::Bash => {
+            insta::assert_snapshot!(output, @r"
+            --revisions
+            --reversed[EOF]
+            ");
+        }
+        Shell::Zsh => {
+            insta::assert_snapshot!(output, @r"
+            --revisions:Which revisions to show
+            --reversed:Show revisions in the opposite order (older revisions first)[EOF]
+            ");
+        }
+        Shell::Fish => {
+            insta::assert_snapshot!(output, @r"
+            --revisions	Which revisions to show
+            --reversed	Show revisions in the opposite order (older revisions first)
+            [EOF]
+            ");
+        }
+        _ => unimplemented!("unexpected shell '{shell}'"),
+    }
 }
 
 #[test]
@@ -381,6 +406,99 @@ fn test_completions_are_generated() {
     ");
 }
 
+#[test_case(Shell::Bash; "bash")]
+#[test_case(Shell::Zsh; "zsh")]
+#[test_case(Shell::Fish; "fish")]
+fn test_default_command_is_resolved(shell: Shell) {
+    let test_env = TestEnvironment::default();
+
+    let output = test_env
+        .complete_at(shell, 1, ["--"])
+        .take_stdout_n_lines(2);
+    match shell {
+        Shell::Bash => {
+            insta::assert_snapshot!(output, @r"
+            --revisions
+            --limit
+            [EOF]
+            ");
+        }
+        Shell::Zsh => {
+            insta::assert_snapshot!(output, @r"
+            --revisions:Which revisions to show
+            --limit:Limit number of revisions to show
+            [EOF]
+            ");
+        }
+        Shell::Fish => {
+            insta::assert_snapshot!(output, @r"
+            --revisions	Which revisions to show
+            --limit	Limit number of revisions to show
+            [EOF]
+            ");
+        }
+        _ => unimplemented!("unexpected shell '{shell}'"),
+    }
+
+    test_env.add_config("ui.default-command = ['abandon']");
+    let output = test_env
+        .complete_at(shell, 1, ["--"])
+        .take_stdout_n_lines(2);
+    match shell {
+        Shell::Bash => {
+            insta::assert_snapshot!(output, @r"
+            --retain-bookmarks
+            --restore-descendants
+            [EOF]
+            ");
+        }
+        Shell::Zsh => {
+            insta::assert_snapshot!(output, @r"
+            --retain-bookmarks:Do not delete bookmarks pointing to the revisions to abandon
+            --restore-descendants:Do not modify the content of the children of the abandoned commits
+            [EOF]
+            ");
+        }
+        Shell::Fish => {
+            insta::assert_snapshot!(output, @r"
+            --retain-bookmarks	Do not delete bookmarks pointing to the revisions to abandon
+            --restore-descendants	Do not modify the content of the children of the abandoned commits
+            [EOF]
+            ");
+        }
+        _ => unimplemented!("unexpected shell '{shell}'"),
+    }
+
+    test_env.add_config("ui.default-command = ['bookmark', 'move']");
+    let output = test_env
+        .complete_at(shell, 1, ["--"])
+        .take_stdout_n_lines(2);
+    match shell {
+        Shell::Bash => {
+            insta::assert_snapshot!(output, @r"
+            --from
+            --to
+            [EOF]
+            ");
+        }
+        Shell::Zsh => {
+            insta::assert_snapshot!(output, @r"
+            --from:Move bookmarks from the given revisions
+            --to:Move bookmarks to this revision
+            [EOF]
+            ");
+        }
+        Shell::Fish => {
+            insta::assert_snapshot!(output, @r"
+            --from	Move bookmarks from the given revisions
+            --to	Move bookmarks to this revision
+            [EOF]
+            ");
+        }
+        _ => unimplemented!("unexpected shell '{shell}'"),
+    }
+}
+
 #[test_case(Shell::Bash; "bash")]
 #[test_case(Shell::Zsh; "zsh")]
 #[test_case(Shell::Fish; "fish")]
@@ -389,9 +507,7 @@ fn test_command_completion(shell: Shell) {
 
     // Command names should be suggested. If the default command were expanded,
     // only "log" would be listed.
-    let output = test_env
-        .complete_at(shell, 1, [""])
-        .normalize_stdout_with(|s| s.split_inclusive('\n').take(2).collect());
+    let output = test_env.complete_at(shell, 1, [""]).take_stdout_n_lines(2);
     match shell {
         Shell::Bash => {
             insta::assert_snapshot!(output, @r"
@@ -419,7 +535,7 @@ fn test_command_completion(shell: Shell) {
 
     let output = test_env
         .complete_at(shell, 2, ["--no-pager", ""])
-        .normalize_stdout_with(|s| s.split_inclusive('\n').take(2).collect());
+        .take_stdout_n_lines(2);
     match shell {
         Shell::Bash => {
             insta::assert_snapshot!(output, @r"

cli/tests/test_global_opts.rs

@@ -37,7 +37,7 @@ fn test_non_utf8_arg() {
     let output = test_env.run_jj_in(".", [&invalid_utf]);
     insta::assert_snapshot!(output, @r"
     ------- stderr -------
-    Error: Non-utf8 argument
+    Error: Non-UTF-8 argument
     [EOF]
     [exit status: 2]
     ");