Refactor LSP custom request handling

2 files changed, 30 insertions, 36 deletions

diff --git a/autoload/ale/lsp.vim b/autoload/ale/lsp.vim
index 1ffdba2e..017096cd 100644
--- a/autoload/ale/lsp.vim
+++ b/autoload/ale/lsp.vim
@@ -5,9 +5,6 @@
 let s:connections = get(s:, 'connections', {})
 let g:ale_lsp_next_message_id = 1
 
-" A Dictionary to track one-shot callbacks for custom LSP requests
-let s:custom_callbacks = get(s:, 'custom_callbacks', {})
-
 " Given an id, which can be an executable or address, and a project path,
 " create a new connection if needed. Return a unique ID for the connection.
 function! ale#lsp#Register(executable_or_address, project, init_options) abort
@@ -299,19 +296,10 @@ function! ale#lsp#HandleMessage(conn_id, message) abort
     " responses.
     if l:conn.initialized
         for l:response in l:response_list
-            if has_key(l:response, 'id') && has_key(s:custom_callbacks, l:response.id)
-                " Response to a custom request, call the registered one-shot handler.
-                try
-                    call s:custom_callbacks[l:response.id](l:response)
-                finally
-                    call remove(s:custom_callbacks, l:response.id)
-                endtry
-            else
-                " Call all of the registered handlers with the response.
-                for l:Callback in l:conn.callback_list
-                    call ale#util#GetFunction(l:Callback)(a:conn_id, l:response)
-                endfor
-            endif
+            " Call all of the registered handlers with the response.
+            for l:Callback in l:conn.callback_list
+                call ale#util#GetFunction(l:Callback)(a:conn_id, l:response)
+            endfor
         endfor
     endif
 endfunction
@@ -537,18 +525,6 @@ function! ale#lsp#Send(conn_id, message) abort
     return l:id == 0 ? -1 : l:id
 endfunction
 
-" Send a custom request to an LSP server.
-" The given callback is called on response.
-function! ale#lsp#SendCustomRequest(conn_id, message, Callback) abort
-    let l:id = ale#lsp#Send(a:conn_id, a:message)
-
-    if l:id > 0
-        let s:custom_callbacks[l:id] = a:Callback
-    endif
-
-    return l:id
-endfunction
-
 " Notify LSP servers or tsserver if a document is opened, if needed.
 " If a document is opened, 1 will be returned, otherwise 0 will be returned.
 function! ale#lsp#OpenDocument(conn_id, buffer, language_id) abort
diff --git a/autoload/ale/lsp_linter.vim b/autoload/ale/lsp_linter.vim
index 82ae60a2..60a93504 100644
--- a/autoload/ale/lsp_linter.vim
+++ b/autoload/ale/lsp_linter.vim
@@ -8,6 +8,9 @@ if !has_key(s:, 'lsp_linter_map')
     let s:lsp_linter_map = {}
 endif
 
+" A Dictionary to track one-shot callbacks for custom LSP requests
+let s:custom_callbacks_map = get(s:, 'custom_callbacks_map', {})
+
 " Check if diagnostics for a particular linter should be ignored.
 function! s:ShouldIgnore(buffer, linter_name) abort
     " Ignore all diagnostics if LSP integration is disabled.
@@ -407,6 +410,7 @@ endfunction
 " Clear LSP linter data for the linting engine.
 function! ale#lsp_linter#ClearLSPData() abort
     let s:lsp_linter_map = {}
+    let s:custom_callbacks_map = {}
 endfunction
 
 " Just for tests.
@@ -414,6 +418,24 @@ function! ale#lsp_linter#SetLSPLinterMap(replacement_map) abort
     let s:lsp_linter_map = a:replacement_map
 endfunction
 
+function! s:HandleLSPResponseToCustomRequests(conn_id, response) abort
+    if has_key(a:response, 'id')
+    \&& has_key(s:custom_callbacks_map, a:response.id)
+        let l:Callback = remove(s:custom_callbacks_map, a:response.id)
+        call l:Callback(a:response)
+    endif
+endfunction
+
+function! s:OnReadyForCustomRequests(message, Callback, linter, lsp_details) abort
+    let l:id = a:lsp_details.connection_id
+    let l:Callback = function('s:HandleLSPResponseToCustomRequests')
+
+    call ale#lsp#RegisterCallback(l:id, l:Callback)
+
+    let l:request_id = ale#lsp#Send(l:id, a:message)
+    let s:custom_callbacks_map[l:request_id] = a:Callback
+endfunction
+
 " Send a custom request to an LSP linter.
 function! ale#lsp_linter#SendRequest(buffer, linter_name, method, parameters, Callback) abort
     let l:filetype = ale#linter#ResolveFiletype(getbufvar(a:buffer, '&filetype'))
@@ -425,17 +447,13 @@ function! ale#lsp_linter#SendRequest(buffer, linter_name, method, parameters, Ca
     endif
 
     let l:linter = l:linter_list[0]
-    let l:executable_or_address = ''
 
-    if l:linter.lsp is# 'socket'
-        let l:executable_or_address = ale#linter#GetAddress(a:buffer, l:linter)
-    else
-        let l:executable_or_address = ale#linter#GetExecutable(a:buffer, l:linter)
+    if empty(l:linter.lsp)
+        throw 'Linter "' . a:linter_name . '" does not support LSP!'
     endif
 
-    let l:root = ale#util#GetFunction(l:linter.project_root)(a:buffer)
-    let l:conn_id = l:executable_or_address . ':' . l:root
     let l:message = [0, a:method, a:parameters]
+    let l:Callback = function('s:OnReadyForCustomRequests', [l:message, a:Callback])
 
-    return ale#lsp#SendCustomRequest(l:conn_id, l:message, a:Callback) == 0
+    return ale#lsp_linter#StartLSP(a:buffer, l:linter, l:Callback)
 endfunction