fix(app): line selection waits on ready

1 files changed, 45 insertions, 8 deletions

diff --git a/packages/ui/src/components/code.tsx b/packages/ui/src/components/code.tsx
index a7687444f..dbf942dbb 100644
--- a/packages/ui/src/components/code.tsx
+++ b/packages/ui/src/components/code.tsx
@@ -128,20 +128,56 @@ export function Code<T>(props: CodeProps<T>) {
     }
   }
 
-  const notifyRendered = () => {
-    if (!local.onRendered) return
+  const lineCount = () => {
+    const text = local.file.contents
+    const total = text.split("\n").length - (text.endsWith("\n") ? 1 : 0)
+    return Math.max(1, total)
+  }
+
+  const applySelection = (range: SelectedLineRange | null) => {
+    const root = getRoot()
+    if (!root) return false
+
+    const lines = lineCount()
+    if (root.querySelectorAll("[data-line]").length < lines) return false
+
+    if (!range) {
+      file().setSelectedLines(null)
+      return true
+    }
+
+    const start = Math.min(range.start, range.end)
+    const end = Math.max(range.start, range.end)
+
+    if (start < 1 || end > lines) {
+      file().setSelectedLines(null)
+      return true
+    }
+
+    if (!root.querySelector(`[data-line="${start}"]`) || !root.querySelector(`[data-line="${end}"]`)) {
+      file().setSelectedLines(null)
+      return true
+    }
+
+    const normalized = (() => {
+      if (range.endSide != null) return { start: range.start, end: range.end }
+      if (range.side !== "deletions") return range
+      if (root.querySelector("[data-deletions]") != null) return range
+      return { start: range.start, end: range.end }
+    })()
 
+    file().setSelectedLines(normalized)
+    return true
+  }
+
+  const notifyRendered = () => {
     observer?.disconnect()
     observer = undefined
     renderToken++
 
     const token = renderToken
 
-    const lines = (() => {
-      const text = local.file.contents
-      const total = text.split("\n").length - (text.endsWith("\n") ? 1 : 0)
-      return Math.max(1, total)
-    })()
+    const lines = lineCount()
 
     const isReady = (root: ShadowRoot) => root.querySelectorAll("[data-line]").length >= lines
 
@@ -152,6 +188,7 @@ export function Code<T>(props: CodeProps<T>) {
       observer = undefined
       requestAnimationFrame(() => {
         if (token !== renderToken) return
+        applySelection(lastSelection)
         local.onRendered?.()
       })
     }
@@ -241,7 +278,7 @@ export function Code<T>(props: CodeProps<T>) {
 
   const setSelectedLines = (range: SelectedLineRange | null) => {
     lastSelection = range
-    file().setSelectedLines(range)
+    applySelection(range)
   }
 
   const scheduleSelectionUpdate = () => {