diff --git a/plugin/src/assistDocument/RequestRunInstructionProvider.tsx b/plugin/src/assistDocument/RequestRunInstructionProvider.tsx
index 6b43661..00829a8 100644
--- a/plugin/src/assistDocument/RequestRunInstructionProvider.tsx
+++ b/plugin/src/assistDocument/RequestRunInstructionProvider.tsx
@@ -2,7 +2,6 @@ import {useCallback, useEffect, useState} from 'react'
 import {ObjectSchemaType, PatchEvent, SanityDocument, unset} from 'sanity'
 
 import {useRunInstruction} from '../assistLayout/RunInstructionProvider'
-import {publicId} from '../helpers/ids'
 
 export interface DraftDelayedTaskArgs<T> {
   documentOnChange: (event: PatchEvent) => void
@@ -19,11 +18,6 @@ export function isDocAssistable(
   return !!(documentSchemaType.liveEdit ? published : draft)
 }
 
-export function getAssistableDocId(documentSchemaType: ObjectSchemaType, documentId: string) {
-  const baseId = publicId(documentId)
-  return documentSchemaType.liveEdit ? baseId : `drafts.${baseId}`
-}
-
 export function useRequestRunInstruction(args: {
   documentOnChange: (event: PatchEvent) => void
   // indicates if the document is a draft or liveEditable currently
diff --git a/plugin/src/assistDocument/components/FieldRefPreview.tsx b/plugin/src/assistDocument/components/FieldRefPreview.tsx
index 38f2f62..443d638 100644
--- a/plugin/src/assistDocument/components/FieldRefPreview.tsx
+++ b/plugin/src/assistDocument/components/FieldRefPreview.tsx
@@ -7,22 +7,20 @@ import {useSelectedField} from '../../assistInspector/helpers'
 import {SelectedFieldContext} from './SelectedFieldContext'
 
 export function FieldRefPreview(props: PreviewProps & {path?: string}) {
+  const {actions} = props
   const documentSchema = useContext(SelectedFieldContext)?.documentSchema
   const path = (useContext(InlineBlockValueContext) as {path?: string})?.path ?? props.path
   const selectedField = useSelectedField(documentSchema, path)
   return (
     <Flex gap={2} align="center" style={{width: '100%'}}>
       <Flex flex={1} gap={2} align="center" paddingY={3} paddingX={1}>
-        {/*<Box>
-          <Text>{selectedField?.icon ? createElement(selectedField?.icon) : <CodeIcon />}</Text>
-        </Box>*/}
         <Box>
           <Text size={1} textOverflow="ellipsis">
             {selectedField?.title ?? 'Select field'}
           </Text>
         </Box>
       </Flex>
-      <>{props.actions}</>
+      {actions as any}
     </Flex>
   )
 }
diff --git a/plugin/src/assistDocument/hooks/useAssistDocumentContextValue.tsx b/plugin/src/assistDocument/hooks/useAssistDocumentContextValue.tsx
index efb856c..12ef4d8 100644
--- a/plugin/src/assistDocument/hooks/useAssistDocumentContextValue.tsx
+++ b/plugin/src/assistDocument/hooks/useAssistDocumentContextValue.tsx
@@ -23,7 +23,7 @@ export function useAssistDocumentContextValue(
   } = useDocumentPane()
 
   // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-  // @ts-ignore this is a valid option available in `corel` - Remove after corel is merged to next
+  // @ts-ignore version is available in `corel` - Remove after corel is merged to next
   const {published, draft, version} = useEditState(
     getPublishedId(documentId),
     documentSchemaType.name,
@@ -36,7 +36,7 @@ export function useAssistDocumentContextValue(
   const assistableDocumentId = version?._id || draft?._id || published?._id || documentId
   const documentIsNew = selectedReleaseId ? !version?._id : !draft?._id && !published?._id
   const documentIsAssistable = selectedReleaseId
-    ? Boolean(version)
+    ? !!version
     : isDocAssistable(documentSchemaType, published, draft)
 
   const {params} = useAiPaneRouter()
diff --git a/plugin/src/assistInspector/AssistInspector.tsx b/plugin/src/assistInspector/AssistInspector.tsx
index 2141343..d273bc4 100644
--- a/plugin/src/assistInspector/AssistInspector.tsx
+++ b/plugin/src/assistInspector/AssistInspector.tsx
@@ -4,7 +4,6 @@ import {useCallback, useMemo, useRef} from 'react'
 import {
   type DocumentInspectorProps,
   PresenceOverlay,
-  useEditState,
   VirtualizerScrollInstanceProvider,
 } from 'sanity'
 import {
@@ -17,12 +16,9 @@ import {styled} from 'styled-components'
 
 import {DocumentForm} from '../_lib/form'
 import {AssistTypeContext} from '../assistDocument/components/AssistTypeContext'
+import {useAssistDocumentContextValue} from '../assistDocument/hooks/useAssistDocumentContextValue'
 import {useStudioAssistDocument} from '../assistDocument/hooks/useStudioAssistDocument'
-import {
-  getAssistableDocId,
-  isDocAssistable,
-  useRequestRunInstruction,
-} from '../assistDocument/RequestRunInstructionProvider'
+import {useRequestRunInstruction} from '../assistDocument/RequestRunInstructionProvider'
 import {useAiAssistanceConfig} from '../assistLayout/AiAssistanceConfigContext'
 import {giveFeedbackUrl, pluginTitle, releaseAnnouncementUrl, salesUrl} from '../constants'
 import {getConditionalMembers} from '../helpers/conditionalMembers'
@@ -209,15 +205,18 @@ export function AssistInspector(props: DocumentInspectorProps) {
     onChange: documentOnChange,
     formState,
   } = documentPane
-  const {published, draft} = useEditState(documentId, documentType, 'low')
+
+  const {assistableDocumentId, documentIsAssistable} = useAssistDocumentContextValue(
+    documentId,
+    schemaType,
+  )
 
   const formStateRef = useRef(formState)
   formStateRef.current = formState
 
-  const assistableDocId = getAssistableDocId(schemaType, documentId)
   const {instructionLoading, requestRunInstruction} = useRequestRunInstruction({
     documentOnChange,
-    isDocAssistable: isDocAssistable(schemaType, published, draft),
+    isDocAssistable: documentIsAssistable,
   })
 
   const typePath = useTypePath(docValue, pathKey ?? '')
@@ -268,14 +267,14 @@ export function AssistInspector(props: DocumentInspectorProps) {
       pathKey &&
       typePath &&
       requestRunInstruction({
-        documentId: assistableDocId,
+        documentId: assistableDocumentId,
         path: pathKey,
         typePath,
         assistDocumentId: assistDocumentId(documentType),
         instruction,
         conditionalMembers: formStateRef.current ? getConditionalMembers(formStateRef.current) : [],
       }),
-    [pathKey, instruction, typePath, documentType, assistableDocId, requestRunInstruction],
+    [pathKey, instruction, typePath, documentType, assistableDocumentId, requestRunInstruction],
   )
 
   const Region = useCallback((_props: any) => {
@@ -380,7 +379,7 @@ export function AssistInspector(props: DocumentInspectorProps) {
           )}
 
           <InstructionTaskHistoryButton
-            documentId={assistableDocId}
+            documentId={assistableDocumentId}
             tasks={tasks}
             instructions={instructions}
             showTitles={!instructionKey}
diff --git a/plugin/src/fieldActions/assistFieldActions.tsx b/plugin/src/fieldActions/assistFieldActions.tsx
index f8e31b6..51ce99b 100644
--- a/plugin/src/fieldActions/assistFieldActions.tsx
+++ b/plugin/src/fieldActions/assistFieldActions.tsx
@@ -13,10 +13,7 @@ import {useDocumentPane} from 'sanity/structure'
 import {useAssistDocumentContext} from '../assistDocument/AssistDocumentContext'
 import {getIcon} from '../assistDocument/components/instruction/appearance/IconInput'
 import {useAssistDocumentContextValue} from '../assistDocument/hooks/useAssistDocumentContextValue'
-import {
-  getAssistableDocId,
-  useRequestRunInstruction,
-} from '../assistDocument/RequestRunInstructionProvider'
+import {useRequestRunInstruction} from '../assistDocument/RequestRunInstructionProvider'
 import {aiInspectorId} from '../assistInspector/constants'
 import {useSelectedField, useTypePath} from '../assistInspector/helpers'
 import {pluginTitleShort} from '../constants'
@@ -50,7 +47,6 @@ export const assistFieldActions: DocumentFieldAction = {
       inspector,
       documentOnChange,
       documentSchemaType,
-      documentId,
       selectedPath,
       assistableDocumentId,
     } =
@@ -72,7 +68,6 @@ export const assistFieldActions: DocumentFieldAction = {
     const typePath = useTypePath(docValue, pathKey)
     const assistDocumentId = assistDocument?._id
 
-    const assistableDocId = getAssistableDocId(documentSchemaType, documentId)
     const {requestRunInstruction} = useRequestRunInstruction({
       documentOnChange,
       isDocAssistable: documentIsAssistable ?? false,
@@ -121,11 +116,11 @@ export const assistFieldActions: DocumentFieldAction = {
 
     const onInstructionAction = useCallback(
       (instruction: StudioInstruction) => {
-        if (!pathKey || !fieldAssistKey || !assistDocumentId || !assistableDocId) {
+        if (!pathKey || !fieldAssistKey || !assistDocumentId || !assistableDocumentId) {
           return
         }
         requestRunInstruction({
-          documentId: assistableDocId,
+          documentId: assistableDocumentId,
           assistDocumentId,
           path: pathKey,
           typePath,
@@ -135,7 +130,14 @@ export const assistFieldActions: DocumentFieldAction = {
             : [],
         })
       },
-      [requestRunInstruction, assistableDocId, pathKey, typePath, assistDocumentId, fieldAssistKey],
+      [
+        requestRunInstruction,
+        assistableDocumentId,
+        pathKey,
+        typePath,
+        assistDocumentId,
+        fieldAssistKey,
+      ],
     )
 
     const privateInstructions = useMemo(