AI Assistant: Fix selected block while using the AI Assistant extension on a nested block

projects/plugins/jetpack/changelog/fix-ai-assistant-inline-extension-nested-input

@@ -0,0 +1,4 @@
+Significance: patch
+Type: other
+
+AI Assistant: Fix selected block while using the AI Assistant extension on a nested block

projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/ai-assistant/index.ts

@@ -54,8 +54,8 @@ unreleasedInlineExtensions.forEach( block => {
 } );
 
 // Since the lists depend on the feature flag, we need to define the types manually.
-export type ExtendedBlockProp = 'core/paragraph' | 'core/list';
-export type ExtendedInlineBlockProp = 'core/heading';
+export type ExtendedBlockProp = 'core/list';
+export type ExtendedInlineBlockProp = 'core/heading' | 'core/paragraph';
 
 type BlockSettingsProps = {
 	supports: {

projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-auto-scroll/index.ts

@@ -7,7 +7,7 @@ import debugFactory from 'debug';
 const debug = debugFactory( 'jetpack-ai-assistant:use-auto-scroll' );
 
 const useAutoScroll = (
-	blockRef: React.MutableRefObject< HTMLElement >,
+	blockRef: React.MutableRefObject< HTMLElement | null >,
 	contentRef?: React.MutableRefObject< HTMLElement >,
 	useBlockAsTarget: boolean = false
 ) => {
@@ -85,7 +85,7 @@ const useAutoScroll = (
 	}, [ blockRef, contentRef, useBlockAsTarget, userScrollHandler ] );
 
 	const getScrollParent = useCallback(
-		( el: HTMLElement | null ): HTMLElement | Document | null => {
+		( el: HTMLElement | null | undefined ): HTMLElement | Document | null => {
 			if ( el == null ) {
 				return null;
 			}

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

@@ -21,6 +21,7 @@ import type { RequestingErrorProps, RequestingStateProp } from '@automattic/jetp
 import type { ReactElement, MouseEvent } from 'react';
 
 export type AiAssistantInputProps = {
+	className?: string;
 	requestingState: RequestingStateProp;
 	requestingError?: RequestingErrorProps;
 	inputRef?: React.MutableRefObject< HTMLInputElement | null >;
@@ -34,12 +35,13 @@ export type AiAssistantInputProps = {
 	tryAgain?: () => void;
 };
 
-const className = classNames(
+const defaultClassNames = classNames(
 	'jetpack-ai-assistant-extension-ai-input',
 	'wp-block' // Some themes, like Twenty Twenty, use this class to set the element's side margins.
 );
 
 export default function AiAssistantInput( {
+	className,
 	requestingState,
 	requestingError,
 	inputRef,
@@ -156,7 +158,7 @@ export default function AiAssistantInput( {
 
 	return (
 		<ExtensionAIControl
-			className={ className }
+			className={ classNames( defaultClassNames, className ) }
 			placeholder={ placeholder }
 			disabled={ disabled }
 			value={ value }

projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-input/style.scss

@@ -1,3 +1,4 @@
-.jetpack-ai-assistant-extension-ai-input {
+.jetpack-components-ai-control__container-wrapper.jetpack-ai-assistant-extension-ai-input {
   z-index: 1;
+  position: sticky;
 }

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

@@ -87,8 +87,9 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 
 			return { postId: getCurrentPostId() };
 		}, [] );
-		// The block's id to find it in the DOM for the positioning adjustments.
-		const { id } = useBlockProps();
+		// The block's id to find it in the DOM for the positioning adjustments
+		// The classname is used by nested blocks to determine which block's toolbar to display when the input is focused.
+		const { id, className } = useBlockProps();
 		// Jetpack AI Assistant feature functions.
 		const { increaseRequestsCount, dequeueAsyncRequest, requireUpgrade } = useAiFeature();
 
@@ -101,6 +102,10 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 			true
 		);
 
+		const focusInput = useCallback( () => {
+			inputRef.current?.focus();
+		}, [] );
+
 		// Data and functions with block-specific implementations.
 		const {
 			onSuggestion: onBlockSuggestion,
@@ -216,9 +221,16 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 			// Make sure the block element has the necessary bottom padding, as it can be replaced or changed
 			setTimeout( () => {
 				adjustBlockPadding();
-				inputRef.current?.focus();
+				focusInput();
 			}, 100 );
-		}, [ disableAutoScroll, onBlockDone, increaseRequestsCount, getContent, adjustBlockPadding ] );
+		}, [
+			disableAutoScroll,
+			onBlockDone,
+			increaseRequestsCount,
+			getContent,
+			adjustBlockPadding,
+			focusInput,
+		] );
 
 		// Called when an error is received.
 		const onError = useCallback(
@@ -305,8 +317,8 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 			disableAutoScroll();
 			stopSuggestion();
 
-			inputRef.current?.focus();
-		}, [ disableAutoScroll, stopSuggestion ] );
+			focusInput();
+		}, [ disableAutoScroll, stopSuggestion, focusInput ] );
 
 		// Called when the user clicks the "Try Again" button in the input error message.
 		const handleTryAgain = useCallback( () => {
@@ -348,9 +360,9 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 				// Save the block's ownerDocument to use it later, as the editor can be in an iframe.
 				ownerDocument.current = inputRef.current.ownerDocument;
 				// Focus the input when the AI Control is displayed.
-				inputRef.current.focus();
+				focusInput();
 			}
-		}, [ showAiControl ] );
+		}, [ showAiControl, focusInput ] );
 
 		// Adjusts the input position in the editor by increasing the block's bottom-padding
 		// and setting the control's margin-top, "wrapping" the input with the block.
@@ -418,6 +430,7 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
 
 				{ showAiControl && (
 					<AiAssistantInput
+						className={ className }
 						requestingState={ requestingState }
 						requestingError={ error }
 						wrapperRef={ controlRef }