feat: add one shot commands to the coder ssh command (#17779)

Closes #2154

> [!WARNING]  
> The tests in this PR were co-authored by AI

cli/ssh.go

@@ -90,15 +90,33 @@ func (r *RootCmd) ssh() *serpent.Command {
 	wsClient := workspacesdk.New(client)
 	cmd := &serpent.Command{
 		Annotations: workspaceCommand,
-		Use:         "ssh <workspace>",
-		Short:       "Start a shell into a workspace",
-		Long:        "This command does not have full parity with the standard SSH command. For users who need the full functionality of SSH, create an ssh configuration with `coder config-ssh`.",
+		Use:         "ssh <workspace> [command]",
+		Short:       "Start a shell into a workspace or run a command",
+		Long: "This command does not have full parity with the standard SSH command. For users who need the full functionality of SSH, create an ssh configuration with `coder config-ssh`.\n\n" +
+			FormatExamples(
+				Example{
+					Description: "Use `--` to separate and pass flags directly to the command executed via SSH.",
+					Command:     "coder ssh <workspace> -- ls -la",
+				},
+			),
 		Middleware: serpent.Chain(
-			serpent.RequireNArgs(1),
+			// Require at least one arg for the workspace name
+			func(next serpent.HandlerFunc) serpent.HandlerFunc {
+				return func(i *serpent.Invocation) error {
+					got := len(i.Args)
+					if got < 1 {
+						return xerrors.New("expected the name of a workspace")
+					}
+
+					return next(i)
+				}
+			},
 			r.InitClient(client),
 			initAppearance(client, &appearanceConfig),
 		),
 		Handler: func(inv *serpent.Invocation) (retErr error) {
+			command := strings.Join(inv.Args[1:], " ")
+
 			// Before dialing the SSH server over TCP, capture Interrupt signals
 			// so that if we are interrupted, we have a chance to tear down the
 			// TCP session cleanly before exiting.  If we don't, then the TCP
@@ -548,40 +566,46 @@ func (r *RootCmd) ssh() *serpent.Command {
 			sshSession.Stdout = inv.Stdout
 			sshSession.Stderr = inv.Stderr
 
-			err = sshSession.Shell()
-			if err != nil {
-				return xerrors.Errorf("start shell: %w", err)
-			}
+			if command != "" {
+				err := sshSession.Run(command)
+				if err != nil {
+					return xerrors.Errorf("run command: %w", err)
+				}
+			} else {
+				err = sshSession.Shell()
+				if err != nil {
+					return xerrors.Errorf("start shell: %w", err)
+				}
 
-			// Put cancel at the top of the defer stack to initiate
-			// shutdown of services.
-			defer cancel()
+				// Put cancel at the top of the defer stack to initiate
+				// shutdown of services.
+				defer cancel()
 
-			if validOut {
-				// Set initial window size.
-				width, height, err := term.GetSize(int(stdoutFile.Fd()))
-				if err == nil {
-					_ = sshSession.WindowChange(height, width)
+				if validOut {
+					// Set initial window size.
+					width, height, err := term.GetSize(int(stdoutFile.Fd()))
+					if err == nil {
+						_ = sshSession.WindowChange(height, width)
+					}
+				}
+
+				err = sshSession.Wait()
+				conn.SendDisconnectedTelemetry()
+				if err != nil {
+					if exitErr := (&gossh.ExitError{}); errors.As(err, &exitErr) {
+						// Clear the error since it's not useful beyond
+						// reporting status.
+						return ExitError(exitErr.ExitStatus(), nil)
+					}
+					// If the connection drops unexpectedly, we get an
+					// ExitMissingError but no other error details, so try to at
+					// least give the user a better message
+					if errors.Is(err, &gossh.ExitMissingError{}) {
+						return ExitError(255, xerrors.New("SSH connection ended unexpectedly"))
+					}
+					return xerrors.Errorf("session ended: %w", err)
 				}
 			}
-
-			err = sshSession.Wait()
-			conn.SendDisconnectedTelemetry()
-			if err != nil {
-				if exitErr := (&gossh.ExitError{}); errors.As(err, &exitErr) {
-					// Clear the error since it's not useful beyond
-					// reporting status.
-					return ExitError(exitErr.ExitStatus(), nil)
-				}
-				// If the connection drops unexpectedly, we get an
-				// ExitMissingError but no other error details, so try to at
-				// least give the user a better message
-				if errors.Is(err, &gossh.ExitMissingError{}) {
-					return ExitError(255, xerrors.New("SSH connection ended unexpectedly"))
-				}
-				return xerrors.Errorf("session ended: %w", err)
-			}
-
 			return nil
 		},
 	}

cli/ssh_test.go

@@ -2200,6 +2200,127 @@ func TestSSH_CoderConnect(t *testing.T) {
 
 		<-cmdDone
 	})
+
+	t.Run("OneShot", func(t *testing.T) {
+		t.Parallel()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+		inv, root := clitest.New(t, "ssh", workspace.Name, "echo 'hello world'")
+		clitest.SetupConfig(t, client, root)
+
+		// Capture command output
+		output := new(bytes.Buffer)
+		inv.Stdout = output
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		cmdDone := tGo(t, func() {
+			err := inv.WithContext(ctx).Run()
+			assert.NoError(t, err)
+		})
+
+		_ = agenttest.New(t, client.URL, agentToken)
+		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+		<-cmdDone
+
+		// Verify command output
+		assert.Contains(t, output.String(), "hello world")
+	})
+
+	t.Run("OneShotExitCode", func(t *testing.T) {
+		t.Parallel()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+
+		// Setup agent first to avoid race conditions
+		_ = agenttest.New(t, client.URL, agentToken)
+		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		// Test successful exit code
+		t.Run("Success", func(t *testing.T) {
+			inv, root := clitest.New(t, "ssh", workspace.Name, "exit 0")
+			clitest.SetupConfig(t, client, root)
+
+			err := inv.WithContext(ctx).Run()
+			assert.NoError(t, err)
+		})
+
+		// Test error exit code
+		t.Run("Error", func(t *testing.T) {
+			inv, root := clitest.New(t, "ssh", workspace.Name, "exit 1")
+			clitest.SetupConfig(t, client, root)
+
+			err := inv.WithContext(ctx).Run()
+			assert.Error(t, err)
+			var exitErr *ssh.ExitError
+			assert.True(t, errors.As(err, &exitErr))
+			assert.Equal(t, 1, exitErr.ExitStatus())
+		})
+	})
+
+	t.Run("OneShotStdio", func(t *testing.T) {
+		t.Parallel()
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+		_, _ = tGoContext(t, func(ctx context.Context) {
+			// Run this async so the SSH command has to wait for
+			// the build and agent to connect!
+			_ = agenttest.New(t, client.URL, agentToken)
+			<-ctx.Done()
+		})
+
+		clientOutput, clientInput := io.Pipe()
+		serverOutput, serverInput := io.Pipe()
+		defer func() {
+			for _, c := range []io.Closer{clientOutput, clientInput, serverOutput, serverInput} {
+				_ = c.Close()
+			}
+		}()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		inv, root := clitest.New(t, "ssh", "--stdio", workspace.Name, "echo 'hello stdio'")
+		clitest.SetupConfig(t, client, root)
+		inv.Stdin = clientOutput
+		inv.Stdout = serverInput
+		inv.Stderr = io.Discard
+
+		cmdDone := tGo(t, func() {
+			err := inv.WithContext(ctx).Run()
+			assert.NoError(t, err)
+		})
+
+		conn, channels, requests, err := ssh.NewClientConn(&testutil.ReaderWriterConn{
+			Reader: serverOutput,
+			Writer: clientInput,
+		}, "", &ssh.ClientConfig{
+			// #nosec
+			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		})
+		require.NoError(t, err)
+		defer conn.Close()
+
+		sshClient := ssh.NewClient(conn, channels, requests)
+		session, err := sshClient.NewSession()
+		require.NoError(t, err)
+		defer session.Close()
+
+		// Capture and verify command output
+		output, err := session.Output("echo 'hello back'")
+		require.NoError(t, err)
+		assert.Contains(t, string(output), "hello back")
+
+		err = sshClient.Close()
+		require.NoError(t, err)
+		_ = clientOutput.Close()
+
+		<-cmdDone
+	})
 }
 
 type fakeCoderConnectDialer struct{}

cli/testdata/coder_--help.golden

@@ -46,7 +46,7 @@ SUBCOMMANDS:
     show              Display details of a workspace's resources and agents
     speedtest         Run upload and download tests from your machine to a
                       workspace
-    ssh               Start a shell into a workspace
+    ssh               Start a shell into a workspace or run a command
     start             Start a workspace
     stat              Show resource usage for the current workspace.
     state             Manually manage Terraform state to fix broken workspaces

cli/testdata/coder_ssh_--help.golden

@@ -1,13 +1,18 @@
 coder v0.0.0-devel
 
 USAGE:
-  coder ssh [flags] <workspace>
+  coder ssh [flags] <workspace> [command]
 
-  Start a shell into a workspace
+  Start a shell into a workspace or run a command
 
   This command does not have full parity with the standard SSH command. For
   users who need the full functionality of SSH, create an ssh configuration with
   `coder config-ssh`.
+  
+    - Use `--` to separate and pass flags directly to the command executed via
+  SSH.:
+  
+       $ coder ssh <workspace> -- ls -la
 
 OPTIONS:
       --disable-autostart bool, $CODER_SSH_DISABLE_AUTOSTART (default: false)

docs/manifest.json

@@ -1460,7 +1460,7 @@
 						},
 						{
 							"title": "ssh",
-							"description": "Start a shell into a workspace",
+							"description": "Start a shell into a workspace or run a command",
 							"path": "reference/cli/ssh.md"
 						},
 						{

docs/reference/cli/index.md

@@ -53,7 +53,7 @@ Coder — A tool for provisioning self-hosted development environments with Terr
 | [<code>schedule</code>](./schedule.md)             | Schedule automated start and stop times for workspaces                                                |
 | [<code>show</code>](./show.md)                     | Display details of a workspace's resources and agents                                                 |
 | [<code>speedtest</code>](./speedtest.md)           | Run upload and download tests from your machine to a workspace                                        |
-| [<code>ssh</code>](./ssh.md)                       | Start a shell into a workspace                                                                        |
+| [<code>ssh</code>](./ssh.md)                       | Start a shell into a workspace or run a command                                                       |
 | [<code>start</code>](./start.md)                   | Start a workspace                                                                                     |
 | [<code>stat</code>](./stat.md)                     | Show resource usage for the current workspace.                                                        |
 | [<code>stop</code>](./stop.md)                     | Stop a workspace                                                                                      |

docs/reference/cli/ssh.md

@@ -1,18 +1,22 @@
 <!-- DO NOT EDIT | GENERATED CONTENT -->
 # ssh
 
-Start a shell into a workspace
+Start a shell into a workspace or run a command
 
 ## Usage
 
 ```console
-coder ssh [flags] <workspace>
+coder ssh [flags] <workspace> [command]
 ```
 
 ## Description
 
 ```console
 This command does not have full parity with the standard SSH command. For users who need the full functionality of SSH, create an ssh configuration with `coder config-ssh`.
+
+  - Use `--` to separate and pass flags directly to the command executed via SSH.:
+
+     $ coder ssh <workspace> -- ls -la
 ```
 
 ## Options