From 948f55b505a363cb09d6a04cec9efc74664991e9 Mon Sep 17 00:00:00 2001
From: Per-Kristian Nordnes <per.kristian.nordnes@gmail.com>
Date: Tue, 31 Oct 2023 13:49:55 +0100
Subject: [PATCH] wip

---
 packages/sanity/package.json                  |   2 +-
 .../inputs/PortableText/PortableTextInput.tsx |   5 +-
 .../inputs/PortableText/object/Annotation.tsx |   1 -
 .../core/form/types/definitionExtensions.ts   |   3 +-
 .../sanity/src/core/form/types/inputProps.ts  |   3 +-
 .../scratchPad/components/editor/Editable.tsx | 217 ++++++++++++++----
 .../scratchPad/components/editor/Input.tsx    | 145 ++++++++++--
 .../components/rendering/renderBlock.tsx      |  65 ------
 .../sanity/src/scratchPad/config/index.ts     |  14 +-
 .../scratchPad/context/ScratchPadProvider.tsx |  28 ++-
 .../sanity/src/scratchPad/scratchPadTool.ts   |  16 +-
 11 files changed, 353 insertions(+), 146 deletions(-)
 delete mode 100644 packages/sanity/src/scratchPad/components/rendering/renderBlock.tsx

diff --git a/packages/sanity/package.json b/packages/sanity/package.json
index 4d86ff71162..6c2754b2658 100644
--- a/packages/sanity/package.json
+++ b/packages/sanity/package.json
@@ -132,7 +132,7 @@
     "build": "pkg-utils build --tsconfig tsconfig.lib.json",
     "postbuild": "run-s check:package",
     "check:package": "pkg-utils --tsconfig tsconfig.lib.json",
-    "clean": "rimraf _internal.js cli.js desk.js lib router.js",
+    "clean": "rimraf _internal.js cli.js desk.js lib router.js scratchPad.js",
     "coverage": "jest --coverage",
     "prepublishOnly": "run-s write:version write:ui-version",
     "test": "pkg-utils --strict && jest",
diff --git a/packages/sanity/src/core/form/inputs/PortableText/PortableTextInput.tsx b/packages/sanity/src/core/form/inputs/PortableText/PortableTextInput.tsx
index c9c1b11e4fa..fb58c9e9276 100644
--- a/packages/sanity/src/core/form/inputs/PortableText/PortableTextInput.tsx
+++ b/packages/sanity/src/core/form/inputs/PortableText/PortableTextInput.tsx
@@ -235,6 +235,9 @@ export function PortableTextInput(props: PortableTextInputProps) {
   // Handle editor changes
   const handleEditorChange = useCallback(
     (change: EditorChange): void => {
+      if (editorRef.current && onEditorChange) {
+        onEditorChange(change, editorRef.current)
+      }
       switch (change.type) {
         case 'mutation':
           onChange(toFormPatches(change.patches))
@@ -279,7 +282,7 @@ export function PortableTextInput(props: PortableTextInputProps) {
         default:
       }
     },
-    [onBlur, onChange, onPathFocus, toast],
+    [onBlur, onChange, onEditorChange, onPathFocus, toast],
   )
 
   useEffect(() => {
diff --git a/packages/sanity/src/core/form/inputs/PortableText/object/Annotation.tsx b/packages/sanity/src/core/form/inputs/PortableText/object/Annotation.tsx
index 569ae9460a0..191415394c2 100644
--- a/packages/sanity/src/core/form/inputs/PortableText/object/Annotation.tsx
+++ b/packages/sanity/src/core/form/inputs/PortableText/object/Annotation.tsx
@@ -278,7 +278,6 @@ export const DefaultAnnotationComponent = (props: BlockAnnotationProps) => {
     }
     return 'default'
   }, [isLink, hasError, hasWarning])
-
   return (
     <Root
       $toneKey={toneKey}
diff --git a/packages/sanity/src/core/form/types/definitionExtensions.ts b/packages/sanity/src/core/form/types/definitionExtensions.ts
index 99326975de6..7ff3999c678 100644
--- a/packages/sanity/src/core/form/types/definitionExtensions.ts
+++ b/packages/sanity/src/core/form/types/definitionExtensions.ts
@@ -23,6 +23,7 @@ import {
   BooleanInputProps,
   NumberInputProps,
   ObjectInputProps,
+  PortableTextInputProps,
   StringInputProps,
 } from './inputProps'
 import {ObjectItem, ObjectItemProps, PrimitiveItemProps} from './itemProps'
@@ -45,7 +46,7 @@ export interface ArrayOfObjectsComponents {
   diff?: ComponentType<any>
   field?: ComponentType<ArrayFieldProps>
   inlineBlock?: ComponentType<BlockProps>
-  input?: ComponentType<ArrayOfObjectsInputProps | PortableTextInputProps>
+  input?: ComponentType<ArrayOfObjectsInputProps>
   item?: ComponentType<ObjectItemProps>
   preview?: ComponentType<PreviewProps>
 }
diff --git a/packages/sanity/src/core/form/types/inputProps.ts b/packages/sanity/src/core/form/types/inputProps.ts
index 7292aee13b4..cdabbf268fd 100644
--- a/packages/sanity/src/core/form/types/inputProps.ts
+++ b/packages/sanity/src/core/form/types/inputProps.ts
@@ -20,6 +20,7 @@ import {
   HotkeyOptions,
   OnCopyFn,
   OnPasteFn,
+  PortableTextEditor,
   RangeDecoration,
 } from '@sanity/portable-text-editor'
 import {FormPatch, PatchEvent} from '../patch'
@@ -497,7 +498,7 @@ export interface PortableTextInputProps
   /**
    * Returns changes from the underlying editor
    */
-  onEditorChange?: (change: EditorChange) => void
+  onEditorChange?: (change: EditorChange, editor: PortableTextEditor) => void
   /**
    * Custom copy function
    */
diff --git a/packages/sanity/src/scratchPad/components/editor/Editable.tsx b/packages/sanity/src/scratchPad/components/editor/Editable.tsx
index 45f80731c8c..0704b339a0a 100644
--- a/packages/sanity/src/scratchPad/components/editor/Editable.tsx
+++ b/packages/sanity/src/scratchPad/components/editor/Editable.tsx
@@ -1,6 +1,8 @@
-import React, {PropsWithChildren, useCallback, useEffect, useMemo, useRef, useState} from 'react'
-import {Stack} from '@sanity/ui'
+import React, {PropsWithChildren, useCallback, useEffect, useMemo, useRef} from 'react'
+import {BoundaryElementProvider, Stack} from '@sanity/ui'
 import {
+  BlockAnnotationRenderProps,
+  BlockChildRenderProps,
   BlockDecoratorRenderProps,
   BlockRenderProps,
   BlockStyleRenderProps,
@@ -12,17 +14,18 @@ import {
 } from '@sanity/portable-text-editor'
 import styled, {css} from 'styled-components'
 import {debounce} from 'lodash'
-import isHotkey from 'is-hotkey'
+import {PortableTextTextBlock} from '@sanity/types'
 import {useScratchPad} from '../../hooks/useScratchPad'
 import {ScratchPadContextValue} from '../../context/ScratchPadProvider'
 import {PortableTextInputProps} from '../../../core'
-import {TextBlock} from '../rendering/renderBlock'
 import {Style} from '../../../core/form/inputs/PortableText/text/Style'
 import {Decorator} from '../../../core/form/inputs/PortableText/text'
+import {Annotation} from '../../../core/form/inputs/PortableText/object/Annotation'
+import {InlineObject} from '../../../core/form/inputs/PortableText/object/InlineObject'
+import {TextBlock} from '../../../core/form/inputs/PortableText/text/TextBlock'
 import {AssistanceRange} from './AssistanceRange'
 
 const INLINE_STYLE: React.CSSProperties = {outline: 'none'}
-const MODIFIER_KEY = 'Alt+Control'
 
 const EditableWrapStack = styled(Stack)(() => {
   return css`
@@ -48,7 +51,6 @@ export function Editable(props: EditableProps) {
   const editableRef = useRef<HTMLDivElement | null>(null)
   const editor = usePortableTextEditor()
   const editorSelection = usePortableTextEditorSelection()
-  const [modifierKeyPressed, setIsModifierKeyPressed] = useState(false)
 
   const {onEditorBeforeInput, onAssistanceRangeSelect, assistanceSelection, editorFocused} =
     useScratchPad()
@@ -68,47 +70,69 @@ export function Editable(props: EditableProps) {
 
   const renderPlaceholder = useCallback(() => <span>{placeholder}</span>, [placeholder])
 
-  const handleSetAssistantRangeKey = useCallback((event: React.KeyboardEvent<HTMLDivElement>) => {
-    if (isHotkey(MODIFIER_KEY, event)) {
-      setIsModifierKeyPressed(true)
-    }
-  }, [])
-
   const handleKeyDown = useCallback(
     (event: React.KeyboardEvent<HTMLDivElement>) => {
-      handleSetAssistantRangeKey(event)
       if (onKeyDown) {
         onKeyDown(event)
       }
     },
-    [handleSetAssistantRangeKey, onKeyDown],
+    [onKeyDown],
   )
 
-  const handleKeyUp = useCallback(() => {
-    setIsModifierKeyPressed(false)
-  }, [])
-
-  // Update the assistance selection when user is done selecting something, and modifier key is pressed
+  // Update the assistance selection when user is 'done' selecting something
   useEffect(() => {
     createAssistSelectionDebounced(editor, onAssistanceRangeSelect)
-  }, [editor, modifierKeyPressed, onAssistanceRangeSelect])
+  }, [editor, onAssistanceRangeSelect, editorSelection])
 
-  // These render functions piggybacks on the PT-input's rendering components
+  // These render functions starting with a underscore piggybacks on the PT-input's rendering components
   const _renderBlock = useCallback(
     (editorRenderProps: BlockRenderProps) => {
-      if (!rootElementRef.current) {
-        return <></>
-      }
       return (
         <TextBlock
-          formProps={formProps}
-          blockProps={editorRenderProps}
-          boundaryElement={rootElementRef.current}
-          scrollElement={rootElementRef.current}
-        />
+          floatingBoundary={editorRenderProps.editorElementRef.current}
+          focused={editorRenderProps.focused}
+          isFullscreen
+          // eslint-disable-next-line react/jsx-handler-names
+          onItemClose={formProps.onItemClose}
+          // eslint-disable-next-line react/jsx-handler-names
+          onItemOpen={formProps.onItemOpen}
+          // eslint-disable-next-line react/jsx-handler-names
+          onItemRemove={formProps.onItemRemove}
+          // eslint-disable-next-line react/jsx-handler-names
+          onPathFocus={formProps.onPathFocus}
+          path={formProps.path.concat(editorRenderProps.path)}
+          readOnly={formProps.readOnly}
+          referenceBoundary={rootElementRef.current}
+          renderAnnotation={formProps.renderAnnotation}
+          renderField={formProps.renderField}
+          renderInlineBlock={formProps.renderInlineBlock}
+          renderInput={formProps.renderInput}
+          renderItem={formProps.renderItem}
+          renderPreview={formProps.renderPreview}
+          renderBlock={formProps.renderBlock}
+          schemaType={editorRenderProps.schemaType}
+          selected={editorRenderProps.selected}
+          value={editorRenderProps.value as PortableTextTextBlock}
+        >
+          {editorRenderProps.children}
+        </TextBlock>
       )
     },
-    [formProps],
+    [
+      formProps.onItemClose,
+      formProps.onItemOpen,
+      formProps.onItemRemove,
+      formProps.onPathFocus,
+      formProps.path,
+      formProps.readOnly,
+      formProps.renderAnnotation,
+      formProps.renderBlock,
+      formProps.renderField,
+      formProps.renderInlineBlock,
+      formProps.renderInput,
+      formProps.renderItem,
+      formProps.renderPreview,
+    ],
   )
 
   const _renderDecorator = useCallback((editorRenderProps: BlockDecoratorRenderProps) => {
@@ -119,24 +143,125 @@ export function Editable(props: EditableProps) {
     return <Style {...editorRenderProps} />
   }, [])
 
+  const _renderAnnotation = useCallback(
+    (editorRenderProps: BlockAnnotationRenderProps) => {
+      return (
+        <Annotation
+          floatingBoundary={editorRenderProps.editorElementRef.current}
+          editorNodeFocused={editorFocused}
+          focused={editorRenderProps.focused}
+          // eslint-disable-next-line react/jsx-handler-names
+          onItemClose={formProps.onItemClose}
+          // eslint-disable-next-line react/jsx-handler-names
+          onItemOpen={formProps.onItemOpen}
+          // eslint-disable-next-line react/jsx-handler-names
+          onPathFocus={formProps.onPathFocus}
+          path={formProps.path.concat(editorRenderProps.path)}
+          referenceBoundary={rootElementRef.current}
+          renderAnnotation={formProps.renderAnnotation}
+          renderField={formProps.renderField}
+          renderInput={formProps.renderInput}
+          renderItem={formProps.renderItem}
+          renderPreview={formProps.renderPreview}
+          selected={editorRenderProps.selected}
+          schemaType={editorRenderProps.schemaType}
+          value={editorRenderProps.value}
+        >
+          {editorRenderProps.children}
+        </Annotation>
+      )
+    },
+    [
+      editorFocused,
+      formProps.onItemClose,
+      formProps.onItemOpen,
+      formProps.onPathFocus,
+      formProps.path,
+      formProps.renderAnnotation,
+      formProps.renderField,
+      formProps.renderInput,
+      formProps.renderItem,
+      formProps.renderPreview,
+    ],
+  )
+
+  const _renderChild = useCallback(
+    (editorRenderProps: BlockChildRenderProps) => {
+      const {
+        children,
+        focused: childFocused,
+        path: childPath,
+        selected,
+        schemaType: childSchemaType,
+        value: child,
+      } = editorRenderProps
+      const isSpan = child._type === editor.schemaTypes.span.name
+      if (isSpan) {
+        return children
+      }
+      return (
+        <InlineObject
+          floatingBoundary={editorRenderProps.editorElementRef.current}
+          focused={childFocused}
+          // eslint-disable-next-line react/jsx-handler-names
+          onItemClose={formProps.onItemClose}
+          // eslint-disable-next-line react/jsx-handler-names
+          onItemOpen={formProps.onItemOpen}
+          // eslint-disable-next-line react/jsx-handler-names
+          onPathFocus={formProps.onPathFocus}
+          path={formProps.path.concat(childPath)}
+          readOnly={formProps.readOnly}
+          referenceBoundary={rootElementRef.current}
+          relativePath={childPath}
+          renderAnnotation={formProps.renderAnnotation}
+          renderBlock={formProps.renderBlock}
+          renderField={formProps.renderField}
+          renderInlineBlock={formProps.renderInlineBlock}
+          renderInput={formProps.renderInput}
+          renderItem={formProps.renderItem}
+          renderPreview={formProps.renderPreview}
+          schemaType={childSchemaType}
+          selected={selected}
+          value={child}
+        />
+      )
+    },
+    [
+      editor.schemaTypes.span.name,
+      formProps.onItemClose,
+      formProps.onItemOpen,
+      formProps.onPathFocus,
+      formProps.path,
+      formProps.readOnly,
+      formProps.renderAnnotation,
+      formProps.renderBlock,
+      formProps.renderField,
+      formProps.renderInlineBlock,
+      formProps.renderInput,
+      formProps.renderItem,
+      formProps.renderPreview,
+    ],
+  )
+
   return (
     <EditableWrapStack ref={rootElementRef} data-ui="EditableWrapStack">
-      <PortableTextEditable
-        onBeforeInput={onEditorBeforeInput}
-        onBlur={onBlur}
-        onFocus={onFocus}
-        onKeyDown={handleKeyDown}
-        onKeyUp={handleKeyUp}
-        ref={editableRef}
-        renderBlock={_renderBlock}
-        renderDecorator={_renderDecorator}
-        // renderChild={renderChild}
-        // renderAnnotation={renderAnnotation}
-        renderStyle={_renderStyle}
-        rangeDecorations={rangeDecorations}
-        renderPlaceholder={renderPlaceholder}
-        style={INLINE_STYLE}
-      />
+      <BoundaryElementProvider element={rootElementRef.current}>
+        <PortableTextEditable
+          onBeforeInput={onEditorBeforeInput}
+          onBlur={onBlur}
+          onFocus={onFocus}
+          onKeyDown={handleKeyDown}
+          ref={editableRef}
+          renderBlock={_renderBlock}
+          renderDecorator={_renderDecorator}
+          renderChild={_renderChild}
+          renderAnnotation={_renderAnnotation}
+          renderStyle={_renderStyle}
+          rangeDecorations={rangeDecorations}
+          renderPlaceholder={renderPlaceholder}
+          style={INLINE_STYLE}
+        />
+      </BoundaryElementProvider>
     </EditableWrapStack>
   )
 }
diff --git a/packages/sanity/src/scratchPad/components/editor/Input.tsx b/packages/sanity/src/scratchPad/components/editor/Input.tsx
index 4e5bedeb9ac..34900180edd 100644
--- a/packages/sanity/src/scratchPad/components/editor/Input.tsx
+++ b/packages/sanity/src/scratchPad/components/editor/Input.tsx
@@ -1,12 +1,28 @@
-import React, {useCallback} from 'react'
+import React, {ReactNode, createRef, startTransition, useCallback, useMemo, useRef} from 'react'
 import {Stack} from '@sanity/ui'
 import {EditorChange, Patch, PortableTextEditor} from '@sanity/portable-text-editor'
-import {PortableTextInputProps, SANITY_PATCH_TYPE} from '../../../core'
+import {isBlockSchemaType} from '@sanity/types'
+import {
+  ArrayOfObjectsItemMember,
+  FIXME,
+  FormInput,
+  InputProps,
+  ObjectFormNode,
+  PortableTextEditorElement,
+  PortableTextInputProps,
+  PortableTextMemberItem,
+  SANITY_PATCH_TYPE,
+  pathToString,
+} from '../../../core'
 import {useScratchPad} from '../../hooks/useScratchPad'
+import {isMemberArrayOfObjects} from '../../../core/form/members/object/fields/asserters'
+import {isArrayOfObjectsFieldMember} from '../../../core/form/inputs/PortableText/_helpers'
+import {PortableTextMemberItemsProvider} from '../../../core/form/inputs/PortableText/contexts/PortableTextMembers'
 import {Editable} from './Editable'
 
-export const ScratchPadInput = function ScratchPadInput(props: PortableTextInputProps) {
-  const {onChange, schemaType, value} = props
+export const ScratchPadInput = function ScratchPadInput(props: InputProps) {
+  const {onChange, schemaType, value, members, onPathFocus} = props as PortableTextInputProps
+  const portableTextMemberItemsRef: React.MutableRefObject<PortableTextMemberItem[]> = useRef([])
 
   const {editorRef, onEditorFocus, onEditorBlur} = useScratchPad()
 
@@ -15,26 +31,127 @@ export const ScratchPadInput = function ScratchPadInput(props: PortableTextInput
       if (change.type === 'mutation') {
         onChange(toFormPatches(change.patches))
       }
-      if (change.type === 'focus') {
+      if (change.type === 'selection') {
         onEditorFocus()
+        startTransition(() => {
+          if (change.selection) {
+            onPathFocus(change.selection.focus.path)
+          }
+        })
       }
       if (change.type === 'blur') {
         onEditorBlur()
       }
     },
-    [onChange, onEditorBlur, onEditorFocus],
+    [onChange, onEditorBlur, onEditorFocus, onPathFocus],
   )
 
+  // Populate the portableTextMembers Map
+  const portableTextMemberItems: PortableTextMemberItem[] = useMemo(() => {
+    const result: {
+      kind: PortableTextMemberItem['kind']
+      member: ArrayOfObjectsItemMember
+      node: ObjectFormNode
+    }[] = []
+
+    for (const member of members) {
+      if (member.kind === 'item') {
+        const isObjectBlock = !isBlockSchemaType(member.item.schemaType)
+        if (isObjectBlock) {
+          result.push({kind: 'objectBlock', member, node: member.item})
+        } else {
+          // Also include regular text blocks with validation, presence or changes.
+          if (
+            member.item.validation.length > 0 ||
+            member.item.changed ||
+            member.item.presence?.length
+          ) {
+            result.push({kind: 'textBlock', member, node: member.item})
+          }
+          // Inline objects
+          const childrenField = member.item.members.find(
+            (f) => f.kind === 'field' && f.name === 'children',
+          )
+          if (
+            childrenField &&
+            childrenField.kind === 'field' &&
+            isMemberArrayOfObjects(childrenField)
+          ) {
+            // eslint-disable-next-line max-depth
+            for (const child of childrenField.field.members) {
+              // eslint-disable-next-line max-depth
+              if (child.kind === 'item' && child.item.schemaType.name !== 'span') {
+                result.push({kind: 'inlineObject', member: child, node: child.item})
+              }
+            }
+          }
+          // Markdefs
+          const markDefArrayMember = member.item.members
+            .filter(isArrayOfObjectsFieldMember)
+            .find((f) => f.name === 'markDefs')
+          if (markDefArrayMember) {
+            // eslint-disable-next-line max-depth
+            for (const child of markDefArrayMember.field.members) {
+              // eslint-disable-next-line max-depth
+              if (child.kind === 'item' && child.item.schemaType.jsonType === 'object') {
+                result.push({
+                  kind: 'annotation',
+                  member: child,
+                  node: child.item,
+                })
+              }
+            }
+          }
+        }
+      }
+    }
+    const items: PortableTextMemberItem[] = result.map((item) => {
+      const key = pathToString(item.node.path)
+      const existingItem = portableTextMemberItemsRef.current.find((ref) => ref.key === key)
+      let input: ReactNode
+
+      if (item.kind !== 'textBlock') {
+        input = <FormInput absolutePath={item.node.path} {...(props as FIXME)} />
+      }
+
+      if (existingItem) {
+        // Only update the input if the node is open or the value has changed
+        // This is a performance optimization.
+        if (item.member.open || existingItem.node.value !== item.node.value) {
+          existingItem.input = input
+        }
+        existingItem.member = item.member
+        existingItem.node = item.node
+        return existingItem
+      }
+
+      return {
+        kind: item.kind,
+        key,
+        member: item.member,
+        node: item.node,
+        elementRef: createRef<PortableTextEditorElement | null>(),
+        input,
+      }
+    })
+
+    portableTextMemberItemsRef.current = items
+
+    return items
+  }, [members, props])
+
   return (
     <Stack data-testid="scratchpad-input">
-      <PortableTextEditor
-        onChange={handleEditorChange}
-        schemaType={schemaType}
-        value={value}
-        ref={editorRef}
-      >
-        <Editable formProps={props} />
-      </PortableTextEditor>
+      <PortableTextMemberItemsProvider memberItems={portableTextMemberItems}>
+        <PortableTextEditor
+          onChange={handleEditorChange}
+          schemaType={schemaType}
+          value={value}
+          ref={editorRef}
+        >
+          <Editable formProps={props as PortableTextInputProps} />
+        </PortableTextEditor>
+      </PortableTextMemberItemsProvider>
     </Stack>
   )
 }
diff --git a/packages/sanity/src/scratchPad/components/rendering/renderBlock.tsx b/packages/sanity/src/scratchPad/components/rendering/renderBlock.tsx
deleted file mode 100644
index 5a79f511012..00000000000
--- a/packages/sanity/src/scratchPad/components/rendering/renderBlock.tsx
+++ /dev/null
@@ -1,65 +0,0 @@
-import {BlockRenderProps} from '@sanity/portable-text-editor'
-import React from 'react'
-import {PortableTextTextBlock} from '@sanity/types'
-import {TextBlock as InputsTextBlock} from '../../../core/form/inputs/PortableText/text/TextBlock'
-import {PortableTextInputProps} from '../../../core'
-
-interface EditorRenderBlockProps {
-  blockProps: BlockRenderProps
-  formProps: PortableTextInputProps
-  boundaryElement: HTMLElement
-  scrollElement: HTMLElement
-}
-
-export const TextBlock = (props: EditorRenderBlockProps) => {
-  const {
-    children,
-    focused: blockFocused,
-    path: blockPath,
-    schemaType: blockSchemaType,
-    selected,
-    value: block,
-  } = props.blockProps
-  const {
-    onItemClose,
-    onItemOpen,
-    onItemRemove,
-    onPathFocus,
-    path,
-    readOnly,
-    renderAnnotation,
-    renderBlock,
-    renderField,
-    renderInlineBlock,
-    renderInput,
-    renderItem,
-    renderPreview,
-  } = props.formProps
-  const {boundaryElement, scrollElement} = props
-  return (
-    <InputsTextBlock
-      floatingBoundary={boundaryElement}
-      focused={blockFocused}
-      isFullscreen={false}
-      onItemClose={onItemClose}
-      onItemOpen={onItemOpen}
-      onItemRemove={onItemRemove}
-      onPathFocus={onPathFocus}
-      path={path.concat(blockPath)}
-      readOnly={readOnly}
-      referenceBoundary={scrollElement}
-      renderAnnotation={renderAnnotation}
-      renderBlock={renderBlock}
-      renderField={renderField}
-      renderInlineBlock={renderInlineBlock}
-      renderInput={renderInput}
-      renderItem={renderItem}
-      renderPreview={renderPreview}
-      schemaType={blockSchemaType}
-      selected={selected}
-      value={block as PortableTextTextBlock}
-    >
-      {children}
-    </InputsTextBlock>
-  )
-}
diff --git a/packages/sanity/src/scratchPad/config/index.ts b/packages/sanity/src/scratchPad/config/index.ts
index 8e3de69184a..2cbad9d53ec 100644
--- a/packages/sanity/src/scratchPad/config/index.ts
+++ b/packages/sanity/src/scratchPad/config/index.ts
@@ -1,4 +1,3 @@
-import {BoldIcon, ItalicIcon, LinkIcon} from '@sanity/icons'
 import {Schema} from '@sanity/schema'
 import {defineField, defineType} from '@sanity/types'
 import {ScratchPadInput} from '../components/editor/Input'
@@ -7,7 +6,20 @@ import {fragmentToAssistantText} from '../utils/toAssistantText'
 const blockType = defineField({
   type: 'block',
   name: 'block',
+  of: [
+    {
+      type: 'object',
+      name: 'test',
+      fields: [
+        {
+          type: 'string',
+          name: 'testString',
+        },
+      ],
+    },
+  ],
   validation: (Rule) =>
+    // Test validation. TODO: remove
     Rule.custom<any>((block) => {
       const text = fragmentToAssistantText(block ? [block] : [])
       return text.length === 1 ? 'Please write a longer paragraph.' : true
diff --git a/packages/sanity/src/scratchPad/context/ScratchPadProvider.tsx b/packages/sanity/src/scratchPad/context/ScratchPadProvider.tsx
index de3499dd37a..12fcf364031 100644
--- a/packages/sanity/src/scratchPad/context/ScratchPadProvider.tsx
+++ b/packages/sanity/src/scratchPad/context/ScratchPadProvider.tsx
@@ -59,6 +59,20 @@ export function ScratchPadProvider(props: ScratchPadProviderProps) {
           },
         ],
       },
+      {
+        _key: '456',
+        _type: 'block',
+        style: 'normal',
+        markDefs: [],
+        children: [
+          {
+            _key: 'def',
+            _type: 'span',
+            text: 'Something something',
+            marks: [],
+          },
+        ],
+      },
     ],
   })
 
@@ -103,17 +117,9 @@ export function ScratchPadProvider(props: ScratchPadProviderProps) {
     // console.log('On editor beforeInput')
   }, [])
 
-  const onAssistanceRangeSelect = useCallback(
-    (selection: EditorSelection): void => {
-      setAssistanceSelection(selection)
-      setTimeout(() => {
-        if (selection) {
-          focusAssistantPrompt()
-        }
-      }, 100)
-    },
-    [focusAssistantPrompt],
-  )
+  const onAssistanceRangeSelect = useCallback((selection: EditorSelection): void => {
+    setAssistanceSelection(selection)
+  }, [])
 
   useEffect(() => {
     if (assistanceSelection && editor) {
diff --git a/packages/sanity/src/scratchPad/scratchPadTool.ts b/packages/sanity/src/scratchPad/scratchPadTool.ts
index 4b11bf27515..bdbecbc0707 100644
--- a/packages/sanity/src/scratchPad/scratchPadTool.ts
+++ b/packages/sanity/src/scratchPad/scratchPadTool.ts
@@ -1,22 +1,30 @@
 import {BulbFilledIcon} from '@sanity/icons'
 import {lazy} from 'react'
-import {definePlugin} from 'sanity'
+import {definePlugin, isArrayOfBlocksSchemaType} from 'sanity'
 
 const ScratchPadRoot = lazy(() => import('./components/Root'))
 
+/**
+ * @alpha
+ */
 export interface ScratchPadToolOptions {
   icon?: React.ComponentType
   name?: string
   title?: string
 }
-
+/**
+ * @alpha
+ */
 export const scratchPadTool = definePlugin<ScratchPadToolOptions | void>((options) => ({
   name: '@sanity/scratchpad-tool',
   form: {
     components: {
-      // No field like field
       field: (fieldProps) => {
-        return fieldProps.children
+        // No 'chrome' for the PT field
+        if (isArrayOfBlocksSchemaType(fieldProps.schemaType)) {
+          return fieldProps.children
+        }
+        return fieldProps.renderDefault(fieldProps)
       },
     },
   },