diff --git a/src/Auth.php b/src/Auth.php
index 791ac98..c6356bd 100644
--- a/src/Auth.php
+++ b/src/Auth.php
@@ -121,11 +121,13 @@ public static function get_token_expiration() {
 	}
 
 	/**
-	 * @param $user
+	 * Retrieves validates user and retrieve signed token
+	 *
+	 * @param User|WP_User $user  Owner of the token.
 	 *
 	 * @return null|string
 	 */
-	protected static function get_signed_token( \WP_User $user, $cap_check = true ) {
+	protected static function get_signed_token( $user, $cap_check = true ) {
 
 		/**
 		 * Only allow the currently signed in user access to a JWT token
diff --git a/src/Login.php b/src/Login.php
index d4fe00f..4ad5db5 100644
--- a/src/Login.php
+++ b/src/Login.php
@@ -1,52 +1,57 @@
 <?php
+/**
+ * Register the login mutation to the WPGraphQL Schema
+ *
+ * @package WPGraphQL\JWT_Authentication
+ * @since 0.0.1
+ */
+
 namespace WPGraphQL\JWT_Authentication;
 
 use GraphQL\Type\Definition\ResolveInfo;
 use WPGraphQL\AppContext;
 
+/**
+ * Class - Login
+ */
 class Login {
-
 	/**
-	 * Register the login mutation to the WPGraphQL Schema
+	 * Register the mutation.
 	 */
 	public static function register_mutation() {
-
-		register_graphql_mutation( 'login', [
-			'description' => __( 'Login a user. Request for an authToken and User details in response', 'wp-graphql-jwt-authentication' ),
-			'inputFields' => [
-				'username' => [
-					'type' => [ 'non_null' => 'String' ],
-					'description' => __( 'The username used for login. Typically a unique or email address depending on specific configuration', 'wp-graphql-jwt-authentication' ),
-				],
-				'password' => [
-					'type' => [ 'non_null' => 'String' ],
-					'description' => __( 'The plain-text password for the user logging in.', 'wp-graphql-jwt-authentication' ),
-				],
-			],
-			'outputFields' => [
-				'authToken' => [
-					'type' => 'String',
-					'description' => __( 'JWT Token that can be used in future requests for Authentication', 'wp-graphql-jwt-authentication' ),
+		register_graphql_mutation(
+			'login',
+			[
+				'description'         => __( 'Login a user. Request for an authToken and User details in response', 'wp-graphql-jwt-authentication' ),
+				'inputFields'         => [
+					'username' => [
+						'type'        => [ 'non_null' => 'String' ],
+						'description' => __( 'The username used for login. Typically a unique or email address depending on specific configuration', 'wp-graphql-jwt-authentication' ),
+					],
+					'password' => [
+						'type'        => [ 'non_null' => 'String' ],
+						'description' => __( 'The plain-text password for the user logging in.', 'wp-graphql-jwt-authentication' ),
+					],
 				],
-				'refreshToken' => [
-					'type' => 'String',
-					'description' => __( 'A JWT token that can be used in future requests to get a refreshed jwtAuthToken. If the refresh token used in a request is revoked or otherwise invalid, a valid Auth token will NOT be issued in the response headers.', 'wp-graphql-jwt-authentication' ),
+				'outputFields'        => [
+					'authToken'    => [
+						'type'        => 'String',
+						'description' => __( 'JWT Token that can be used in future requests for Authentication', 'wp-graphql-jwt-authentication' ),
+					],
+					'refreshToken' => [
+						'type'        => 'String',
+						'description' => __( 'A JWT token that can be used in future requests to get a refreshed jwtAuthToken. If the refresh token used in a request is revoked or otherwise invalid, a valid Auth token will NOT be issued in the response headers.', 'wp-graphql-jwt-authentication' ),
+					],
+					'user'         => [
+						'type'        => 'User',
+						'description' => __( 'The user that was logged in', 'wp-graphql-jwt-authentication' ),
+					],
 				],
-				'user' => [
-					'type' => 'User',
-					'description' => __( 'The user that was logged in', 'wp-graphql-jwt-authentication' ),
-				],
-			],
-			'mutateAndGetPayload' => function( $input, AppContext $context, ResolveInfo $info ) {
-
-				/**
-				 * Login the user in and get an authToken and user in response
-				 */
-				return Auth::login_and_get_token( sanitize_user( $input['username'] ), trim( $input['password'] ) );
-
-			},
-		]);
-
+				'mutateAndGetPayload' => function( $input, AppContext $context, ResolveInfo $info ) {
+					// Login the user in and get an authToken and user in response.
+					return Auth::login_and_get_token( sanitize_user( $input['username'] ), trim( $input['password'] ) );
+				},
+			]
+		);
 	}
-
 }
diff --git a/src/ManageTokens.php b/src/ManageTokens.php
index e57a378..e77052d 100644
--- a/src/ManageTokens.php
+++ b/src/ManageTokens.php
@@ -1,56 +1,39 @@
 <?php
+/**
+ * Defines functionality for managing JWTs.
+ *
+ * @package WPGraphQL\JWT_Authentication
+ * @since 0.0.1
+ */
 
 namespace WPGraphQL\JWT_Authentication;
 
 use GraphQL\Error\UserError;
+use GraphQL\Type\Definition\Type;
 use WPGraphQL\Model\User;
 use WPGraphQL\Types;
 
+/**
+ * Class - ManageToken
+ */
 class ManageTokens {
-
 	/**
 	 * Initialize the funcionality for managing tokens
 	 */
 	public static function init() {
+		// Register JWT fields.
+		add_action( 'graphql_register_types', [ __CLASS__, 'add_jwt_fields' ], 10 );
 
-		/**
-		 * Filter the User type to have a jwtUserSecret and jwtAuthToken field
-		 */
-		add_filter( 'graphql_user_fields', [
-			'\WPGraphQL\JWT_Authentication\ManageTokens',
-			'add_user_fields'
-		] );
+		// Add fields to the input for user mutations.
+		add_filter( 'graphql_user_mutation_input_fields', [ __CLASS__, 'add_user_mutation_input_fields' ] );
+		add_action( 'graphql_user_object_mutation_update_additional_data', [ __CLASS__, 'update_jwt_fields_during_mutation' ], 10, 3 );
 
-		/**
-		 * Add fields to the input for user mutations
-		 */
-		add_filter( 'graphql_user_mutation_input_fields', [
-			'\WPGraphQL\JWT_Authentication\ManageTokens',
-			'add_user_mutation_input_fields'
-		] );
+		// Filter the signed token, preventing it from returning if the user has had their JWT Secret revoked.
+		add_filter( 'graphql_jwt_auth_signed_token', [ __CLASS__, 'prevent_token_from_returning_if_revoked' ], 10, 2 );
 
-		add_action( 'graphql_user_object_mutation_update_additional_data', [
-			'\WPGraphQL\JWT_Authentication\ManageTokens',
-			'update_jwt_fields_during_mutation'
-		], 10, 3 );
-
-		/**
-		 * Filter the signed token, preventing it from returning if the user has had their JWT Secret revoked
-		 */
-		add_filter( 'graphql_jwt_auth_signed_token', [
-			'\WPGraphQL\JWT_Authentication\ManageTokens',
-			'prevent_token_from_returning_if_revoked'
-		], 10, 2 );
-
-		add_filter( 'graphql_response_headers_to_send', [
-			'\WPGraphQL\JWT_Authentication\ManageTokens',
-			'add_tokens_to_graphql_response_headers'
-		] );
-
-		add_filter( 'graphql_response_headers_to_send', [
-			'\WPGraphQL\JWT_Authentication\ManageTokens',
-			'add_auth_headers_to_response'
-		] );
+		// Filter the GraphQL response headers.
+		add_filter( 'graphql_response_headers_to_send', [ __CLASS__, 'add_tokens_to_graphql_response_headers' ] );
+		add_filter( 'graphql_response_headers_to_send', [ __CLASS__, 'add_auth_headers_to_response' ] );
 
 		/**
 		 * Add Auth Headers to REST REQUEST responses
@@ -61,137 +44,121 @@ public static function init() {
 		 * so folks can refresh their tokens after each REST API
 		 * request.
 		 */
-		add_filter( 'rest_request_after_callbacks', [
-			'\WPGraphQL\JWT_Authentication\ManageTokens',
-			'add_auth_headers_to_rest_response'
-		], 10, 3 );
-
-		add_filter( 'graphql_access_control_allow_headers', [
-			'\WPGraphQL\JWT_Authentication\ManageTokens',
-			'add_jwt_allowed_headers'
-		] );
+		add_filter( 'rest_request_after_callbacks', [ __CLASS__, 'add_auth_headers_to_rest_response' ], 10, 3 );
+
+		// Filter allowed HTTP headers.
+		add_filter( 'graphql_access_control_allow_headers', [ __CLASS__, 'add_jwt_allowed_headers' ] );
 	}
 
 	/**
-	 * Filters the User type in the GraphQL Schema to provide fields for querying for user's
-	 * jwtAuthToken and jwtUserSecret
+	 * Registers JWT fields to the GraphQL schema on all targeted types.
 	 *
-	 * @param array $fields The fields for the User type in the GraphQL Schema
-	 *
-	 * @return array $fields
-	 * @throws \Exception
+	 * Type must provided an root object with an User ID assigned to the "ID" field.
+	 * Ex. $source->ID
 	 */
-	public static function add_user_fields( $fields ) {
-
-		$fields['jwtAuthToken'] = [
-			'type'        => Types::string(),
-			'description' => __( 'A JWT token that can be used in future requests for authentication/authorization', 'wp-graphql-jwt-authentication' ),
-			'resolve'     => function ( User $user ) {
-
-				$user = get_user_by( 'id', $user->userId );
-
-				/**
-				 * Get the token for the user
-				 */
-				$token = Auth::get_token( $user );
-
-				/**
-				 * If the token cannot be returned, throw an error
-				 */
-				if ( empty( $token ) || is_wp_error( $token ) ) {
-					throw new UserError( __( 'The JWT token could not be returned', 'wp-graphql-jwt-authentication' ) );
-				}
-
-				return ! empty( $token ) ? $token : null;
-			},
-		];
-
-		$fields['jwtRefreshToken'] = [
-			'type'        => Types::string(),
-			'description' => __( 'A JWT token that can be used in future requests to get a refreshed jwtAuthToken. If the refresh token used in a request is revoked or otherwise invalid, a valid Auth token will NOT be issued in the response headers.', 'wp-graphql-jwt-authentication' ),
-			'resolve'     => function ( User $user ) {
-
-				$user = get_user_by( 'id', $user->userId );
-
-				/**
-				 * Get the token for the user
-				 */
-				$token = Auth::get_refresh_token( $user );
-
-				/**
-				 * If the token cannot be returned, throw an error
-				 */
-				if ( empty( $token ) || is_wp_error( $token ) ) {
-					throw new UserError( __( 'The JWT token could not be returned', 'wp-graphql-jwt-authentication' ) );
-				}
-
-				return ! empty( $token ) ? $token : null;
-			},
-		];
-
-		$fields['jwtUserSecret'] = [
-			'type'        => Types::string(),
-			'description' => __( 'A unique secret tied to the users JWT token that can be revoked or refreshed. Revoking the secret prevents JWT tokens from being issued to the user. Refreshing the token invalidates previously issued tokens, but allows new tokens to be issued.', 'wp-graphql' ),
-			'resolve'     => function ( User $user ) {
-
-				/**
-				 * Get the user's JWT Secret
-				 */
-				$secret = Auth::get_user_jwt_secret( $user->userId );
-
-				/**
-				 * If the secret cannot be returned, throw an error
-				 */
-				if ( is_wp_error( $secret ) ) {
-					throw new UserError( __( 'The user secret could not be returned', 'wp-graphql-jwt-authentication' ) );
-				}
-
-				/**
-				 * Return the secret
-				 */
-				return ! empty( $secret ) ? $secret : null;
-			}
-		];
-
-		$fields['jwtAuthExpiration'] = [
-			'type'        => Types::string(),
-			'description' => __( 'The expiration for the JWT Token for the user. If not set custom for the user, it will use the default sitewide expiration setting', 'wp-graphql-jwt-authentication' ),
-			'resolve'     => function () {
-				$expiration = Auth::get_token_expiration();
-
-				return ! empty( $expiration ) ? $expiration : null;
-			}
-		];
-
-		$fields['isJwtAuthSecretRevoked'] = [
-			'type'        => Types::non_null( Types::boolean() ),
-			'description' => __( 'Whether the JWT User secret has been revoked. If the secret has been revoked, auth tokens will not be issued until an admin, or user with proper capabilities re-issues a secret for the user.', 'wp-graphql-jwt-authentication' ),
-			'resolve'     => function ( User $user ) {
-				$revoked = Auth::is_jwt_secret_revoked( $user->ID );
-
-				return true == $revoked ? true : false;
-			}
-		];
-
+	public static function add_jwt_fields() {
+		$types = apply_filters( 'graphql_jwt_user_types', [ 'User' ] );
+		foreach ( $types as $type ) {
+			self::register_jwt_fields_to( $type );
+		}
+	}
 
-		return $fields;
+	/**
+	 * Adds the JWT fields to the provided type.
+	 *
+	 * @param string $type  Type for the fields to be registered to.
+	 *
+	 * @throws UserError  Invalid token/Token not found.
+	 */
+	public static function register_jwt_fields_to( $type ) {
+		register_graphql_fields(
+			$type,
+			[
+				'jwtAuthToken'           => [
+					'type'        => 'String',
+					'description' => __( 'A JWT token that can be used in future requests for authentication/authorization', 'wp-graphql-jwt-authentication' ),
+					'resolve'     => function ( $user ) {
+						$user = get_user_by( 'id', $user->ID );
+
+						// Get the token for the user.
+						$token = Auth::get_token( $user );
+
+						// If the token cannot be returned, throw an error.
+						if ( empty( $token ) || is_wp_error( $token ) ) {
+							throw new UserError( __( 'The JWT token could not be returned', 'wp-graphql-jwt-authentication' ) );
+						}
+
+						return ! empty( $token ) ? $token : null;
+					},
+				],
+				'jwtRefreshToken'        => [
+					'type'        => 'String',
+					'description' => __( 'A JWT token that can be used in future requests to get a refreshed jwtAuthToken. If the refresh token used in a request is revoked or otherwise invalid, a valid Auth token will NOT be issued in the response headers.', 'wp-graphql-jwt-authentication' ),
+					'resolve'     => function ( $user ) {
+						$user = get_user_by( 'id', $user->ID );
+
+						// Get the token for the user.
+						$token = Auth::get_refresh_token( $user );
+
+						// If the token cannot be returned, throw an error.
+						if ( empty( $token ) || is_wp_error( $token ) ) {
+							throw new UserError( __( 'The JWT token could not be returned', 'wp-graphql-jwt-authentication' ) );
+						}
+
+						return ! empty( $token ) ? $token : null;
+					},
+				],
+				'jwtUserSecret'          => [
+					'type'        => 'String',
+					'description' => __( 'A unique secret tied to the users JWT token that can be revoked or refreshed. Revoking the secret prevents JWT tokens from being issued to the user. Refreshing the token invalidates previously issued tokens, but allows new tokens to be issued.', 'wp-graphql' ),
+					'resolve'     => function ( $user ) {
+						// Get the user's JWT Secret.
+						$secret = Auth::get_user_jwt_secret( $user->ID );
+
+						// If the secret cannot be returned, throw an error.
+						if ( is_wp_error( $secret ) ) {
+							throw new UserError( __( 'The user secret could not be returned', 'wp-graphql-jwt-authentication' ) );
+						}
+
+						// Return the secret.
+						return ! empty( $secret ) ? $secret : null;
+					},
+				],
+				'jwtAuthExpiration'      => [
+					'type'        => 'String',
+					'description' => __( 'The expiration for the JWT Token for the user. If not set custom for the user, it will use the default sitewide expiration setting', 'wp-graphql-jwt-authentication' ),
+					'resolve'     => function () {
+						$expiration = Auth::get_token_expiration();
+
+						return ! empty( $expiration ) ? $expiration : null;
+					},
+				],
+				'isJwtAuthSecretRevoked' => [
+					'type'        => [ 'non_null' => 'Boolean' ],
+					'description' => __( 'Whether the JWT User secret has been revoked. If the secret has been revoked, auth tokens will not be issued until an admin, or user with proper capabilities re-issues a secret for the user.', 'wp-graphql-jwt-authentication' ),
+					'resolve'     => function ( $user ) {
+						$revoked = Auth::is_jwt_secret_revoked( $user->ID );
+
+						return true === $revoked ? true : false;
+					},
+				],
+			]
+		);
 
 	}
 
 	/**
 	 * Given an array of fields, this returns an array with the new fields added
 	 *
-	 * @param array $fields The input fields for user mutations
+	 * @param array $fields The input fields for user mutations.
 	 *
 	 * @return array
 	 */
 	public static function add_user_mutation_input_fields( array $fields ) {
-
-		$fields['revokeJwtUserSecret'] = [
+		$fields['revokeJwtUserSecret']  = [
 			'type'        => Types::boolean(),
 			'description' => __( 'If true, this will revoke the users JWT secret. If false, this will unrevoke the JWT secret AND issue a new one. To revoke, the user must have proper capabilities to edit users JWT secrets.', 'wp-graphql-jwt-authentication' ),
 		];
-
 		$fields['refreshJwtUserSecret'] = [
 			'type'        => Types::boolean(),
 			'description' => __( 'If true, this will refresh the users JWT secret.' ),
@@ -201,12 +168,13 @@ public static function add_user_mutation_input_fields( array $fields ) {
 	}
 
 	/**
-	 * @param int    $user_id       The ID of the user being mutated
-	 * @param array  $input         The input args of the GraphQL mutation request
-	 * @param string $mutation_name The name of the mutation
+	 * Processes JWT Authentication-related parameters in User mutations.
+	 *
+	 * @param int    $user_id       The ID of the user being mutated.
+	 * @param array  $input         The input args of the GraphQL mutation request.
+	 * @param string $mutation_name The name of the mutation.
 	 */
 	public static function update_jwt_fields_during_mutation( $user_id, array $input, $mutation_name ) {
-
 		/**
 		 * If there was input to revokeJwtUserSecret, check the value for true or false, and
 		 * revoke or unRevoke the token accordingly
@@ -219,9 +187,7 @@ public static function update_jwt_fields_during_mutation( $user_id, array $input
 			}
 		}
 
-		/**
-		 * If refreshJwtUserSecret is true.
-		 */
+		// If refreshJwtUserSecret is true.
 		if ( isset( $input['refreshJwtUserSecret'] ) ) {
 			if ( true === $input['refreshJwtUserSecret'] ) {
 				Auth::issue_new_user_secret( $user_id );
@@ -233,21 +199,18 @@ public static function update_jwt_fields_during_mutation( $user_id, array $input
 	/**
 	 * This filters the token to prevent it from being issued if it has been revoked.
 	 *
-	 * @param string $token
-	 * @param int    $user_id
+	 * @param string $token    Token.
+	 * @param int    $user_id  User associated with token.
 	 *
+	 * @throws UserError  Revoked token.
 	 * @return string $token
 	 */
 	public static function prevent_token_from_returning_if_revoked( $token, $user_id ) {
 
-		/**
-		 * Check to see if the user's auth secret has been revoked.
-		 */
+		// Check to see if the user's auth secret has been revoked.
 		$revoked = Auth::is_jwt_secret_revoked( $user_id );
 
-		/**
-		 * If the token has been revoked, prevent it from being returned
-		 */
+		// If the token has been revoked, prevent it from being returned.
 		if ( true === $revoked ) {
 			throw new UserError( __( 'The JWT token cannot be issued for this user', 'wp-graphql-jwt-authentication' ) );
 		}
@@ -259,13 +222,11 @@ public static function prevent_token_from_returning_if_revoked( $token, $user_id
 	/**
 	 * Returns tokens in the response headers
 	 *
-	 * @param $headers
+	 * @param array $headers  GraphQL HTTP response headers.
 	 *
-	 * @return mixed
-	 * @throws \Exception
+	 * @return array
 	 */
 	public static function add_tokens_to_graphql_response_headers( $headers ) {
-
 		/**
 		 * If the request _is_ SSL, or GRAPHQL_DEBUG is defined, return the tokens
 		 * otherwise do not return them.
@@ -274,28 +235,19 @@ public static function add_tokens_to_graphql_response_headers( $headers ) {
 			return $headers;
 		}
 
-		/**
-		 * If there's a Refresh-Authorization token in the request headers, validate it
-		 */
+		// If there's a Refresh-Authorization token in the request headers, validate it.
 		$validate_refresh_header = Auth::validate_token( Auth::get_refresh_header(), true );
 
-		/**
-		 * If the refresh token in the request headers is valid, return a JWT Auth token that can be used for future requests
-		 */
+		// If the refresh token in the request headers is valid, return a JWT Auth token that can be used for future requests.
 		if ( ! is_wp_error( $validate_refresh_header ) && ! empty( $validate_refresh_header->data->user->id ) ) {
 
-			/**
-			 * Get an auth token and refresh token to return
-			 */
+			// Get an auth token and refresh token to return.
 			$auth_token = Auth::get_token( new \WP_User( $validate_refresh_header->data->user->id ), false );
 
-			/**
-			 * If the tokens can be generated (not revoked, etc), return them
-			 */
+			// If the tokens can be generated (not revoked, etc), return them.
 			if ( ! empty( $auth_token ) && ! is_wp_error( $auth_token ) ) {
 				$headers['X-JWT-Auth'] = $auth_token;
 			}
-
 		}
 
 		$validate_auth_header = Auth::validate_token( null, false );
@@ -307,11 +259,9 @@ public static function add_tokens_to_graphql_response_headers( $headers ) {
 			if ( ! empty( $refresh_token ) && ! is_wp_error( $refresh_token ) ) {
 				$headers['X-JWT-Refresh'] = $refresh_token;
 			}
-
 		}
 
 		return $headers;
-
 	}
 
 	/**
@@ -320,12 +270,12 @@ public static function add_tokens_to_graphql_response_headers( $headers ) {
 	 * This allows clients the ability to Authenticate with WPGraphQL, use the token
 	 * with REST API Requests, but get new refresh tokens from the REST API Headers
 	 *
+	 * @param WP_HTTP_Response $response  Response object.
+	 *
 	 * @return \WP_HTTP_Response
-	 * @throws \Exception
 	 */
-	public static function add_auth_headers_to_rest_response( $response, $handler, $request ) {
-
-		if( ! $response instanceof \WP_HTTP_Response ) {
+	public static function add_auth_headers_to_rest_response( $response ) {
+		if ( ! $response instanceof \WP_HTTP_Response ) {
 			return $response;
 		}
 
@@ -345,9 +295,9 @@ public static function add_auth_headers_to_rest_response( $response, $handler, $
 		 *
 		 * Might need a patch to core to allow for individual filtering.
 		 */
-		$response->set_headers( [
-			'Access-Control-Expose-Headers' => 'X-WP-Total, X-WP-TotalPages, X-JWT-Refresh',
-		] );
+		$response->set_headers(
+			[ 'Access-Control-Expose-Headers' => 'X-WP-Total, X-WP-TotalPages, X-JWT-Refresh' ]
+		);
 
 		$refresh_token = null;
 
@@ -360,9 +310,7 @@ public static function add_auth_headers_to_rest_response( $response, $handler, $
 		}
 
 		if ( $refresh_token ) {
-			$response->set_headers( [
-				'X-JWT-Refresh' => $refresh_token,
-			] );
+			$response->set_headers( [ 'X-JWT-Refresh' => $refresh_token ] );
 		}
 
 		return $response;
@@ -372,7 +320,7 @@ public static function add_auth_headers_to_rest_response( $response, $handler, $
 	 * Expose the X-JWT-Refresh tokens in the response headers. This allows
 	 * folks to grab new refresh tokens from authenticated requests for subsequent use.
 	 *
-	 * @param array $headers The existing response headers
+	 * @param array $headers The existing response headers.
 	 *
 	 * @return array
 	 */
@@ -385,7 +333,7 @@ public static function add_auth_headers_to_response( array $headers ) {
 	/**
 	 * Expose the X-JWT-Auth and X-JWT-Refresh as allowed headers in GraphQL responses
 	 *
-	 * @param array $allowed_headers The existing allowed headers
+	 * @param array $allowed_headers The existing allowed headers.
 	 *
 	 * @return array
 	 */
@@ -395,5 +343,4 @@ public static function add_jwt_allowed_headers( array $allowed_headers ) {
 
 		return $allowed_headers;
 	}
-
 }
diff --git a/src/RefreshToken.php b/src/RefreshToken.php
index d2de17f..fecf3c0 100644
--- a/src/RefreshToken.php
+++ b/src/RefreshToken.php
@@ -1,50 +1,40 @@
 <?php
+/**
+ * Registers the "refreshToken" mutation to the WPGraphQL Schema.
+ *
+ * @package WPGraphQL\JWT_Authentication
+ * @since 0.0.1
+ */
+
 namespace WPGraphQL\JWT_Authentication;
 
 use GraphQL\Error\UserError;
-use GraphQLRelay\Relay;
-use WPGraphQL\Types;
 
+/**
+ * Class - RefreshToken
+ */
 class RefreshToken {
-
-	private static $mutation;
-
 	/**
-	 * Takes an array of fields from the RootMutation and returns the fields
-	 * with the "login" mutation field added
-	 *
-	 * @param array $fields The fields in the RootMutation of the Schema
-	 *
-	 * @return array $fields
+	 * Registers the mutation.
 	 */
-	public static function root_mutation_fields( $fields ) {
-		$fields['refreshJwtAuthToken'] = self::mutation();
-
-		return $fields;
-	}
-
-	public static function mutation() {
-
-		if ( empty( self::$mutation ) ) {
-
-			self::$mutation = Relay::mutationWithClientMutationId([
-				'name' => 'RefreshJwtAuthToken',
-				'isPrivate' => false,
-				'description' => __( 'Use a valid JWT Refresh token to retrieve a new JWT Auth Token', 'wp-graphql-jwt-authentication' ),
-				'inputFields' => [
+	public static function register_mutation() {
+		register_graphql_mutation(
+			'refreshJwtAuthToken',
+			[
+				'description'         => __( 'Use a valid JWT Refresh token to retrieve a new JWT Auth Token', 'wp-graphql-jwt-authentication' ),
+				'inputFields'         => [
 					'jwtRefreshToken' => [
-						'type' => Types::non_null( Types::string() ),
+						'type'        => [ 'non_null' => 'String' ],
 						'description' => __( 'A valid, previously issued JWT refresh token. If valid a new Auth token will be provided. If invalid, expired, revoked or otherwise invalid, a new AuthToken will not be provided.', 'wp-graphql-jwt-authentication' ),
 					],
 				],
-				'outputFields' => [
+				'outputFields'        => [
 					'authToken' => [
-						'type' => Types::string(),
+						'type'        => 'String',
 						'description' => __( 'JWT Token that can be used in future requests for Authentication', 'wp-graphql-jwt-authentication' ),
 					],
 				],
 				'mutateAndGetPayload' => function( $input ) {
-
 					$refresh_token = ! empty( $input['jwtRefreshToken'] ) ? Auth::validate_token( $input['jwtRefreshToken'] ) : null;
 
 					$id = isset( $refresh_token->data->user->id ) || 0 === $refresh_token->data->user->id ? absint( $refresh_token->data->user->id ) : 0;
@@ -55,17 +45,9 @@ public static function mutation() {
 					$user = new \WP_User( $id );
 					$auth_token = Auth::get_token( $user, false );
 
-					return [
-						'authToken' => $auth_token,
-					];
-
+					return [ 'authToken' => $auth_token ];
 				},
-			]);
-
-		}
-
-		return ( ! empty( self::$mutation ) ) ? self::$mutation : null;
-
+			]
+		);
 	}
-
 }
diff --git a/wp-graphql-jwt-authentication.php b/wp-graphql-jwt-authentication.php
index 82a71c4..1a4b1c3 100644
--- a/wp-graphql-jwt-authentication.php
+++ b/wp-graphql-jwt-authentication.php
@@ -35,8 +35,10 @@
 
 if ( ! class_exists( '\WPGraphQL\JWT_Authentication' ) ) :
 
+	/**
+	 * Class - JWT_Authentication
+	 */
 	final class JWT_Authentication {
-
 		/**
 		 * Stores the instance of the JWT_Authentication class
 		 *
@@ -54,7 +56,6 @@ final class JWT_Authentication {
 		 * @access public
 		 */
 		public static function instance() {
-
 			if ( ! isset( self::$instance ) && ! ( self::$instance instanceof JWT_Authentication ) ) {
 				self::$instance = new JWT_Authentication;
 				self::$instance->setup_constants();
@@ -86,10 +87,8 @@ public static function instance() {
 		 * @return void
 		 */
 		public function __clone() {
-
 			// Cloning instances of the class is forbidden.
 			_doing_it_wrong( __FUNCTION__, esc_html__( 'The Init_JWT_Authentication class should not be cloned.', 'wp-graphql-jwt-authentication' ), '0.0.1' );
-
 		}
 
 		/**
@@ -100,10 +99,8 @@ public function __clone() {
 		 * @return void
 		 */
 		public function __wakeup() {
-
 			// De-serializing instances of the class is forbidden.
 			_doing_it_wrong( __FUNCTION__, esc_html__( 'De-serializing instances of the WPGraphQL class is not allowed', 'wp-graphql-jwt-authentication' ), '0.0.1' );
-
 		}
 
 		/**
@@ -114,7 +111,6 @@ public function __wakeup() {
 		 * @return void
 		 */
 		private function setup_constants() {
-
 			// Plugin version.
 			if ( ! defined( 'WPGRAPHQL_JWT_AUTHENTICATION_VERSION' ) ) {
 				define( 'WPGRAPHQL_JWT_AUTHENTICATION_VERSION', '0.3.4' );
@@ -135,11 +131,10 @@ private function setup_constants() {
 				define( 'WPGRAPHQL_JWT_AUTHENTICATION_PLUGIN_FILE', __FILE__ );
 			}
 
-			// Whether to autoload the files or not
+			// Whether to autoload the files or not.
 			if ( ! defined( 'WPGRAPHQL_JWT_AUTHENTICATION_AUTOLOAD' ) ) {
 				define( 'WPGRAPHQL_JWT_AUTHENTICATION_AUTOLOAD', true );
 			}
-
 		}
 
 		/**
@@ -151,9 +146,8 @@ private function setup_constants() {
 		 * @return void
 		 */
 		private function includes() {
-
-			// Autoload Required Classes
-			if ( defined( 'WPGRAPHQL_JWT_AUTHENTICATION_AUTOLOAD' ) && true == WPGRAPHQL_JWT_AUTHENTICATION_AUTOLOAD ) {
+			// Autoload Required Classes.
+			if ( defined( 'WPGRAPHQL_JWT_AUTHENTICATION_AUTOLOAD' ) && true === WPGRAPHQL_JWT_AUTHENTICATION_AUTOLOAD ) {
 				require_once( WPGRAPHQL_JWT_AUTHENTICATION_PLUGIN_DIR . 'vendor/autoload.php' );
 			}
 		}
@@ -162,43 +156,38 @@ private function includes() {
 		 * Initialize the plugin
 		 */
 		private static function init() {
-
-			/**
-			 * Initialize the GraphQL fields for managing tokens
-			 */
+			// Initialize the GraphQL fields for managing tokens.
 			ManageTokens::init();
 
-			/**
-			 * Filter how WordPress determines the current user
-			 */
-			add_filter( 'determine_current_user', [
-				'\WPGraphQL\JWT_Authentication\Auth',
-				'filter_determine_current_user'
-			], 99, 1 );
-
-			/**
-			 * Register the login mutation to the Schema
-			 */
-			add_action( 'graphql_register_types', [
-				'\WPGraphQL\JWT_Authentication\Login',
-				'register_mutation'
-			], 10 );
-
-			add_filter( 'graphql_rootMutation_fields', [
-				'\WPGraphQL\JWT_Authentication\RefreshToken',
-				'root_mutation_fields'
-			], 10, 1 );
-
-
-
+			// Filter how WordPress determines the current user.
+			add_filter(
+				'determine_current_user',
+				[ '\WPGraphQL\JWT_Authentication\Auth', 'filter_determine_current_user' ],
+				99
+			);
+
+			// Register the "login" mutation to the Schema.
+			add_action(
+				'graphql_register_types',
+				[ '\WPGraphQL\JWT_Authentication\Login', 'register_mutation' ],
+				10
+			);
+
+			// Register the "refreshToken" mutation to the Schema.
+			add_filter(
+				'graphql_register_types',
+				[ '\WPGraphQL\JWT_Authentication\RefreshToken', 'register_mutation' ],
+				10
+			);
 		}
-
 	}
 
 endif;
 
+/**
+ * Start JWT_Authentication.
+ */
 function init() {
 	return JWT_Authentication::instance();
 }
-
 add_action( 'plugins_loaded', '\WPGraphQL\JWT_Authentication\init', 1 );