Update/optimize title to ability

projects/plugins/jetpack/changelog/add-title-optimization-ability-path

@@ -0,0 +1,4 @@
+Significance: patch
+Type: enhancement
+
+AI Assistant: Route title optimization UI directly through the wp-orchestrator agent endpoint, removing the custom REST proxy endpoint.

projects/plugins/jetpack/extensions/blocks/ai-assistant/ai-assistant.php

@@ -18,6 +18,9 @@
 	exit( 0 );
 }
 
+require_once __DIR__ . '/class-title-optimization-ability.php';
+Title_Optimization_Ability::init();
+
 /**
  * Registers our block for use in Gutenberg
  * This is done via an action so that we can disable

projects/plugins/jetpack/extensions/blocks/ai-assistant/class-title-optimization-ability.php

@@ -0,0 +1,184 @@
+<?php
+/**
+ * Title optimization ability registration and execution.
+ *
+ * The UI calls the wp-orchestrator endpoint directly (POST /wpcom/v2/ai/agent)
+ * with agent_id "wp-orchestrator" and ability "wpcom/optimize-title".  This
+ * class registers the ability with the WordPress Abilities API so the
+ * orchestrator can discover it, and provides an execute callback for any
+ * server-side (non-orchestrator) invocation.
+ *
+ * @package automattic/jetpack
+ */
+
+// @phan-file-suppress PhanUndeclaredFunction, PhanUndeclaredClassMethod @phan-suppress-current-line UnusedSuppression -- Ability API added in WP 6.9, but we keep compatibility with older WP in CI.
+
+namespace Automattic\Jetpack\Extensions\AIAssistant;
+
+use WP_Error;
+
+/**
+ * Registers and executes the title optimization ability.
+ */
+class Title_Optimization_Ability {
+	/**
+	 * Ability category slug.
+	 *
+	 * @var string
+	 */
+	const CATEGORY_SLUG = 'jetpack-ai';
+
+	/**
+	 * Ability name.
+	 *
+	 * @var string
+	 */
+	const ABILITY_NAME = 'wpcom/optimize-title';
+
+	/**
+	 * Feature name used by the orchestrator.
+	 *
+	 * @var string
+	 */
+	const FEATURE_NAME = 'jetpack-ai-title-optimization';
+
+	/**
+	 * Initialize ability registration.
+	 *
+	 * @return void
+	 */
+	public static function init() {
+		if ( did_action( 'wp_abilities_api_categories_init' ) ) {
+			self::register_category();
+		} else {
+			add_action( 'wp_abilities_api_categories_init', array( __CLASS__, 'register_category' ) );
+		}
+
+		if ( did_action( 'wp_abilities_api_init' ) ) {
+			self::register_ability();
+		} else {
+			add_action( 'wp_abilities_api_init', array( __CLASS__, 'register_ability' ) );
+		}
+	}
+
+	/**
+	 * Register the ability category.
+	 *
+	 * @return void
+	 */
+	public static function register_category() {
+		if ( ! function_exists( 'wp_register_ability_category' ) ) {
+			return;
+		}
+
+		wp_register_ability_category(
+			self::CATEGORY_SLUG,
+			array(
+				// "Jetpack AI" is a product name and should not be translated.
+				'label'       => 'Jetpack AI',
+				'description' => __( 'Abilities for Jetpack AI features.', 'jetpack' ),
+			)
+		);
+	}
+
+	/**
+	 * Register the title optimization ability.
+	 *
+	 * @return void
+	 */
+	public static function register_ability() {
+		if ( ! function_exists( 'wp_register_ability' ) ) {
+			return;
+		}
+
+		wp_register_ability(
+			self::ABILITY_NAME,
+			array(
+				'label'               => __( 'Optimize title', 'jetpack' ),
+				'description'         => __( 'Generate optimized title options based on post content and optional keywords.', 'jetpack' ),
+				'category'            => self::CATEGORY_SLUG,
+				'input_schema'        => array(
+					'type'                 => 'object',
+					'required'             => array( 'content' ),
+					'properties'           => array(
+						'content'  => array(
+							'type'        => 'string',
+							'description' => __( 'The post content used to generate title suggestions.', 'jetpack' ),
+						),
+						'keywords' => array(
+							'type'        => 'string',
+							'description' => __( 'Optional keywords to include in title optimization.', 'jetpack' ),
+						),
+						'post_id'  => array(
+							'type'        => 'integer',
+							'description' => __( 'Optional post ID associated with the request.', 'jetpack' ),
+						),
+					),
+					'additionalProperties' => false,
+				),
+				'execute_callback'    => array( __CLASS__, 'execute' ),
+				'permission_callback' => array( __CLASS__, 'permission_callback' ),
+				'meta'                => array(
+					'annotations'  => array(
+						'readonly'    => true,
+						'destructive' => false,
+						'idempotent'  => false,
+					),
+					'show_in_rest' => true,
+				),
+			)
+		);
+	}
+
+	/**
+	 * Permission check callback for ability execution.
+	 *
+	 * @return bool
+	 */
+	public static function permission_callback() {
+		if ( ! class_exists( 'Jetpack_AI_Helper' ) ) {
+			require_once JETPACK__PLUGIN_DIR . '_inc/lib/class-jetpack-ai-helper.php';
+		}
+
+		return \Jetpack_AI_Helper::get_status_permission_check();
+	}
+
+	/**
+	 * Execute the title optimization ability.
+	 *
+	 * This is the execute callback invoked when the ability is called
+	 * server-side (e.g. via the WordPress Abilities API).  It validates
+	 * and sanitises the input, then returns the structured message array
+	 * that the orchestrator uses to generate title suggestions.
+	 *
+	 * @param array $input Ability input.
+	 * @return array|WP_Error
+	 */
+	public static function execute( $input ) {
+		$content = isset( $input['content'] ) ? sanitize_textarea_field( $input['content'] ) : '';
+		if ( '' === trim( $content ) ) {
+			return new WP_Error(
+				'jetpack_ai_title_optimization_invalid_content',
+				__( 'Content is required to optimize a title.', 'jetpack' )
+			);
+		}
+
+		$keywords = isset( $input['keywords'] ) ? sanitize_text_field( $input['keywords'] ) : '';
+		$post_id  = isset( $input['post_id'] ) ? absint( $input['post_id'] ) : 0;
+
+		$result = array(
+			'role'    => 'jetpack-ai',
+			'context' => array(
+				'type'     => 'title-optimization',
+				'content'  => $content,
+				'keywords' => $keywords,
+			),
+		);
+
+		if ( $post_id > 0 ) {
+			$result['context']['post_id'] = $post_id;
+		}
+
+		return $result;
+	}
+}

projects/plugins/jetpack/extensions/plugins/ai-assistant-plugin/components/title-optimization/index.tsx

@@ -2,7 +2,7 @@
  * External dependencies
  */
 import {
-	useAiSuggestions,
+	ERROR_CONTEXT_TOO_LARGE,
 	RequestingErrorProps,
 	ERROR_QUOTA_EXCEEDED,
 	ERROR_NETWORK,
@@ -11,10 +11,11 @@
 	QuotaExceededMessage,
 	usePostContent,
 	AiAssistantModal,
+	requestJwt,
 } from '@automattic/jetpack-ai-client';
 import { useAnalytics, useAutosaveAndRedirect } from '@automattic/jetpack-shared-extension-utils';
 import { Button, Spinner, ExternalLink, Notice } from '@wordpress/components';
-import { useDispatch } from '@wordpress/data';
+import { useDispatch, useSelect } from '@wordpress/data';
 import { useState, useCallback } from '@wordpress/element';
 import { __ } from '@wordpress/i18n';
 /**
@@ -41,13 +42,62 @@
 );
 
 const ERROR_JSON_PARSE = 'json-parse-error';
+
+/**
+ * wp-orchestrator constants matching the ability registered in
+ * class-title-optimization-ability.php.
+ */
+const ORCHESTRATOR_URL = 'https://public-api.wordpress.com/wpcom/v2/ai/agent';
+const ORCHESTRATOR_AGENT_ID = 'wp-orchestrator';
+const ABILITY_NAME = 'wpcom/optimize-title';
+const FEATURE_NAME = 'jetpack-ai-title-optimization';
 type TitleOptimizationJSONError = {
 	code: typeof ERROR_JSON_PARSE;
 	message: string;
 };
 
 type TitleOptimizationError = RequestingErrorProps | TitleOptimizationJSONError;
 
+function getErrorFromApiResponse( status?: number ): RequestingErrorProps {
+	if ( status === 429 ) {
+		return {
+			code: ERROR_QUOTA_EXCEEDED,
+			message: '',
+			severity: 'info',
+		};
+	}
+
+	if ( status === 503 ) {
+		return {
+			code: ERROR_SERVICE_UNAVAILABLE,
+			message: genericErrorMessage,
+			severity: 'info',
+		};
+	}
+
+	if ( status === 413 ) {
+		return {
+			code: ERROR_CONTEXT_TOO_LARGE,
+			message: genericErrorMessage,
+			severity: 'info',
+		};
+	}
+
+	if ( status === 422 ) {
+		return {
+			code: ERROR_UNCLEAR_PROMPT,
+			message: genericErrorMessage,
+			severity: 'info',
+		};
+	}
+
+	return {
+		code: ERROR_NETWORK,
+		message: genericErrorMessage,
+		severity: 'info',
+	};
+}
+
 const TitleOptimizationErrorMessage = ( { error }: { error: TitleOptimizationError } ) => {
 	if ( error.code === ERROR_QUOTA_EXCEEDED ) {
 		return (
@@ -105,6 +155,7 @@
 	const [ error, setError ] = useState< TitleOptimizationError | null >( null );
 	const [ optimizationKeywords, setOptimizationKeywords ] = useState( '' );
 	const { editPost } = useDispatch( 'core/editor' );
+	const postId = useSelect( select => select( 'core/editor' ).getCurrentPostId(), [] );
 	const { autosave } = useAutosaveAndRedirect();
 	const { increaseAiAssistantRequestsCount } = useDispatch( 'wordpress-com/plans' );
 	const { tracks } = useAnalytics();
@@ -134,16 +185,8 @@
 		[ increaseAiAssistantRequestsCount ]
 	);
 
-	const { request, stopSuggestion } = useAiSuggestions( {
-		onDone: handleDone,
-		onError: ( e: RequestingErrorProps ) => {
-			setError( e );
-			setGenerating( false );
-		},
-	} );
-
 	const handleRequest = useCallback(
-		( isRetry: boolean = false ) => {
+		async ( isRetry: boolean = false ) => {
 			// track the generate title optimization options
 			recordEvent( 'jetpack_ai_title_optimization_generate', {
 				placement,
@@ -152,21 +195,59 @@
 			} );
 
 			setGenerating( true );
-			// Message to request a backend prompt for this feature
-			const messages = [
-				{
-					role: 'jetpack-ai' as const,
-					context: {
-						type: 'title-optimization',
-						content: getPostContent(),
-						keywords: optimizationKeywords,
+			setError( null );
+			try {
+				const { token, blogId } = await requestJwt();
+
+				const content = getPostContent();
+				const requestBody: Record< string, unknown > = {
+					agent_id: ORCHESTRATOR_AGENT_ID,
+					ability: ABILITY_NAME,
+					feature: FEATURE_NAME,
+					stream: false,
+					site_id: Number( blogId ),
+					input: {
+						messages: [
+							{
+								role: 'jetpack-ai',
+								context: {
+									type: 'title-optimization',
+									content,
+									keywords: optimizationKeywords,
+								},
+							},
+						],
+						...( postId ? { post_id: postId } : {} ),
+					},
+				};
+
+				const res = await fetch( ORCHESTRATOR_URL, {
+					method: 'POST',
+					headers: {
+						Authorization: `Bearer ${ token }`,
+						'Content-Type': 'application/json',
 					},
-				},
-			];
+					body: JSON.stringify( requestBody ),
+				} );
+
+				if ( ! res.ok ) {
+					setError( getErrorFromApiResponse( res.status ) );
+					setGenerating( false );
+					return;
+				}
 
-			request( messages, { feature: 'jetpack-ai-title-optimization' } );
+				const data = ( await res.json() ) as {
+					choices?: Array< { message?: { content?: string } } >;
+				};
+
+				handleDone( data?.choices?.[ 0 ]?.message?.content ?? '' );
+			} catch ( e ) {
+				const requestError = e as { status?: number; data?: { status?: number } };
+				setError( getErrorFromApiResponse( requestError?.data?.status || requestError?.status ) );
+				setGenerating( false );
+			}
 		},
-		[ recordEvent, placement, getPostContent, optimizationKeywords, request ]
+		[ recordEvent, placement, getPostContent, optimizationKeywords, postId, handleDone ]
 	);
 
 	const handleTitleOptimization = useCallback( () => {
@@ -214,8 +295,7 @@
 		setError( null );
 		toggleTitleOptimizationModal();
 		setOptimizationKeywords( '' );
-		stopSuggestion();
-	}, [ stopSuggestion, toggleTitleOptimizationModal ] );
+	}, [ toggleTitleOptimizationModal ] );
 
 	// When can we retry?
 	const showTryAgainButton =