Add optional command extensions

src/bin/dosh-client.rs

@@ -212,11 +212,7 @@ async fn main() -> Result<()> {
         .unwrap_or_default();
     let raw_server = host.ssh.clone().unwrap_or_else(|| requested_server.clone());
     let server = ssh_with_user(&raw_server, host.user.as_deref());
-    let startup_command = startup_command(&args.command).or_else(|| {
-        host.default_command
-            .as_deref()
-            .map(startup_command_from_string)
-    });
+    let startup_command = resolved_startup_command(&args.command, &config, &host);
     let local_forwards = parse_local_forwards(&args.local_forward)?;
     let remote_forwards = parse_remote_forwards(&args.remote_forward)?;
     let dynamic_forwards = parse_dynamic_forwards(&args.dynamic_forward)?;
@@ -776,6 +772,73 @@ fn startup_command(args: &[String]) -> Option<Vec<u8>> {
     Some(command.into_bytes())
 }
 
+fn resolved_startup_command(
+    args: &[String],
+    config: &dosh::config::ClientConfig,
+    host: &dosh::config::HostConfig,
+) -> Option<Vec<u8>> {
+    if args.is_empty() {
+        return host
+            .default_command
+            .as_deref()
+            .map(startup_command_from_string);
+    }
+    match resolve_command_extension(args, config, host) {
+        ExtensionResolution::Command(command) => Some(startup_command_from_string(&command)),
+        ExtensionResolution::Disabled => None,
+        ExtensionResolution::NoMatch => startup_command(args),
+    }
+}
+
+enum ExtensionResolution {
+    Command(String),
+    Disabled,
+    NoMatch,
+}
+
+fn resolve_command_extension(
+    args: &[String],
+    config: &dosh::config::ClientConfig,
+    host: &dosh::config::HostConfig,
+) -> ExtensionResolution {
+    let Some(name) = args.first() else {
+        return ExtensionResolution::NoMatch;
+    };
+    if let Some(extension) = host.extensions.get(name) {
+        if extension.disabled {
+            return ExtensionResolution::Disabled;
+        }
+        if let Some(command) = extension.command.as_deref() {
+            return ExtensionResolution::Command(expand_command_extension(command, &args[1..]));
+        }
+    }
+    if let Some(extension) = config.extensions.get(name) {
+        if extension.disabled {
+            return ExtensionResolution::Disabled;
+        }
+        if let Some(command) = extension.command.as_deref() {
+            return ExtensionResolution::Command(expand_command_extension(command, &args[1..]));
+        }
+    }
+    ExtensionResolution::NoMatch
+}
+
+fn expand_command_extension(template: &str, args: &[String]) -> String {
+    let quoted_args = args
+        .iter()
+        .map(|arg| shell_word(arg))
+        .collect::<Vec<_>>()
+        .join(" ");
+    if template.contains("{args}") {
+        return template.replace("{args}", &quoted_args);
+    }
+    if args.is_empty() {
+        template.to_string()
+    } else {
+        format!("{template} {quoted_args}")
+    }
+}
+
 fn startup_command_from_string(command: &str) -> Vec<u8> {
     let mut command = command.to_string();
     command.push('\n');
@@ -4226,10 +4289,11 @@ mod tests {
         RemoteForward, SshConfig, StatusAction, auth_allows,
         load_first_native_identity_with_prompt, parse_dynamic_forward, parse_local_forward,
         parse_remote_forward, parse_ssh_config, raw_contains_host_table, recv_response_until,
-        render_status_clear, render_status_overlay, requested_env, ssh_destination_host,
-        ssh_username, ssh_with_user, startup_command, toml_bare_key_or_quoted,
+        render_status_clear, render_status_overlay, requested_env, resolved_startup_command,
+        ssh_destination_host, ssh_username, ssh_with_user, startup_command,
+        toml_bare_key_or_quoted,
     };
-    use dosh::config::{ClientConfig, HostConfig};
+    use dosh::config::{ClientConfig, CommandExtension, HostConfig};
     use dosh::native::EnvVar;
     use dosh::protocol::{self, Frame, PacketKind};
     use ed25519_dalek::{SigningKey, VerifyingKey};
@@ -4775,6 +4839,100 @@ mod tests {
         );
     }
 
+    #[test]
+    fn resolved_startup_command_expands_global_extension() {
+        let mut config = ClientConfig::default();
+        config.extensions.insert(
+            "tm".to_string(),
+            CommandExtension {
+                command: Some("tm {args}".to_string()),
+                description: Some("server tmux dashboard".to_string()),
+                disabled: false,
+            },
+        );
+        let host = HostConfig::default();
+
+        assert_eq!(
+            resolved_startup_command(&["tm".into(), "work space".into()], &config, &host),
+            Some(b"tm 'work space'\n".to_vec())
+        );
+    }
+
+    #[test]
+    fn resolved_startup_command_uses_host_extension_override() {
+        let mut config = ClientConfig::default();
+        config.extensions.insert(
+            "tm".to_string(),
+            CommandExtension {
+                command: Some("tm {args}".to_string()),
+                description: None,
+                disabled: false,
+            },
+        );
+        let mut host = HostConfig::default();
+        host.extensions.insert(
+            "tm".to_string(),
+            CommandExtension {
+                command: Some("/opt/tm/bin/tm --remote {args}".to_string()),
+                description: None,
+                disabled: false,
+            },
+        );
+
+        assert_eq!(
+            resolved_startup_command(&["tm".into(), "dosh".into()], &config, &host),
+            Some(b"/opt/tm/bin/tm --remote 'dosh'\n".to_vec())
+        );
+    }
+
+    #[test]
+    fn resolved_startup_command_host_can_disable_global_extension() {
+        let mut config = ClientConfig::default();
+        config.extensions.insert(
+            "tm".to_string(),
+            CommandExtension {
+                command: Some("tm {args}".to_string()),
+                description: None,
+                disabled: false,
+            },
+        );
+        let mut host = HostConfig::default();
+        host.extensions.insert(
+            "tm".to_string(),
+            CommandExtension {
+                command: None,
+                description: None,
+                disabled: true,
+            },
+        );
+
+        assert_eq!(
+            resolved_startup_command(&["tm".into(), "dosh".into()], &config, &host),
+            None
+        );
+    }
+
+    #[test]
+    fn resolved_startup_command_keeps_raw_command_fallback_and_default_command() {
+        let config = ClientConfig::default();
+        let mut host = HostConfig {
+            default_command: Some("tm".to_string()),
+            ..HostConfig::default()
+        };
+
+        assert_eq!(
+            resolved_startup_command(&["echo".into(), "ok".into()], &config, &host),
+            Some(b"echo ok\n".to_vec())
+        );
+        assert_eq!(
+            resolved_startup_command(&[], &config, &host),
+            Some(b"tm\n".to_vec())
+        );
+
+        host.default_command = None;
+        assert_eq!(resolved_startup_command(&[], &config, &host), None);
+    }
+
     #[test]
     fn parses_local_forward_with_default_bind() {
         assert_eq!(