Refactor api limit reached block logic.

Author: rahulsprajapati

includes/Classifai/Features/APIUsageTracking.php

@@ -156,8 +156,48 @@ public function feature_setup() {
 
 		add_action( self::FORCE_CRON_HOOK, [ $this, 'run_usage_force_refresh' ] );
 		add_action( self::CRON_HOOK, [ $this, 'run_usage_refresh' ] );
+
+		add_filter( 'classifai_pre_fetch_feature_response', [ $this, 'pre_fetch_feature_response' ], 10, 2 );
 	}
 
+	/**
+	 * Pre-fetch the feature response to check if the feature is allowed to run based on the usage tracking data.
+	 *
+	 * @param mixed                         $response Response to return.
+	 * @param \Classifai\Providers\Provider $provider_instance provider used.
+	 *
+	 * @return mixed Response to return.
+	 */
+	public function pre_fetch_feature_response( $response, $provider_instance ) {
+		$usage_tracking_provider = $this->get_feature_provider_instance();
+		$provider_ids            = [];
+
+		if ( ! empty( $usage_tracking_provider ) && $usage_tracking_provider instanceof UsageTrackingProvider ) {
+			$provider_ids = $usage_tracking_provider->get_provider_ids();
+		}
+
+		if (
+			empty( $provider_ids )
+			|| empty( $provider_instance::ID )
+			|| ! in_array( $provider_instance::ID, $provider_ids, true )
+		) {
+			return $response;
+		}
+
+		$limit_reached = get_option( self::HARD_LIMIT_REACHED_KEY, false );
+
+		if ( $limit_reached ) {
+			return new WP_Error(
+				'classifai_hard_limit_reached',
+				__( 'Usage has reached the configured hard limit. Re-enable in ClassifAI / Pricing / Usage settings.', 'classifai' ),
+				[
+					'status' => 403,
+				]
+			);
+		}
+
+		return $response;
+	}
 
 	/**
 	 * Register any needed endpoints.

includes/Classifai/Features/Feature.php

@@ -1399,6 +1399,26 @@ public function run( ...$args ) {
 		$provider_id       = $settings['provider'];
 		$provider_instance = $this->get_feature_provider_instance( $provider_id );
 
+		/**
+		 * Filter the response before the feature api endpoint is called.
+		 *
+		 * @since x.x.x
+		 *
+		 * @hook classifai_pre_fetch_feature_response
+		 *
+		 * @param mixed                         $response Response to return.
+		 * @param \Classifai\Providers\Provider $provider_instance provider used.
+		 * @param mixed                         $args Arguments used by the feature.
+		 * @param \Classifai\Features\Feature   $this Current feature class.
+		 *
+		 * @return mixed Response to return.
+		 */
+		$pre_fetch_response = apply_filters( 'classifai_pre_fetch_feature_response', null, $provider_instance, $args, $this );
+
+		if ( ! is_null( $pre_fetch_response ) ) {
+			return $pre_fetch_response;
+		}
+
 		if ( ! is_callable( [ $provider_instance, 'rest_endpoint_callback' ] ) ) {
 			return new WP_Error( 'invalid_route', esc_html__( 'The selected provider does not have a valid callback in place.', 'classifai' ) );
 		}

includes/Classifai/Providers/OpenAI/APIRequest.php

@@ -2,8 +2,6 @@
 
 namespace Classifai\Providers\OpenAI;
 
-use Classifai\Features\APIUsageTracking;
-use Classifai\Providers\OpenAI\UsageTracking;
 use Classifai\Providers\Provider;
 use WP_Error;
 
@@ -75,13 +73,6 @@ public function __construct( string $api_key = '', string $feature = '', ?Provid
 	 * @return array|WP_Error
 	 */
 	public function get( string $url, array $options = [] ) {
-		if ( ! $this->is_request_allowed() ) {
-			return new WP_Error(
-				'classifai_hard_limit',
-				__( 'Usage has reached the configured hard limit. Re-enable in ClassifAI / Pricing / Usage settings.', 'classifai' )
-			);
-		}
-
 		/**
 		 * Filter the URL for the get request.
 		 *
@@ -142,13 +133,6 @@ public function get( string $url, array $options = [] ) {
 	 * @return array|WP_Error
 	 */
 	public function post( string $url = '', array $options = [] ) {
-		if ( ! $this->is_request_allowed() ) {
-			return new WP_Error(
-				'classifai_hard_limit',
-				__( 'Usage has reached the configured hard limit. Re-enable in ClassifAI / Pricing / Usage settings.', 'classifai' )
-			);
-		}
-
 		$options = wp_parse_args(
 			$options,
 			[
@@ -216,13 +200,6 @@ public function post( string $url = '', array $options = [] ) {
 	 * @return array|WP_Error
 	 */
 	public function post_form( string $url = '', array $body = [] ) {
-		if ( ! $this->is_request_allowed() ) {
-			return new WP_Error(
-				'classifai_hard_limit',
-				__( 'Usage has reached the configured hard limit. Re-enable in ClassifAI / Pricing / Usage settings.', 'classifai' )
-			);
-		}
-
 		/**
 		 * Filter the URL for the post form request.
 		 *
@@ -391,35 +368,4 @@ public function get_api_key() {
 
 		return $this->api_key ?? '';
 	}
-
-	/**
-	 * Whether API requests are allowed (not blocked by hard usage limit).
-	 *
-	 * @return bool
-	 */
-	private function is_request_allowed(): bool {
-		if ( empty( $this->provider::ID ) || ! in_array( $this->provider::ID, UsageTracking::get_provider_ids(), true ) ) {
-			return true;
-		}
-
-		/**
-		 * Filter whether API requests are allowed.
-		 * When the hard usage limit is reached and not overridden, this is set to false.
-		 *
-		 * @since x.x.x
-		 * @hook classifai_can_make_request
-		 * @param bool $allowed Whether the request is allowed. Default true.
-		 * @param string $feature The feature name.
-		 * @param Provider $provider The provider instance.
-		 *
-		 * @return bool
-		 */
-		$allowed = (bool) apply_filters( 'classifai_can_make_request', true, $this->feature, $this->provider );
-
-		if ( ! $allowed ) {
-			return false;
-		}
-
-		return ! (bool) get_option( APIUsageTracking::HARD_LIMIT_REACHED_KEY, false );
-	}
 }