diff --git a/projects/plugins/jetpack/changelog/update-jetpack-ai-remove-old-extensions b/projects/plugins/jetpack/changelog/update-jetpack-ai-remove-old-extensions
new file mode 100644
index 0000000000000..1c5e5afef06b1
--- /dev/null
+++ b/projects/plugins/jetpack/changelog/update-jetpack-ai-remove-old-extensions
@@ -0,0 +1,4 @@
+Significance: patch
+Type: other
+
+AI Assistant: Remove dead code from transformative extensions
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/ai-assistant.php b/projects/plugins/jetpack/extensions/blocks/ai-assistant/ai-assistant.php
index e3a9f6dfebf90..9cae364a75c57 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/ai-assistant.php
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/ai-assistant.php
@@ -138,18 +138,6 @@ function () {
 	}
 );
 
-/**
- * Register the `ai-assistant-extensions-support` extension.
- */
-add_action(
-	'jetpack_register_gutenberg_extensions',
-	function () {
-		if ( apply_filters( 'jetpack_ai_enabled', true ) && apply_filters( 'jetpack_inline_extensions_enabled', false ) ) {
-			\Jetpack_Gutenberg::set_extension_available( 'ai-assistant-extensions-support' );
-		}
-	}
-);
-
 /**
  * Register the `ai-assistant-experimental-image-generation-support` extension.
  */
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/components/ai-assistant-toolbar-dropdown/dropdown-content.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/components/ai-assistant-toolbar-dropdown/dropdown-content.tsx
index 37a612b506eba..7cd5133c4d3c0 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/components/ai-assistant-toolbar-dropdown/dropdown-content.tsx
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/components/ai-assistant-toolbar-dropdown/dropdown-content.tsx
@@ -9,7 +9,7 @@ import React from 'react';
 /**
  * Internal dependencies
  */
-import { EXTENDED_INLINE_BLOCKS } from '../../extensions/ai-assistant';
+import { EXTENDED_BLOCKS } from '../../extensions/constants';
 import {
 	PROMPT_TYPE_CHANGE_TONE,
 	PROMPT_TYPE_CORRECT_SPELLING,
@@ -27,7 +27,7 @@ import './style.scss';
 /**
  * Types and constants
  */
-import type { ExtendedBlockProp, ExtendedInlineBlockProp } from '../../extensions/ai-assistant';
+import type { ExtendedBlockProp } from '../../extensions/constants';
 import type { PromptTypeProp } from '../../lib/prompt';
 import type { ToneProp } from '../tone-dropdown-control';
 import type { ReactElement } from 'react';
@@ -113,7 +113,7 @@ const quickActionsList: {
 			icon: postContent,
 		},
 	],
-	'core/list': EXTENDED_INLINE_BLOCKS.includes( 'core/list' )
+	'core/list': EXTENDED_BLOCKS.includes( 'core/list' )
 		? [
 				{
 					name: __( 'Simplify', 'jetpack' ),
@@ -172,7 +172,7 @@ export type OnRequestSuggestion = (
 ) => void;
 
 type AiAssistantToolbarDropdownContentProps = {
-	blockType: ExtendedBlockProp | ExtendedInlineBlockProp;
+	blockType: ExtendedBlockProp;
 	disabled?: boolean;
 	onAskAiAssistant: () => void;
 	onRequestSuggestion: OnRequestSuggestion;
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/components/ai-assistant-toolbar-dropdown/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/components/ai-assistant-toolbar-dropdown/index.tsx
index b006bc86575a1..b21e4768532e3 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/components/ai-assistant-toolbar-dropdown/index.tsx
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/components/ai-assistant-toolbar-dropdown/index.tsx
@@ -12,7 +12,7 @@ import React from 'react';
 /**
  * Internal dependencies
  */
-import { getStoreBlockId } from '../../extensions/ai-assistant/with-ai-assistant';
+import { getStoreBlockId } from '../../hooks/use-transform-to-assistant';
 import { getBlocksContent, getRawTextFromHTML } from '../../lib/utils/block-content';
 import { transformToAIAssistantBlock } from '../../transforms';
 import AiAssistantToolbarDropdownContent from './dropdown-content';
@@ -21,7 +21,7 @@ import './style.scss';
  * Types and constants
  */
 import type { AiAssistantDropdownOnChangeOptionsArgProps } from './dropdown-content';
-import type { ExtendedBlockProp } from '../../extensions/ai-assistant';
+import type { ExtendedBlockProp } from '../../extensions/constants';
 import type { PromptTypeProp } from '../../lib/prompt';
 import type { ReactElement } from 'react';
 
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/edit.js b/projects/plugins/jetpack/extensions/blocks/ai-assistant/edit.js
index 5f950854ae5d3..38343aa3c6a90 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/edit.js
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/edit.js
@@ -31,11 +31,11 @@ import ConnectPrompt from './components/connect-prompt';
 import FeedbackControl from './components/feedback-control';
 import QuotaExceededMessage, { FairUsageNotice } from './components/quota-exceeded-message';
 import ToolbarControls from './components/toolbar-controls';
-import { getStoreBlockId } from './extensions/ai-assistant/with-ai-assistant';
 import useAIAssistant from './hooks/use-ai-assistant';
 import useAICheckout from './hooks/use-ai-checkout';
 import useAiFeature from './hooks/use-ai-feature';
 import useAiProductPage from './hooks/use-ai-product-page';
+import { getStoreBlockId } from './hooks/use-transform-to-assistant';
 import { isUserConnected } from './lib/connection';
 import './editor.scss';
 
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/editor.js b/projects/plugins/jetpack/extensions/blocks/ai-assistant/editor.js
index 67e56bc01b045..9f5e73872dcef 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/editor.js
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/editor.js
@@ -6,12 +6,9 @@ import transforms from './transforms';
 import './editor.scss';
 
 /**
- * Supports and extensions
+ * Extensions
  */
-import './supports';
-import './extensions/ai-assistant';
-import './extensions/jetpack-contact-form';
-import './inline-extensions/with-ai-extension';
+import './extensions/with-ai-extension';
 
 registerJetpackBlockFromMetadata( metadata, {
 	edit,
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/ai-assistant/index.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/ai-assistant/index.ts
deleted file mode 100644
index 9ab83d3cb5b52..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/ai-assistant/index.ts
+++ /dev/null
@@ -1,186 +0,0 @@
-/*
- * External dependencies
- */
-import { getBlockType } from '@wordpress/blocks';
-import { select } from '@wordpress/data';
-import { addFilter } from '@wordpress/hooks';
-/*
- * Internal dependencies
- */
-import metadata from '../../block.json';
-import { isUserConnected } from '../../lib/connection';
-import { getFeatureAvailability } from '../../lib/utils/get-feature-availability';
-
-// We have two types of block extensions for now, transformative and inline.
-// The transformative blocks are transformed into an AI Assistant block when a request is made.
-// The inline blocks are updated in place.
-// Once all transformative blocks are converted to inline blocks, we can remove the distinction, but for now, we need to keep it.
-
-export const AI_ASSISTANT_SUPPORT_NAME = 'ai-assistant-support';
-export const AI_ASSISTANT_EXTENSIONS_SUPPORT_NAME = 'ai-assistant-extensions-support';
-
-// Check if the AI Assistant support is enabled.
-export const isAiAssistantSupportEnabled = getFeatureAvailability( AI_ASSISTANT_SUPPORT_NAME );
-// Check if the AI Assistant inline extensions support is enabled.
-export const isAiAssistantExtensionsSupportEnabled = getFeatureAvailability(
-	AI_ASSISTANT_EXTENSIONS_SUPPORT_NAME
-);
-
-// All Jetpack Form blocks to extend
-export const JETPACK_FORM_CHILDREN_BLOCKS = [
-	'jetpack/field-name',
-	'jetpack/field-email',
-	'jetpack/field-text',
-	'jetpack/field-textarea',
-	'jetpack/field-checkbox',
-	'jetpack/field-date',
-	'jetpack/field-telephone',
-	'jetpack/field-url',
-	'jetpack/field-checkbox-multiple',
-	'jetpack/field-radio',
-	'jetpack/field-select',
-	'jetpack/field-consent',
-	'jetpack/button',
-] as const;
-
-// The list of all extended blocks before the inline extensions were released. Does not include the list-item block.
-export const ALL_EXTENDED_BLOCKS = [
-	'core/paragraph',
-	'core/list',
-	'core/heading',
-	'jetpack/contact-form',
-	...JETPACK_FORM_CHILDREN_BLOCKS,
-];
-
-// The blocks will be converted one by one to inline blocks, so we update the lists accordingly, under the feature flag.
-export let EXTENDED_TRANSFORMATIVE_BLOCKS: string[] = [ ...ALL_EXTENDED_BLOCKS ];
-export const EXTENDED_INLINE_BLOCKS: string[] = [];
-
-// Temporarily keep track of inline extensions that have been released to production.
-const releasedInlineExtensions = [
-	'core/heading',
-	'core/paragraph',
-	'core/list-item',
-	'core/list',
-	'jetpack/contact-form',
-	...JETPACK_FORM_CHILDREN_BLOCKS,
-];
-
-// Temporarily keep track of inline extensions that are being worked on.
-const unreleasedInlineExtensions = [];
-
-releasedInlineExtensions.forEach( block => {
-	// Add the released inline extension to the inline list...
-	EXTENDED_INLINE_BLOCKS.push( block );
-	// ...and remove it from the transformative list.
-	EXTENDED_TRANSFORMATIVE_BLOCKS = EXTENDED_TRANSFORMATIVE_BLOCKS.filter( b => b !== block );
-} );
-
-unreleasedInlineExtensions.forEach( block => {
-	if ( isAiAssistantExtensionsSupportEnabled ) {
-		// Add the unreleased inline extension to the inline list...
-		EXTENDED_INLINE_BLOCKS.push( block );
-		// ...and remove it from the transformative list.
-		EXTENDED_TRANSFORMATIVE_BLOCKS = EXTENDED_TRANSFORMATIVE_BLOCKS.filter( b => b !== block );
-	}
-} );
-
-// Since the lists depend on the feature flag, we need to define the types manually.
-export type ExtendedBlockProp = string;
-export type ExtendedInlineBlockProp =
-	| 'core/heading'
-	| 'core/paragraph'
-	| 'core/list-item'
-	| 'core/list'
-	| 'jetpack/contact-form'
-	| ( typeof JETPACK_FORM_CHILDREN_BLOCKS )[ number ];
-
-type BlockSettingsProps = {
-	supports: {
-		'jetpack/ai': {
-			assistant: boolean;
-		};
-	};
-};
-
-/**
- * Check if it is possible to extend the block.
- *
- * @return {boolean} True if it is possible to extend the block.
- */
-export function isPossibleToExtendBlock(): boolean {
-	const isBlockRegistered = getBlockType( metadata.name );
-	if ( ! isBlockRegistered ) {
-		return false;
-	}
-
-	// Check Jetpack extension is enabled.
-	if ( ! isAiAssistantSupportEnabled ) {
-		return false;
-	}
-
-	// Do not extend the block if the site is not connected.
-	const connected = isUserConnected();
-	if ( ! connected ) {
-		return false;
-	}
-
-	// Do not extend if there is an error getting the feature.
-	const { errorCode } = select( 'wordpress-com/plans' )?.getAiAssistantFeature?.() || {};
-	if ( errorCode ) {
-		return false;
-	}
-
-	/*
-	 * Do not extend if the AI Assistant block is hidden
-	 * ToDo: the `editPostStore` is undefined for P2 sites.
-	 * Let's find a way to check if the block is hidden.
-	 */
-	const { getHiddenBlockTypes } = select( 'core/edit-post' ) || {};
-	const hiddenBlocks = getHiddenBlockTypes?.() || []; // It will extend the block if the function is undefined.
-	if ( hiddenBlocks.includes( metadata.name ) ) {
-		return false;
-	}
-
-	return true;
-}
-
-/**
- * Add jetpack/ai support to the extended blocks.
- *
- * @param {BlockSettingsProps} settings - Block settings.
- * @param {ExtendedBlockProp}  name     - Block name.
- * @return {BlockSettingsProps}          Block settings.
- */
-function addJetpackAISupport(
-	settings: BlockSettingsProps,
-	name: ExtendedBlockProp
-): BlockSettingsProps {
-	// Only extend the blocks in the list.
-	if ( ! EXTENDED_TRANSFORMATIVE_BLOCKS.includes( name ) ) {
-		return settings;
-	}
-
-	// Do not extend Form blocks, as they are handled differently.
-	const formBlocks = [ 'jetpack/contact-form', ...JETPACK_FORM_CHILDREN_BLOCKS ];
-	if ( formBlocks.includes( name ) ) {
-		return settings;
-	}
-
-	if ( ! isPossibleToExtendBlock() ) {
-		return settings;
-	}
-
-	return {
-		...settings,
-		supports: {
-			...settings.supports,
-			'jetpack/ai': {
-				assistant: true,
-			},
-		},
-	};
-}
-
-// Extend BlockType.
-addFilter( 'blocks.registerBlockType', 'jetpack/ai-assistant-support', addJetpackAISupport, 100 );
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/ai-assistant/with-ai-assistant.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/ai-assistant/with-ai-assistant.tsx
deleted file mode 100644
index 8ae77eaa2db51..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/ai-assistant/with-ai-assistant.tsx
+++ /dev/null
@@ -1,46 +0,0 @@
-/**
- * External dependencies
- */
-import { BlockControls } from '@wordpress/block-editor';
-import { createHigherOrderComponent } from '@wordpress/compose';
-import React from 'react';
-/**
- * Internal dependencies
- */
-import AiAssistantBlockToolbarDropdown from '../../components/ai-assistant-toolbar-dropdown';
-
-export function getStoreBlockId( clientId ) {
-	return `ai-assistant-block-${ clientId }`;
-}
-
-/*
- * Extend the withAIAssistant function of the block
- * to implement multiple blocks edition
- */
-export const withAIAssistant = createHigherOrderComponent(
-	BlockEdit => props => {
-		const { name: blockType } = props;
-
-		const blockControlProps = {
-			group: 'block' as const,
-		};
-
-		/*
-		 * CAUTION: code added before this line will be executed for all extended blocks,
-		 * defined by the EXTENDED_BLOCKS constant in ../, not just the selected blocks.
-		 * Code added above this line should be carefully evaluated for its impact on performance.
-		 */
-		return (
-			<>
-				<BlockEdit { ...props } />
-
-				<BlockControls { ...blockControlProps }>
-					<AiAssistantBlockToolbarDropdown blockType={ blockType } />
-				</BlockControls>
-			</>
-		);
-	},
-	'withAIAssistant'
-);
-
-export default withAIAssistant;
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/block-handler.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/block-handler.tsx
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/block-handler.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/block-handler.tsx
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-input/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/components/ai-assistant-input/index.tsx
similarity index 97%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-input/index.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/components/ai-assistant-input/index.tsx
index 90356810fd123..e3e2fa37ecfe8 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-input/index.tsx
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/components/ai-assistant-input/index.tsx
@@ -16,7 +16,7 @@ import './style.scss';
 /*
  * Types
  */
-import type { ExtendedInlineBlockProp } from '../../../extensions/ai-assistant';
+import type { ExtendedBlockProp } from '../../constants';
 import type { RequestingErrorProps, RequestingStateProp } from '@automattic/jetpack-ai-client';
 import type { ReactElement } from 'react';
 
@@ -27,7 +27,7 @@ export type AiAssistantInputProps = {
 	inputRef?: React.MutableRefObject< HTMLInputElement | null >;
 	wrapperRef?: React.MutableRefObject< HTMLDivElement | null >;
 	action?: string;
-	blockType: ExtendedInlineBlockProp;
+	blockType: ExtendedBlockProp;
 	feature: string;
 	request: ( question: string ) => void;
 	stopSuggestion?: () => void;
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-input/style.scss b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/components/ai-assistant-input/style.scss
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-input/style.scss
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/components/ai-assistant-input/style.scss
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/extensions/components/ai-assistant-toolbar-dropdown/index.tsx
similarity index 97%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/components/ai-assistant-toolbar-dropdown/index.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/components/ai-assistant-toolbar-dropdown/index.tsx
index 6909cae078180..5cfc20b6ede84 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/extensions/components/ai-assistant-toolbar-dropdown/index.tsx
@@ -19,13 +19,13 @@ import type {
 	AiAssistantDropdownOnChangeOptionsArgProps,
 	OnRequestSuggestion,
 } from '../../../components/ai-assistant-toolbar-dropdown/dropdown-content';
-import type { ExtendedInlineBlockProp } from '../../../extensions/ai-assistant';
 import type { PromptTypeProp } from '../../../lib/prompt';
+import type { ExtendedBlockProp } from '../../constants';
 import type { BlockBehavior } from '../../types';
 import type { ReactElement } from 'react';
 
 type AiAssistantExtensionToolbarDropdownContentProps = {
-	blockType: ExtendedInlineBlockProp;
+	blockType: ExtendedBlockProp;
 	onClose: () => void;
 	onAskAiAssistant: () => void;
 	onRequestSuggestion: OnRequestSuggestion;
@@ -101,7 +101,7 @@ function AiAssistantExtensionToolbarDropdownContent( {
 
 type AiAssistantExtensionToolbarDropdownProps = {
 	behavior: BlockBehavior;
-	blockType: ExtendedInlineBlockProp;
+	blockType: ExtendedBlockProp;
 	label?: string;
 	onAskAiAssistant: () => void;
 	onRequestSuggestion: OnRequestSuggestion;
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/constants.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/constants.ts
new file mode 100644
index 0000000000000..65475feca7f65
--- /dev/null
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/constants.ts
@@ -0,0 +1,33 @@
+// All Jetpack Form blocks to extend
+export const JETPACK_FORM_CHILDREN_BLOCKS = [
+	'jetpack/field-name',
+	'jetpack/field-email',
+	'jetpack/field-text',
+	'jetpack/field-textarea',
+	'jetpack/field-checkbox',
+	'jetpack/field-date',
+	'jetpack/field-telephone',
+	'jetpack/field-url',
+	'jetpack/field-checkbox-multiple',
+	'jetpack/field-radio',
+	'jetpack/field-select',
+	'jetpack/field-consent',
+	'jetpack/button',
+] as const;
+
+export const EXTENDED_BLOCKS = [
+	'core/heading',
+	'core/paragraph',
+	'core/list-item',
+	'core/list',
+	'jetpack/contact-form',
+	...JETPACK_FORM_CHILDREN_BLOCKS,
+];
+
+export type ExtendedBlockProp =
+	| 'core/heading'
+	| 'core/paragraph'
+	| 'core/list-item'
+	| 'core/list'
+	| 'jetpack/contact-form'
+	| ( typeof JETPACK_FORM_CHILDREN_BLOCKS )[ number ];
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/get-block-handler.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/get-block-handler.tsx
similarity index 80%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/get-block-handler.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/get-block-handler.tsx
index 529ebdd77b73a..8708c9d11400d 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/get-block-handler.tsx
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/get-block-handler.tsx
@@ -6,11 +6,8 @@ import debugFactory from 'debug';
 /**
  * Internal dependencies
  */
-import {
-	JETPACK_FORM_CHILDREN_BLOCKS,
-	type ExtendedInlineBlockProp,
-} from '../extensions/ai-assistant';
 import { BlockHandler } from './block-handler';
+import { ExtendedBlockProp, JETPACK_FORM_CHILDREN_BLOCKS } from './constants';
 import { HeadingHandler } from './heading';
 import { JetpackFormHandler, JetpackChildrenFormHandler } from './jetpack-form';
 import { ListHandler } from './list';
@@ -43,14 +40,11 @@ export const InlineExtensionsContext = createContext( {} );
 /**
  * Gets the block handler based on the block type.
  * The block handler is used to handle the request suggestions.
- * @param {ExtendedInlineBlockProp} blockType - The block type.
- * @param {string}                  clientId  - The block client ID.
- * @return {IBlockHandler}                     The block handler.
+ * @param {ExtendedBlockProp} blockType - The block type.
+ * @param {string}            clientId  - The block client ID.
+ * @return {IBlockHandler}                The block handler.
  */
-export function getBlockHandler(
-	blockType: ExtendedInlineBlockProp,
-	clientId: string
-): IBlockHandler {
+export function getBlockHandler( blockType: ExtendedBlockProp, clientId: string ): IBlockHandler {
 	let HandlerClass = handlers[ blockType ];
 
 	if ( ! HandlerClass ) {
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/heading/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/heading/index.tsx
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/heading/index.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/heading/index.tsx
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-bar/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-bar/index.tsx
deleted file mode 100644
index 75fbb13963429..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-bar/index.tsx
+++ /dev/null
@@ -1,297 +0,0 @@
-/**
- * External dependencies
- */
-import { useAiContext, BlockAIControl, ERROR_QUOTA_EXCEEDED } from '@automattic/jetpack-ai-client';
-import { useAnalytics } from '@automattic/jetpack-shared-extension-utils';
-import { serialize } from '@wordpress/blocks';
-import { KeyboardShortcuts } from '@wordpress/components';
-import { useViewportMatch } from '@wordpress/compose';
-import { select } from '@wordpress/data';
-import { useDispatch } from '@wordpress/data';
-import {
-	useContext,
-	useCallback,
-	useRef,
-	useState,
-	useEffect,
-	createPortal,
-} from '@wordpress/element';
-import { __ } from '@wordpress/i18n';
-import { store as noticesStore } from '@wordpress/notices';
-import clsx from 'clsx';
-import React from 'react';
-/**
- * Internal dependencies
- */
-import ConnectPrompt from '../../../../components/connect-prompt';
-import QuotaExceededMessage from '../../../../components/quota-exceeded-message';
-import useAiFeature from '../../../../hooks/use-ai-feature';
-import { isUserConnected } from '../../../../lib/connection';
-import { getJetpackFormCustomPrompt } from '../../../../lib/prompt';
-import { AiAssistantUiContext } from '../../ui-handler/context';
-import { AI_ASSISTANT_JETPACK_FORM_NOTICE_ID } from '../../ui-handler/with-ui-handler-data-provider';
-import './style.scss';
-
-/*
- * Core viewport breakpoints.
- * @see https://github.com/WordPress/gutenberg/blob/d5d8533cf2cc04bb005bda147114cf00782d6c38/packages/base-styles/_breakpoints.scss#L5-L14
- */
-const BREAKPOINTS = {
-	huge: 1440,
-	wide: 1280,
-	large: 960,
-	medium: 782,
-	small: 600,
-	mobile: 480,
-};
-
-/**
- * Return the serialized content from the childrens block.
- *
- * @param {string} clientId - The block client ID.
- * @return {string}          The serialized content.
- */
-function getSerializedContentFromBlock( clientId: string ): string {
-	if ( ! clientId?.length ) {
-		return '';
-	}
-
-	const block = select( 'core/block-editor' ).getBlock( clientId );
-	if ( ! block ) {
-		return '';
-	}
-
-	const { innerBlocks } = block;
-	if ( ! innerBlocks?.length ) {
-		return '';
-	}
-
-	return innerBlocks.reduce( ( acc, innerBlock ) => {
-		return acc + serialize( innerBlock ) + '\n\n';
-	}, '' );
-}
-
-export default function AiAssistantBar( {
-	clientId,
-	className = '',
-}: {
-	clientId: string;
-	className?: string;
-} ) {
-	const wrapperRef = useRef< HTMLDivElement >( null );
-	const inputRef = useRef< HTMLInputElement >( null );
-	const { tracks } = useAnalytics();
-
-	const connected = isUserConnected();
-
-	const { inputValue, setInputValue, isVisible, assistantAnchor } =
-		useContext( AiAssistantUiContext );
-
-	const { dequeueAiAssistantFeatureAsyncRequest } = useDispatch( 'wordpress-com/plans' );
-
-	const focusOnPrompt = () => {
-		// Small delay to avoid focus crash
-		setTimeout( () => {
-			inputRef.current?.focus?.();
-		}, 100 );
-	};
-
-	const { requestSuggestion, requestingState, stopSuggestion, requestingError } = useAiContext( {
-		onDone: focusOnPrompt,
-	} );
-
-	const { requireUpgrade } = useAiFeature();
-
-	const siteRequireUpgrade = requestingError?.code === ERROR_QUOTA_EXCEEDED || requireUpgrade;
-
-	const isLoading = requestingState === 'requesting' || requestingState === 'suggesting';
-
-	const showGuideLine = requestingState === 'suggesting' || requestingState === 'done';
-
-	const placeholder = __( 'Ask Jetpack AI to create your form', 'jetpack' );
-
-	const loadingPlaceholder = __( 'Creating your form. Please wait a few moments.', 'jetpack' );
-
-	const { removeNotice } = useDispatch( noticesStore );
-
-	const handleSend = useCallback( () => {
-		// Do not send the request if the input value is empty.
-		if ( ! inputValue?.length ) {
-			return;
-		}
-
-		// Remove previous error notice.
-		removeNotice( AI_ASSISTANT_JETPACK_FORM_NOTICE_ID );
-
-		const prompt = getJetpackFormCustomPrompt( {
-			request: inputValue,
-			content: getSerializedContentFromBlock( clientId ),
-		} );
-
-		/*
-		 * Always dequeue/cancel the AI Assistant feature async request,
-		 * in case there is one pending,
-		 * when performing a new AI suggestion request.
-		 */
-		dequeueAiAssistantFeatureAsyncRequest();
-
-		requestSuggestion( prompt, { feature: 'jetpack-form-ai-extension' } );
-		tracks.recordEvent( 'jetpack_ai_assistant_block_generate', {
-			feature: 'jetpack-form-ai-extension',
-		} );
-		wrapperRef?.current?.focus();
-	}, [
-		clientId,
-		dequeueAiAssistantFeatureAsyncRequest,
-		inputValue,
-		removeNotice,
-		requestSuggestion,
-		tracks,
-	] );
-
-	const handleStopSuggestion = useCallback( () => {
-		stopSuggestion();
-		focusOnPrompt();
-		tracks.recordEvent( 'jetpack_ai_assistant_block_stop', {
-			feature: 'jetpack-form-ai-extension',
-		} );
-	}, [ stopSuggestion, tracks ] );
-
-	/*
-	 * Fix the assistant bar when the viewport is mobile,
-	 * and the Assistant anchor exists.
-	 */
-	const isMobileViewport = useViewportMatch( 'medium', '<' );
-	const isAssistantBarFixed = isMobileViewport && assistantAnchor;
-
-	/*
-	 * Auto-mobile-switching mode.
-	 * Update the bar layout depending on the bar component width.
-	 */
-	const observerRef = useRef< ResizeObserver | null >( null );
-	const isMobileModeRef = useRef( isMobileViewport );
-	const [ isMobileMode, setMobileMode ] = useState( isMobileViewport );
-
-	useEffect( () => {
-		// Get the Assistant bar DOM element.
-		const barElement = wrapperRef?.current;
-		if ( ! barElement ) {
-			return;
-		}
-
-		// Only create a new observer if there isn't one already
-		if ( ! observerRef?.current ) {
-			observerRef.current = new ResizeObserver( entries => {
-				if ( ! isVisible ) {
-					return;
-				}
-
-				if ( isAssistantBarFixed ) {
-					return;
-				}
-
-				const barWidth = entries[ 0 ].contentRect.width;
-				const isMobileModeBasedOnWidth = barWidth < BREAKPOINTS.mobile;
-
-				// Only update the state if the mode has changed.
-				if ( isMobileModeBasedOnWidth !== isMobileModeRef.current ) {
-					isMobileModeRef.current = isMobileModeBasedOnWidth; // Update the ref to be able to compare later.
-					setMobileMode( isMobileModeBasedOnWidth ); // Update the state (and re-render)
-				}
-			} );
-		}
-
-		// Start observing the Assistant bar element.
-		observerRef.current.observe( barElement );
-
-		return () => {
-			// Disconnect the observer when the component is unmounted.
-			observerRef?.current?.disconnect();
-		};
-	}, [ isAssistantBarFixed, isVisible ] );
-
-	useEffect( () => {
-		if ( isVisible ) {
-			tracks.recordEvent( 'jetpack_ai_assistant_prompt_show', {
-				block_type: 'jetpack/contact-form',
-			} );
-		}
-	}, [ isVisible, tracks ] );
-
-	// focus input on first render only (for a11y reasons, toggling on/off should not focus the input)
-	useEffect( () => {
-		/*
-		 * Only focus the input when the Assistant bar is visible.
-		 * Also, add a small delay to avoid focus when the Assistant bar is toggled off.
-		 */
-		const timeId = setTimeout( () => {
-			if ( ! isVisible ) {
-				return;
-			}
-
-			if ( ! inputRef?.current ) {
-				return;
-			}
-
-			inputRef.current.focus();
-		}, 300 );
-
-		return function () {
-			clearTimeout( timeId );
-		};
-	}, [] ); // eslint-disable-line react-hooks/exhaustive-deps -- only run on first render
-
-	if ( ! isVisible ) {
-		return null;
-	}
-
-	// Assistant bar component.
-	const AiAssistantBarComponent = (
-		<div className="jetpack-ai-assistant__bar-wrapper">
-			<KeyboardShortcuts
-				bindGlobal
-				shortcuts={ {
-					esc: () => {
-						if ( [ 'requesting', 'suggesting' ].includes( requestingState ) ) {
-							handleStopSuggestion();
-						}
-					},
-				} }
-			>
-				<div
-					ref={ wrapperRef }
-					className={ clsx( 'jetpack-ai-assistant__bar', {
-						[ className ]: className,
-						'is-fixed': isAssistantBarFixed,
-						'is-mobile-mode': isMobileMode,
-					} ) }
-					tabIndex={ -1 }
-				>
-					{ siteRequireUpgrade && <QuotaExceededMessage placement="jetpack-form-block" /> }
-					{ ! connected && <ConnectPrompt /> }
-					<BlockAIControl
-						ref={ inputRef }
-						disabled={ siteRequireUpgrade || ! connected }
-						value={ inputValue }
-						placeholder={ isLoading ? loadingPlaceholder : placeholder }
-						onChange={ setInputValue }
-						onSend={ handleSend }
-						onStop={ handleStopSuggestion }
-						state={ requestingState }
-						isTransparent={ siteRequireUpgrade || ! connected }
-						showButtonLabels={ ! isMobileMode }
-						showGuideLine={ showGuideLine }
-					/>
-				</div>
-			</KeyboardShortcuts>
-		</div>
-	);
-
-	// Check if the Assistant bar should be rendered in the Assistant anchor (fixed mode)
-	if ( isAssistantBarFixed ) {
-		return createPortal( AiAssistantBarComponent, assistantAnchor );
-	}
-
-	// Render in the editor canvas.
-	return AiAssistantBarComponent;
-}
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-bar/style.scss b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-bar/style.scss
deleted file mode 100644
index 9b6453879e6e2..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-bar/style.scss
+++ /dev/null
@@ -1,43 +0,0 @@
-.jetpack-ai-assistant__bar {
-
-    &:not(.is-fixed) {
-        position: relative;
-        z-index: 30;
-        margin-top: 10px;
-    }
-
-    &.is-fixed {
-		.jetpack-components-ai-control__container-wrapper {
-			position: relative;
-			bottom: 0;
-		}
-
-		.jetpack-components-ai-control__container {
-			border-radius: 0;
-			box-shadow: none;
-		}
-
-		.jetpack-components-ai-control__wrapper {
-			padding: 7px 14px;
-		}
-	}
-}
-
-.jetpack-ai-assistant-bar__slot {
-	border: none;
-	border-radius: 0;
-	display: block;
-	position: sticky;
-	z-index: 31;
-
-	.jetpack-ai-assistant__bar-wrapper {
-		position: relative;
-		bottom: 0;
-	}
-}
-
-.jetpack-ai-assistant__bar-wrapper {
-	position: sticky;
-	z-index: 30;
-	bottom: 16px;
-}
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-toolbar-button/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-toolbar-button/index.tsx
deleted file mode 100644
index ff386d97faf9c..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/components/ai-assistant-toolbar-button/index.tsx
+++ /dev/null
@@ -1,137 +0,0 @@
-/*
- * External dependencies
- */
-import { aiAssistantIcon, useAiContext } from '@automattic/jetpack-ai-client';
-import { KeyboardShortcuts, ToolbarButton } from '@wordpress/components';
-import { useViewportMatch } from '@wordpress/compose';
-import { useEffect, useContext, useRef, useCallback } from '@wordpress/element';
-import { __ } from '@wordpress/i18n';
-/*
- * Internal dependencies
- */
-import { AiAssistantUiContext } from '../../ui-handler/context';
-import { selectFormBlock } from '../../ui-handler/with-ui-handler-data-provider';
-import type { ReactElement } from 'react';
-
-const AI_ASSISTANT_BAR_SLOT_CLASS = 'jetpack-ai-assistant-bar__slot';
-
-/**
- * The toolbar button that toggles the Assistant Bar.
- * Also, it creates a slot just after the contextual toolbar
- * to be used as the anchor for the Assistant Bar.
- *
- * @param {object} props                     - The component props.
- * @param {string} props.jetpackFormClientId - The Jetpack Form block client ID.
- * @return {ReactElement}               The toolbar button.
- */
-export default function AiAssistantToolbarButton( {
-	jetpackFormClientId,
-}: {
-	jetpackFormClientId?: string;
-} ): ReactElement {
-	const { isVisible, toggle, setAnchor, assistantAnchor } = useContext( AiAssistantUiContext );
-	const { requestingState } = useAiContext();
-
-	const toolbarButtonRef = useRef< HTMLElement | null >( null );
-
-	const isMobileViewport = useViewportMatch( 'medium', '<' );
-
-	/*
-	 * When the toolbar button is rendered, we need to find the
-	 * contextual toolbar and create a slot just after it.
-	 * This slot will be used as the anchor for the Assistant Bar.
-	 */
-	useEffect( () => {
-		if ( ! toolbarButtonRef.current ) {
-			return;
-		}
-
-		const toolbar = toolbarButtonRef.current.closest(
-			'.block-editor-block-contextual-toolbar'
-		) as HTMLElement;
-		if ( ! toolbar ) {
-			return;
-		}
-
-		/*
-		 * AI Assistant bar slot element.
-		 * When the viewport is in mobile mode,
-		 * create an element just after the contextual toolbar
-		 * to be used as the anchor for the Assistant Bar.
-		 */
-
-		/*
-		 * Check if the slot already exists,
-		 * quering from the block-toolbar parent element.
-		 * It should not happend, since the slot removes
-		 * when the viewport is not mobile.
-		 */
-		let slot = toolbar.parentElement?.querySelector(
-			`.${ AI_ASSISTANT_BAR_SLOT_CLASS }`
-		) as HTMLElement;
-
-		if ( slot ) {
-			// always move the slot right after the toolbar.
-			toolbar.after( slot );
-			return setAnchor( slot );
-		}
-
-		// Slot not found - create it.
-		slot = document.createElement( 'div' );
-
-		// Set role="toolbar" and Aria label
-		slot.setAttribute( 'role', 'toolbar' );
-		slot.setAttribute( 'aria-label', __( 'AI Assistant', 'jetpack' ) );
-		slot.setAttribute( 'aria-orientation', 'horizontal' );
-		slot.className = AI_ASSISTANT_BAR_SLOT_CLASS;
-
-		// Set the top position based on the toolbar height.
-		const toolbarHeight = toolbar.offsetHeight;
-		slot.style.top = `${ toolbarHeight }px`;
-		toolbar.after( slot );
-
-		// Set the anchor where the Assistant Bar will be rendered.
-		setAnchor( slot );
-	}, [ setAnchor ] );
-
-	// Remove the slot when the view is not mobile.
-	useEffect( () => {
-		if ( isMobileViewport ) {
-			return;
-		}
-
-		assistantAnchor?.remove();
-	}, [ isMobileViewport, assistantAnchor ] );
-
-	const toggleFromToolbar = useCallback( () => {
-		if ( ! jetpackFormClientId ) {
-			return toggle();
-		}
-
-		selectFormBlock( jetpackFormClientId, toggle );
-	}, [ jetpackFormClientId, toggle ] );
-
-	const isDisabled = requestingState === 'requesting' || requestingState === 'suggesting';
-
-	return (
-		<>
-			<KeyboardShortcuts
-				bindGlobal
-				shortcuts={ {
-					'mod+/': toggle,
-				} }
-			/>
-
-			<ToolbarButton
-				ref={ toolbarButtonRef }
-				showTooltip
-				onClick={ toggleFromToolbar }
-				aria-haspopup="true"
-				aria-expanded={ isVisible }
-				label={ __( 'Ask AI Assistant', 'jetpack' ) }
-				icon={ aiAssistantIcon }
-				disabled={ isDisabled }
-			/>
-		</>
-	);
-}
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/constants.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/constants.ts
deleted file mode 100644
index 0563fb73d8dfd..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/constants.ts
+++ /dev/null
@@ -1,22 +0,0 @@
-export const JETPACK_FORM_AI_COMPOSITION_EXTENSION = 'ai-assistant-form-support' as const;
-
-export const isJetpackFromBlockAiCompositionAvailable =
-	window?.Jetpack_Editor_Initial_State?.available_blocks?.[ JETPACK_FORM_AI_COMPOSITION_EXTENSION ]
-		?.available;
-
-// All blocks to extend
-export const JETPACK_FORM_CHILDREN_BLOCKS = [
-	'jetpack/field-name',
-	'jetpack/field-email',
-	'jetpack/field-text',
-	'jetpack/field-textarea',
-	'jetpack/field-checkbox',
-	'jetpack/field-date',
-	'jetpack/field-telephone',
-	'jetpack/field-url',
-	'jetpack/field-checkbox-multiple',
-	'jetpack/field-radio',
-	'jetpack/field-select',
-	'jetpack/field-consent',
-	'jetpack/button',
-];
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/index.tsx
deleted file mode 100644
index 0895c729cc4be..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/index.tsx
+++ /dev/null
@@ -1,305 +0,0 @@
-/*
- * External dependencies
- */
-import { useAiContext, withAiDataProvider } from '@automattic/jetpack-ai-client';
-import { useModuleStatus } from '@automattic/jetpack-shared-extension-utils';
-import { BlockControls } from '@wordpress/block-editor';
-import { getBlockType } from '@wordpress/blocks';
-import { createHigherOrderComponent } from '@wordpress/compose';
-import { select, useDispatch, useSelect } from '@wordpress/data';
-import { useEffect, useCallback } from '@wordpress/element';
-import { addFilter } from '@wordpress/hooks';
-import React from 'react';
-/*
- * Internal dependencies
- */
-import { getFeatureAvailability } from '../../lib/utils/get-feature-availability';
-import {
-	AI_ASSISTANT_EXTENSIONS_SUPPORT_NAME,
-	EXTENDED_TRANSFORMATIVE_BLOCKS,
-} from '../ai-assistant';
-import AiAssistantBar from './components/ai-assistant-bar';
-import AiAssistantToolbarButton from './components/ai-assistant-toolbar-button';
-import { isJetpackFromBlockAiCompositionAvailable } from './constants';
-import { JETPACK_FORM_CHILDREN_BLOCKS } from './constants';
-import withUiHandlerDataProvider from './ui-handler/with-ui-handler-data-provider';
-
-type IsPossibleToExtendJetpackFormBlockProps = {
-	checkChildrenBlocks?: boolean;
-	clientId: string;
-};
-
-/**
- * Check if it is possible to extend the block.
- *
- * @param {string}  blockName           - The block name.
- * @param {boolean} checkChildrenBlocks - Check if the block is a child of a Jetpack Form block.
- * @return {boolean}                     True if it is possible to extend the block.
- */
-export function useIsPossibleToExtendJetpackFormBlock(
-	blockName: string | undefined,
-	{ checkChildrenBlocks = false, clientId }: IsPossibleToExtendJetpackFormBlockProps = {
-		clientId: '',
-	}
-): boolean {
-	// Check if the AI Assistant block is registered.
-	const isBlockRegistered = getBlockType( 'jetpack/ai-assistant' );
-	const { isModuleActive } = useModuleStatus( 'contact-form' );
-
-	if ( ! isModuleActive ) {
-		return false;
-	}
-
-	if ( ! isBlockRegistered ) {
-		return false;
-	}
-
-	// Check if there is a block name.
-	if ( typeof blockName !== 'string' ) {
-		return false;
-	}
-
-	// Only extend the blocks in the allowed list.
-	if ( ! EXTENDED_TRANSFORMATIVE_BLOCKS.includes( blockName ) ) {
-		return false;
-	}
-
-	// Check if Jetpack extension is enabled.
-	if ( ! isJetpackFromBlockAiCompositionAvailable ) {
-		return false;
-	}
-
-	// clientId is required
-	if ( ! clientId?.length ) {
-		return false;
-	}
-
-	// Only extend allowed blocks.
-	if ( checkChildrenBlocks ) {
-		// First, check if it should check for children blocks. (false by default)
-		if ( ! JETPACK_FORM_CHILDREN_BLOCKS.includes( blockName ) ) {
-			return false;
-		}
-	} else if ( blockName !== 'jetpack/contact-form' ) {
-		// If it is not a child block, check if it is the Jetpack Form block.
-		return false;
-	}
-
-	/*
-	 * Do not extend if the AI Assistant block is hidden
-	 * Todo: Do we want to make the extension depend on the block visibility?
-	 * ToDo: the `editPostStore` is undefined for P2 sites.
-	 * Let's find a way to check if the block is hidden.
-	 */
-	const { getHiddenBlockTypes } = select( 'core/edit-post' ) || {};
-	const hiddenBlocks = getHiddenBlockTypes?.() || []; // It will extend the block if the function is undefined.
-	if ( hiddenBlocks.includes( 'jetpack/ai-assistant' ) ) {
-		return false;
-	}
-
-	return true;
-}
-
-/**
- * HOC to populate the Jetpack Form edit component
- * with the AI Assistant bar and button.
- */
-const jetpackFormEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
-	return props => {
-		const possibleToExtendJetpackFormBlock = useIsPossibleToExtendJetpackFormBlock( props?.name, {
-			clientId: props.clientId,
-		} );
-
-		const { increaseAiAssistantRequestsCount } = useDispatch( 'wordpress-com/plans' );
-
-		const { eventSource } = useAiContext( {
-			onDone: useCallback( () => {
-				/*
-				 * Increase the AI Suggestion counter.
-				 * @todo: move this at store level.
-				 */
-				increaseAiAssistantRequestsCount();
-			}, [ increaseAiAssistantRequestsCount ] ),
-			onError: useCallback(
-				error => {
-					/*
-					 * Incrses AI Suggestion counter
-					 * only for valid errors.
-					 * @todo: move this at store level.
-					 */
-					if ( error.code === 'error_network' || error.code === 'error_quota_exceeded' ) {
-						return;
-					}
-
-					// Increase the AI Suggestion counter.
-					increaseAiAssistantRequestsCount();
-				},
-				[ increaseAiAssistantRequestsCount ]
-			),
-		} );
-
-		const stopSuggestion = useCallback( () => {
-			if ( ! eventSource ) {
-				return;
-			}
-			eventSource?.close();
-		}, [ eventSource ] );
-
-		useEffect( () => {
-			/*
-			 * Cleanup function to remove the event listeners
-			 * and close the event source.
-			 */
-			return () => {
-				// Only stop when the parent block is unmounted.
-				if ( props?.name !== 'jetpack/contact-form' ) {
-					return;
-				}
-
-				stopSuggestion();
-			};
-		}, [ stopSuggestion, props?.name ] );
-
-		// Only extend Jetpack Form block (children not included).
-		if ( ! possibleToExtendJetpackFormBlock ) {
-			return <BlockEdit { ...props } />;
-		}
-
-		const blockControlsProps = {
-			group: 'block' as const,
-		};
-
-		return (
-			<>
-				<BlockEdit { ...props } />
-
-				<AiAssistantBar clientId={ props.clientId } />
-
-				<BlockControls { ...blockControlsProps }>
-					<AiAssistantToolbarButton />
-				</BlockControls>
-			</>
-		);
-	};
-}, 'jetpackFormEditWithAiComponents' );
-
-/**
- * Function used to extend the registerBlockType settings.
- *
- * - Populate the Jetpack Form edit component
- * with the AI Assistant bar and button (jetpackFormEditWithAiComponents).
- * - Add the UI Handler data provider (withUiHandlerDataProvider).
- * - Add the AI Assistant data provider (withAiDataProvider).
- *
- * @param {object} settings - The block settings.
- * @param {string} name     - The block name.
- * @return {object}          The block settings.
- */
-function jetpackFormWithAiSupport( settings, name: string ) {
-	// Only extend Jetpack Form block type.
-	if ( name !== 'jetpack/contact-form' ) {
-		return settings;
-	}
-
-	// Only extend the blocks in the allowed list.
-	if ( ! EXTENDED_TRANSFORMATIVE_BLOCKS.includes( name ) ) {
-		return settings;
-	}
-
-	// Disable if Inline Extension is enabled
-	if ( getFeatureAvailability( AI_ASSISTANT_EXTENSIONS_SUPPORT_NAME ) ) {
-		return settings;
-	}
-
-	return {
-		...settings,
-		edit: withAiDataProvider(
-			withUiHandlerDataProvider( jetpackFormEditWithAiComponents( settings.edit ) )
-		),
-	};
-}
-
-addFilter(
-	'blocks.registerBlockType',
-	'jetpack/ai-assistant-support',
-	jetpackFormWithAiSupport,
-	100
-);
-
-/**
- * HOC to populate the Jetpack Form children blocks edit components:
- * - AI Assistant toolbar button.
- *
- * This HOC must be used only for children blocks of the Jetpack Form block.
- */
-const jetpackFormChildrenEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
-	return props => {
-		// Get clientId of the parent block.
-		const parentClientId = useSelect(
-			selectData => {
-				const blockEditorSelectData: {
-					getBlockParentsByBlockName: ( clientId: string, blockName: string ) => string[];
-				} = selectData( 'core/block-editor' );
-				const { getBlockParentsByBlockName } = blockEditorSelectData;
-
-				return getBlockParentsByBlockName( props.clientId, 'jetpack/contact-form' )?.[ 0 ];
-			},
-			[ props.clientId ]
-		);
-
-		const possibleToExtendJetpackFormBlock = useIsPossibleToExtendJetpackFormBlock( props?.name, {
-			checkChildrenBlocks: true,
-			clientId: parentClientId,
-		} );
-
-		if ( ! possibleToExtendJetpackFormBlock ) {
-			return <BlockEdit { ...props } />;
-		}
-
-		const blockControlsProps = {
-			group: 'parent' as const,
-		};
-
-		return (
-			<>
-				<BlockEdit { ...props } />
-
-				<BlockControls { ...blockControlsProps }>
-					<AiAssistantToolbarButton jetpackFormClientId={ parentClientId } />
-				</BlockControls>
-			</>
-		);
-	};
-}, 'jetpackFormChildrenEditWithAiComponents' );
-
-/*
- * Extend children blocks of Jetpack Form block
- * with the AI Assistant components.
- */
-function jetpackFormChildrenEditWithAiSupport( settings, name ) {
-	// Only extend the blocks in the allowed list.
-	if ( ! EXTENDED_TRANSFORMATIVE_BLOCKS.includes( name ) ) {
-		return settings;
-	}
-
-	// Only extend allowed blocks (Jetpack form and its children)
-	if ( ! JETPACK_FORM_CHILDREN_BLOCKS.includes( name ) ) {
-		return settings;
-	}
-
-	// Disable if Inline Extension is enabled
-	if ( getFeatureAvailability( AI_ASSISTANT_EXTENSIONS_SUPPORT_NAME ) ) {
-		return settings;
-	}
-
-	return {
-		...settings,
-		edit: jetpackFormChildrenEditWithAiComponents( settings.edit ),
-	};
-}
-
-addFilter(
-	'blocks.registerBlockType',
-	'jetpack/ai-assistant-support',
-	jetpackFormChildrenEditWithAiSupport
-);
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/context.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/context.tsx
deleted file mode 100644
index d80532be85bb2..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/context.tsx
+++ /dev/null
@@ -1,47 +0,0 @@
-/**
- * External dependencies
- */
-import { createContext } from '@wordpress/element';
-
-export type AiAssistantUiContextProps = {
-	inputValue: string;
-
-	isVisible: boolean;
-
-	isFixed: boolean;
-
-	assistantAnchor: HTMLElement | null;
-	setAnchor: ( anchor: HTMLElement | null ) => void;
-
-	setInputValue: ( value: string ) => void;
-
-	show: () => void;
-	hide: () => void;
-	toggle: () => void;
-
-	setAssistantFixed: ( isFixed: boolean ) => void;
-};
-
-type AiAssistantUiContextProviderProps = {
-	/*
-	 * Open the AI Assistant
-	 */
-	value: AiAssistantUiContextProps;
-
-	/*
-	 * Children
-	 */
-	children: React.ReactNode;
-};
-
-/**
- * Ai Assistant Context
- */
-export const AiAssistantUiContext = createContext( {} as AiAssistantUiContextProps );
-
-export const AiAssistantUiContextProvider = ( {
-	value,
-	children,
-}: AiAssistantUiContextProviderProps ) => (
-	<AiAssistantUiContext.Provider value={ value } children={ children } />
-);
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/with-ui-handler-data-provider.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/with-ui-handler-data-provider.tsx
deleted file mode 100644
index 10bd22ede26b2..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/with-ui-handler-data-provider.tsx
+++ /dev/null
@@ -1,307 +0,0 @@
-/**
- * External dependencies
- */
-import { ERROR_QUOTA_EXCEEDED, useAiContext } from '@automattic/jetpack-ai-client';
-import { createBlock } from '@wordpress/blocks';
-import { parse } from '@wordpress/blocks';
-import { KeyboardShortcuts } from '@wordpress/components';
-import { createHigherOrderComponent } from '@wordpress/compose';
-import { useDispatch, useSelect, dispatch } from '@wordpress/data';
-import { useState, useMemo, useCallback, useEffect, useRef } from '@wordpress/element';
-import { __ } from '@wordpress/i18n';
-import { store as noticesStore } from '@wordpress/notices';
-/**
- * Internal dependencies
- */
-import { useIsPossibleToExtendJetpackFormBlock } from '..';
-import { compareBlocks } from '../../../lib/utils/compare-blocks';
-import { fixIncompleteHTML } from '../../../lib/utils/fix-incomplete-html';
-import { AiAssistantUiContextProvider } from './context';
-/**
- * Types
- */
-import type { Block } from '@automattic/jetpack-ai-client';
-import type { RequestingErrorProps } from '@automattic/jetpack-ai-client';
-
-// An identifier to use on the extension error notices,
-export const AI_ASSISTANT_JETPACK_FORM_NOTICE_ID = 'ai-assistant';
-
-type BlockEditorSelect = {
-	getBlock: ( clientId: string ) => Block;
-};
-
-type CoreEditorSelect = {
-	getCurrentPostId: () => number;
-};
-
-/**
- * Select the Jetpack Form block,
- * based on the block client ID.
- * Then, run the function passed as parameter (optional).
- *
- * @param {string}   clientId - The block client ID.
- * @param {Function} fn       - The function to run after selecting the block.
- * @return {void}
- */
-export function selectFormBlock( clientId: string, fn: () => void ): void {
-	const blockEditorDispatch = dispatch( 'core/block-editor' );
-	blockEditorDispatch.selectBlock( clientId );
-	fn?.();
-}
-
-const withUiHandlerDataProvider = createHigherOrderComponent( BlockListBlock => {
-	return props => {
-		const { clientId, isSelected } = props;
-
-		const { replaceInnerBlocks } = useDispatch( 'core/block-editor' );
-		const coreEditorSelect = useSelect( select => select( 'core/editor' ), [] ) as CoreEditorSelect;
-		const blockEditorSelect = useSelect(
-			select => select( 'core/block-editor' ),
-			[]
-		) as BlockEditorSelect;
-		const postId = coreEditorSelect.getCurrentPostId();
-		const blockFromPostId = blockEditorSelect.getBlock( clientId );
-
-		// AI Assistant input value
-		const [ inputValue, setInputValue ] = useState( '' );
-
-		// AI Assistant component visibility
-		// Only show input automatically for empty forms
-		const shouldStartVisible =
-			blockFromPostId?.name === 'jetpack/contact-form' &&
-			blockFromPostId?.innerBlocks?.length === 0;
-		const [ isVisible, setAssistantVisibility ] = useState( shouldStartVisible );
-
-		// AI Assistant component is-fixed state
-		const [ isFixed, setAssistantFixed ] = useState( false );
-
-		const [ assistantAnchor, setAssistantAnchor ] = useState< HTMLElement | null >( null );
-
-		// Keep track of the current list of valid blocks between renders.
-		const currentListOfValidBlocks = useRef< Array< Block > >( [] );
-
-		/**
-		 * Show the AI Assistant
-		 *
-		 * @return {void}
-		 */
-		const show = useCallback( () => {
-			setAssistantVisibility( true );
-		}, [] );
-
-		/**
-		 * Hide the AI Assistant
-		 *
-		 * @return {void}
-		 */
-		const hide = useCallback( () => {
-			setAssistantVisibility( false );
-		}, [] );
-
-		/**
-		 * Toggle the AI Assistant visibility
-		 *
-		 * @return {void}
-		 */
-		const toggle = useCallback( () => {
-			setAssistantVisibility( ! isVisible );
-		}, [ isVisible ] );
-
-		/**
-		 * Set the AI Assistant anchor
-		 *
-		 * @param {HTMLElement} anchor
-		 */
-		const setAnchor = useCallback( ( anchor: HTMLElement | null ) => {
-			setAssistantAnchor( anchor );
-		}, [] );
-
-		const { createNotice } = useDispatch( noticesStore );
-
-		/**
-		 * Show the error notice
-		 *
-		 * @param {RequestingErrorProps} suggestionError
-		 * @return {void}
-		 */
-		const showSuggestionError = useCallback(
-			( { severity, message, code }: RequestingErrorProps ) => {
-				if ( code !== ERROR_QUOTA_EXCEEDED ) {
-					createNotice( severity, message, {
-						isDismissible: true,
-						id: AI_ASSISTANT_JETPACK_FORM_NOTICE_ID,
-					} );
-				}
-			},
-			[ createNotice ]
-		);
-
-		// Show/hide the assistant based on the block selection.
-		useEffect( () => {
-			if ( isSelected ) {
-				return;
-			}
-			hide();
-		}, [ isSelected, hide ] );
-
-		// Build the context value to pass to the provider.
-		const contextValue = useMemo(
-			() => ( {
-				// Value of the input element.
-				inputValue,
-				setInputValue,
-
-				// Assistant bar visibility.
-				isVisible,
-				show,
-				hide,
-				toggle,
-
-				// Assistant bar position and size.
-				isFixed,
-				setAssistantFixed,
-
-				// Assistant bar anchor.
-				assistantAnchor,
-				setAnchor,
-			} ),
-			[ inputValue, isVisible, show, hide, toggle, isFixed, assistantAnchor, setAnchor ]
-		);
-
-		const setContent = useCallback(
-			( newContent: string, isRequestDone = false ) => {
-				// Remove the Jetpack Form block from the content.
-				const processedContent = newContent.replace(
-					/<!-- (\/)*wp:jetpack\/(contact-)*form ({[^}]*} )*(\/)*-->/g,
-					''
-				);
-
-				// Fix HTML tags that are not closed properly.
-				const fixedContent = fixIncompleteHTML( processedContent );
-
-				const newContentBlocks = parse( fixedContent );
-
-				// Check if the generated blocks are valid.
-				const validBlocks = newContentBlocks.filter( block => {
-					return (
-						block.isValid &&
-						! [ 'core/freeform', 'core/missing', 'core/html' ].includes( block.name )
-					);
-				} );
-
-				let lastBlockUpdated = false;
-
-				// While streaming, the last block can go from valid to invalid and back as new children are added token by token.
-				if ( validBlocks.length < currentListOfValidBlocks.current.length ) {
-					// The last block is temporarily invalid, so we use the last valid state.
-					validBlocks.push(
-						currentListOfValidBlocks.current[ currentListOfValidBlocks.current.length - 1 ]
-					);
-				} else if (
-					validBlocks.length === currentListOfValidBlocks.current.length &&
-					validBlocks.length > 0
-				) {
-					// Update the last valid block with the new content if it is different.
-					const lastBlock = validBlocks[ validBlocks.length - 1 ];
-					const lastBlockFromCurrentList =
-						currentListOfValidBlocks.current[ validBlocks.length - 1 ];
-
-					lastBlockUpdated = ! compareBlocks( lastBlock, lastBlockFromCurrentList );
-				}
-
-				if (
-					// Only update the blocks when there are valid blocks, to avoid having no children and triggering the empty state.
-					validBlocks.length > 0 &&
-					// Only update the blocks when the valid list changed, meaning a new block arrived or the last block was updated.
-					( validBlocks.length !== currentListOfValidBlocks.current.length || lastBlockUpdated )
-				) {
-					// Only update the valid blocks
-					replaceInnerBlocks( clientId, validBlocks );
-
-					// Update the list of current valid blocks
-					currentListOfValidBlocks.current = validBlocks;
-				}
-
-				// Final form adjustments (only when the request is done)
-				if ( isRequestDone ) {
-					/*
-					 * Inspect generated blocks list,
-					 * checking if the jetpack/button block:
-					 * - if it exists twice or more, remove the first one.
-					 * - if it does not exist, create one.
-					 */
-					const allButtonBlocks = validBlocks.filter( block => block.name === 'jetpack/button' );
-					currentListOfValidBlocks.current = currentListOfValidBlocks.current || [];
-					if ( allButtonBlocks.length > 1 ) {
-						// Remove all button blocks, less the last one.
-						let buttonCounter = 0;
-						currentListOfValidBlocks.current = currentListOfValidBlocks.current.filter( block => {
-							if ( block.name !== 'jetpack/button' ) {
-								return true;
-							}
-
-							buttonCounter++;
-							if ( buttonCounter === allButtonBlocks.length ) {
-								return true;
-							}
-							return false;
-						} );
-
-						replaceInnerBlocks( clientId, currentListOfValidBlocks.current );
-					} else if ( allButtonBlocks.length === 0 ) {
-						// One button block is required.
-						replaceInnerBlocks( clientId, [
-							...currentListOfValidBlocks.current,
-							createBlock( 'jetpack/button', {
-								label: __( 'Submit', 'jetpack' ),
-								element: 'button',
-								text: __( 'Submit', 'jetpack' ),
-								borderRadius: 8,
-								lock: {
-									remove: true,
-								},
-							} ),
-						] );
-					}
-
-					// Reset the list of valid blocks after the request is done.
-					currentListOfValidBlocks.current = [];
-				}
-			},
-			[ clientId, replaceInnerBlocks ]
-		);
-
-		useAiContext( {
-			askQuestionOptions: { postId },
-			onDone: finalContent => {
-				setContent( finalContent, true );
-				setInputValue( '' );
-			},
-			onSuggestion: setContent,
-			onError: showSuggestionError,
-		} );
-
-		/*
-		 * Ensure to provide data context,
-		 * and the AI Assistant component (popover)
-		 * only if is't possible to extend the block.
-		 */
-		if ( ! useIsPossibleToExtendJetpackFormBlock( props.name, { clientId: props.clientId } ) ) {
-			return <BlockListBlock { ...props } />;
-		}
-
-		return (
-			<AiAssistantUiContextProvider value={ contextValue }>
-				<KeyboardShortcuts
-					shortcuts={ {
-						'mod+/': () => selectFormBlock( clientId, show ),
-					} }
-				>
-					<BlockListBlock { ...props } />
-				</KeyboardShortcuts>
-			</AiAssistantUiContextProvider>
-		);
-	};
-}, 'withUiHandlerDataProvider' );
-
-export default withUiHandlerDataProvider;
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/jetpack-form/children.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-form/children.tsx
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/jetpack-form/children.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-form/children.tsx
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/jetpack-form/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-form/index.tsx
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/jetpack-form/index.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-form/index.tsx
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/lib/is-possible-to-extend-block.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/lib/is-possible-to-extend-block.ts
similarity index 71%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/lib/is-possible-to-extend-block.ts
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/lib/is-possible-to-extend-block.ts
index 9fbc5c12e0069..2239bafa6f025 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/lib/is-possible-to-extend-block.ts
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/lib/is-possible-to-extend-block.ts
@@ -6,12 +6,18 @@ import { select } from '@wordpress/data';
 /*
  * Internal dependencies
  */
-import { EXTENDED_INLINE_BLOCKS, isAiAssistantSupportEnabled } from '../../extensions/ai-assistant';
+import { getFeatureAvailability } from '../../lib/utils/get-feature-availability';
+import { EXTENDED_BLOCKS } from '../constants';
+
+export const AI_ASSISTANT_SUPPORT_NAME = 'ai-assistant-support';
+
+// Check if the AI Assistant support is enabled.
+export const isAiAssistantSupportEnabled = getFeatureAvailability( AI_ASSISTANT_SUPPORT_NAME );
 
 /**
  * Check if it is possible to extend the block as an inline extension.
  * @param {string} blockName - The block name.
- * @return {boolean}          Whether it is possible to extend the block.
+ * @return {boolean}           Whether it is possible to extend the block.
  */
 export function isPossibleToExtendBlock( blockName: string ): boolean {
 	// Check if the AI Assistant block is registered. If not, we understand that Jetpack AI is not active.
@@ -27,7 +33,7 @@ export function isPossibleToExtendBlock( blockName: string ): boolean {
 	}
 
 	// Only extend the blocks in the inline blocks list
-	if ( ! EXTENDED_INLINE_BLOCKS.includes( blockName ) ) {
+	if ( ! EXTENDED_BLOCKS.includes( blockName ) ) {
 		return false;
 	}
 
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/list-item/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/list-item/index.tsx
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/list-item/index.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/list-item/index.tsx
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/list/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/list/index.tsx
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/list/index.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/list/index.tsx
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/paragraph/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/paragraph/index.tsx
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/paragraph/index.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/paragraph/index.tsx
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/types.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/types.ts
similarity index 100%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/types.ts
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/types.ts
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/with-ai-extension.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/with-ai-extension.tsx
similarity index 98%
rename from projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/with-ai-extension.tsx
rename to projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/with-ai-extension.tsx
index 6f11f273c17f1..f0cdb68296901 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/inline-extensions/with-ai-extension.tsx
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/with-ai-extension.tsx
@@ -27,11 +27,11 @@ import { isPossibleToExtendBlock } from './lib/is-possible-to-extend-block';
 /*
  * Types
  */
+import type { ExtendedBlockProp } from './constants';
 import type {
 	AiAssistantDropdownOnChangeOptionsArgProps,
 	OnRequestSuggestion,
 } from '../components/ai-assistant-toolbar-dropdown/dropdown-content';
-import type { ExtendedInlineBlockProp } from '../extensions/ai-assistant';
 import type { PromptTypeProp } from '../lib/prompt';
 import type {
 	PromptMessagesProp,
@@ -547,7 +547,7 @@ const blockEditWithAiComponents = createHigherOrderComponent( BlockEdit => {
  * @param {string} name     - The block name.
  * @return {object}          The extended block settings.
  */
-function blockWithInlineExtension( settings, name: ExtendedInlineBlockProp ) {
+function blockWithInlineExtension( settings, name: ExtendedBlockProp ) {
 	// Only extend the allowed block types and when AI is enabled
 	const possibleToExtendBlock = isPossibleToExtendBlock( name );
 
@@ -558,6 +558,12 @@ function blockWithInlineExtension( settings, name: ExtendedInlineBlockProp ) {
 	return {
 		...settings,
 		edit: blockEditWithAiComponents( settings.edit ),
+		supports: {
+			...settings.supports,
+			'jetpack/ai': {
+				assistant: true,
+			},
+		},
 	};
 }
 
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-transform-to-assistant/index.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-transform-to-assistant/index.ts
index c04e0d811fa1e..1b6c826aa2cfb 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-transform-to-assistant/index.ts
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-transform-to-assistant/index.ts
@@ -7,14 +7,13 @@ import { useCallback } from 'react';
 /*
  * Internal dependencies
  */
-import { ALL_EXTENDED_BLOCKS } from '../../extensions/ai-assistant';
-import { getStoreBlockId } from '../../extensions/ai-assistant/with-ai-assistant';
 import { getBlocksContent } from '../../lib/utils/block-content';
-import { transformToAIAssistantBlock } from '../../transforms';
+import { TRANSFORMABLE_BLOCKS, transformToAIAssistantBlock } from '../../transforms';
 /*
  * Types
  */
 import type { AiAssistantDropdownOnChangeOptionsArgProps } from '../../components/ai-assistant-toolbar-dropdown/dropdown-content';
+import type { ExtendedBlockProp } from '../../extensions/constants';
 import type { PromptTypeProp } from '../../lib/prompt';
 import type { Block } from '@automattic/jetpack-ai-client';
 
@@ -31,6 +30,10 @@ type CoreBlockEditorSelect = {
 	getBlockParents: ( clientId: string ) => string[];
 };
 
+export function getStoreBlockId( clientId ) {
+	return `ai-assistant-block-${ clientId }`;
+}
+
 const useTransformToAssistant = () => {
 	const { replaceBlock, removeBlocks } = useDispatch(
 		'core/block-editor'
@@ -53,8 +56,7 @@ const useTransformToAssistant = () => {
 				return false;
 			}
 
-			// The block must be an extended block
-			if ( ! ALL_EXTENDED_BLOCKS.includes( blockName ) ) {
+			if ( ! TRANSFORMABLE_BLOCKS.includes( blockName ) ) {
 				return false;
 			}
 
@@ -97,7 +99,7 @@ const useTransformToAssistant = () => {
 			};
 
 			const newAIAssistantBlock = transformToAIAssistantBlock(
-				firstBlock.name as string,
+				firstBlock.name as ExtendedBlockProp,
 				extendedBlockAttributes
 			);
 
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/supports/Readme.md b/projects/plugins/jetpack/extensions/blocks/ai-assistant/supports/Readme.md
deleted file mode 100644
index 81f72bfd30751..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/supports/Readme.md
+++ /dev/null
@@ -1,16 +0,0 @@
-# `jetpack/ai` supports
-
-It allows blocks to add Jetpack AI features supports
-
-## assistant
-
-```json
-"supports": {
-	"jetpack/ai": {
-		"assistance": true
-	}
-}
-```
-
-It will enhance the block with the assistant features,
-such as populating the block sidebar with the AI Assistant panel.
\ No newline at end of file
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/supports/index.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/supports/index.ts
deleted file mode 100644
index b33d515e1dce2..0000000000000
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/supports/index.ts
+++ /dev/null
@@ -1,40 +0,0 @@
-/**
- * External dependencies
- */
-import { getBlockSupport } from '@wordpress/blocks';
-import { addFilter } from '@wordpress/hooks';
-/**
- * Internal dependencies
- */
-import { isPossibleToExtendBlock } from '../extensions/ai-assistant';
-import withAIAssistant from '../extensions/ai-assistant/with-ai-assistant';
-
-export const SUPPORT_NAME = 'jetpack/ai';
-
-function handleJetpackAISupports( settings ) {
-	const jetpackAISupports = getBlockSupport( settings, SUPPORT_NAME );
-	if ( ! jetpackAISupports ) {
-		return settings;
-	}
-
-	// Check specific for the `assistant` support.
-	if ( ! jetpackAISupports?.assistant ) {
-		return settings;
-	}
-
-	if ( ! isPossibleToExtendBlock() ) {
-		return settings;
-	}
-
-	return {
-		...settings,
-		edit: withAIAssistant( settings.edit ),
-	};
-}
-
-addFilter(
-	'blocks.registerBlockType',
-	'jetpack/handle-jetpack-ai-supports',
-	handleJetpackAISupports,
-	100
-);
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/transforms/index.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/transforms/index.tsx
index 468158dbdc40c..19d33a4bfe8f1 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/transforms/index.tsx
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/transforms/index.tsx
@@ -7,16 +7,14 @@ import { createBlock, getSaveContent } from '@wordpress/blocks';
  * Internal dependencies
  */
 import metadata from '../block.json';
-import {
-	EXTENDED_TRANSFORMATIVE_BLOCKS,
-	isPossibleToExtendBlock,
-} from '../extensions/ai-assistant';
 /**
  * Types
  */
-import type { ExtendedBlockProp, ExtendedInlineBlockProp } from '../extensions/ai-assistant';
+import type { ExtendedBlockProp } from '../extensions/constants';
 import type { PromptItemProps } from '../lib/prompt';
 
+export const TRANSFORMABLE_BLOCKS = [ 'core/heading', 'core/paragraph', 'core/list' ];
+
 const from: unknown[] = [];
 
 /**
@@ -26,10 +24,7 @@ const from: unknown[] = [];
  * @param {object}            attrs     - Block attributes.
  * @return {object}                      AI Assistant block instance.
  */
-export function transformToAIAssistantBlock(
-	blockType: ExtendedBlockProp | ExtendedInlineBlockProp,
-	attrs
-) {
+export function transformToAIAssistantBlock( blockType: ExtendedBlockProp, attrs ) {
 	const { content, ...restAttrs } = attrs;
 	let htmlContent = content;
 
@@ -65,11 +60,11 @@ export function transformToAIAssistantBlock(
 /*
  * Create individual transform handler for each block type.
  */
-for ( const blockType of EXTENDED_TRANSFORMATIVE_BLOCKS ) {
+for ( const blockType of TRANSFORMABLE_BLOCKS ) {
 	from.push( {
 		type: 'block',
 		blocks: [ blockType ],
-		isMatch: () => isPossibleToExtendBlock(),
+		isMatch: () => TRANSFORMABLE_BLOCKS.includes( blockType ),
 		transform: ( attrs, innerBlocks ) => {
 			const content = getSaveContent( blockType, attrs, innerBlocks );
 			return transformToAIAssistantBlock( blockType as ExtendedBlockProp, { ...attrs, content } );
diff --git a/projects/plugins/jetpack/extensions/index.json b/projects/plugins/jetpack/extensions/index.json
index 36eae8e4272fe..41b3eea98d413 100644
--- a/projects/plugins/jetpack/extensions/index.json
+++ b/projects/plugins/jetpack/extensions/index.json
@@ -79,7 +79,6 @@
 		"v6-video-frame-poster",
 		"videopress/video-chapters",
 		"ai-assistant-backend-prompts",
-		"ai-assistant-extensions-support",
 		"ai-title-optimization-keywords-support"
 	],
 	"experimental": [ "ai-image", "ai-paragraph" ],
diff --git a/projects/plugins/jetpack/extensions/plugins/ai-content-lens/editor.js b/projects/plugins/jetpack/extensions/plugins/ai-content-lens/editor.js
index fa5ab1b5ffa8c..53658e93e720e 100644
--- a/projects/plugins/jetpack/extensions/plugins/ai-content-lens/editor.js
+++ b/projects/plugins/jetpack/extensions/plugins/ai-content-lens/editor.js
@@ -4,7 +4,7 @@ import { store as editorStore } from '@wordpress/editor';
 import { addFilter } from '@wordpress/hooks';
 import debugFactory from 'debug';
 import metadata from '../../blocks/ai-assistant/block.json';
-import { isPossibleToExtendBlock } from '../../blocks/ai-assistant/extensions/ai-assistant';
+import { isPossibleToExtendBlock } from '../../blocks/ai-assistant/extensions/lib/is-possible-to-extend-block';
 import { aiExcerptPluginName, aiExcerptPluginSettings } from '.';
 
 const debug = debugFactory( 'jetpack-ai-content-lens:registration' );