From 9a877c1f9ff6b4af5d621426c9d92aeb3b8f4d67 Mon Sep 17 00:00:00 2001
From: rdunk <rupert@rupertdunk.com>
Date: Thu, 22 Feb 2024 17:03:18 +0000
Subject: [PATCH] fix: allow encoding draft ids

---
 packages/presentation/src/loader/LoaderQueries.tsx       | 5 +----
 packages/presentation/src/useParams.ts                   | 4 ++++
 .../src/csm/transformSanityNodeData.ts                   | 4 ++--
 packages/visual-editing/src/ui/Overlays.tsx              | 9 ++++++++-
 packages/visual-editing/src/ui/elementsReducer.ts        | 4 +++-
 5 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/packages/presentation/src/loader/LoaderQueries.tsx b/packages/presentation/src/loader/LoaderQueries.tsx
index 4efea2bce..7323f1c00 100644
--- a/packages/presentation/src/loader/LoaderQueries.tsx
+++ b/packages/presentation/src/loader/LoaderQueries.tsx
@@ -430,10 +430,7 @@ export function turboChargeResultIfSourceMap<T = unknown>(
         return undefined
       }
       // If there's a displayed document, always prefer it
-      if (
-        liveDocument?._id &&
-        getPublishedId(liveDocument._id) === getPublishedId(sourceDocument._id)
-      ) {
+      if (liveDocument?._id === sourceDocument._id) {
         return liveDocument
       }
       // Fallback to general documents cache
diff --git a/packages/presentation/src/useParams.ts b/packages/presentation/src/useParams.ts
index 152ff2a42..6bd27d288 100644
--- a/packages/presentation/src/useParams.ts
+++ b/packages/presentation/src/useParams.ts
@@ -11,6 +11,7 @@ import {
   PresentationSearchParams,
   PresentationStateParams,
 } from './types'
+import { getPublishedId } from 'sanity'
 
 function pruneObject<T extends RouterState | PresentationParams>(obj: T): T {
   return Object.fromEntries(
@@ -104,6 +105,9 @@ export function useParams({
     () =>
       debounce<PresentationNavigate>(
         (nextState, nextSearchState = {}, forceReplace) => {
+          // Force navigation to use published IDs only
+          if (nextState.id) nextState.id = getPublishedId(nextState.id)
+
           // Extract type, id and path as 'routerState'
           const { _searchParams: routerSearchParams, ...routerState } =
             routerStateRef.current
diff --git a/packages/visual-editing-helpers/src/csm/transformSanityNodeData.ts b/packages/visual-editing-helpers/src/csm/transformSanityNodeData.ts
index 31d0b11b2..3944a2bdb 100644
--- a/packages/visual-editing-helpers/src/csm/transformSanityNodeData.ts
+++ b/packages/visual-editing-helpers/src/csm/transformSanityNodeData.ts
@@ -1,4 +1,4 @@
-import { getPublishedId, studioPath } from '@sanity/client/csm'
+import { studioPath } from '@sanity/client/csm'
 import {
   is,
   minLength,
@@ -65,7 +65,7 @@ export function encodeSanityNodeData(node: SanityNode): string | undefined {
   }
 
   const parts = [
-    ['id', getPublishedId(_id)],
+    ['id', _id],
     ['type', type],
     ['path', pathToUrlString(studioPath.fromString(path))],
     ['base', encodeURIComponent(baseUrl)],
diff --git a/packages/visual-editing/src/ui/Overlays.tsx b/packages/visual-editing/src/ui/Overlays.tsx
index 9ee2d917c..a99269b2f 100644
--- a/packages/visual-editing/src/ui/Overlays.tsx
+++ b/packages/visual-editing/src/ui/Overlays.tsx
@@ -30,6 +30,7 @@ import type {
 import { ElementOverlay } from './ElementOverlay'
 import { overlayStateReducer } from './overlayStateReducer'
 import { useController } from './useController'
+import { getPublishedId } from '@sanity/client/csm'
 
 const Root = styled.div<{
   $zIndex?: string | number
@@ -151,7 +152,13 @@ export const Overlays: FunctionComponent<{
   const overlayEventHandler: OverlayEventHandler = useCallback(
     (message) => {
       if (message.type === 'element/click') {
-        channel.send('overlay/focus', message.sanity)
+        const { sanity } = message
+        if ('id' in sanity) {
+          channel.send('overlay/focus', {
+            ...sanity,
+            id: getPublishedId(sanity.id),
+          })
+        }
       } else if (message.type === 'overlay/activate') {
         channel.send('overlay/toggle', { enabled: true })
       } else if (message.type === 'overlay/deactivate') {
diff --git a/packages/visual-editing/src/ui/elementsReducer.ts b/packages/visual-editing/src/ui/elementsReducer.ts
index d2906ac5f..cd4595f9d 100644
--- a/packages/visual-editing/src/ui/elementsReducer.ts
+++ b/packages/visual-editing/src/ui/elementsReducer.ts
@@ -1,4 +1,5 @@
 import { PresentationMsg } from '@sanity/visual-editing-helpers'
+import { getPublishedId } from '@sanity/client/csm'
 
 import { ElementState, OverlayMsg } from '../types'
 
@@ -86,7 +87,8 @@ export const elementsReducer = (
         // We want to focus any element which matches the id and path...
         const focused =
           'path' in e.sanity &&
-          e.sanity.id === message.data.id &&
+          // The encoded ID might relate to a draft, but Presentation tool will send the published ID
+          getPublishedId(e.sanity.id) === message.data.id &&
           e.sanity.path === message.data.path
         return {
           ...e,