diff --git a/lib/src/ambient_widget.dart b/lib/src/ambient_widget.dart
index c70d667..b8c73a2 100644
--- a/lib/src/ambient_widget.dart
+++ b/lib/src/ambient_widget.dart
@@ -23,7 +23,7 @@ class AmbientMode extends StatefulWidget {
   final Widget? child;
 
   /// Called each time the the wear device triggers an ambient update request.
-  final VoidCallback? onUpdate;
+  final Function(WearMode)? onUpdate;
 
   /// Get current [WearMode].
   static WearMode wearModeOf(BuildContext context) {
@@ -72,6 +72,7 @@ class _AmbientModeState extends State<AmbientMode> with AmbientCallback {
   void _updateMode(bool isAmbient) {
     if (mounted) {
       setState(() => _ambientMode = isAmbient ? WearMode.ambient : WearMode.active);
+      if(widget.onUpdate != null) widget.onUpdate!(_ambientMode);
     }
   }
 
@@ -84,7 +85,6 @@ class _AmbientModeState extends State<AmbientMode> with AmbientCallback {
   @override
   void onUpdateAmbient() {
     _updateMode(true);
-    widget.onUpdate?.call();
   }
 }