fix: improve exit codes for agent/agentssh and cli/ssh (#10850)

cli/root.go

@@ -136,14 +136,22 @@ func (r *RootCmd) RunMain(subcommands []*clibase.Cmd) {
 	}
 	err = cmd.Invoke().WithOS().Run()
 	if err != nil {
+		code := 1
+		var exitErr *exitError
+		if errors.As(err, &exitErr) {
+			code = exitErr.code
+			err = exitErr.err
+		}
 		if errors.Is(err, cliui.Canceled) {
 			//nolint:revive
-			os.Exit(1)
+			os.Exit(code)
 		}
 		f := prettyErrorFormatter{w: os.Stderr, verbose: r.verbose}
-		f.format(err)
+		if err != nil {
+			f.format(err)
+		}
 		//nolint:revive
-		os.Exit(1)
+		os.Exit(code)
 	}
 }
 
@@ -953,6 +961,30 @@ func DumpHandler(ctx context.Context) {
 	}
 }
 
+type exitError struct {
+	code int
+	err  error
+}
+
+var _ error = (*exitError)(nil)
+
+func (e *exitError) Error() string {
+	if e.err != nil {
+		return fmt.Sprintf("exit code %d: %v", e.code, e.err)
+	}
+	return fmt.Sprintf("exit code %d", e.code)
+}
+
+func (e *exitError) Unwrap() error {
+	return e.err
+}
+
+// ExitError returns an error that will cause the CLI to exit with the given
+// exit code. If err is non-nil, it will be wrapped by the returned error.
+func ExitError(code int, err error) error {
+	return &exitError{code: code, err: err}
+}
+
 // IiConnectionErr is a convenience function for checking if the source of an
 // error is due to a 'connection refused', 'no such host', etc.
 func isConnectionError(err error) bool {

cli/ssh.go

@@ -379,11 +379,16 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 
 			err = sshSession.Wait()
 			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 xerrors.New("SSH connection ended unexpectedly")
+					return ExitError(255, xerrors.New("SSH connection ended unexpectedly"))
 				}
 				return xerrors.Errorf("session ended: %w", err)
 			}