diff --git a/example/lib/simple_example.dart b/example/lib/simple_example.dart
index 7f21acb..8e9f251 100644
--- a/example/lib/simple_example.dart
+++ b/example/lib/simple_example.dart
@@ -22,7 +22,7 @@ class _SimpleExampleState extends State<SimpleExample> {
         GestureDetector(
           onTap: () {
             ScaffoldMessenger.of(context).showSnackBar(SnackBar(content: Text('Tapped a container')));
-            controller.focusOnChild(context, i);
+            controller.focusOnChild(i);
           },
           child: InfoContainer(color: Colors.primaries[i % Colors.primaries.length]),
         ),
diff --git a/lib/core/controller/controller.dart b/lib/core/controller/controller.dart
index 78781b7..1ef89e4 100644
--- a/lib/core/controller/controller.dart
+++ b/lib/core/controller/controller.dart
@@ -25,6 +25,7 @@ class LazyCanvasController with ChangeNotifier {
   bool _init = false;
   late final SpatialHashing<int> _spatialHash;
   TickerProvider? _ticker;
+  late BuildContext _context;
 
   bool debug;
 
@@ -71,6 +72,10 @@ class LazyCanvasController with ChangeNotifier {
     _ticker = ticker;
   }
 
+  void setBuildContext(BuildContext context) {
+    _context = context;
+  }
+
   // ==================== Child Management ====================
 
   /// Add a child at a given position with a widget. Returns the child ID.
@@ -174,7 +179,7 @@ class LazyCanvasController with ChangeNotifier {
   // ==================== Positioning Logic ====================
 
   /// Currently rendered widgets with their position info
-  List<ChildInfo> widgetsWithScreenPositions(BuildContext context, {bool forceRebuild = false}) {
+  List<ChildInfo> widgetsWithScreenPositions({bool forceRebuild = false}) {
     if (!_init) return [];
 
     final idsToBuild = _childrenWithinBuildArea(_gsCenter, _buildExtent);
@@ -218,7 +223,6 @@ class LazyCanvasController with ChangeNotifier {
   /// an offstage rendering will be used ( double render )
   /// Preferred horizontal margin used for [ScalingMode.fitInViewport].
   void focusOnChild(
-    BuildContext context,
     int id, {
     ScalingMode scalingMode = ScalingMode.keepScale,
     bool animate = true,
@@ -236,7 +240,7 @@ class LazyCanvasController with ChangeNotifier {
     // else do an offstage render
     childSize ??= childInfo.lastRenderedSize != null && !forceRedraw
         ? childInfo.lastRenderedSize
-        : measureWidgetSize(context, (_) => childInfo.widget);
+        : measureWidgetSize(_context, (_) => childInfo.widget);
 
     /*
     margin is symmatric on ltrb so
diff --git a/lib/core/render.dart b/lib/core/render.dart
index ccf15b6..3f723f5 100644
--- a/lib/core/render.dart
+++ b/lib/core/render.dart
@@ -33,6 +33,7 @@ class _LazyCanvasState extends State<LazyCanvas> with TickerProviderStateMixin<L
 
   @override
   Widget build(BuildContext context) {
+    widget.controller.setBuildContext(context);
     return GestureDetector(
       behavior: HitTestBehavior.translucent,
       onScaleUpdate: widget.controller.onScaleUpdate,
@@ -40,7 +41,7 @@ class _LazyCanvasState extends State<LazyCanvas> with TickerProviderStateMixin<L
       child: ListenableBuilder(
         listenable: widget.controller,
         builder: (_, _) {
-          final childrenWithPositions = widget.controller.widgetsWithScreenPositions(context);
+          final childrenWithPositions = widget.controller.widgetsWithScreenPositions();
           final ssPositions = childrenWithPositions.map((e) => e.ssPosition).toList();
           final childrenIds = childrenWithPositions.map((e) => e.id).toList();
           final children = childrenWithPositions.map((e) => e.child).toList();
diff --git a/test/canvas_view_test.dart b/test/canvas_view_test.dart
index b805e72..e28f8a0 100644
--- a/test/canvas_view_test.dart
+++ b/test/canvas_view_test.dart
@@ -55,7 +55,9 @@ void main() {
       ),
     );
     await tester.pumpAndSettle();
+
     final context = tester.element(find.byType(LazyCanvas));
+    controller.setBuildContext(context);
 
     // Initial size check
     final finder = find.byType(TestChild);
@@ -66,7 +68,7 @@ void main() {
     controller.onScaleUpdate(ScaleUpdateDetails(focalPoint: const Offset(0, 0), scale: 2.0));
     await tester.pumpAndSettle();
     expect(controller.scale, 2.0);
-    final ssPositions = controller.widgetsWithScreenPositions(context).map((e) => e.ssPosition).toList();
+    final ssPositions = controller.widgetsWithScreenPositions().map((e) => e.ssPosition).toList();
     expect(ssPositions, [Offset.zero, const Offset(200, 200)]);
   });
 }