fix(marionette_mcp): time out connect when the VM service stops answering

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Unreleased
 
 - Fix `record-video --width/--height` recording only the top-left corner of the screen at full scale instead of the whole screen scaled down to that size. Web recordings were not affected ([#134](https://github.com/leancodepl/marionette_mcp/issues/134))
+- Fix `connect` waiting indefinitely when the app's VM service accepts the connection but stops answering, for example while the device is asleep: it now fails after 30 seconds with an error that says what to check, and stops waiting for a previously connected app that no longer responds to close its connection ([#131](https://github.com/leancodepl/marionette_mcp/issues/131))
 
 # 0.7.0
 

packages/marionette_mcp/lib/src/vm_service/vm_service_connector.dart

@@ -155,20 +155,43 @@ class VmServiceConnector {
   /// Returns true if currently connected to a VM service.
   bool get isConnected => _service != null && _isolateId != null;
 
+  /// How long [connect] waits for the app's VM service by default.
+  static const defaultConnectTimeout = Duration(seconds: 30);
+
+  /// How long [disconnect] waits for the old VM service to close.
+  static const _disposeTimeout = Duration(seconds: 5);
+
   /// Connects to a VM service at the given URI.
   ///
+  /// Gives up with a [TimeoutException] when the VM service doesn't finish
+  /// the handshake within [timeout], for example because the app is paused
+  /// or its device went to sleep, and leaves the connector disconnected.
+  ///
   /// Throws an exception if connection fails.
-  Future<void> connect(String uri) async {
+  Future<void> connect(
+    String uri, {
+    Duration timeout = defaultConnectTimeout,
+  }) async {
     if (isConnected) {
       _logger.warning('Already connected, disconnecting first');
       await disconnect();
     }
 
     _logger.info('Connecting to VM service at $uri');
 
-    try {
-      _service = await vmServiceConnectUri(uri);
-      _serviceEventSubscription = _service!.onServiceEvent.listen((e) {
+    // A handshake that times out keeps running in the background, so it
+    // must neither publish its service nor leak it once it finally returns.
+    var abandoned = false;
+    VmService? service;
+
+    Future<String> handshake() async {
+      final connected = await vmServiceConnectUri(uri);
+      if (abandoned) {
+        await connected.dispose();
+        throw StateError('Abandoned connection to $uri');
+      }
+      service = _service = connected;
+      _serviceEventSubscription = connected.onServiceEvent.listen((e) {
         switch (e.kind) {
           case EventKind.kServiceRegistered:
             final serviceName = e.service!;
@@ -188,33 +211,68 @@ class VmServiceConnector {
         }
       });
       try {
-        await _service!
+        await connected
             .streamListen(EventStreams.kService)
             .timeout(const Duration(milliseconds: 500));
       } catch (e) {
         _logger.info('streamListen(EventStreams.kService) skipped: $e');
       }
 
-      _isolateId = await _findIsolateWithMarionetteExtensions();
+      return _findIsolateWithMarionetteExtensions();
+    }
+
+    try {
+      final isolateId = await handshake().timeout(
+        timeout,
+        onTimeout: () {
+          abandoned = true;
+          throw TimeoutException(
+            'Timed out connecting to the VM service at $uri after '
+            '${timeout.inSeconds}s. Make sure the app is still running, '
+            'is not paused in the debugger, and its device is awake, '
+            'then call connect again.',
+            timeout,
+          );
+        },
+      );
+      _isolateId = isolateId;
       _logger.info('Connected to isolate: $_isolateId');
     } catch (err) {
       _logger.severe('Failed to connect to VM service', err);
-      await disconnect();
+      // Only tear down the service this attempt opened: a later connect may
+      // already have replaced it.
+      if (identical(_service, service)) {
+        await disconnect();
+      }
       rethrow;
     }
   }
 
   /// Disconnects from the current VM service.
+  ///
+  /// Doesn't wait more than a few seconds for an unresponsive VM service to
+  /// close its connection, so a hung app can't block the next [connect].
   Future<void> disconnect() async {
-    if (_service != null) {
+    final service = _service;
+    if (service != null) {
       _logger.info('Disconnecting from VM service');
-      await _serviceEventSubscription?.cancel();
-      _serviceEventSubscription = null;
-      await _service!.dispose();
+      final subscription = _serviceEventSubscription;
       _service = null;
       _isolateId = null;
+      _serviceEventSubscription = null;
       _registeredServices.clear();
       _pendingServiceRequests.clear();
+      try {
+        await Future.wait([
+          if (subscription != null) subscription.cancel(),
+          service.dispose(),
+        ]).timeout(_disposeTimeout);
+      } on TimeoutException {
+        _logger.warning(
+          'VM service did not close within ${_disposeTimeout.inSeconds}s; '
+          'dropping the connection',
+        );
+      }
       _logger.fine('Disconnected');
     }
   }

packages/marionette_mcp/test/vm_service/connect_timeout_test.dart

@@ -0,0 +1,203 @@
+import 'dart:async';
+import 'dart:convert';
+import 'dart:io';
+
+import 'package:marionette_mcp/src/vm_service/vm_service_connector.dart';
+import 'package:test/test.dart';
+
+/// A local stand-in for an app's VM service whose replies the test controls.
+class _FakeVmService {
+  _FakeVmService._(this._server);
+
+  final HttpServer _server;
+  final _sockets = <WebSocket>[];
+
+  /// Whether the server completes the WebSocket upgrade at all.
+  bool acceptUpgrade = true;
+
+  /// Whether the server answers JSON-RPC requests once connected.
+  bool answerRequests = true;
+
+  /// Extensions the single fake isolate reports.
+  List<String> extensionRPCs = ['ext.flutter.marionette.getLogs'];
+
+  static Future<_FakeVmService> start() async {
+    final server = await HttpServer.bind(InternetAddress.loopbackIPv4, 0);
+    final fake = _FakeVmService._(server);
+    server.listen(fake._handle);
+    return fake;
+  }
+
+  String get uri => 'ws://127.0.0.1:${_server.port}/ws';
+
+  Future<void> _handle(HttpRequest request) async {
+    // Leaving the request open simulates a VM service that accepted the
+    // TCP connection but never finishes the handshake.
+    if (!acceptUpgrade) return;
+    final socket = await WebSocketTransformer.upgrade(request);
+    _sockets.add(socket);
+    socket.listen((data) {
+      if (!answerRequests) return;
+      final request = jsonDecode(data as String) as Map<String, dynamic>;
+      socket.add(jsonEncode({
+        'jsonrpc': '2.0',
+        'id': request['id'],
+        'result': _result(request['method'] as String),
+      }));
+    });
+  }
+
+  Map<String, dynamic> _result(String method) {
+    const isolateRef = {
+      'type': '@Isolate',
+      'id': 'isolates/1',
+      'number': '1',
+      'name': 'main',
+      'isSystemIsolate': false,
+    };
+    return switch (method) {
+      'getVM' => {
+          'type': 'VM',
+          'name': 'vm',
+          'architectureBits': 64,
+          'hostCPU': '',
+          'operatingSystem': '',
+          'targetCPU': '',
+          'version': '',
+          'pid': 1,
+          'startTime': 0,
+          'isolates': [isolateRef],
+          'isolateGroups': <Object>[],
+          'systemIsolates': <Object>[],
+          'systemIsolateGroups': <Object>[],
+        },
+      'getIsolate' => {
+          ...isolateRef,
+          'type': 'Isolate',
+          'isolateFlags': <Object>[],
+          'startTime': 0,
+          'runnable': true,
+          'livePorts': 0,
+          'pauseOnExit': false,
+          'pauseEvent': {'type': 'Event', 'kind': 'Resume', 'timestamp': 0},
+          'libraries': <Object>[],
+          'breakpoints': <Object>[],
+          'exceptionPauseMode': 'None',
+          'extensionRPCs': extensionRPCs,
+        },
+      _ => {'type': 'Success'},
+    };
+  }
+
+  Future<void> close() async {
+    for (final socket in _sockets) {
+      await socket.close();
+    }
+    await _server.close(force: true);
+  }
+}
+
+void main() {
+  late _FakeVmService vmService;
+  late VmServiceConnector connector;
+
+  setUp(() async {
+    vmService = await _FakeVmService.start();
+    connector = VmServiceConnector();
+  });
+
+  tearDown(() async {
+    await connector.disconnect();
+    await vmService.close();
+  });
+
+  test('connects to a responsive VM service', () async {
+    await connector.connect(vmService.uri);
+
+    expect(connector.isConnected, isTrue);
+  });
+
+  test('times out when the VM service stops answering requests', () async {
+    vmService.answerRequests = false;
+
+    final stopwatch = Stopwatch()..start();
+    await expectLater(
+      connector.connect(
+        vmService.uri,
+        timeout: const Duration(milliseconds: 300),
+      ),
+      throwsA(
+        isA<TimeoutException>().having(
+          (e) => e.message,
+          'message',
+          contains(
+              'Timed out connecting to the VM service at ${vmService.uri}'),
+        ),
+      ),
+    );
+
+    expect(stopwatch.elapsed, lessThan(const Duration(seconds: 5)));
+    expect(connector.isConnected, isFalse);
+  });
+
+  test('times out when the WebSocket handshake never completes', () async {
+    vmService.acceptUpgrade = false;
+
+    await expectLater(
+      connector.connect(
+        vmService.uri,
+        timeout: const Duration(milliseconds: 300),
+      ),
+      throwsA(isA<TimeoutException>()),
+    );
+
+    expect(connector.isConnected, isFalse);
+  });
+
+  test('can connect again after a timed-out attempt', () async {
+    vmService.answerRequests = false;
+    await expectLater(
+      connector.connect(
+        vmService.uri,
+        timeout: const Duration(milliseconds: 300),
+      ),
+      throwsA(isA<TimeoutException>()),
+    );
+
+    vmService.answerRequests = true;
+    await connector.connect(vmService.uri);
+
+    expect(connector.isConnected, isTrue);
+  });
+
+  test('still reports a missing marionette binding without waiting', () async {
+    vmService.extensionRPCs = [];
+
+    await expectLater(
+      connector.connect(vmService.uri),
+      throwsA(
+        isA<Exception>().having(
+          (e) => e.toString(),
+          'message',
+          contains('No isolate found'),
+        ),
+      ),
+    );
+
+    expect(connector.isConnected, isFalse);
+  });
+
+  test('reconnects when the previous VM service stopped answering', () async {
+    await connector.connect(vmService.uri);
+    vmService.answerRequests = false;
+
+    final other = await _FakeVmService.start();
+    addTearDown(other.close);
+
+    final stopwatch = Stopwatch()..start();
+    await connector.connect(other.uri);
+
+    expect(connector.isConnected, isTrue);
+    expect(stopwatch.elapsed, lessThan(const Duration(seconds: 10)));
+  });
+}