[flutter_tools] Generalize waitForExtension (#77220)

packages/flutter_tools/lib/src/resident_devtools_handler.dart

@@ -7,7 +7,6 @@
 import 'dart:async';
 
 import 'package:meta/meta.dart';
-import 'package:vm_service/vm_service.dart' as vm_service;
 
 import 'base/logger.dart';
 import 'resident_runner.dart';
@@ -66,12 +65,12 @@ class FlutterResidentDevtoolsHandler implements ResidentDevtoolsHandler {
       // report their URLs yet. Do so now.
       _residentRunner.printDebuggerList(includeObservatory: false);
     }
-    await _waitForExtensions(flutterDevices);
+    final List<FlutterDevice> devicesWithExtension = await _devicesWithExtensions(flutterDevices);
     await _maybeCallDevToolsUriServiceExtension(
-      flutterDevices,
+      devicesWithExtension,
     );
     await _callConnectedVmServiceUriExtension(
-      flutterDevices,
+      devicesWithExtension,
     );
   }
 
@@ -107,12 +106,28 @@ class FlutterResidentDevtoolsHandler implements ResidentDevtoolsHandler {
     }
   }
 
-  Future<void> _waitForExtensions(List<FlutterDevice> flutterDevices) async {
-    await Future.wait(<Future<void>>[
+  Future<List<FlutterDevice>> _devicesWithExtensions(List<FlutterDevice> flutterDevices) async {
+    final List<FlutterDevice> devices = await Future.wait(<Future<FlutterDevice>>[
       for (final FlutterDevice device in flutterDevices)
-        if (device.vmService != null)
-          waitForExtension(device.vmService.service, 'ext.flutter.connectedVmServiceUri'),
+        _waitForExtensionsForDevice(device)
     ]);
+    return devices.where((FlutterDevice device) => device != null).toList();
+  }
+
+  /// Returns null if the service extension cannot be found on the device.
+  Future<FlutterDevice> _waitForExtensionsForDevice(FlutterDevice flutterDevice) async {
+    const String extension = 'ext.flutter.connectedVmServiceUri';
+    try {
+      await flutterDevice.vmService?.findExtensionIsolate(extension);
+      return flutterDevice;
+    } on VmServiceDisappearedException {
+      _logger.printTrace(
+        'The VM Service for ${flutterDevice.device} disappeared while trying to'
+        ' find the $extension service extension. Skipping subsequent DevTools '
+        'setup for this device.',
+      );
+      return null;
+    }
   }
 
   Future<void> _callConnectedVmServiceUriExtension(List<FlutterDevice> flutterDevices) async {
@@ -164,10 +179,10 @@ class FlutterResidentDevtoolsHandler implements ResidentDevtoolsHandler {
 
   @override
   Future<void> hotRestart(List<FlutterDevice> flutterDevices) async {
-    await _waitForExtensions(flutterDevices);
+    final List<FlutterDevice> devicesWithExtension = await _devicesWithExtensions(flutterDevices);
     await Future.wait(<Future<void>>[
-      _maybeCallDevToolsUriServiceExtension(flutterDevices),
-      _callConnectedVmServiceUriExtension(flutterDevices),
+      _maybeCallDevToolsUriServiceExtension(devicesWithExtension),
+      _callConnectedVmServiceUriExtension(devicesWithExtension),
     ]);
   }
 
@@ -181,37 +196,6 @@ class FlutterResidentDevtoolsHandler implements ResidentDevtoolsHandler {
   }
 }
 
-
-@visibleForTesting
-Future<void> waitForExtension(vm_service.VmService vmService, String extension) async {
-  final Completer<void> completer = Completer<void>();
-  try {
-    await vmService.streamListen(vm_service.EventStreams.kExtension);
-  } on Exception {
-    // do nothing
-  }
-  StreamSubscription<vm_service.Event> extensionStream;
-  extensionStream = vmService.onExtensionEvent.listen((vm_service.Event event) {
-    if (event.json['extensionKind'] == 'Flutter.FrameworkInitialization') {
-      // The 'Flutter.FrameworkInitialization' event is sent on hot restart
-      // as well, so make sure we don't try to complete this twice.
-      if (!completer.isCompleted) {
-        completer.complete();
-        extensionStream.cancel();
-      }
-    }
-  });
-  final vm_service.VM vm = await vmService.getVM();
-  if (vm.isolates.isNotEmpty) {
-    final vm_service.IsolateRef isolateRef = vm.isolates.first;
-    final vm_service.Isolate isolate = await vmService.getIsolate(isolateRef.id);
-    if (isolate.extensionRPCs.contains(extension)) {
-      return;
-    }
-  }
-  await completer.future;
-}
-
 @visibleForTesting
 NoOpDevtoolsHandler createNoOpHandler(DevtoolsLauncher launcher, ResidentRunner runner, Logger logger) {
   return NoOpDevtoolsHandler();

packages/flutter_tools/lib/src/vmservice.dart

@@ -4,6 +4,8 @@
 
 // @dart = 2.8
 
+import 'dart:async';
+
 import 'package:file/file.dart';
 import 'package:meta/meta.dart' show required;
 import 'package:vm_service/vm_service.dart' as vm_service;
@@ -479,7 +481,7 @@ class FlutterVmService {
     @required Uri assetsDirectory,
   }) async {
     try {
-      await service.streamListen('Isolate');
+      await service.streamListen(vm_service.EventStreams.kIsolate);
     } on vm_service.RPCError {
       // Do nothing, since the tool is already subscribed.
     }
@@ -784,6 +786,58 @@ class FlutterVmService {
     }
   }
 
+  /// Waits for a signal from the VM service that [extensionName] is registered.
+  ///
+  /// Looks at the list of loaded extensions for first Flutter view, as well as
+  /// the stream of added extensions to avoid races.
+  ///
+  /// Throws a [VmServiceDisappearedException] should the VM Service disappear
+  /// while making calls to it.
+  Future<vm_service.IsolateRef> findExtensionIsolate(String extensionName) async {
+    try {
+      await service.streamListen(vm_service.EventStreams.kIsolate);
+    } on vm_service.RPCError {
+      // Do nothing, since the tool is already subscribed.
+    }
+
+    final Completer<vm_service.IsolateRef> extensionAdded = Completer<vm_service.IsolateRef>();
+    StreamSubscription<vm_service.Event> isolateEvents;
+    isolateEvents = service.onIsolateEvent.listen((vm_service.Event event) {
+      if (event.kind == vm_service.EventKind.kServiceExtensionAdded
+          && event.extensionRPC == extensionName) {
+        isolateEvents.cancel();
+        extensionAdded.complete(event.isolate);
+      }
+    });
+
+    try {
+      final List<FlutterView> flutterViews = await getFlutterViews();
+      if (flutterViews.isEmpty) {
+        throw VmServiceDisappearedException();
+      }
+
+      for (final FlutterView flutterView in flutterViews) {
+        final vm_service.IsolateRef isolateRef = flutterView.uiIsolate;
+        if (isolateRef == null) {
+          continue;
+        }
+
+        final vm_service.Isolate isolate = await service.getIsolate(isolateRef.id);
+        if (isolate.extensionRPCs.contains(extensionName)) {
+          return isolateRef;
+        }
+      }
+      return await extensionAdded.future;
+    } finally {
+      await isolateEvents.cancel();
+      try {
+        await service.streamCancel(vm_service.EventStreams.kIsolate);
+      } on vm_service.RPCError {
+        // It's ok for cleanup to fail, such as when the service disappears.
+      }
+    }
+  }
+
   /// Attempt to retrieve the isolate with id [isolateId], or `null` if it has
   /// been collected.
   Future<vm_service.Isolate> getIsolateOrNull(String isolateId) {
@@ -845,6 +899,9 @@ class FlutterVmService {
   }
 }
 
+/// Thrown when the VM Service disappears while calls are being made to it.
+class VmServiceDisappearedException implements Exception {}
+
 /// Whether the event attached to an [Isolate.pauseEvent] should be considered
 /// a "pause" event.
 bool isPauseEvent(String kind) {

packages/flutter_tools/test/general.shard/resident_devtools_handler_test.dart

@@ -5,6 +5,7 @@
 // @dart = 2.8
 
 import 'package:flutter_tools/src/base/platform.dart';
+import 'package:flutter_tools/src/device.dart';
 import 'package:flutter_tools/src/devtools_launcher.dart';
 import 'package:flutter_tools/src/vmservice.dart';
 import 'package:vm_service/vm_service.dart' as vm_service;
@@ -17,7 +18,7 @@ import 'package:test/fake.dart';
 import '../src/common.dart';
 import '../src/context.dart';
 
- final vm_service.Isolate isolate = vm_service.Isolate(
+final vm_service.Isolate isolate = vm_service.Isolate(
   id: '1',
   pauseEvent: vm_service.Event(
     kind: vm_service.EventKind.kResume,
@@ -40,60 +41,17 @@ import '../src/context.dart';
   startTime: 0,
   isSystemIsolate: false,
   isolateFlags: <vm_service.IsolateFlag>[],
-  extensionRPCs: <String>['foo']
-);
-
-final vm_service.Isolate fakeUnpausedIsolate = vm_service.Isolate(
-  id: '1',
-  pauseEvent: vm_service.Event(
-    kind: vm_service.EventKind.kResume,
-    timestamp: 0
-  ),
-  breakpoints: <vm_service.Breakpoint>[],
-  exceptionPauseMode: null,
-  extensionRPCs: <String>[],
-  libraries: <vm_service.LibraryRef>[
-    vm_service.LibraryRef(
-      id: '1',
-      uri: 'file:///hello_world/main.dart',
-      name: '',
-    ),
-  ],
-  livePorts: 0,
-  name: 'test',
-  number: '1',
-  pauseOnExit: false,
-  runnable: true,
-  startTime: 0,
-  isSystemIsolate: false,
-  isolateFlags: <vm_service.IsolateFlag>[],
-);
-
-final vm_service.VM fakeVM = vm_service.VM(
-  isolates: <vm_service.IsolateRef>[fakeUnpausedIsolate],
-  pid: 1,
-  hostCPU: '',
-  isolateGroups: <vm_service.IsolateGroupRef>[],
-  targetCPU: '',
-  startTime: 0,
-  name: 'dart',
-  architectureBits: 64,
-  operatingSystem: '',
-  version: '',
-  systemIsolateGroups: <vm_service.IsolateGroupRef>[],
-  systemIsolates: <vm_service.IsolateRef>[],
-);
-
-final FlutterView fakeFlutterView = FlutterView(
-  id: 'a',
-  uiIsolate: fakeUnpausedIsolate,
+  extensionRPCs: <String>['ext.flutter.connectedVmServiceUri'],
 );
 
 final FakeVmServiceRequest listViews = FakeVmServiceRequest(
   method: kListViewsMethod,
   jsonResponse: <String, Object>{
     'views': <Object>[
-      fakeFlutterView.toJson(),
+      FlutterView(
+        id: 'a',
+        uiIsolate: isolate,
+      ).toJson()
     ],
   },
 );
@@ -173,10 +131,10 @@ void main() {
       const FakeVmServiceRequest(
         method: 'streamListen',
         args: <String, Object>{
-          'streamId': 'Extension',
+          'streamId': 'Isolate',
         }
       ),
-      FakeVmServiceRequest(method: 'getVM', jsonResponse: fakeVM.toJson()),
+      listViews,
       FakeVmServiceRequest(
         method: 'getIsolate',
         jsonResponse: isolate.toJson(),
@@ -184,13 +142,11 @@ void main() {
           'isolateId': '1',
         },
       ),
-      FakeVmServiceStreamResponse(
-        streamId: 'Extension',
-        event: vm_service.Event(
-          timestamp: 0,
-          extensionKind: 'Flutter.FrameworkInitialization',
-          kind: 'test',
-        ),
+      const FakeVmServiceRequest(
+        method: 'streamCancel',
+        args: <String, Object>{
+          'streamId': 'Isolate',
+        },
       ),
       listViews,
       const FakeVmServiceRequest(
@@ -218,15 +174,55 @@ void main() {
     );
   });
 
-  testWithoutContext('wait for extension handles an immediate extension', () {
+  testWithoutContext('serveAndAnnounceDevTools with skips calling service extensions when VM service disappears', () async {
+    final ResidentDevtoolsHandler handler = FlutterResidentDevtoolsHandler(
+      FakeDevtoolsLauncher()..activeDevToolsServer = DevToolsServerAddress('localhost', 8080),
+      FakeResidentRunner(),
+      BufferLogger.test(),
+    );
     final FakeVmServiceHost fakeVmServiceHost = FakeVmServiceHost(requests: <VmServiceExpectation>[
       const FakeVmServiceRequest(
         method: 'streamListen',
         args: <String, Object>{
-          'streamId': 'Extension',
-        }
+          'streamId': 'Isolate',
+        },
       ),
-      FakeVmServiceRequest(method: 'getVM', jsonResponse: fakeVM.toJson()),
+      const FakeVmServiceRequest(
+        method: kListViewsMethod,
+        errorCode: RPCErrorCodes.kServiceDisappeared,
+      ),
+      const FakeVmServiceRequest(
+        method: 'streamCancel',
+        args: <String, Object>{
+          'streamId': 'Isolate',
+        },
+        errorCode: RPCErrorCodes.kServiceDisappeared,
+      ),
+    ], httpAddress: Uri.parse('http://localhost:1234'));
+
+    final FakeFlutterDevice device = FakeFlutterDevice()
+      ..vmService = fakeVmServiceHost.vmService;
+
+    await handler.serveAndAnnounceDevTools(
+      flutterDevices: <FlutterDevice>[device],
+    );
+  });
+
+  testWithoutContext('serveAndAnnounceDevTools with multiple devices and VM service disappears on one', () async {
+    final ResidentDevtoolsHandler handler = FlutterResidentDevtoolsHandler(
+      FakeDevtoolsLauncher()..activeDevToolsServer = DevToolsServerAddress('localhost', 8080),
+      FakeResidentRunner(),
+      BufferLogger.test(),
+    );
+
+    final FakeVmServiceHost vmServiceHost = FakeVmServiceHost(requests: <VmServiceExpectation>[
+      const FakeVmServiceRequest(
+        method: 'streamListen',
+        args: <String, Object>{
+          'streamId': 'Isolate',
+        },
+      ),
+      listViews,
       FakeVmServiceRequest(
         method: 'getIsolate',
         jsonResponse: isolate.toJson(),
@@ -234,46 +230,61 @@ void main() {
           'isolateId': '1',
         },
       ),
-    ]);
-    waitForExtension(fakeVmServiceHost.vmService.service, 'foo');
-  });
+      const FakeVmServiceRequest(
+        method: 'streamCancel',
+        args: <String, Object>{
+          'streamId': 'Isolate',
+        },
+      ),
+      listViews,
+      const FakeVmServiceRequest(
+        method: 'ext.flutter.activeDevToolsServerAddress',
+        args: <String, Object>{
+          'isolateId': '1',
+          'value': 'http://localhost:8080',
+        },
+      ),
+      listViews,
+      const FakeVmServiceRequest(
+        method: 'ext.flutter.connectedVmServiceUri',
+        args: <String, Object>{
+          'isolateId': '1',
+          'value': 'http://localhost:1234',
+        },
+      ),
+    ], httpAddress: Uri.parse('http://localhost:1234'));
 
-  testWithoutContext('wait for extension handles no isolates', () {
-    final FakeVmServiceHost fakeVmServiceHost = FakeVmServiceHost(requests: <VmServiceExpectation>[
+    final FakeVmServiceHost vmServiceHostThatDisappears = FakeVmServiceHost(requests: <VmServiceExpectation>[
       const FakeVmServiceRequest(
         method: 'streamListen',
         args: <String, Object>{
-          'streamId': 'Extension',
-        }
+          'streamId': 'Isolate',
+        },
       ),
-      FakeVmServiceRequest(method: 'getVM', jsonResponse: vm_service.VM(
-        isolates: <vm_service.IsolateRef>[],
-        pid: 1,
-        hostCPU: '',
-        isolateGroups: <vm_service.IsolateGroupRef>[],
-        targetCPU: '',
-        startTime: 0,
-        name: 'dart',
-        architectureBits: 64,
-        operatingSystem: '',
-        version: '',
-        systemIsolateGroups: <vm_service.IsolateGroupRef>[],
-        systemIsolates: <vm_service.IsolateRef>[],
-      ).toJson()),
-      FakeVmServiceStreamResponse(
-        streamId: 'Extension',
-        event: vm_service.Event(
-          timestamp: 0,
-          extensionKind: 'Flutter.FrameworkInitialization',
-          kind: 'test',
-        ),
+      const FakeVmServiceRequest(
+        method: kListViewsMethod,
+        errorCode: RPCErrorCodes.kServiceDisappeared,
       ),
-    ]);
-    waitForExtension(fakeVmServiceHost.vmService.service, 'foo');
+      const FakeVmServiceRequest(
+        method: 'streamCancel',
+        args: <String, Object>{
+          'streamId': 'Isolate',
+        },
+        errorCode: RPCErrorCodes.kServiceDisappeared,
+      ),
+    ], httpAddress: Uri.parse('http://localhost:5678'));
+
+    await handler.serveAndAnnounceDevTools(
+      flutterDevices: <FlutterDevice>[
+        FakeFlutterDevice()
+          ..vmService = vmServiceHostThatDisappears.vmService,
+        FakeFlutterDevice()
+          ..vmService = vmServiceHost.vmService,
+      ],
+    );
   });
 }
 
-
 class FakeDevtoolsLauncher extends Fake implements DevtoolsLauncher {
   @override
   DevToolsServerAddress activeDevToolsServer;
@@ -296,6 +307,11 @@ class FakeResidentRunner extends Fake implements ResidentRunner {
 }
 
 class FakeFlutterDevice extends Fake implements FlutterDevice {
+  @override
+  final Device device = FakeDevice();
+
   @override
   FlutterVmService vmService;
 }
+
+class FakeDevice extends Fake implements Device {}

packages/flutter_tools/test/general.shard/vmservice_test.dart

@@ -45,52 +45,47 @@ final Map<String, Object> vm = <String, dynamic>{
   ],
 };
 
-final vm_service.Isolate isolate = vm_service.Isolate.parse(
-  <String, dynamic>{
-    'type': 'Isolate',
-    'fixedId': true,
-    'id': 'isolates/242098474',
-    'name': 'main.dart:main()',
-    'number': 242098474,
-    '_originNumber': 242098474,
-    'startTime': 1540488745340,
-    '_heaps': <String, dynamic>{
-      'new': <String, dynamic>{
-        'used': 0,
-        'capacity': 0,
-        'external': 0,
-        'collections': 0,
-        'time': 0.0,
-        'avgCollectionPeriodMillis': 0.0,
-      },
-      'old': <String, dynamic>{
-        'used': 0,
-        'capacity': 0,
-        'external': 0,
-        'collections': 0,
-        'time': 0.0,
-        'avgCollectionPeriodMillis': 0.0,
-      },
-    },
-  }
+const String kExtensionName = 'ext.flutter.test.interestingExtension';
+
+final vm_service.Isolate isolate = vm_service.Isolate(
+  id: '1',
+  pauseEvent: vm_service.Event(
+    kind: vm_service.EventKind.kResume,
+    timestamp: 0
+  ),
+  breakpoints: <vm_service.Breakpoint>[],
+  exceptionPauseMode: null,
+  libraries: <vm_service.LibraryRef>[
+    vm_service.LibraryRef(
+      id: '1',
+      uri: 'file:///hello_world/main.dart',
+      name: '',
+    ),
+  ],
+  livePorts: 0,
+  name: 'test',
+  number: '1',
+  pauseOnExit: false,
+  runnable: true,
+  startTime: 0,
+  isSystemIsolate: false,
+  isolateFlags: <vm_service.IsolateFlag>[],
+  extensionRPCs: <String>[kExtensionName],
 );
 
-final Map<String, Object> listViews = <String, dynamic>{
-  'type': 'FlutterViewList',
-  'views': <dynamic>[
-    <String, dynamic>{
-      'type': 'FlutterView',
-      'id': '_flutterView/0x4a4c1f8',
-      'isolate': <String, dynamic>{
-        'type': '@Isolate',
-        'fixedId': true,
-        'id': 'isolates/242098474',
-        'name': 'main.dart:main()',
-        'number': 242098474,
-      },
-    },
-  ]
-};
+final FlutterView fakeFlutterView = FlutterView(
+  id: 'a',
+  uiIsolate: isolate,
+);
+
+final FakeVmServiceRequest listViewsRequest = FakeVmServiceRequest(
+  method: kListViewsMethod,
+  jsonResponse: <String, Object>{
+    'views': <Object>[
+      fakeFlutterView.toJson(),
+    ],
+  },
+);
 
 typedef ServiceCallback = Future<Map<String, dynamic>> Function(Map<String, Object>);
 
@@ -408,17 +403,7 @@ void main() {
             'views': <Object>[],
           },
         ),
-        const FakeVmServiceRequest(
-          method: kListViewsMethod,
-          jsonResponse: <String, Object>{
-            'views': <Object>[
-              <String, Object>{
-                'id': 'a',
-                'isolate': <String, Object>{},
-              },
-            ],
-          },
-        ),
+        listViewsRequest,
       ]
     );
 
@@ -452,6 +437,187 @@ void main() {
     expect(fakeVmServiceHost.hasRemainingExpectations, false);
   });
 
+  group('findExtensionIsolate', () {
+
+    testWithoutContext('returns an isolate with the registered extensionRPC', () async {
+      final FakeVmServiceHost fakeVmServiceHost = FakeVmServiceHost(requests: <VmServiceExpectation>[
+        const FakeVmServiceRequest(
+          method: 'streamListen',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+        ),
+        listViewsRequest,
+        FakeVmServiceRequest(
+          method: 'getIsolate',
+          jsonResponse: isolate.toJson(),
+          args: <String, Object>{
+            'isolateId': '1',
+          },
+        ),
+        const FakeVmServiceRequest(
+          method: 'streamCancel',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+        ),
+      ]);
+
+      final vm_service.IsolateRef isolateRef = await fakeVmServiceHost.vmService.findExtensionIsolate(kExtensionName);
+      expect(isolateRef.id, '1');
+    });
+
+    testWithoutContext('returns the isolate with the registered extensionRPC when there are multiple FlutterViews', () async {
+      const String otherExtensionName = 'ext.flutter.test.otherExtension';
+
+      // Copy the other isolate and change a few fields.
+      final vm_service.Isolate isolate2 = vm_service.Isolate.parse(
+        isolate.toJson()
+          ..['id'] = '2'
+          ..['extensionRPCs'] = <String>[otherExtensionName],
+      );
+
+      final FlutterView fakeFlutterView2 = FlutterView(
+        id: '2',
+        uiIsolate: isolate2,
+      );
+
+      final FakeVmServiceHost fakeVmServiceHost = FakeVmServiceHost(requests: <VmServiceExpectation>[
+        const FakeVmServiceRequest(
+          method: 'streamListen',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+        ),
+        FakeVmServiceRequest(
+          method: kListViewsMethod,
+          jsonResponse: <String, Object>{
+            'views': <Object>[
+              fakeFlutterView.toJson(),
+              fakeFlutterView2.toJson(),
+            ],
+          },
+        ),
+        FakeVmServiceRequest(
+          method: 'getIsolate',
+          jsonResponse: isolate.toJson(),
+          args: <String, Object>{
+            'isolateId': '1',
+          },
+        ),
+        FakeVmServiceRequest(
+          method: 'getIsolate',
+          jsonResponse: isolate2.toJson(),
+          args: <String, Object>{
+            'isolateId': '2',
+          },
+        ),
+        const FakeVmServiceRequest(
+          method: 'streamCancel',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+        ),
+      ]);
+
+      final vm_service.IsolateRef isolateRef = await fakeVmServiceHost.vmService.findExtensionIsolate(otherExtensionName);
+      expect(isolateRef.id, '2');
+    });
+
+    testWithoutContext('when the isolate stream is already subscribed, returns an isolate with the registered extensionRPC', () async {
+      final FakeVmServiceHost fakeVmServiceHost = FakeVmServiceHost(requests: <VmServiceExpectation>[
+        const FakeVmServiceRequest(
+          method: 'streamListen',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+          // Stream already subscribed - https://github.com/dart-lang/sdk/blob/master/runtime/vm/service/service.md#streamlisten
+          errorCode: 103,
+        ),
+        listViewsRequest,
+        FakeVmServiceRequest(
+          method: 'getIsolate',
+          jsonResponse: isolate.toJson()..['extensionRPCs'] = <String>[kExtensionName],
+          args: <String, Object>{
+            'isolateId': '1',
+          },
+        ),
+        const FakeVmServiceRequest(
+          method: 'streamCancel',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+        ),
+      ]);
+
+      final vm_service.IsolateRef isolateRef = await fakeVmServiceHost.vmService.findExtensionIsolate(kExtensionName);
+      expect(isolateRef.id, '1');
+    });
+
+    testWithoutContext('returns an isolate with a extensionRPC that is registered later', () async {
+      final FakeVmServiceHost fakeVmServiceHost = FakeVmServiceHost(requests: <VmServiceExpectation>[
+        const FakeVmServiceRequest(
+          method: 'streamListen',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+        ),
+        listViewsRequest,
+        FakeVmServiceRequest(
+          method: 'getIsolate',
+          jsonResponse: isolate.toJson(),
+          args: <String, Object>{
+            'isolateId': '1',
+          },
+        ),
+        FakeVmServiceStreamResponse(
+          streamId: 'Isolate',
+          event: vm_service.Event(
+            kind: vm_service.EventKind.kServiceExtensionAdded,
+            extensionRPC: kExtensionName,
+            timestamp: 1,
+          ),
+        ),
+        const FakeVmServiceRequest(
+          method: 'streamCancel',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+        ),
+      ]);
+
+      final vm_service.IsolateRef isolateRef = await fakeVmServiceHost.vmService.findExtensionIsolate(kExtensionName);
+      expect(isolateRef.id, '1');
+    });
+
+    testWithoutContext('throws when the service disappears', () async {
+      final FakeVmServiceHost fakeVmServiceHost = FakeVmServiceHost(requests: <VmServiceExpectation>[
+        const FakeVmServiceRequest(
+          method: 'streamListen',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+        ),
+        const FakeVmServiceRequest(
+          method: kListViewsMethod,
+          errorCode: RPCErrorCodes.kServiceDisappeared,
+        ),
+        const FakeVmServiceRequest(
+          method: 'streamCancel',
+          args: <String, Object>{
+            'streamId': 'Isolate',
+          },
+          errorCode: RPCErrorCodes.kServiceDisappeared,
+        ),
+      ]);
+
+      expect(
+        () => fakeVmServiceHost.vmService.findExtensionIsolate(kExtensionName),
+        throwsA(isA<VmServiceDisappearedException>()),
+      );
+    });
+  });
+
   testWithoutContext('Can process log events from the vm service', () {
     final vm_service.Event event = vm_service.Event(
       bytes: base64.encode(utf8.encode('Hello There\n')),