3 files changed, 48 insertions, 27 deletions
diff --git a/src/components/Menu/context.tsx b/src/components/Menu/context.tsx
index 1ddcd583f..908ad352e 100644
--- a/src/components/Menu/context.tsx
+++ b/src/components/Menu/context.tsx
@@ -10,3 +10,23 @@ export const Context = React.createContext<ContextType>({
 export const ItemContext = React.createContext<ItemContextType>({
   disabled: false,
 })
+
+export function useMenuContext() {
+  const context = React.useContext(Context)
+
+  if (!context) {
+    throw new Error('useMenuContext must be used within a Context.Provider')
+  }
+
+  return context
+}
+
+export function useMenuItemContext() {
+  const context = React.useContext(ItemContext)
+
+  if (!context) {
+    throw new Error('useMenuItemContext must be used within a Context.Provider')
+  }
+
+  return context
+}
diff --git a/src/components/Menu/index.tsx b/src/components/Menu/index.tsx
index a171e9772..73eb9da52 100644
--- a/src/components/Menu/index.tsx
+++ b/src/components/Menu/index.tsx
@@ -9,7 +9,12 @@ import {atoms as a, useTheme} from '#/alf'
 import {Button, ButtonText} from '#/components/Button'
 import * as Dialog from '#/components/Dialog'
 import {useInteractionState} from '#/components/hooks/useInteractionState'
-import {Context, ItemContext} from '#/components/Menu/context'
+import {
+  Context,
+  ItemContext,
+  useMenuContext,
+  useMenuItemContext,
+} from '#/components/Menu/context'
 import {
   ContextType,
   GroupProps,
@@ -25,10 +30,6 @@ export {
   useDialogControl as useMenuControl,
 } from '#/components/Dialog'
 
-export function useMemoControlContext() {
-  return React.useContext(Context)
-}
-
 export function Root({
   children,
   control,
@@ -47,7 +48,7 @@ export function Root({
 }
 
 export function Trigger({children, label, role = 'button'}: TriggerProps) {
-  const {control} = React.useContext(Context)
+  const context = useMenuContext()
   const {state: focused, onIn: onFocus, onOut: onBlur} = useInteractionState()
   const {
     state: pressed,
@@ -57,14 +58,14 @@ export function Trigger({children, label, role = 'button'}: TriggerProps) {
 
   return children({
     isNative: true,
-    control,
+    control: context.control,
     state: {
       hovered: false,
       focused,
       pressed,
     },
     props: {
-      onPress: control.open,
+      onPress: context.control.open,
       onFocus,
       onBlur,
       onPressIn,
@@ -82,7 +83,7 @@ export function Outer({
   showCancel?: boolean
   style?: StyleProp<ViewStyle>
 }>) {
-  const context = React.useContext(Context)
+  const context = useMenuContext()
   const {_} = useLingui()
 
   return (
@@ -105,7 +106,7 @@ export function Outer({
 
 export function Item({children, label, style, onPress, ...rest}: ItemProps) {
   const t = useTheme()
-  const {control} = React.useContext(Context)
+  const context = useMenuContext()
   const {state: focused, onIn: onFocus, onOut: onBlur} = useInteractionState()
   const {
     state: pressed,
@@ -121,10 +122,9 @@ export function Item({children, label, style, onPress, ...rest}: ItemProps) {
       onFocus={onFocus}
       onBlur={onBlur}
       onPress={async e => {
-        await onPress(e)
-        if (!e.defaultPrevented) {
-          control?.close()
-        }
+        context.control.close(() => {
+          onPress?.(e)
+        })
       }}
       onPressIn={e => {
         onPressIn()
@@ -156,7 +156,7 @@ export function Item({children, label, style, onPress, ...rest}: ItemProps) {
 
 export function ItemText({children, style}: ItemTextProps) {
   const t = useTheme()
-  const {disabled} = React.useContext(ItemContext)
+  const {disabled} = useMenuItemContext()
   return (
     <Text
       numberOfLines={1}
@@ -177,7 +177,7 @@ export function ItemText({children, style}: ItemTextProps) {
 
 export function ItemIcon({icon: Comp}: ItemIconProps) {
   const t = useTheme()
-  const {disabled} = React.useContext(ItemContext)
+  const {disabled} = useMenuItemContext()
   return (
     <Comp
       size="lg"
@@ -223,7 +223,7 @@ export function Group({children, style}: GroupProps) {
 
 function Cancel() {
   const {_} = useLingui()
-  const {control} = React.useContext(Context)
+  const context = useMenuContext()
 
   return (
     <Button
@@ -231,7 +231,7 @@ function Cancel() {
       size="small"
       variant="ghost"
       color="secondary"
-      onPress={() => control.close()}>
+      onPress={() => context.control.close()}>
       <ButtonText>
         <Trans>Cancel</Trans>
       </ButtonText>
diff --git a/src/components/Menu/index.web.tsx b/src/components/Menu/index.web.tsx
index 37ad67e29..ab0c9d20a 100644
--- a/src/components/Menu/index.web.tsx
+++ b/src/components/Menu/index.web.tsx
@@ -7,7 +7,12 @@ import * as DropdownMenu from '@radix-ui/react-dropdown-menu'
 import {atoms as a, flatten, useTheme, web} from '#/alf'
 import * as Dialog from '#/components/Dialog'
 import {useInteractionState} from '#/components/hooks/useInteractionState'
-import {Context, ItemContext} from '#/components/Menu/context'
+import {
+  Context,
+  ItemContext,
+  useMenuContext,
+  useMenuItemContext,
+} from '#/components/Menu/context'
 import {
   ContextType,
   GroupProps,
@@ -40,10 +45,6 @@ export function useMenuControl(): Dialog.DialogControlProps {
   )
 }
 
-export function useMemoControlContext() {
-  return React.useContext(Context)
-}
-
 export function Root({
   children,
   control,
@@ -110,7 +111,7 @@ const RadixTriggerPassThrough = React.forwardRef(
 RadixTriggerPassThrough.displayName = 'RadixTriggerPassThrough'
 
 export function Trigger({children, label, role = 'button'}: TriggerProps) {
-  const {control} = React.useContext(Context)
+  const {control} = useMenuContext()
   const {
     state: hovered,
     onIn: onMouseEnter,
@@ -203,7 +204,7 @@ export function Outer({
 
 export function Item({children, label, onPress, ...rest}: ItemProps) {
   const t = useTheme()
-  const {control} = React.useContext(Context)
+  const {control} = useMenuContext()
   const {
     state: hovered,
     onIn: onMouseEnter,
@@ -262,7 +263,7 @@ export function Item({children, label, onPress, ...rest}: ItemProps) {
 
 export function ItemText({children, style}: ItemTextProps) {
   const t = useTheme()
-  const {disabled} = React.useContext(ItemContext)
+  const {disabled} = useMenuItemContext()
   return (
     <Text
       style={[
@@ -279,7 +280,7 @@ export function ItemText({children, style}: ItemTextProps) {
 
 export function ItemIcon({icon: Comp, position = 'left'}: ItemIconProps) {
   const t = useTheme()
-  const {disabled} = React.useContext(ItemContext)
+  const {disabled} = useMenuItemContext()
   return (
     <View
       style={[