fix dialog root complexity

1 files changed, 95 insertions, 58 deletions

diff --git a/packages/ui/src/context/dialog.tsx b/packages/ui/src/context/dialog.tsx
index af5da06f9..fae0c57b4 100644
--- a/packages/ui/src/context/dialog.tsx
+++ b/packages/ui/src/context/dialog.tsx
@@ -1,79 +1,116 @@
-import { For, Show, type JSX } from "solid-js"
-import { createStore } from "solid-js/store"
-import { createSimpleContext } from "@opencode-ai/ui/context"
+import {
+  createContext,
+  createMemo,
+  createSignal,
+  getOwner,
+  Owner,
+  ParentProps,
+  runWithOwner,
+  Show,
+  useContext,
+  type JSX,
+} from "solid-js"
+import { Dialog as Kobalte } from "@kobalte/core/dialog"
 
-type DialogElement = JSX.Element | (() => JSX.Element)
+type DialogElement = () => JSX.Element
 
-export const { use: useDialog, provider: DialogProvider } = createSimpleContext({
-  name: "Dialog",
-  init: () => {
-    const [store, setStore] = createStore({
-      stack: [] as {
-        element: DialogElement
-        onClose?: () => void
-      }[],
-    })
+const Context = createContext<ReturnType<typeof init>>()
 
-    return {
-      get stack() {
-        return store.stack
-      },
-      push(element: DialogElement, onClose?: () => void) {
-        setStore("stack", (s) => [...s, { element, onClose }])
-      },
-      pop() {
-        const current = store.stack.at(-1)
-        current?.onClose?.()
-        setStore("stack", store.stack.slice(0, -1))
-      },
-      replace(element: DialogElement, onClose?: () => void) {
-        for (const item of store.stack) {
-          item.onClose?.()
-        }
-        setStore("stack", [{ element, onClose }])
-      },
-      clear() {
-        for (const item of store.stack) {
-          item.onClose?.()
-        }
-        setStore("stack", [])
-      },
-    }
-  },
-})
+function init() {
+  const [store, setStore] = createSignal<
+    {
+      element: DialogElement
+      onClose?: () => void
+      owner: Owner
+    }[]
+  >([])
 
-import { Dialog as Kobalte } from "@kobalte/core/dialog"
+  return {
+    get stack() {
+      return store()
+    },
+    push(element: DialogElement, owner: Owner, onClose?: () => void) {
+      setStore((s) => [...s, { element, onClose, owner }])
+    },
+    pop() {
+      const current = store().at(-1)
+      current?.onClose?.()
+      setStore((stack) => stack.slice(0, -1))
+    },
+    replace(element: DialogElement, owner: Owner, onClose?: () => void) {
+      for (const item of store()) {
+        item.onClose?.()
+      }
+      setStore([{ element, onClose, owner }])
+    },
+    clear() {
+      for (const item of store()) {
+        item.onClose?.()
+      }
+      setStore([])
+    },
+  }
+}
 
-export function DialogRoot(props: { children?: JSX.Element }) {
-  const dialog = useDialog()
+export function DialogProvider(props: ParentProps) {
+  const ctx = init()
+  const last = createMemo(() => ctx.stack.at(-1))
   return (
-    <>
+    <Context.Provider value={ctx}>
       {props.children}
-      <Show when={dialog.stack.length > 0}>
-        <div data-component="dialog-stack">
-          <For each={dialog.stack}>
-            {(item, index) => (
-              <Show when={index() === dialog.stack.length - 1}>
+      <div data-component="dialog-stack">
+        <Show when={last()}>
+          {(item) =>
+            runWithOwner(item().owner, () => {
+              return (
                 <Kobalte
                   modal
                   defaultOpen
                   onOpenChange={(open) => {
                     if (!open) {
-                      item.onClose?.()
-                      dialog.pop()
+                      item().onClose?.()
+                      ctx.pop()
                     }
                   }}
                 >
                   <Kobalte.Portal>
                     <Kobalte.Overlay data-component="dialog-overlay" />
-                    {typeof item.element === "function" ? item.element() : item.element}
+                    {item().element()}
                   </Kobalte.Portal>
                 </Kobalte>
-              </Show>
-            )}
-          </For>
-        </div>
-      </Show>
-    </>
+              )
+            })
+          }
+        </Show>
+      </div>
+    </Context.Provider>
   )
 }
+
+export function useDialog() {
+  const ctx = useContext(Context)
+  const owner = getOwner()
+  if (!owner) {
+    throw new Error("useDialog must be used within a DialogProvider")
+  }
+  if (!ctx) {
+    throw new Error("useDialog must be used within a DialogProvider")
+  }
+  return {
+    get stack() {
+      return ctx.stack
+    },
+    replace(element: DialogElement, onClose?: () => void) {
+      ctx.replace(element, owner, onClose)
+    },
+    push(element: DialogElement, onClose?: () => void) {
+      ctx.push(element, owner, onClose)
+    },
+    pop() {
+      ctx.pop()
+    },
+    clear() {
+      ctx.clear()
+    },
+  }
+}