diff --git a/examples/src/App.vue b/examples/src/App.vue
index 670cd47..6f3a951 100644
--- a/examples/src/App.vue
+++ b/examples/src/App.vue
@@ -1,21 +1,24 @@
 <template>
   <div>
-    <TestButton @click="handleClick">
-      Button
-    </TestButton>
-    <TestButton disabled @click="handleClick">
-      Button
-    </TestButton>
+    <DialogDisclosure v-bind="dialog" :as="TestButton">Open Dialog</DialogDisclosure>
+    <Dialog v-bind="dialog">Hello Dialog</Dialog>
   </div>
 </template>
 
 <script>
+import { markRaw } from 'vue'
+import { Dialog, DialogDisclosure, useDialogState } from 'vue-ari'
 import TestButton from './MyButton.vue'
 
 export default {
   name: 'App',
   components: {
-    TestButton,
+    Dialog,
+    DialogDisclosure,
+  },
+  setup() {
+    const dialog = useDialogState()
+    return { dialog, TestButton: markRaw(TestButton) }
   },
   methods: {
     handleClick() {
diff --git a/package.json b/package.json
index 89ee21e..0d921b6 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "vue-ari",
-  "version": "0.0.0-alpha.4",
+  "version": "0.0.0-alpha.5",
   "license": "MIT",
   "repository": "https://github.com/visualjerk/ari",
   "main": "lib/index.js",
diff --git a/src/Dialog/Dialog.ts b/src/Dialog/Dialog.ts
index 503f9a8..e8daf9c 100644
--- a/src/Dialog/Dialog.ts
+++ b/src/Dialog/Dialog.ts
@@ -101,12 +101,23 @@ function useHandleToggleFocus(props: DialogProps, ref: Ref<HTMLElement>) {
     (visible) => {
       if (visible) {
         focusFirstFocusable(ref.value)
-      } else if (focusIsWithin(ref.value)) {
+      }
+    }
+  )
+  // Needs to be a 'sync' watcher, so we can check
+  // if focus was within the dialog before it is closed
+  watch(
+    () => props.visible,
+    (visible) => {
+      if (!visible && focusIsWithin(ref.value)) {
         const disclosure: HTMLElement = document.querySelector(
           `[aria-controls="${props.baseId}"]`
         )
         disclosure.focus()
       }
+    },
+    {
+      flush: 'sync',
     }
   )
 }