diff --git a/projects/plugins/jetpack/changelog/add-ai-inline-extensions-custom-behavior b/projects/plugins/jetpack/changelog/add-ai-inline-extensions-custom-behavior
new file mode 100644
index 0000000000000..29d83674ddacf
--- /dev/null
+++ b/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
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/block-handler.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/block-handler.tsx
index 417ef7f65e63d..2ecad93603ab6 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/block-handler.tsx
+++ b/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 {
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-toolbar-dropdown/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-toolbar-dropdown/index.tsx
index 033076b9f8afc..e9591e20e2840 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-toolbar-dropdown/index.tsx
+++ b/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();
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/get-block-handler.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/get-block-handler.tsx
index fcbf9caf0a907..bf0420dc1ca51 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/get-block-handler.tsx
+++ b/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,
};
}
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/list-item/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/list-item/index.tsx
index 4203e51033546..b04e42072794e 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/list-item/index.tsx
+++ b/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;
}
}
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/types.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/types.ts
index 6652b26a1d9cb..093b61071dd0e 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/types.ts
+++ b/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 = {
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/with-ai-extension.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/with-ai-extension.tsx
index 36cd3f77268ff..11abd3a99da8d 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/with-ai-extension.tsx
+++ b/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 = (
<>
@@ -462,6 +463,20 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
>
);
+
+ if ( isChildBlock ) {
+ return aiInlineExtensionContent;
+ }
+
+ const ProviderProps = {
+ value: { [ blockName ]: { handleAskAiAssistant, handleRequestSuggestion } },
+ };
+
+ return (
+
+ { aiInlineExtensionContent }
+
+ );
};
}, 'blockEditWithAiComponents' );