diff --git a/src/lib/components/description/DescriptionProvider.svelte b/src/lib/components/description/DescriptionProvider.svelte
index 38b2b5e..9d4881d 100644
--- a/src/lib/components/description/DescriptionProvider.svelte
+++ b/src/lib/components/description/DescriptionProvider.svelte
@@ -18,7 +18,7 @@
   import { getContext, setContext } from "svelte";
   import { writable, Writable } from "svelte/store";
   export let name: string;
-  let descriptionIds = [];
+  let descriptionIds: string[] = [];
   let contextStore: Writable<DescriptionContext> = writable({
     name,
     register,
diff --git a/src/lib/components/dialog/Dialog.svelte b/src/lib/components/dialog/Dialog.svelte
index 8ed7517..ba6eacb 100644
--- a/src/lib/components/dialog/Dialog.svelte
+++ b/src/lib/components/dialog/Dialog.svelte
@@ -15,9 +15,9 @@
   export interface StateDefinition {
     dialogState: DialogStates;
 
-    titleId: string | null;
+    titleId?: string;
 
-    setTitleId(id: string | null): void;
+    setTitleId(id?: string): void;
 
     close(): void;
   }
@@ -26,9 +26,9 @@
 
   export function useDialogContext(
     component: string
-  ): Writable<StateDefinition | undefined> {
+  ): Writable<StateDefinition> {
     let context = getContext(DIALOG_CONTEXT_NAME) as
-      | Writable<StateDefinition | undefined>
+      | Writable<StateDefinition>
       | undefined;
     if (context === undefined) {
       throw new Error(
@@ -112,14 +112,12 @@
     }
   });
 
-  let titleId: StateDefinition["titleId"] = null;
+  let titleId: StateDefinition["titleId"];
 
-  let api: Writable<StateDefinition | undefined> = writable();
-  setContext(DIALOG_CONTEXT_NAME, api);
-  $: api.set({
+  let api: Writable<StateDefinition> = writable({
     titleId,
     dialogState,
-    setTitleId(id: string | null) {
+    setTitleId(id?: string) {
       if (titleId === id) return;
       titleId = id;
     },
@@ -127,6 +125,14 @@
       dispatch("close", false);
     },
   });
+  setContext(DIALOG_CONTEXT_NAME, api);
+  $: api.update((obj) => {
+    return {
+      ...obj,
+      titleId,
+      dialogState,
+    };
+  });
 
   // Handle outside click
   async function handleWindowMousedown(event: MouseEvent) {
@@ -136,7 +142,7 @@
     if (containers.size !== 1) return;
     if (contains(containers, target)) return;
 
-    $api.close();
+    $api?.close();
     await tick();
     target?.focus();
   }
@@ -148,7 +154,7 @@
     if (containers.size > 1) return; // 1 is myself, otherwise other elements in the Stack
     event.preventDefault();
     event.stopPropagation();
-    $api.close();
+    $api?.close();
   }
 
   let mounted = false;
@@ -196,7 +202,7 @@
           entry.boundingClientRect.width === 0 &&
           entry.boundingClientRect.height === 0
         ) {
-          $api.close();
+          $api?.close();
         }
       }
     });
diff --git a/src/lib/components/dialog/DialogOverlay.svelte b/src/lib/components/dialog/DialogOverlay.svelte
index 74de6b3..6ef9100 100644
--- a/src/lib/components/dialog/DialogOverlay.svelte
+++ b/src/lib/components/dialog/DialogOverlay.svelte
@@ -7,7 +7,7 @@
     if (event.target !== event.currentTarget) return;
     event.preventDefault();
     event.stopPropagation();
-    $api.close();
+    $api?.close();
   }
   $: propsWeControl = {
     id,
@@ -16,5 +16,5 @@
 </script>
 
 <div {...{ ...$$restProps, ...propsWeControl }} on:click={handleClick}>
-  <slot open={$api.dialogState === DialogStates.Open} />
+  <slot open={$api?.dialogState === DialogStates.Open} />
 </div>
diff --git a/src/lib/components/disclosure/Disclosure.svelte b/src/lib/components/disclosure/Disclosure.svelte
index 480e74a..24f5096 100644
--- a/src/lib/components/disclosure/Disclosure.svelte
+++ b/src/lib/components/disclosure/Disclosure.svelte
@@ -26,8 +26,8 @@
 
   export function useDisclosureContext(
     component: string
-  ): Writable<StateDefinition | undefined> {
-    let context: Writable<StateDefinition | undefined> | undefined = getContext(
+  ): Writable<StateDefinition> {
+    let context: Writable<StateDefinition> | undefined = getContext(
       DISCLOSURE_CONTEXT_NAME
     );
 
@@ -55,10 +55,7 @@
   let panelStore: StateDefinition["panelStore"] = writable(null);
   let buttonStore: StateDefinition["buttonStore"] = writable(null);
 
-  let api: Writable<StateDefinition | undefined> = writable();
-  setContext(DISCLOSURE_CONTEXT_NAME, api);
-
-  $: api.set({
+  let api: Writable<StateDefinition> = writable({
     buttonId,
     panelId,
     disclosureState,
@@ -87,6 +84,14 @@
       restoreElement?.focus();
     },
   });
+  setContext(DISCLOSURE_CONTEXT_NAME, api);
+
+  $: api.update((obj) => {
+    return {
+      ...obj,
+      disclosureState,
+    };
+  });
 
   let openClosedState: Writable<State> | undefined = writable();
   setContext("OpenClosed", openClosedState);
diff --git a/src/lib/components/label/LabelProvider.svelte b/src/lib/components/label/LabelProvider.svelte
index 6df97f6..1b91721 100644
--- a/src/lib/components/label/LabelProvider.svelte
+++ b/src/lib/components/label/LabelProvider.svelte
@@ -16,7 +16,7 @@
   import { getContext, setContext } from "svelte";
   import { writable, Writable } from "svelte/store";
   export let name: string;
-  let labelIds = [];
+  let labelIds: string[] = [];
   let contextStore: Writable<LabelContext> = writable({
     name,
     register,
diff --git a/src/lib/components/listbox/Listbox.svelte b/src/lib/components/listbox/Listbox.svelte
index 2aebe20..b31bfc7 100644
--- a/src/lib/components/listbox/Listbox.svelte
+++ b/src/lib/components/listbox/Listbox.svelte
@@ -38,8 +38,8 @@
   const LISTBOX_CONTEXT_NAME = "ListboxContext";
   export function useListboxContext(
     component: string
-  ): Writable<StateDefinition | undefined> {
-    let context: Writable<StateDefinition | undefined> | undefined =
+  ): Writable<StateDefinition> {
+    let context: Writable<StateDefinition> | undefined =
       getContext(LISTBOX_CONTEXT_NAME);
 
     if (context === undefined) {
@@ -61,6 +61,12 @@
   import { writable, Writable } from "svelte/store";
   import { match } from "$lib/utils/match";
   import { State, useOpenClosedProvider } from "$lib/internal/open-closed";
+  export let disabled = false;
+  export let horizontal = false;
+  export let value: any;
+  $: orientation = (horizontal ? "horizontal" : "vertical") as
+    | "horizontal"
+    | "vertical";
 
   const dispatch = createEventDispatcher();
 
@@ -68,41 +74,23 @@
   let labelStore = writable(null);
   setContext("labelStore", labelStore);
   $: labelRef = $labelStore;
-  let buttonStore = writable(null);
-  setContext("buttonStore", buttonStore);
-  $: buttonRef = $buttonStore;
-  let optionsStore = writable(null);
-  setContext("optionsStore", optionsStore);
-  $: optionsRef = $optionsStore;
-  let options = [];
-  let searchQuery = "";
-  let activeOptionIndex = null;
 
-  let api: Writable<StateDefinition | undefined> = writable();
-  setContext(LISTBOX_CONTEXT_NAME, api);
+  let buttonRef: Writable<StateDefinition["buttonRef"]> = writable(null);
+  setContext("buttonStore", buttonRef);
 
-  let openClosedState = writable(State.Closed);
-  useOpenClosedProvider(openClosedState);
-  $: openClosedState.set(
-    match(listboxState, {
-      [ListboxStates.Open]: State.Open,
-      [ListboxStates.Closed]: State.Closed,
-    })
-  );
+  let optionsRef: Writable<StateDefinition["optionsRef"]> = writable(null);
+  setContext("optionsStore", optionsRef);
 
-  export let disabled = false;
-  export let horizontal = false;
-  export let value: any;
+  let options: StateDefinition["options"] = [];
+  let searchQuery: StateDefinition["searchQuery"] = "";
+  let activeOptionIndex: StateDefinition["activeOptionIndex"] = null;
 
-  $: orientation = (horizontal ? "horizontal" : "vertical") as
-    | "horizontal"
-    | "vertical";
-  $: api.set({
+  let api: Writable<StateDefinition> = writable({
     listboxState,
     labelRef,
     value,
-    buttonRef,
-    optionsRef,
+    buttonRef: $buttonRef,
+    optionsRef: $optionsRef,
     options,
     searchQuery,
     activeOptionIndex,
@@ -131,7 +119,7 @@
           resolveItems: () => options,
           resolveActiveIndex: () => activeOptionIndex,
           resolveId: (option) => option.id,
-          resolveDisabled: (option) => option.disabled,
+          resolveDisabled: (option) => option.dataRef.disabled,
         }
       );
 
@@ -147,7 +135,9 @@
       searchQuery += value.toLowerCase();
 
       let match = options.findIndex(
-        (option) => !option.disabled && option.textValue.startsWith(searchQuery)
+        (option) =>
+          !option.dataRef.disabled &&
+          option.dataRef.textValue.startsWith(searchQuery)
       );
 
       if (match === -1 || match === activeOptionIndex) return;
@@ -184,18 +174,44 @@
       dispatch("updateValue", { value });
     },
   });
+  setContext(LISTBOX_CONTEXT_NAME, api);
+
+  let openClosedState = writable(State.Closed);
+  useOpenClosedProvider(openClosedState);
+  $: openClosedState.set(
+    match(listboxState, {
+      [ListboxStates.Open]: State.Open,
+      [ListboxStates.Closed]: State.Closed,
+    })
+  );
+
+  $: api.update((obj) => {
+    return {
+      ...obj,
+      listboxState,
+      labelRef,
+      value,
+      buttonRef: $buttonRef,
+      optionsRef: $optionsRef,
+      options,
+      searchQuery,
+      activeOptionIndex,
+      disabled,
+      orientation,
+    };
+  });
 
   function handleMousedown(event: MouseEvent) {
     let target = event.target as HTMLElement;
     let active = document.activeElement;
 
     if (listboxState !== ListboxStates.Open) return;
-    if (buttonRef?.contains(target)) return;
+    if ($buttonRef?.contains(target)) return;
 
-    if (!optionsRef?.contains(target)) $api.closeListbox();
+    if (!$optionsRef?.contains(target)) $api.closeListbox();
     if (active !== document.body && active?.contains(target)) return; // Keep focus on newly clicked/focused element
     if (!event.defaultPrevented) {
-      buttonRef?.focus({ preventScroll: true });
+      $buttonRef?.focus({ preventScroll: true });
     }
   }
 </script>
diff --git a/src/lib/components/switch/Switch.svelte b/src/lib/components/switch/Switch.svelte
index 834e78c..9bcd4ce 100644
--- a/src/lib/components/switch/Switch.svelte
+++ b/src/lib/components/switch/Switch.svelte
@@ -19,7 +19,6 @@
   let descriptionContext = useDescriptionContext();
   let id = `headlessui-switch-${useId()}`;
   $: switchStore = $api?.switchStore;
-  let internalSwitchRef = null;
 
   function toggle() {
     dispatch("updateValue", !checked);
@@ -63,7 +62,6 @@
 {:else}
   <button
     {...{ ...$$restProps, ...propsWeControl }}
-    bind:this={$internalSwitchRef}
     on:click={handleClick}
     on:keyup={handleKeyUp}
     on:keypress={handleKeyPress}
diff --git a/src/lib/internal/StackContextProvider.svelte b/src/lib/internal/StackContextProvider.svelte
index 329abf2..145687d 100644
--- a/src/lib/internal/StackContextProvider.svelte
+++ b/src/lib/internal/StackContextProvider.svelte
@@ -8,7 +8,7 @@
 <script lang="ts">
   import { getContext, onDestroy, setContext } from "svelte";
   import { writable, Writable } from "svelte/store";
-  type OnUpdate = (message: StackMessage, element: HTMLElement | null) => void;
+  type OnUpdate = (message: StackMessage, element: HTMLElement) => void;
 
   export let onUpdate: OnUpdate | undefined;
   export let element: HTMLElement | null;
@@ -31,8 +31,9 @@
       _cleanup();
     }
     if (!element) return null;
-    $notifyStore(StackMessage.Add, element);
-    return () => $notifyStore(StackMessage.Remove, element);
+    let savedElement = element;
+    $notifyStore(StackMessage.Add, savedElement);
+    return () => $notifyStore(StackMessage.Remove, savedElement);
   })();
 
   onDestroy(() => {
diff --git a/src/routes/dialog/index.svelte b/src/routes/dialog/index.svelte
index b3840cf..a6be1c1 100644
--- a/src/routes/dialog/index.svelte
+++ b/src/routes/dialog/index.svelte
@@ -19,11 +19,17 @@
     modifiers: [{ name: "offset", options: { offset: [0, 10] } }],
   };
 
-  function classNames(...classes) {
+  function classNames(...classes: unknown[]) {
     return classes.filter(Boolean).join(" ");
   }
 
-  function resolveClass({ active, disabled }) {
+  function resolveClass({
+    active,
+    disabled,
+  }: {
+    active: boolean;
+    disabled: boolean;
+  }) {
     return classNames(
       "flex justify-between w-full px-4 py-2 text-sm leading-5 text-left",
       active ? "bg-gray-100 text-gray-900" : "text-gray-700",