From c9626833f3727d55d889cd2c8cdf67a6b69027ed Mon Sep 17 00:00:00 2001
From: Elliott Brooks <21270878+elliette@users.noreply.github.com>
Date: Tue, 28 May 2024 13:05:31 -0700
Subject: [PATCH] Add a reload handler along with a restart handler

---
 dwds/lib/src/dwds_vm_client.dart | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/dwds/lib/src/dwds_vm_client.dart b/dwds/lib/src/dwds_vm_client.dart
index b5a89f579..0a66c7385 100644
--- a/dwds/lib/src/dwds_vm_client.dart
+++ b/dwds/lib/src/dwds_vm_client.dart
@@ -28,6 +28,7 @@ typedef VmResponse = Map<String, Object?>;
 
 enum _NamespacedServiceExtension {
   extDwdsEmitEvent(method: 'ext.dwds.emitEvent'),
+  extDwdsReload(method: 'ext.dwds.reload'),
   extDwdsRestart(method: 'ext.dwds.restart'),
   extDwdsScreenshot(method: 'ext.dwds.screenshot'),
   extDwdsSendEvent(method: 'ext.dwds.sendEvent'),
@@ -77,6 +78,8 @@ class DwdsVmClient {
     final requestSink = requestController.sink;
     final requestStream = requestController.stream;
 
+    final clientCompleter = Completer<VmService>();
+
     _setUpVmServerConnection(
       chromeProxyService: chromeProxyService,
       debugService: debugService,
@@ -85,6 +88,7 @@ class DwdsVmClient {
       requestStream: requestStream,
       requestSink: requestSink,
       dwdsStats: dwdsStats,
+      clientFuture: clientCompleter.future,
     );
 
     final client = ddsUri == null
@@ -97,6 +101,10 @@ class DwdsVmClient {
             ddsUri: ddsUri,
           );
 
+    if (!clientCompleter.isCompleted) {
+      clientCompleter.complete(client);
+    }
+
     final dwdsVmClient =
         DwdsVmClient(client, requestController, responseController);
 
@@ -158,12 +166,14 @@ class DwdsVmClient {
     required StreamSink<VmResponse> responseSink,
     required Stream<VmRequest> requestStream,
     required StreamSink<VmRequest> requestSink,
+    required Future<VmService> clientFuture,
   }) {
     responseStream.listen((request) async {
       final response = await _maybeHandleServiceExtensionRequest(
         request,
         chromeProxyService: chromeProxyService,
         dwdsStats: dwdsStats,
+        clientFuture: clientFuture,
       );
       if (response != null) {
         requestSink.add(response);
@@ -187,6 +197,7 @@ class DwdsVmClient {
     VmResponse request, {
     required ChromeProxyService chromeProxyService,
     required DwdsStats dwdsStats,
+    required Future<VmService> clientFuture,
   }) async {
     VmRequest? response;
     final method = request['method'];
@@ -194,8 +205,11 @@ class DwdsVmClient {
       response = await _flutterListViewsHandler(chromeProxyService);
     } else if (method == _NamespacedServiceExtension.extDwdsEmitEvent.method) {
       response = _extDwdsEmitEventHandler(request);
+    } else if (method == _NamespacedServiceExtension.extDwdsReload.method) {
+      response = await _extDwdsReloadHandler(chromeProxyService);
     } else if (method == _NamespacedServiceExtension.extDwdsRestart.method) {
-      response = await _extDwdsRestartHandler(chromeProxyService);
+      final client = await clientFuture;
+      response = await _extDwdsRestartHandler(chromeProxyService, client);
     } else if (method == _NamespacedServiceExtension.extDwdsSendEvent.method) {
       response = await _extDwdsSendEventHandler(request, dwdsStats);
     } else if (method == _NamespacedServiceExtension.extDwdsScreenshot.method) {
@@ -265,13 +279,21 @@ class DwdsVmClient {
     return {'result': Success().toJson()};
   }
 
-  static Future<Map<String, Object>> _extDwdsRestartHandler(
+  static Future<Map<String, Object>> _extDwdsReloadHandler(
     ChromeProxyService chromeProxyService,
   ) async {
     await _fullReload(chromeProxyService);
     return {'result': Success().toJson()};
   }
 
+  static Future<Map<String, Object>> _extDwdsRestartHandler(
+    ChromeProxyService chromeProxyService,
+    VmService client,
+  ) async {
+    await _hotRestart(chromeProxyService, client);
+    return {'result': Success().toJson()};
+  }
+
   static Future<void> _registerServiceExtensions({
     required VmService client,
     required ChromeProxyService chromeProxyService,