[ Widget Preview ] Fix WidgetInspectorService override (#176550)

When registering the `WidgetPreviewScaffoldInspectorService`, we were
originally setting it after the bindings were initialized. This meant
that the widget inspector service extensions were registered with the
original `WidgetInspectorService` and were not taking the custom
codepaths in the override.

This change moves the `WidgetPreviewScaffoldInspectorService`
initialization to before the bindings are initialized.

packages/flutter_tools/templates/template_manifest.json

@@ -341,6 +341,7 @@
 
     "templates/widget_preview_scaffold/lib/main.dart.tmpl",
     "templates/widget_preview_scaffold/lib/src/widget_preview.dart.tmpl",
+    "templates/widget_preview_scaffold/lib/src/widget_preview_inspector_service.dart.tmpl",
     "templates/widget_preview_scaffold/lib/src/widget_preview_rendering.dart.tmpl",
     "templates/widget_preview_scaffold/lib/src/widget_preview_scaffold_controller.dart.tmpl",
     "templates/widget_preview_scaffold/lib/src/controls.dart.tmpl",

packages/flutter_tools/templates/widget_preview_scaffold/lib/src/widget_preview_inspector_service.dart.tmpl

@@ -0,0 +1,87 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'package:flutter/widgets.dart';
+import 'package:widget_preview_scaffold/src/dtd/dtd_services.dart';
+import 'package:widget_preview_scaffold/src/dtd/editor_service.dart';
+import 'package:widget_preview_scaffold/src/widget_preview_rendering.dart';
+
+/// A custom [WidgetInspectorService] responsible for routing navigation events
+/// to the IDE.
+///
+/// IMPORTANT NOTE: this **must** be called before WidgetsFlutterBinding.ensureInitialized()
+/// is called, otherwise the inspector service extensions will be registered against
+/// the default WidgetInspectorService, causing overrides to not be invoked.
+class WidgetPreviewScaffoldInspectorService with WidgetInspectorService {
+  WidgetPreviewScaffoldInspectorService({required this.dtdServices}) {
+    WidgetInspectorService.instance = this;
+  }
+
+  /// The DTD services instance used to communicate with the tool.
+  final WidgetPreviewScaffoldDtdServices dtdServices;
+
+  // Keys used to specify the creation location of a widget when serializing a
+  // DiagnosticsNode to JSON. This location is used by the widget inspector
+  // to jump to the creation location of a selected widget.
+  static const kFile = 'fileUri';
+  static const kLine = 'line';
+  static const kColumn = 'column';
+
+  CodeLocation? _nextNavigationLocation;
+
+  @protected
+  @override
+  bool setSelection(Object? object, [String? groupName]) {
+    // The next navigation event sent to `postEvent` will be for this selection.
+    // Save the location of preview annotation applications so we can override
+    // the navigation target in `postEvent`.
+    if (object is PreviewWidgetElement) {
+      final previewData = (object.widget as PreviewWidget).preview;
+      _nextNavigationLocation = CodeLocation(
+        uri: previewData.scriptUri,
+        line: previewData.line,
+        column: previewData.column,
+      );
+    }
+    final result = super.setSelection(object, groupName);
+    _nextNavigationLocation = null;
+    return result;
+  }
+
+  @override
+  void postEvent(
+    String eventKind,
+    Map<Object, Object?> eventData, {
+    String stream = 'Extension',
+  }) {
+    // It's unlikely that the widget previewer will be connected to directly by
+    // an IDE via the VM service, so we forward navigation events via the
+    // Editor DTD service.
+    if (eventKind == 'navigate') {
+      CodeLocation? location = _nextNavigationLocation;
+      if (eventData case {
+        kFile: final String file,
+        kLine: final int line,
+        kColumn: final int column,
+      } when location == null) {
+        location = CodeLocation(uri: file, line: line, column: column);
+      } else if (location != null) {
+        // If a [PreviewWidgetElement] was selected, we're not navigating to the
+        // creation location of the widget. Override the location details in the
+        // event data, just in case an IDE is attached and listening for
+        // navigation events through the VM service.
+        // TODO(bkonyi): determine if this is necessary
+        eventData.addAll(<String, Object>{
+          kFile: location.uri,
+          kLine: location.line!,
+          kColumn: location.column!,
+        });
+      }
+      if (location != null) {
+        dtdServices.navigateToCode(location);
+      }
+    }
+    super.postEvent(eventKind, eventData, stream: stream);
+  }
+}

packages/flutter_tools/templates/widget_preview_scaffold/lib/src/widget_preview_rendering.dart.tmpl

@@ -16,6 +16,7 @@ import 'package:google_fonts/google_fonts.dart';
 import 'package:stack_trace/stack_trace.dart';
 import 'package:url_launcher/url_launcher.dart';
 import 'package:widget_preview_scaffold/src/dtd/editor_service.dart';
+import 'package:widget_preview_scaffold/src/widget_preview_inspector_service.dart';
 
 import 'controls.dart';
 import 'generated_preview.dart';
@@ -928,14 +929,20 @@ class PreviewAssetBundle extends PlatformAssetBundle {
 /// the preview scaffold project which prevents us from being able to use hot
 /// restart to iterate on this file.
 Future<void> mainImpl() async {
+  final controller = WidgetPreviewScaffoldController(previews: previews);
+  await controller.initialize();
+  // WARNING: do not move this line. This constructor sets
+  // [WidgetInspectorService.instance] to the custom service for the widget
+  // previewer. If [WidgetsFlutterBinding.ensureInitialized()] is invoked before
+  // the custom service is set, inspector service extensions will be registered
+  // against the wrong service.
+  WidgetPreviewScaffoldInspectorService(dtdServices: controller.dtdServices);
   final WidgetsBinding binding = WidgetsFlutterBinding.ensureInitialized();
   // Disable the injection of [WidgetInspector] into the widget tree built by
   // [WidgetsApp]. [WidgetInspector] instances will be created for each
   // individual preview so the widget inspector won't allow for users to select
   // widgets that make up the widget preview scaffolding.
   binding.debugExcludeRootWidgetInspector = true;
-  final controller = WidgetPreviewScaffoldController(previews: previews);
-  await controller.initialize();
   runWidget(
     DisableWidgetInspectorScope(
       child: binding.wrapWithDefaultView(

packages/flutter_tools/templates/widget_preview_scaffold/lib/src/widget_preview_scaffold_controller.dart.tmpl

@@ -3,9 +3,7 @@
 // found in the LICENSE file.
 
 import 'package:flutter/foundation.dart';
-import 'package:flutter/widgets.dart';
 import 'package:path/path.dart' as path;
-import 'package:widget_preview_scaffold/src/dtd/editor_service.dart';
 import 'package:widget_preview_scaffold/src/widget_preview_rendering.dart';
 import 'dtd/dtd_services.dart';
 import 'widget_preview.dart';
@@ -24,11 +22,7 @@ class WidgetPreviewScaffoldController {
     required PreviewsCallback previews,
     @visibleForTesting WidgetPreviewScaffoldDtdServices? dtdServicesOverride,
   }) : _previews = previews,
-       dtdServices = dtdServicesOverride ?? WidgetPreviewScaffoldDtdServices() {
-    // Overrides the default WidgetInspectorService instance to handle selection
-    // events.
-    _WidgetPreviewScaffoldInspectorService(dtdServices: dtdServices);
-  }
+       dtdServices = dtdServicesOverride ?? WidgetPreviewScaffoldDtdServices();
 
   @visibleForTesting
   static const kFilterBySelectedFilePreference = 'filterBySelectedFile';
@@ -170,77 +164,3 @@ class WidgetPreviewScaffoldController {
     }
   }
 }
-
-/// A custom [WidgetInspectorService] responsible for routing navigation events
-/// to the IDE.
-class _WidgetPreviewScaffoldInspectorService with WidgetInspectorService {
-  _WidgetPreviewScaffoldInspectorService({required this.dtdServices}) {
-    WidgetInspectorService.instance = this;
-  }
-
-  final WidgetPreviewScaffoldDtdServices dtdServices;
-
-  // Keys used to specify the creation location of a widget when serializing a
-  // DiagnosticsNode to JSON. This location is used by the widget inspector
-  // to jump to the creation location of a selected widget.
-  static const kFile = 'fileUri';
-  static const kLine = 'line';
-  static const kColumn = 'column';
-
-  CodeLocation? _nextNavigationLocation;
-
-  @protected
-  @override
-  bool setSelection(Object? object, [String? groupName]) {
-    // The next navigation event sent to `postEvent` will be for this selection.
-    // Save the location of preview annotation applications so we can override
-    // the navigation target in `postEvent`.
-    if (object is PreviewWidgetElement) {
-      final previewData = (object.widget as PreviewWidget).preview;
-      _nextNavigationLocation = CodeLocation(
-        uri: previewData.scriptUri,
-        line: previewData.line,
-        column: previewData.column,
-      );
-    }
-    final result = super.setSelection(object, groupName);
-    _nextNavigationLocation = null;
-    return result;
-  }
-
-  @override
-  void postEvent(
-    String eventKind,
-    Map<Object, Object?> eventData, {
-    String stream = 'Extension',
-  }) {
-    // It's unlikely that the widget previewer will be connected to directly by
-    // an IDE via the VM service, so we forward navigation events via the
-    // Editor DTD service.
-    if (eventKind == 'navigate') {
-      CodeLocation? location = _nextNavigationLocation;
-      if (eventData case {
-        kFile: final String file,
-        kLine: final int line,
-        kColumn: final int column,
-      } when location == null) {
-        location = CodeLocation(uri: file, line: line, column: column);
-      } else if (location != null) {
-        // If a [PreviewWidgetElement] was selected, we're not navigating to the
-        // creation location of the widget. Override the location details in the
-        // event data, just in case an IDE is attached and listening for
-        // navigation events through the VM service.
-        // TODO(bkonyi): determine if this is necessary
-        eventData.addAll(<String, Object>{
-          kFile: location.uri,
-          kLine: location.line!,
-          kColumn: location.column!,
-        });
-      }
-      if (location != null) {
-        dtdServices.navigateToCode(location);
-      }
-    }
-    super.postEvent(eventKind, eventData, stream: stream);
-  }
-}

packages/flutter_tools/test/widget_preview_scaffold.shard/widget_preview_scaffold/lib/src/widget_preview_inspector_service.dart

@@ -0,0 +1,87 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'package:flutter/widgets.dart';
+import 'package:widget_preview_scaffold/src/dtd/dtd_services.dart';
+import 'package:widget_preview_scaffold/src/dtd/editor_service.dart';
+import 'package:widget_preview_scaffold/src/widget_preview_rendering.dart';
+
+/// A custom [WidgetInspectorService] responsible for routing navigation events
+/// to the IDE.
+///
+/// IMPORTANT NOTE: this **must** be called before WidgetsFlutterBinding.ensureInitialized()
+/// is called, otherwise the inspector service extensions will be registered against
+/// the default WidgetInspectorService, causing overrides to not be invoked.
+class WidgetPreviewScaffoldInspectorService with WidgetInspectorService {
+  WidgetPreviewScaffoldInspectorService({required this.dtdServices}) {
+    WidgetInspectorService.instance = this;
+  }
+
+  /// The DTD services instance used to communicate with the tool.
+  final WidgetPreviewScaffoldDtdServices dtdServices;
+
+  // Keys used to specify the creation location of a widget when serializing a
+  // DiagnosticsNode to JSON. This location is used by the widget inspector
+  // to jump to the creation location of a selected widget.
+  static const kFile = 'fileUri';
+  static const kLine = 'line';
+  static const kColumn = 'column';
+
+  CodeLocation? _nextNavigationLocation;
+
+  @protected
+  @override
+  bool setSelection(Object? object, [String? groupName]) {
+    // The next navigation event sent to `postEvent` will be for this selection.
+    // Save the location of preview annotation applications so we can override
+    // the navigation target in `postEvent`.
+    if (object is PreviewWidgetElement) {
+      final previewData = (object.widget as PreviewWidget).preview;
+      _nextNavigationLocation = CodeLocation(
+        uri: previewData.scriptUri,
+        line: previewData.line,
+        column: previewData.column,
+      );
+    }
+    final result = super.setSelection(object, groupName);
+    _nextNavigationLocation = null;
+    return result;
+  }
+
+  @override
+  void postEvent(
+    String eventKind,
+    Map<Object, Object?> eventData, {
+    String stream = 'Extension',
+  }) {
+    // It's unlikely that the widget previewer will be connected to directly by
+    // an IDE via the VM service, so we forward navigation events via the
+    // Editor DTD service.
+    if (eventKind == 'navigate') {
+      CodeLocation? location = _nextNavigationLocation;
+      if (eventData case {
+        kFile: final String file,
+        kLine: final int line,
+        kColumn: final int column,
+      } when location == null) {
+        location = CodeLocation(uri: file, line: line, column: column);
+      } else if (location != null) {
+        // If a [PreviewWidgetElement] was selected, we're not navigating to the
+        // creation location of the widget. Override the location details in the
+        // event data, just in case an IDE is attached and listening for
+        // navigation events through the VM service.
+        // TODO(bkonyi): determine if this is necessary
+        eventData.addAll(<String, Object>{
+          kFile: location.uri,
+          kLine: location.line!,
+          kColumn: location.column!,
+        });
+      }
+      if (location != null) {
+        dtdServices.navigateToCode(location);
+      }
+    }
+    super.postEvent(eventKind, eventData, stream: stream);
+  }
+}

packages/flutter_tools/test/widget_preview_scaffold.shard/widget_preview_scaffold/lib/src/widget_preview_rendering.dart

@@ -16,6 +16,7 @@ import 'package:google_fonts/google_fonts.dart';
 import 'package:stack_trace/stack_trace.dart';
 import 'package:url_launcher/url_launcher.dart';
 import 'package:widget_preview_scaffold/src/dtd/editor_service.dart';
+import 'package:widget_preview_scaffold/src/widget_preview_inspector_service.dart';
 
 import 'controls.dart';
 import 'generated_preview.dart';
@@ -928,14 +929,20 @@ class PreviewAssetBundle extends PlatformAssetBundle {
 /// the preview scaffold project which prevents us from being able to use hot
 /// restart to iterate on this file.
 Future<void> mainImpl() async {
+  final controller = WidgetPreviewScaffoldController(previews: previews);
+  await controller.initialize();
+  // WARNING: do not move this line. This constructor sets
+  // [WidgetInspectorService.instance] to the custom service for the widget
+  // previewer. If [WidgetsFlutterBinding.ensureInitialized()] is invoked before
+  // the custom service is set, inspector service extensions will be registered
+  // against the wrong service.
+  WidgetPreviewScaffoldInspectorService(dtdServices: controller.dtdServices);
   final WidgetsBinding binding = WidgetsFlutterBinding.ensureInitialized();
   // Disable the injection of [WidgetInspector] into the widget tree built by
   // [WidgetsApp]. [WidgetInspector] instances will be created for each
   // individual preview so the widget inspector won't allow for users to select
   // widgets that make up the widget preview scaffolding.
   binding.debugExcludeRootWidgetInspector = true;
-  final controller = WidgetPreviewScaffoldController(previews: previews);
-  await controller.initialize();
   runWidget(
     DisableWidgetInspectorScope(
       child: binding.wrapWithDefaultView(

packages/flutter_tools/test/widget_preview_scaffold.shard/widget_preview_scaffold/lib/src/widget_preview_scaffold_controller.dart

@@ -3,9 +3,7 @@
 // found in the LICENSE file.
 
 import 'package:flutter/foundation.dart';
-import 'package:flutter/widgets.dart';
 import 'package:path/path.dart' as path;
-import 'package:widget_preview_scaffold/src/dtd/editor_service.dart';
 import 'package:widget_preview_scaffold/src/widget_preview_rendering.dart';
 import 'dtd/dtd_services.dart';
 import 'widget_preview.dart';
@@ -24,11 +22,7 @@ class WidgetPreviewScaffoldController {
     required PreviewsCallback previews,
     @visibleForTesting WidgetPreviewScaffoldDtdServices? dtdServicesOverride,
   }) : _previews = previews,
-       dtdServices = dtdServicesOverride ?? WidgetPreviewScaffoldDtdServices() {
-    // Overrides the default WidgetInspectorService instance to handle selection
-    // events.
-    _WidgetPreviewScaffoldInspectorService(dtdServices: dtdServices);
-  }
+       dtdServices = dtdServicesOverride ?? WidgetPreviewScaffoldDtdServices();
 
   @visibleForTesting
   static const kFilterBySelectedFilePreference = 'filterBySelectedFile';
@@ -170,77 +164,3 @@ class WidgetPreviewScaffoldController {
     }
   }
 }
-
-/// A custom [WidgetInspectorService] responsible for routing navigation events
-/// to the IDE.
-class _WidgetPreviewScaffoldInspectorService with WidgetInspectorService {
-  _WidgetPreviewScaffoldInspectorService({required this.dtdServices}) {
-    WidgetInspectorService.instance = this;
-  }
-
-  final WidgetPreviewScaffoldDtdServices dtdServices;
-
-  // Keys used to specify the creation location of a widget when serializing a
-  // DiagnosticsNode to JSON. This location is used by the widget inspector
-  // to jump to the creation location of a selected widget.
-  static const kFile = 'fileUri';
-  static const kLine = 'line';
-  static const kColumn = 'column';
-
-  CodeLocation? _nextNavigationLocation;
-
-  @protected
-  @override
-  bool setSelection(Object? object, [String? groupName]) {
-    // The next navigation event sent to `postEvent` will be for this selection.
-    // Save the location of preview annotation applications so we can override
-    // the navigation target in `postEvent`.
-    if (object is PreviewWidgetElement) {
-      final previewData = (object.widget as PreviewWidget).preview;
-      _nextNavigationLocation = CodeLocation(
-        uri: previewData.scriptUri,
-        line: previewData.line,
-        column: previewData.column,
-      );
-    }
-    final result = super.setSelection(object, groupName);
-    _nextNavigationLocation = null;
-    return result;
-  }
-
-  @override
-  void postEvent(
-    String eventKind,
-    Map<Object, Object?> eventData, {
-    String stream = 'Extension',
-  }) {
-    // It's unlikely that the widget previewer will be connected to directly by
-    // an IDE via the VM service, so we forward navigation events via the
-    // Editor DTD service.
-    if (eventKind == 'navigate') {
-      CodeLocation? location = _nextNavigationLocation;
-      if (eventData case {
-        kFile: final String file,
-        kLine: final int line,
-        kColumn: final int column,
-      } when location == null) {
-        location = CodeLocation(uri: file, line: line, column: column);
-      } else if (location != null) {
-        // If a [PreviewWidgetElement] was selected, we're not navigating to the
-        // creation location of the widget. Override the location details in the
-        // event data, just in case an IDE is attached and listening for
-        // navigation events through the VM service.
-        // TODO(bkonyi): determine if this is necessary
-        eventData.addAll(<String, Object>{
-          kFile: location.uri,
-          kLine: location.line!,
-          kColumn: location.column!,
-        });
-      }
-      if (location != null) {
-        dtdServices.navigateToCode(location);
-      }
-    }
-    super.postEvent(eventKind, eventData, stream: stream);
-  }
-}

packages/flutter_tools/test/widget_preview_scaffold.shard/widget_preview_scaffold/test/widget_inspector_service_override_test.dart

@@ -0,0 +1,56 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'package:flutter/widgets.dart';
+import 'package:flutter_test/flutter_test.dart';
+import 'package:widget_preview_scaffold/src/widget_preview_inspector_service.dart';
+
+import 'utils/widget_preview_scaffold_test_utils.dart';
+
+void main() {
+  test(
+    'Ensure $WidgetPreviewScaffoldInspectorService is initialized correctly',
+    () async {
+      expect(
+        WidgetInspectorService.instance,
+        isNot(isA<WidgetPreviewScaffoldInspectorService>()),
+      );
+      final dtdServices = FakeWidgetPreviewScaffoldDtdServices();
+      // Override the original WidgetInspectorService with our custom version.
+      TestWidgetPreviewScaffoldInspectorService(dtdServices: dtdServices);
+      expect(
+        WidgetInspectorService.instance,
+        isA<WidgetPreviewScaffoldInspectorService>(),
+      );
+      // Initialize the bindings and verify that the inspector service hasn't been
+      // changed and that initServiceExtensions has been invoked. This indicates
+      // that the inspector service extensions have been initialized with the custom
+      // inspector service instance and will be routed through the overridden methods.
+      WidgetsFlutterBinding.ensureInitialized();
+      expect(
+        WidgetInspectorService.instance,
+        isA<WidgetPreviewScaffoldInspectorService>(),
+      );
+      expect(
+        TestWidgetPreviewScaffoldInspectorService.serviceExtensionsRegistered,
+        true,
+      );
+    },
+  );
+}
+
+final class TestWidgetPreviewScaffoldInspectorService
+    extends WidgetPreviewScaffoldInspectorService {
+  TestWidgetPreviewScaffoldInspectorService({required super.dtdServices});
+
+  static bool serviceExtensionsRegistered = false;
+
+  @override
+  void initServiceExtensions(
+    RegisterServiceExtensionCallback registerExtension,
+  ) {
+    super.initServiceExtensions(registerExtension);
+    serviceExtensionsRegistered = true;
+  }
+}

packages/flutter_tools/test/widget_preview_scaffold.shard/widget_preview_scaffold/test/widget_inspector_test.dart

@@ -6,6 +6,7 @@ import 'package:flutter/widget_previews.dart';
 import 'package:flutter/widgets.dart';
 import 'package:flutter_test/flutter_test.dart';
 import 'package:widget_preview_scaffold/src/widget_preview.dart';
+import 'package:widget_preview_scaffold/src/widget_preview_inspector_service.dart';
 import 'package:widget_preview_scaffold/src/widget_preview_rendering.dart';
 
 import 'utils/widget_preview_scaffold_test_utils.dart';
@@ -89,9 +90,14 @@ void main() {
   testWidgets('WidgetInspector navigates to Preview application location', (
     tester,
   ) async {
-    final controller = FakeWidgetPreviewScaffoldController();
+    final dtd = FakeWidgetPreviewScaffoldDtdServices();
+    // Install the WidgetInspectorService override (this is done in the
+    // constructor body).
+    WidgetPreviewScaffoldInspectorService(dtdServices: dtd);
+    final controller = FakeWidgetPreviewScaffoldController(
+      dtdServicesOverride: dtd,
+    );
     await controller.initialize();
-    final dtd = controller.dtdServices as FakeWidgetPreviewScaffoldDtdServices;
 
     final service = WidgetInspectorService.instance;
     service.isSelectMode = true;