Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
41.67%
15 / 36
33.33%
1 / 3
 CRAP
0.00%
0 / 1
REST_AI
41.67%
15 / 36
33.33%
1 / 3
 16.73
0.00%
0 / 1
 __construct
91.67%
11 / 12
0.00%
0 / 1
 3.01
 is_rest_endpoint_registered
100.00%
4 / 4
100.00%
1 / 1
 1
 get_openai_jwt
0.00%
0 / 20
0.00%
0 / 1
 12
1<?php
2/**
3 * Sets up the AI REST API endpoints.
4 *
5 * @package automattic/my-jetpack
6 */
7
8namespace Automattic\Jetpack\My_Jetpack;
9
10use Automattic\Jetpack\Connection\Client;
11use Automattic\Jetpack\Connection\Manager as Connection_Manager;
12use Jetpack_Options;
13use WP_Error;
14
15/**
16 * Registers the REST routes for AI.
17 *
18 * @phan-constructor-used-for-side-effects
19 */
20class REST_AI {
21    /**
22     * Constructor.
23     */
24    public function __construct() {
25        /*
26         * Check if the `jetpack/v4/jetpack-ai-jwt` endpoint is registered
27         * by the Jetpack plugin to avoid registering it again.
28         * In case it's not registered, register it
29         * to make it available for Jetpack products that depend on it.
30         */
31        if ( ! self::is_rest_endpoint_registered( 'jetpack/v4', '/jetpack-ai-jwt' ) ) {
32            register_rest_route(
33                'jetpack/v4',
34                'jetpack-ai-jwt',
35                array(
36                    'methods'             => \WP_REST_Server::EDITABLE,
37                    'callback'            => __CLASS__ . '::get_openai_jwt',
38                    'permission_callback' => function () {
39                        return ( new Connection_Manager( 'jetpack' ) )->is_user_connected() && current_user_can( 'edit_posts' );
40                    },
41                )
42            );
43        }
44    }
45
46    /**
47     * Check if a specific REST endpoint is registered.
48     *
49     * @param string $namespace - The namespace of the endpoint.
50     * @param string $route     - The route of the endpoint.
51     * @return bool               True if the endpoint is registered, false otherwise.
52     */
53    public static function is_rest_endpoint_registered( $namespace, $route ) {
54        $server        = rest_get_server();
55        $routes        = $server->get_routes();
56        $full_endpoint = '/' . trim( $namespace, '/' ) . $route;
57        return isset( $routes[ $full_endpoint ] );
58    }
59
60    /**
61     * Ask WPCOM for a JWT token to use for OpenAI completion.
62     */
63    public static function get_openai_jwt() {
64        $blog_id = Jetpack_Options::get_option( 'id' );
65
66        $response = Client::wpcom_json_api_request_as_user(
67            "/sites/$blog_id/jetpack-openai-query/jwt",
68            '2',
69            array(
70                'method'  => 'POST',
71                'headers' => array( 'Content-Type' => 'application/json; charset=utf-8' ),
72            ),
73            wp_json_encode( array(), JSON_UNESCAPED_SLASHES | JSON_UNESCAPED_UNICODE ),
74            'wpcom'
75        );
76
77        if ( is_wp_error( $response ) ) {
78            return $response;
79        }
80
81        $json = json_decode( wp_remote_retrieve_body( $response ) );
82
83        if ( ! isset( $json->token ) ) {
84            return new WP_Error( 'no-token', 'No token returned from WPCOM' );
85        }
86
87        return array(
88            'token'   => $json->token,
89            'blog_id' => $blog_id,
90        );
91    }
92}