AI Inline Extensions: Support custom behavior (#37634)

* AI Inline Extensions: Support Custom Behavior

* AI Inline Extensions: Expose actions through Context

* changelog

* AI Inline Extensions: Mark list-item as children

* AI Inline Extensions: Update child block name and context render

projects/plugins/jetpack/changelog/add-ai-inline-extensions-custom-behavior

@@ -0,0 +1,4 @@
+Significance: minor
+Type: other
+
+Add support for custom behavior on AI Inline Extensions

projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/block-handler.tsx

@@ -23,15 +23,18 @@ export class BlockHandler {
 	public renderRules: RenderHTMLRules = [];
 	public firstUpdate: boolean = true;
 	public behavior: BlockBehavior = 'dropdown' as const;
+	public isChildBlock: boolean = false;
 
 	constructor(
 		clientId: string,
 		renderRules: RenderHTMLRules = [],
-		behavior: BlockBehavior = 'dropdown'
+		behavior: BlockBehavior = 'dropdown',
+		isChildBlock: boolean = false
 	) {
 		this.clientId = clientId;
 		this.renderRules = renderRules;
 		this.behavior = behavior;
+		this.isChildBlock = isChildBlock;
 	}
 
 	public getBlock(): Block {

projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-toolbar-dropdown/index.tsx

@@ -4,13 +4,14 @@
 import { aiAssistantIcon } from '@automattic/jetpack-ai-client';
 import { useAnalytics } from '@automattic/jetpack-shared-extension-utils';
 import { ToolbarButton, Dropdown } from '@wordpress/components';
+import React, { useCallback, useContext } from '@wordpress/element';
 import { __ } from '@wordpress/i18n';
-import React, { useCallback } from 'react';
 /*
  * Internal dependencies
  */
 import AiAssistantToolbarDropdownContent from '../../../components/ai-assistant-toolbar-dropdown/dropdown-content';
 import useTransformToAssistant from '../../../hooks/use-transform-to-assistant';
+import { InlineExtensionsContext } from '../../get-block-handler';
 /*
  * Types
  */
@@ -114,6 +115,7 @@ export default function AiAssistantExtensionToolbarDropdown( {
 	onRequestSuggestion,
 }: AiAssistantExtensionToolbarDropdownProps ): ReactElement {
 	const { tracks } = useAnalytics();
+	const inlineExtensionsContext = useContext( InlineExtensionsContext );
 
 	const toggleHandler = useCallback(
 		( isOpen: boolean ) => {
@@ -153,6 +155,11 @@ export default function AiAssistantExtensionToolbarDropdown( {
 			} }
 			renderToggle={ ( { isOpen, onToggle } ) => {
 				const handleClick = () => {
+					if ( typeof behavior === 'function' ) {
+						behavior( { onToggle, onAskAiAssistant, context: inlineExtensionsContext } );
+						return;
+					}
+
 					switch ( behavior ) {
 						case 'action':
 							handleAskAiAssistant();

projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/get-block-handler.tsx

@@ -1,6 +1,7 @@
 /**
  * External dependencies
  */
+import { createContext } from '@wordpress/element';
 import debugFactory from 'debug';
 /**
  * Internal dependencies
@@ -25,6 +26,8 @@ const handlers = {
 	'core/list': ListHandler,
 };
 
+export const InlineExtensionsContext = createContext( {} );
+
 /**
  * Gets the block handler based on the block type.
  * The block handler is used to handle the request suggestions.
@@ -50,5 +53,6 @@ export function getBlockHandler(
 		onDone: handler.onDone.bind( handler ),
 		getContent: handler.getContent.bind( handler ),
 		behavior: handler.behavior,
+		isChildBlock: handler.isChildBlock,
 	};
 }

projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/list-item/index.tsx

@@ -6,5 +6,6 @@ import { BlockHandler } from '../block-handler';
 export class ListItemHandler extends BlockHandler {
 	constructor( clientId: string ) {
 		super( clientId, [ 'listItem' ] );
+		this.isChildBlock = true;
 	}
 }

projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/types.ts

@@ -9,13 +9,15 @@ import type { Block } from '@automattic/jetpack-ai-client';
 
 export type OnSuggestion = ( suggestion: string ) => void;
 
-export type BlockBehavior = 'dropdown' | 'action';
+type CustomBlockBehavior = ( { onToggle, onAskAiAssistant, context } ) => void;
+export type BlockBehavior = 'dropdown' | 'action' | CustomBlockBehavior;
 
 export interface IBlockHandler {
 	onSuggestion: OnSuggestion;
 	onDone: ( suggestion: string ) => void;
 	getContent: () => string;
 	behavior: BlockBehavior;
+	isChildBlock?: boolean;
 }
 
 export type BlockEditorSelect = {

projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/with-ai-extension.tsx

@@ -21,7 +21,7 @@ import useAutoScroll from '../hooks/use-auto-scroll';
 import { mapInternalPromptTypeToBackendPromptType } from '../lib/prompt/backend-prompt';
 import AiAssistantInput from './components/ai-assistant-input';
 import AiAssistantExtensionToolbarDropdown from './components/ai-assistant-toolbar-dropdown';
-import { getBlockHandler } from './get-block-handler';
+import { getBlockHandler, InlineExtensionsContext } from './get-block-handler';
 import { isPossibleToExtendBlock } from './lib/is-possible-to-extend-block';
 /*
  * Types
@@ -114,6 +114,7 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 			onDone: onBlockDone,
 			getContent,
 			behavior,
+			isChildBlock,
 		} = useMemo( () => getBlockHandler( blockName, clientId ), [ blockName, clientId ] );
 
 		// Called when the user clicks the "Ask AI Assistant" button.
@@ -431,7 +432,7 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 			};
 		}, [ adjustBlockPadding, clientId, controlObserver, id, showAiControl ] );
 
-		return (
+		const aiInlineExtensionContent = (
 			<>
 				<BlockEdit { ...props } />
 
@@ -462,6 +463,20 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 				</BlockControls>
 			</>
 		);
+
+		if ( isChildBlock ) {
+			return aiInlineExtensionContent;
+		}
+
+		const ProviderProps = {
+			value: { [ blockName ]: { handleAskAiAssistant, handleRequestSuggestion } },
+		};
+
+		return (
+			<InlineExtensionsContext.Provider { ...ProviderProps }>
+				{ aiInlineExtensionContent }
+			</InlineExtensionsContext.Provider>
+		);
 	};
 }, 'blockEditWithAiComponents' );