fix(lsp): make spec complient JSON-RPC response and properly implement definition link suport

compiler/src/language_server/goto.re

@@ -27,22 +27,40 @@ let send_no_result = (~id: Protocol.message_id) => {
   Protocol.response(~id, `Null);
 };
 
-let send_location_link =
+let send_location_response =
     (
       ~id: Protocol.message_id,
-      ~origin_range: Protocol.range,
+      goto_request_type: goto_request_type,
+      ~origin_loc: Location.t,
       ~target_uri: Protocol.uri,
-      ~target_range: Protocol.range,
+      ~target_loc: Location.t,
     ) => {
-  Protocol.response(
-    ~id,
-    Protocol.location_link_to_yojson({
-      origin_selection_range: origin_range,
-      target_uri,
-      target_range,
-      target_selection_range: target_range,
-    }),
-  );
+  let client_supports_links =
+    switch (goto_request_type) {
+    | Definition => Initialize.client_definition_link_support^
+    | TypeDefinition => Initialize.client_type_definition_link_support^
+    };
+  if (client_supports_links) {
+    let origin_range = Utils.loc_to_range(origin_loc);
+    let target_range = Utils.loc_to_range(target_loc);
+    Protocol.response(
+      ~id,
+      Protocol.location_link_to_yojson({
+        origin_selection_range: origin_range,
+        target_uri,
+        target_range,
+        target_selection_range: target_range,
+      }),
+    );
+  } else {
+    Protocol.response(
+      ~id,
+      Protocol.location_to_yojson({
+        uri: target_uri,
+        range: Utils.loc_to_range(target_loc),
+      }),
+    );
+  };
 };
 
 type check_position =
@@ -151,11 +169,12 @@ let process =
     switch (result) {
     | None => send_no_result(~id)
     | Some((origin_loc, target_loc, target_uri)) =>
-      send_location_link(
+      send_location_response(
         ~id,
-        ~origin_range=Utils.loc_to_range(origin_loc),
+        goto_request_type,
+        ~origin_loc,
         ~target_uri,
-        ~target_range=Utils.loc_to_range(target_loc),
+        ~target_loc,
       )
     };
   };

compiler/src/language_server/initialize.re

@@ -1,5 +1,30 @@
 open Grain_typed;
 
+[@deriving yojson({strict: false})]
+type link_support_capability = {
+  [@key "linkSupport"] [@default false]
+  link_support: bool,
+};
+
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_definition
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_typeDefinition
+[@deriving yojson({strict: false})]
+type text_document_client_capability = {
+  [@key "definition"] [@default None]
+  definition: option(link_support_capability),
+  [@key "typeDefinition"] [@default None]
+  type_definition: option(link_support_capability),
+};
+
+[@deriving yojson({strict: false})]
+type client_capabilities = {
+  [@key "textDocument"] [@default None]
+  text_document: option(text_document_client_capability),
+};
+
+let client_definition_link_support = ref(false);
+let client_type_definition_link_support = ref(false);
+
 // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#initializeParams
 module RequestParams = {
   [@deriving yojson({strict: false})]
@@ -22,9 +47,27 @@ module RequestParams = {
     root_uri: option(Protocol.uri),
     [@default "off"]
     trace: Protocol.trace_value,
+    [@key "capabilities"] [@default None]
+    capabilities: option(client_capabilities),
   };
 };
 
+// (definition linkSupport, typeDefinition linkSupport)
+let take_text_document_link_support = (params: RequestParams.t) =>
+  switch (params.capabilities) {
+  | Some({text_document: Some({definition, type_definition})}) => (
+      switch (definition) {
+      | None => false
+      | Some({link_support}) => link_support
+      },
+      switch (type_definition) {
+      | None => false
+      | Some({link_support}) => link_support
+      },
+    )
+  | _ => (false, false)
+  };
+
 // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#initializeResult
 module ResponseResult = {
   [@deriving yojson]
@@ -42,7 +85,7 @@ module ResponseResult = {
     [@key "hoverProvider"]
     hover_provider: bool,
     [@key "definitionProvider"]
-    definition_provider: Protocol.definition_client_capabilities,
+    definition_provider: bool,
     [@key "typeDefinitionProvider"]
     type_definition_provider: bool,
     [@key "referencesProvider"]
@@ -69,9 +112,7 @@ module ResponseResult = {
     document_formatting_provider: true,
     text_document_sync: Full,
     hover_provider: true,
-    definition_provider: {
-      link_support: true,
-    },
+    definition_provider: true,
     type_definition_provider: true,
     references_provider: false,
     document_symbol_provider: true,
@@ -95,6 +136,10 @@ let process =
       ~documents: Hashtbl.t(Protocol.uri, string),
       params: RequestParams.t,
     ) => {
+  let (definition_ls, type_definition_ls) =
+    take_text_document_link_support(params);
+  client_definition_link_support := definition_ls;
+  client_type_definition_link_support := type_definition_ls;
   // The initialize request can set up the initial trace level
   Trace.set_level(params.trace);
   Protocol.response(

compiler/src/language_server/initialize.rei

@@ -1,5 +1,11 @@
 open Grain_typed;
 
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_definition
+let client_definition_link_support: ref(bool);
+
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_typeDefinition
+let client_type_definition_link_support: ref(bool);
+
 // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#initializeParams
 module RequestParams: {
   [@deriving yojson({strict: false})]

compiler/src/language_server/protocol.re

@@ -130,13 +130,6 @@ type text_document_sync_kind =
   | Full
   | Incremental;
 
-//https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#definitionClientCapabilities
-[@deriving yojson({strict: false})]
-type definition_client_capabilities = {
-  [@key "linkSupport"]
-  link_support: bool,
-};
-
 let text_document_sync_kind_to_yojson = kind =>
   text_document_sync_kind_to_enum(kind) |> [%to_yojson: int];
 let text_document_sync_kind_of_yojson = json =>
@@ -207,7 +200,6 @@ type request_message = {
 };
 
 // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#responseMessage
-[@deriving yojson({strict: false})]
 type response_message = {
   jsonrpc: version,
   id: option(message_id),
@@ -279,18 +271,42 @@ let request = (): result(request_message, string) => {
   };
 };
 
+// Response optionally contains an error field however many clients (such as Neovim) get confused when that value is null so we omit it.
+// Additionally, according to the jsonrpc spec, the error field must not exist if there is no error.
+// https://www.jsonrpc.org/specification#response_object
+let jsonrpc_response_success =
+    (~id: option(message_id), result: Yojson.Safe.t) => {
+  let fields =
+    switch (id) {
+    | None => [("jsonrpc", `String(version)), ("result", result)]
+    | Some(id_val) => [
+        ("jsonrpc", `String(version)),
+        ("id", `Int(id_val)),
+        ("result", result),
+      ]
+    };
+  `Assoc(fields);
+};
+
+// On an error, respose must not contain the result field.
+// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#responseMessage
+let jsonrpc_response_error = (~id: option(message_id), err: response_error) => {
+  let err_json = response_error_to_yojson(err);
+  let fields =
+    switch (id) {
+    | None => [("jsonrpc", `String(version)), ("error", err_json)]
+    | Some(id_val) => [
+        ("jsonrpc", `String(version)),
+        ("id", `Int(id_val)),
+        ("error", err_json),
+      ]
+    };
+  `Assoc(fields);
+};
+
 let response = (~id=?, result) => {
-  let response_message = {
-    jsonrpc: version,
-    id,
-    result: Some(result),
-    error: None,
-  };
   let content =
-    Yojson.Safe.to_string(
-      ~std=true,
-      response_message_to_yojson(response_message),
-    );
+    Yojson.Safe.to_string(~std=true, jsonrpc_response_success(~id, result));
   let length = String.length(content);
 
   let len = string_of_int(length);
@@ -303,16 +319,10 @@ let response = (~id=?, result) => {
 };
 
 let empty_response = id => {
-  let response_message = {
-    jsonrpc: version,
-    id: Some(id),
-    result: None,
-    error: None,
-  };
   let content =
     Yojson.Safe.to_string(
       ~std=true,
-      response_message_to_yojson(response_message),
+      jsonrpc_response_success(~id=Some(id), `Null),
     );
   let length = String.length(content);
 
@@ -323,17 +333,8 @@ let empty_response = id => {
 };
 
 let error = (~id=?, error) => {
-  let response_message = {
-    jsonrpc: version,
-    id,
-    result: None,
-    error: Some(error),
-  };
   let content =
-    Yojson.Safe.to_string(
-      ~std=true,
-      response_message_to_yojson(response_message),
-    );
+    Yojson.Safe.to_string(~std=true, jsonrpc_response_error(~id, error));
   let length = String.length(content);
   let len = string_of_int(length);
   let msg = header_prefix ++ len ++ sep ++ content;

compiler/src/language_server/protocol.rei

@@ -161,7 +161,6 @@ type request_message = {
 };
 
 // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#responseMessage
-[@deriving yojson({strict: false})]
 type response_message = {
   jsonrpc: version,
   id: option(message_id),
@@ -200,13 +199,6 @@ type workspace_edit = {
   document_changes: list(text_document_edit),
 };
 
-//https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#definitionClientCapabilities
-[@deriving yojson({strict: false})]
-type definition_client_capabilities = {
-  [@key "linkSupport"]
-  link_support: bool,
-};
-
 let request: unit => result(request_message, string);
 
 let response: (~id: message_id=?, Yojson.Safe.t) => unit;

compiler/test/runner.re

@@ -695,18 +695,27 @@ let lsp_success_response = result => {
     ("jsonrpc", `String("2.0")),
     ("id", `Int(1)),
     ("result", result),
-    ("error", `Null),
   ]);
 };
 
-let lsp_setup_teardown_requests = (code_uri, code) => {
-  let init_request =
-    lsp_input(
-      "initialize",
-      Yojson.Safe.from_string(
-        {|{"processId":1,"clientInfo":null,"locale":null,"rootUri":null,"trace":"off"}|},
-      ),
-    );
+let lsp_default_initialize_params =
+  Yojson.Safe.from_string(
+    {|{"processId":1,"clientInfo":null,"locale":null,"rootUri":null,"trace":"off","capabilities":{"textDocument":{"definition":{"linkSupport":true},"typeDefinition":{"linkSupport":true}}}}|},
+  );
+
+let lsp_initialize_params_without_link_support =
+  Yojson.Safe.from_string(
+    {|{"processId":1,"clientInfo":null,"locale":null,"rootUri":null,"trace":"off"}|},
+  );
+
+let lsp_initialize_params_definition_plain_type_link =
+  Yojson.Safe.from_string(
+    {|{"processId":1,"clientInfo":null,"locale":null,"rootUri":null,"trace":"off","capabilities":{"textDocument":{"definition":{"linkSupport":false},"typeDefinition":{"linkSupport":true}}}}|},
+  );
+
+let lsp_setup_teardown_requests =
+    (~initialize_params=lsp_default_initialize_params, code_uri, code) => {
+  let init_request = lsp_input("initialize", initialize_params);
 
   let open_request =
     lsp_input(
@@ -742,7 +751,7 @@ let assert_lsp_responses =
     lsp_expected_response(
       lsp_success_response(
         Yojson.Safe.from_string(
-          {|{"capabilities":{"documentFormattingProvider":true,"textDocumentSync":1,"hoverProvider":true,"definitionProvider":{"linkSupport":true},"typeDefinitionProvider":true,"referencesProvider":false,"documentSymbolProvider":true,"codeActionProvider":true,"codeLensProvider":{"resolveProvider":true},"documentHighlightProvider":false,"documentRangeFormattingProvider":false,"renameProvider":false,"inlayHintProvider":{"resolveProvider":false}}}|},
+          {|{"capabilities":{"documentFormattingProvider":true,"textDocumentSync":1,"hoverProvider":true,"definitionProvider":true,"typeDefinitionProvider":true,"referencesProvider":false,"documentSymbolProvider":true,"codeActionProvider":true,"codeLensProvider":{"resolveProvider":true},"documentHighlightProvider":false,"documentRangeFormattingProvider":false,"renameProvider":false,"inlayHintProvider":{"resolveProvider":false}}}|},
         ),
       ),
     );
@@ -772,12 +781,20 @@ let assert_lsp_responses =
 };
 
 let makeLspRunner =
-    (test, name, code_uri, code, request_params, expected_output) => {
+    (
+      test,
+      ~initialize_params=lsp_default_initialize_params,
+      name,
+      code_uri,
+      code,
+      request_params,
+      expected_output,
+    ) => {
   test(
     name,
     ({expect}) => {
       let (setup_request, teardown_request) =
-        lsp_setup_teardown_requests(code_uri, code);
+        lsp_setup_teardown_requests(~initialize_params, code_uri, code);
 
       let (result, code) =
         lsp(