Connect queries from macro through to host. (#22)

* Connect up augmentations from macro output through to host.

* Connect queries from macro through to host.

Add union types to JSON generation, use them for requests and responses.

* Rebase.

* Address review comments.

pkgs/_macro_client/lib/macro_client.dart

@@ -2,31 +2,35 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
+import 'dart:async';
 import 'dart:convert';
 import 'dart:io';
 
+import 'package:dart_model/dart_model.dart';
 import 'package:macro/macro.dart';
 import 'package:macro_service/macro_service.dart';
 
-/// Runs macros, connecting them to a macro host.
+/// Local macro client which runs macros as directed by requests from a remote
+/// macro host.
 ///
 /// TODO(davidmorgan): handle shutdown and dispose.
 /// TODO(davidmorgan): split to multpile implementations depending on
 /// transport used to connect to host.
 class MacroClient {
-  final RemoteMacroHost host = RemoteMacroHost();
   final Iterable<Macro> macros;
   final Socket socket;
+  late final RemoteMacroHost _host;
+  Completer<Response>? _responseCompleter;
 
   MacroClient._(this.macros, this.socket) {
+    _host = RemoteMacroHost(this);
+
     // TODO(davidmorgan): negotiation about protocol version goes here.
 
     // Tell the host which macros are in this bundle.
     for (final macro in macros) {
-      final request = MacroStartedRequest(macroDescription: macro.description);
-      // TODO(davidmorgan): currently this is JSON with one request per line,
-      // switch to binary.
-      socket.writeln(json.encode(request.node));
+      _sendRequest(MacroRequest.macroStartedRequest(
+          MacroStartedRequest(macroDescription: macro.description)));
     }
 
     const Utf8Decoder()
@@ -42,21 +46,77 @@ class MacroClient {
     return MacroClient._(macros, socket);
   }
 
-  void _handleRequest(String request) async {
-    // TODO(davidmorgan): support more than one request type.
-    final augmentRequest =
-        AugmentRequest.fromJson(json.decode(request) as Map<String, Object?>);
-    // TODO(davidmorgan): support multiple macros.
-    final response = await macros.single.augment(host, augmentRequest);
-    _send(response.node);
+  void _sendRequest(MacroRequest request) {
+    // TODO(davidmorgan): currently this is JSON with one request per line,
+    // switch to binary.
+    socket.writeln(json.encode(request.node));
   }
 
-  void _send(Map<String, Object?> node) {
+  void _sendResponse(Response response) {
     // TODO(davidmorgan): currently this is JSON with one request per line,
     // switch to binary.
-    socket.writeln(json.encode(node));
+    socket.writeln(json.encode(response.node));
+  }
+
+  void _handleRequest(String request) async {
+    final jsonData = json.decode(request) as Map<String, Object?>;
+    final hostRequest = HostRequest.fromJson(jsonData);
+    switch (hostRequest.type) {
+      case HostRequestType.augmentRequest:
+        _sendResponse(Response.augmentResponse(
+            await macros.single.augment(_host, hostRequest.asAugmentRequest)));
+      default:
+      // Ignore unknown request.
+      // TODO(davidmorgan): make handling of unknown request types a designed
+      // part of the protocol+code, update implementation here and below.
+    }
+    final response = Response.fromJson(jsonData);
+    switch (response.type) {
+      case ResponseType.unknown:
+        // Ignore unknown response.
+        break;
+      default:
+        // TODO(davidmorgan): track requests and responses properly.
+        if (_responseCompleter != null) {
+          _responseCompleter!.complete(response);
+          _responseCompleter = null;
+        }
+    }
   }
 }
 
 /// [Host] that is connected to a remote macro host.
-class RemoteMacroHost implements Host {}
+///
+/// Wraps `MacroClient` exposing just what should be available to the macro.
+///
+/// This gets passed into user-written macro code, so fields and methods here
+/// can be accessed by the macro code if they are public, even if they are not
+/// on `Host`, via dynamic dispatch.
+///
+/// TODO(language/issues/3951): follow up on security implications.
+///
+class RemoteMacroHost implements Host {
+  final MacroClient _client;
+
+  RemoteMacroHost(this._client);
+
+  @override
+  Future<Model> query(Query query) async {
+    _client._sendRequest(MacroRequest.queryRequest(QueryRequest(query: query)));
+    // TODO(davidmorgan): this is needed because the constructor doesn't wait
+    // for responses to `MacroStartedRequest`, so we need to discard the
+    // responses. Properly track requests and responses.
+    while (true) {
+      final nextResponse = await _nextResponse();
+      if (nextResponse.type == ResponseType.macroStartedResponse) {
+        continue;
+      }
+      return nextResponse.asQueryResponse.model;
+    }
+  }
+
+  Future<Response> _nextResponse() async {
+    _client._responseCompleter = Completer<Response>();
+    return await _client._responseCompleter!.future;
+  }
+}

pkgs/_macro_client/pubspec.yaml

@@ -9,6 +9,7 @@ environment:
   sdk: ^3.6.0-48.0.dev
 
 dependencies:
+  dart_model: any
   macro: any
   macro_service: any
 

pkgs/_macro_client/test/macro_client_test.dart

@@ -8,7 +8,9 @@ import 'dart:io';
 
 import 'package:_macro_client/macro_client.dart';
 import 'package:_test_macros/declare_x_macro.dart';
+import 'package:_test_macros/query_class.dart';
 import 'package:async/async.dart';
+import 'package:dart_model/dart_model.dart';
 import 'package:macro_service/macro_service.dart';
 import 'package:test/test.dart';
 
@@ -26,7 +28,7 @@ void main() {
           serverSocket.first.timeout(const Duration(seconds: 10)), completes);
     });
 
-    test('sends requests to and from macros', () async {
+    test('sends augmentation requests to macros, sends reponse', () async {
       final serverSocket = await ServerSocket.bind('localhost', 0);
 
       unawaited(MacroClient.run(
@@ -38,11 +40,56 @@ void main() {
       final responses = StreamQueue(
           const Utf8Decoder().bind(socket).transform(const LineSplitter()));
       final descriptionResponse = await responses.next;
-      expect(descriptionResponse, '{"macroDescription":{"runsInPhases":[2]}}');
+      expect(
+          descriptionResponse,
+          '{"type":"MacroStartedRequest","value":'
+          '{"macroDescription":{"runsInPhases":[2]}}}');
 
-      socket.writeln(json.encode(AugmentRequest(phase: 2)));
+      socket.writeln(
+          json.encode(HostRequest.augmentRequest(AugmentRequest(phase: 2))));
       final augmentResponse = await responses.next;
-      expect(augmentResponse, '{"augmentations":[{"code":"int get x => 3;"}]}');
+      expect(
+          augmentResponse,
+          '{"type":"AugmentResponse","value":'
+          '{"augmentations":[{"code":"int get x => 3;"}]}}');
+    });
+
+    test('sends query requests to host, sends reponse', () async {
+      final serverSocket = await ServerSocket.bind('localhost', 0);
+
+      unawaited(MacroClient.run(
+          endpoint: HostEndpoint(port: serverSocket.port),
+          macros: [QueryClassImplementation()]));
+
+      final socket = await serverSocket.first;
+
+      final responses = StreamQueue(socket
+          .cast<List<int>>()
+          .transform(const Utf8Decoder())
+          .transform(const LineSplitter()));
+      final descriptionResponse = await responses.next;
+      expect(
+          descriptionResponse,
+          '{"type":"MacroStartedRequest","value":'
+          '{"macroDescription":{"runsInPhases":[3]}}}');
+
+      socket.writeln(
+          json.encode(HostRequest.augmentRequest(AugmentRequest(phase: 3))));
+      final queryRequest = await responses.next;
+      expect(
+        queryRequest,
+        '{"type":"QueryRequest","value":{"query":{}}}',
+      );
+
+      socket.writeln(json.encode(Response.queryResponse(QueryResponse(
+          model: Model(uris: {'package:foo/foo.dart': Library()})))));
+
+      final augmentRequest = await responses.next;
+      expect(
+        augmentRequest,
+        '{"type":"AugmentResponse","value":'
+        '{"augmentations":[{"code":"// {uris: {package:foo/foo.dart: {}}}"}]}}',
+      );
     });
   });
 }

pkgs/_macro_host/lib/macro_host.dart

@@ -13,8 +13,8 @@ import 'package:macro_service/macro_service.dart';
 /// Hosts macros: builds them, runs them, serves the macro service.
 ///
 /// Tools that want to support macros, such as the Analyzer and the CFE, can
-/// do so by running a `MacroHost` and providing their own `MacroService`.
-class MacroHost implements MacroService {
+/// do so by running a `MacroHost` and providing their own `HostService`.
+class MacroHost implements HostService {
   final MacroServer macroServer;
   final ListOfServices services;
   final MacroBuilder macroBuilder = MacroBuilder();
@@ -24,10 +24,6 @@ class MacroHost implements MacroService {
   // lifecycle state.
   Completer<Set<int>>? _macroPhases;
 
-  // TODO(davidmorgan): actually match up requests and responses instead of
-  // this hack.
-  Completer<AugmentResponse>? _responseCompleter;
-
   MacroHost._(this.macroServer, this.services) {
     services.services.insert(0, this);
   }
@@ -39,7 +35,7 @@ class MacroHost implements MacroService {
   ///
   /// TODO(davidmorgan): make this split clearer, it should be in the protocol
   /// definition somewhere which requests the host handles.
-  static Future<MacroHost> serve({required MacroService service}) async {
+  static Future<MacroHost> serve({required HostService service}) async {
     final listOfServices = ListOfServices();
     listOfServices.services.add(service);
     final server = await MacroServer.serve(service: listOfServices);
@@ -70,44 +66,34 @@ class MacroHost implements MacroService {
       QualifiedName name, AugmentRequest request) async {
     // TODO(davidmorgan): this just assumes the macro is running, actually
     // track macro lifecycle.
-    macroServer.sendToMacro(name, request);
-    if (_responseCompleter != null) {
-      throw StateError('request is already pending');
-    }
-    _responseCompleter = Completer<AugmentResponse>();
-    return await _responseCompleter!.future;
+    final response = await macroServer.sendToMacro(
+        name, HostRequest.augmentRequest(request));
+    return response.asAugmentResponse;
   }
 
   /// Handle requests that are for the host.
   @override
-  Future<Object?> handle(Object request) async {
-    // TODO(davidmorgan): differentiate requests and responses, rather than
-    // handling as requests.
-    if (_responseCompleter != null) {
-      final augmentResponse =
-          AugmentResponse.fromJson(request as Map<String, Object?>);
-      _responseCompleter!.complete(augmentResponse);
-      _responseCompleter = null;
-      return Object();
+  Future<Response?> handle(MacroRequest request) async {
+    switch (request.type) {
+      case MacroRequestType.macroStartedRequest:
+        _macroPhases!.complete(request
+            .asMacroStartedRequest.macroDescription.runsInPhases
+            .toSet());
+        return Response.macroStartedResponse(MacroStartedResponse());
+      default:
+        return null;
     }
-    // TODO(davidmorgan): don't assume the type. Return `null` for types
-    // that should be passed through to the service that was passed in.
-    final macroStartedRequest =
-        MacroStartedRequest.fromJson(request as Map<String, Object?>);
-    _macroPhases!
-        .complete(macroStartedRequest.macroDescription.runsInPhases.toSet());
-    return MacroStartedResponse();
   }
 }
 
 // TODO(davidmorgan): this is used to handle some requests in the host while
 // letting some fall through to the passed in service. Differentiate in a
 // better way.
-class ListOfServices implements MacroService {
-  List<MacroService> services = [];
+class ListOfServices implements HostService {
+  List<HostService> services = [];
 
   @override
-  Future<Object> handle(Object request) async {
+  Future<Response> handle(MacroRequest request) async {
     for (final service in services) {
       final result = await service.handle(request);
       if (result != null) return result;

pkgs/_macro_host/test/macro_host_test.dart

@@ -11,7 +11,7 @@ import 'package:test/test.dart';
 
 void main() {
   group(MacroHost, () {
-    test('hosts a macro', () async {
+    test('hosts a macro, receives augmentations', () async {
       final service = TestMacroHostService();
       final host = await MacroHost.serve(service: service);
 
@@ -27,12 +27,34 @@ void main() {
           AugmentResponse(
               augmentations: [Augmentation(code: 'int get x => 3;')]));
     });
+
+    test('hosts a macro, responds to queries', () async {
+      final service = TestMacroHostService();
+      final host = await MacroHost.serve(service: service);
+
+      final macroName = QualifiedName(
+          'package:_test_macros/query_class.dart#QueryClassImplementation');
+      final packageConfig = Isolate.packageConfigSync!;
+
+      expect(host.isMacro(packageConfig, macroName), true);
+      expect(await host.queryMacroPhases(packageConfig, macroName), {3});
+
+      expect(
+          await host.augment(macroName, AugmentRequest(phase: 2)),
+          AugmentResponse(augmentations: [
+            Augmentation(code: '// {uris: {package:foo/foo.dart: {}}}')
+          ]));
+    });
   });
 }
 
-class TestMacroHostService implements MacroService {
+class TestMacroHostService implements HostService {
   @override
-  Future<Object> handle(Object request) async {
-    return Object();
+  Future<Response?> handle(MacroRequest request) async {
+    if (request.type == MacroRequestType.queryRequest) {
+      return Response.queryResponse(QueryResponse(
+          model: Model(uris: {'package:foo/foo.dart': Library()})));
+    }
+    return null;
   }
 }

pkgs/_macro_runner/lib/macro_runner.dart

@@ -15,6 +15,11 @@ class MacroRunner {
   /// Starts [macroBundle] connected to [endpoint].
   void start(
       {required BuiltMacroBundle macroBundle, required HostEndpoint endpoint}) {
-    Process.run(macroBundle.executablePath, [json.encode(endpoint)]);
+    Process.run(macroBundle.executablePath, [json.encode(endpoint)])
+        .then((result) {
+      if (result.exitCode != 0) {
+        print('Macro process exited with error: ${result.stderr}');
+      }
+    });
   }
 }