Close #5112 - Support newer Pyright and other LSPs

Add support for dynamic capability registration for the diagnostic
pull model to support Pyright >= 1.1.407 and other language servers.
This is a rather complex and intricate change tested with Pyright and
gopls, and may need further tweaking if something breaks with another
server.

4 files changed, 485 insertions, 9 deletions

diff --git a/autoload/ale/lsp.vim b/autoload/ale/lsp.vim
index f49f3ae00..7103879ba 100644
--- a/autoload/ale/lsp.vim
+++ b/autoload/ale/lsp.vim
@@ -23,6 +23,7 @@ function! ale#lsp#Register(executable_or_address, project, language, init_option
         " config: Configuration settings to send to the server.
         " callback_list: A list of callbacks for handling LSP responses.
         " capabilities_queue: The list of callbacks to call with capabilities.
+        " dynamic_registrations: A map of dynamically registered capabilities.
         " capabilities: Features the server supports.
         let s:connections[l:conn_id] = {
         \   'id': l:conn_id,
@@ -37,6 +38,8 @@ function! ale#lsp#Register(executable_or_address, project, language, init_option
         \   'config': {},
         \   'callback_list': [],
         \   'init_queue': [],
+        \   'dynamic_registrations': {},
+        \   'static_pull_model': 0,
         \   'capabilities': {
         \       'hover': 0,
         \       'rename': 0,
@@ -196,6 +199,27 @@ function! ale#lsp#ReadMessageData(data) abort
     return [l:remainder, l:response_list]
 endfunction
 
+function! s:UpdatePullModelCapability(conn) abort
+    let a:conn.capabilities.pull_model = get(a:conn, 'static_pull_model', 0)
+
+    if a:conn.capabilities.pull_model
+        return
+    endif
+
+    for l:registration in values(get(a:conn, 'dynamic_registrations', {}))
+        if get(l:registration, 'method', '') is# 'textDocument/diagnostic'
+            let l:options = get(l:registration, 'registerOptions', {})
+
+            if type(l:options) is v:t_dict
+            \&& type(get(l:options, 'interFileDependencies')) is v:t_bool
+                let a:conn.capabilities.pull_model = 1
+
+                return
+            endif
+        endif
+    endfor
+endfunction
+
 " Update capabilities from the server, so we know which features the server
 " supports.
 function! ale#lsp#UpdateCapabilities(conn_id, capabilities) abort
@@ -280,7 +304,8 @@ function! ale#lsp#UpdateCapabilities(conn_id, capabilities) abort
     " Check if the language server supports pull model diagnostics.
     if type(get(a:capabilities, 'diagnosticProvider')) is v:t_dict
         if type(get(a:capabilities.diagnosticProvider, 'interFileDependencies')) is v:t_bool
-            let l:conn.capabilities.pull_model = 1
+            let l:conn.static_pull_model = 1
+            call s:UpdatePullModelCapability(l:conn)
         endif
     endif
 
@@ -311,6 +336,121 @@ function! ale#lsp#UpdateCapabilities(conn_id, capabilities) abort
     endif
 endfunction
 
+" Update capabilities registered dynamically with client/registerCapability.
+" Returns 1 when pull diagnostics were registered.
+function! ale#lsp#RegisterCapabilities(conn_id, registrations) abort
+    let l:conn = get(s:connections, a:conn_id, {})
+
+    if empty(l:conn) || type(a:registrations) isnot v:t_list
+        return 0
+    endif
+
+    if !has_key(l:conn, 'dynamic_registrations')
+        let l:conn.dynamic_registrations = {}
+    endif
+
+    let l:registered_pull_diagnostics = 0
+
+    for l:registration in a:registrations
+        if type(l:registration) isnot v:t_dict
+            continue
+        endif
+
+        let l:id = get(l:registration, 'id', '')
+
+        if empty(l:id)
+            continue
+        endif
+
+        let l:conn.dynamic_registrations[l:id] = l:registration
+
+        if get(l:registration, 'method', '') is# 'textDocument/diagnostic'
+            let l:options = get(l:registration, 'registerOptions', {})
+
+            if type(l:options) is v:t_dict
+            \&& type(get(l:options, 'interFileDependencies')) is v:t_bool
+                let l:registered_pull_diagnostics = 1
+            endif
+        endif
+    endfor
+
+    call s:UpdatePullModelCapability(l:conn)
+
+    return l:registered_pull_diagnostics
+endfunction
+
+" Update capabilities removed dynamically with client/unregisterCapability.
+" The LSP spec names the list "unregisterations".
+function! ale#lsp#UnregisterCapabilities(conn_id, unregisterations) abort
+    let l:conn = get(s:connections, a:conn_id, {})
+
+    if empty(l:conn) || type(a:unregisterations) isnot v:t_list
+        return 0
+    endif
+
+    if !has_key(l:conn, 'dynamic_registrations')
+        let l:conn.dynamic_registrations = {}
+    endif
+
+    let l:unregistered_pull_diagnostics = 0
+
+    for l:unregistration in a:unregisterations
+        if type(l:unregistration) isnot v:t_dict
+            continue
+        endif
+
+        let l:id = get(l:unregistration, 'id', '')
+        let l:method = get(l:unregistration, 'method', '')
+
+        if empty(l:id) || !has_key(l:conn.dynamic_registrations, l:id)
+            continue
+        endif
+
+        let l:registration = l:conn.dynamic_registrations[l:id]
+
+        if get(l:registration, 'method', '') isnot# l:method
+            continue
+        endif
+
+        if l:method is# 'textDocument/diagnostic'
+            let l:unregistered_pull_diagnostics = 1
+        endif
+
+        call remove(l:conn.dynamic_registrations, l:id)
+    endfor
+
+    call s:UpdatePullModelCapability(l:conn)
+
+    return l:unregistered_pull_diagnostics
+endfunction
+
+" Send textDocument/diagnostic requests for all open documents on a connection.
+" Returns a list of request details so callers can map responses back to URIs.
+function! ale#lsp#SendDiagnosticsForOpenDocuments(conn_id) abort
+    let l:conn = get(s:connections, a:conn_id, {})
+    let l:request_list = []
+
+    if empty(l:conn) || !ale#lsp#HasCapability(a:conn_id, 'pull_model')
+        return l:request_list
+    endif
+
+    for l:buffer_string in sort(keys(l:conn.open_documents))
+        let l:buffer = str2nr(l:buffer_string)
+        let l:message = ale#lsp#message#Diagnostic(l:buffer)
+        let l:request_id = ale#lsp#Send(a:conn_id, l:message)
+
+        if l:request_id > 0
+            call add(l:request_list, {
+            \   'id': l:request_id,
+            \   'buffer': l:buffer,
+            \   'uri': l:message[2].textDocument.uri,
+            \})
+        endif
+    endfor
+
+    return l:request_list
+endfunction
+
 " Update a connection's configuration dictionary and notify LSP servers
 " of any changes since the last update. Returns 1 if a configuration
 " update was sent; otherwise 0 will be returned.
diff --git a/autoload/ale/lsp_linter.vim b/autoload/ale/lsp_linter.vim
index 6e33f0af1..22f739a79 100644
--- a/autoload/ale/lsp_linter.vim
+++ b/autoload/ale/lsp_linter.vim
@@ -210,6 +210,24 @@ function! s:HandleLSPErrorMessage(linter, response) abort
     call ale#lsp_linter#AddErrorMessage(a:linter.name, l:message)
 endfunction
 
+function! s:SendPullDiagnosticsForOpenDocuments(conn_id) abort
+    let l:linter = get(s:lsp_linter_map, a:conn_id)
+
+    if empty(l:linter)
+        return
+    endif
+
+    for l:request in ale#lsp#SendDiagnosticsForOpenDocuments(a:conn_id)
+        let l:info = get(g:ale_buffer_info, l:request.buffer, {})
+
+        if !empty(l:info)
+            call ale#engine#MarkLinterActive(l:info, l:linter)
+        endif
+
+        let s:diagnostic_uri_map[l:request.id] = l:request.uri
+    endfor
+endfunction
+
 function! ale#lsp_linter#AddErrorMessage(linter_name, message) abort
     " This global variable is set here so we don't load the debugging.vim file
     " until someone uses :ALEInfo.
@@ -228,20 +246,39 @@ function! ale#lsp_linter#HandleLSPResponse(conn_id, response) abort
     if get(a:response, 'jsonrpc', '') is# '2.0' && has_key(a:response, 'error')
         let l:linter = get(s:lsp_linter_map, a:conn_id, {})
 
+        " Clean up values in the map on error.
+        if empty(l:method) && has_key(s:diagnostic_uri_map, get(a:response, 'id'))
+            call remove(s:diagnostic_uri_map, a:response.id)
+        endif
+
         call s:HandleLSPErrorMessage(l:linter, a:response)
     elseif l:method is# 'textDocument/publishDiagnostics'
         let l:uri = a:response.params.uri
         let l:diagnostics = a:response.params.diagnostics
 
         call ale#lsp_linter#HandleLSPDiagnostics(a:conn_id, l:uri, l:diagnostics)
-    elseif has_key(s:diagnostic_uri_map, get(a:response, 'id'))
-        let l:uri = remove(s:diagnostic_uri_map, a:response.id)
-        let l:diagnostics = a:response.result.kind is# 'unchanged'
-        \   ? 'unchanged'
-        \   : a:response.result.items
-
-        call ale#lsp_linter#HandleLSPDiagnostics(a:conn_id, l:uri, l:diagnostics)
+    elseif l:method is# 'workspace/diagnostic/refresh'
+        call ale#lsp#SendResponse(a:conn_id, a:response.id, v:null)
+        call s:SendPullDiagnosticsForOpenDocuments(a:conn_id)
     elseif l:method is# 'client/registerCapability'
+        let l:registered_pull_diagnostics = ale#lsp#RegisterCapabilities(
+        \   a:conn_id,
+        \   get(get(a:response, 'params', {}), 'registrations', []),
+        \)
+
+        call ale#lsp#SendResponse(a:conn_id, a:response.id, v:null)
+
+        if l:registered_pull_diagnostics
+            call s:SendPullDiagnosticsForOpenDocuments(a:conn_id)
+        endif
+    elseif l:method is# 'client/unregisterCapability'
+        let l:unregisterations = get(
+        \   get(a:response, 'params', {}),
+        \   'unregisterations',
+        \   get(get(a:response, 'params', {}), 'unregistrations', []),
+        \)
+
+        call ale#lsp#UnregisterCapabilities(a:conn_id, l:unregisterations)
         call ale#lsp#SendResponse(a:conn_id, a:response.id, v:null)
     elseif l:method is# 'workspace/configuration'
         let l:items = get(get(a:response, 'params', {}), 'items', [])
@@ -253,6 +290,18 @@ function! ale#lsp_linter#HandleLSPResponse(conn_id, response) abort
         \   g:ale_lsp_show_message_format,
         \   a:response.params
         \)
+    elseif empty(l:method) && has_key(s:diagnostic_uri_map, get(a:response, 'id'))
+        let l:uri = remove(s:diagnostic_uri_map, a:response.id)
+        let l:diagnostics = 'unchanged'
+
+        if type(get(a:response, 'result')) is v:t_dict
+            if get(a:response.result, 'kind', '') isnot# 'unchanged'
+            \&& type(get(a:response.result, 'items')) is v:t_list
+                let l:diagnostics = a:response.result.items
+            endif
+        endif
+
+        call ale#lsp_linter#HandleLSPDiagnostics(a:conn_id, l:uri, l:diagnostics)
     elseif get(a:response, 'type', '') is# 'event'
     \&& get(a:response, 'event', '') is# 'semanticDiag'
         call s:HandleTSServerDiagnostics(a:response, 'semantic')
diff --git a/test/lsp/test_engine_lsp_response_handling.vader b/test/lsp/test_engine_lsp_response_handling.vader
index 226051c44..b99409f89 100644
--- a/test/lsp/test_engine_lsp_response_handling.vader
+++ b/test/lsp/test_engine_lsp_response_handling.vader
@@ -1,4 +1,6 @@
 Before:
+  runtime autoload/ale/lsp.vim
+
   Save g:ale_set_lists_synchronously
   Save g:ale_buffer_info
   Save g:ale_lsp_error_messages
@@ -40,6 +42,7 @@ After:
   call ale#test#RestoreDirectory()
   call ale#linter#Reset()
   call ale#lsp_linter#ClearLSPData()
+  call ale#lsp_linter#ClearDiagnosticURIMap()
 
 Given foobar(An empty file):
 Execute(tsserver syntax error responses should be handled correctly):
@@ -534,6 +537,60 @@ Execute(LSP pull model diagnostic responses that are 'unchanged' should be handl
   \ g:ale_buffer_info[bufnr('')].loclist
   AssertEqual [], g:ale_buffer_info[bufnr('')].active_linter_list
 
+Execute(LSP pull model diagnostic errors should clear pending requests):
+  let b:ale_linters = ['eclipselsp']
+  runtime ale_linters/java/eclipselsp.vim
+
+  if has('win32')
+    call ale#test#SetFilename('filename,[]^$.ts')
+  else
+    call ale#test#SetFilename('filename*?,{}[]^$.java')
+  endif
+
+  call ale#engine#InitBufferInfo(bufnr(''))
+  let g:ale_buffer_info[bufnr('')].loclist = [
+  \ {
+  \   'lnum': 1,
+  \   'bufnr': bufnr(''),
+  \   'col': 1,
+  \   'pattern': '',
+  \   'valid': 1,
+  \   'vcol': 0,
+  \   'nr': -1,
+  \   'type': 'W',
+  \   'text': 'Missing JRE 1-8'
+  \ },
+  \]
+
+  call ale#lsp_linter#SetLSPLinterMap({'1': {'name': 'eclipselsp', 'aliases': [], 'lsp': 'stdio'}})
+  call ale#lsp_linter#SetDiagnosticURIMap({'347': ale#util#ToURI(expand('%:p'))})
+
+  call ale#lsp_linter#HandleLSPResponse(1, {
+  \ 'jsonrpc': '2.0',
+  \ 'id': 347,
+  \ 'error': {
+  \   'code': -32802,
+  \   'message': 'server cancelled',
+  \ },
+  \})
+
+  AssertEqual {}, ale#lsp_linter#GetDiagnosticURIMap()
+  AssertEqual
+  \ [
+  \   {
+  \     'lnum': 1,
+  \     'bufnr': bufnr(''),
+  \     'col': 1,
+  \     'pattern': '',
+  \     'valid': 1,
+  \     'vcol': 0,
+  \     'nr': -1,
+  \     'type': 'W',
+  \     'text': 'Missing JRE 1-8'
+  \    }
+  \ ],
+  \ g:ale_buffer_info[bufnr('')].loclist
+
 Execute(workspace/configuration requests should be answered with the connection config):
   let g:sent_responses = []
 
@@ -564,26 +621,134 @@ Execute(workspace/configuration requests should be answered with the connection
 
 Execute(client/registerCapability requests should be acknowledged):
   let g:sent_responses = []
+  let g:registered_capabilities = []
+  let g:sent_pull_diagnostics = []
 
   function! ale#lsp#SendResponse(conn_id, id, result) abort
     call add(g:sent_responses, [a:conn_id, a:id, a:result])
   endfunction
 
+  function! ale#lsp#RegisterCapabilities(conn_id, registrations) abort
+    call add(g:registered_capabilities, [a:conn_id, a:registrations])
+
+    return 1
+  endfunction
+
+  function! ale#lsp#SendDiagnosticsForOpenDocuments(conn_id) abort
+    call add(g:sent_pull_diagnostics, a:conn_id)
+
+    return []
+  endfunction
+
   call ale#lsp_linter#SetLSPLinterMap({'1': {'name': 'expert', 'aliases': [], 'lsp': 'stdio'}})
   call ale#lsp_linter#HandleLSPResponse(1, {
   \ 'jsonrpc': '2.0',
   \ 'id': 12,
   \ 'method': 'client/registerCapability',
   \ 'params': {
-  \   'registrations': [{'id': 'abc', 'method': 'textDocument/didSave'}],
+  \   'registrations': [
+  \     {
+  \       'id': 'abc',
+  \       'method': 'textDocument/diagnostic',
+  \       'registerOptions': {'interFileDependencies': v:true},
+  \     },
+  \   ],
   \ },
   \})
 
   AssertEqual
   \ [[1, 12, v:null]],
   \ g:sent_responses
+  AssertEqual
+  \ [[1, [
+  \   {
+  \     'id': 'abc',
+  \     'method': 'textDocument/diagnostic',
+  \     'registerOptions': {'interFileDependencies': v:true},
+  \   },
+  \ ]]],
+  \ g:registered_capabilities
+  AssertEqual [1], g:sent_pull_diagnostics
+
+  unlet! g:sent_responses
+  unlet! g:registered_capabilities
+  unlet! g:sent_pull_diagnostics
+  runtime autoload/ale/lsp.vim
+
+Execute(workspace/diagnostic/refresh requests should pull diagnostics):
+  let g:sent_responses = []
+  let g:sent_pull_diagnostics = []
+
+  function! ale#lsp#SendResponse(conn_id, id, result) abort
+    call add(g:sent_responses, [a:conn_id, a:id, a:result])
+  endfunction
+
+  function! ale#lsp#SendDiagnosticsForOpenDocuments(conn_id) abort
+    call add(g:sent_pull_diagnostics, a:conn_id)
+
+    return []
+  endfunction
+
+  call ale#lsp_linter#SetLSPLinterMap({'1': {'name': 'expert', 'aliases': [], 'lsp': 'stdio'}})
+  call ale#lsp_linter#SetDiagnosticURIMap({'13': 'file://foo.py'})
+
+  call ale#lsp_linter#HandleLSPResponse(1, {
+  \ 'jsonrpc': '2.0',
+  \ 'id': 13,
+  \ 'method': 'workspace/diagnostic/refresh',
+  \})
+
+  AssertEqual
+  \ [[1, 13, v:null]],
+  \ g:sent_responses
+  AssertEqual [1], g:sent_pull_diagnostics
+  AssertEqual {'13': 'file://foo.py'}, ale#lsp_linter#GetDiagnosticURIMap()
+
+  unlet! g:sent_responses
+  unlet! g:sent_pull_diagnostics
+  runtime autoload/ale/lsp.vim
+
+Execute(client/unregisterCapability requests should unregister and be acknowledged):
+  let g:sent_responses = []
+  let g:unregistered_capabilities = []
+
+  function! ale#lsp#SendResponse(conn_id, id, result) abort
+    call add(g:sent_responses, [a:conn_id, a:id, a:result])
+  endfunction
+
+  function! ale#lsp#UnregisterCapabilities(conn_id, unregisterations) abort
+    call add(g:unregistered_capabilities, [a:conn_id, a:unregisterations])
+
+    return 1
+  endfunction
+
+  call ale#lsp_linter#SetLSPLinterMap({'1': {'name': 'expert', 'aliases': [], 'lsp': 'stdio'}})
+  call ale#lsp_linter#SetDiagnosticURIMap({'14': 'file://foo.py'})
+
+  call ale#lsp_linter#HandleLSPResponse(1, {
+  \ 'jsonrpc': '2.0',
+  \ 'id': 14,
+  \ 'method': 'client/unregisterCapability',
+  \ 'params': {
+  \   'unregisterations': [
+  \     {
+  \       'id': 'abc',
+  \       'method': 'textDocument/diagnostic',
+  \     },
+  \   ],
+  \ },
+  \})
+
+  AssertEqual
+  \ [[1, 14, v:null]],
+  \ g:sent_responses
+  AssertEqual
+  \ [[1, [{'id': 'abc', 'method': 'textDocument/diagnostic'}]]],
+  \ g:unregistered_capabilities
+  AssertEqual {'14': 'file://foo.py'}, ale#lsp_linter#GetDiagnosticURIMap()
 
   unlet! g:sent_responses
+  unlet! g:unregistered_capabilities
   runtime autoload/ale/lsp.vim
 
 Execute(LSP errors should be logged in the history):
diff --git a/test/lsp/test_other_initialize_message_handling.vader b/test/lsp/test_other_initialize_message_handling.vader
index 4a6d89e03..2d1ffbe67 100644
--- a/test/lsp/test_other_initialize_message_handling.vader
+++ b/test/lsp/test_other_initialize_message_handling.vader
@@ -214,6 +214,128 @@ Execute(Capabilities should be enabled when sent as Dictionaries):
   \ b:conn.capabilities
   AssertEqual [[1, 'initialized', {}]], g:message_list
 
+Execute(Pull model diagnostics should be enabled when registered dynamically):
+  AssertEqual 0, b:conn.capabilities.pull_model
+
+  AssertEqual
+  \ 1,
+  \ ale#lsp#RegisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'pyright-diagnostics',
+  \     'method': 'textDocument/diagnostic',
+  \     'registerOptions': {
+  \       'interFileDependencies': v:true,
+  \       'workspaceDiagnostics': v:false,
+  \       'documentSelector': v:null,
+  \       'identifier': 'Pyright',
+  \     },
+  \   },
+  \ ])
+
+  AssertEqual 1, b:conn.capabilities.pull_model
+
+Execute(Pull model diagnostics should be disabled when unregistered dynamically):
+  call ale#lsp#RegisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'pyright-diagnostics',
+  \     'method': 'textDocument/diagnostic',
+  \     'registerOptions': {
+  \       'interFileDependencies': v:true,
+  \     },
+  \   },
+  \ ])
+
+  AssertEqual 1, b:conn.capabilities.pull_model
+
+  AssertEqual
+  \ 1,
+  \ ale#lsp#UnregisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'pyright-diagnostics',
+  \     'method': 'textDocument/diagnostic',
+  \   },
+  \ ])
+
+  AssertEqual 0, b:conn.capabilities.pull_model
+
+Execute(Pull model diagnostics should stay enabled while another registration exists):
+  call ale#lsp#RegisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'pyright-diagnostics-one',
+  \     'method': 'textDocument/diagnostic',
+  \     'registerOptions': {
+  \       'interFileDependencies': v:true,
+  \     },
+  \   },
+  \   {
+  \     'id': 'pyright-diagnostics-two',
+  \     'method': 'textDocument/diagnostic',
+  \     'registerOptions': {
+  \       'interFileDependencies': v:true,
+  \     },
+  \   },
+  \ ])
+
+  call ale#lsp#UnregisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'pyright-diagnostics-one',
+  \     'method': 'textDocument/diagnostic',
+  \   },
+  \ ])
+
+  AssertEqual 1, b:conn.capabilities.pull_model
+
+Execute(Unregistering non-diagnostic registrations should not change diagnostics):
+  call ale#lsp#RegisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'pyright-diagnostics',
+  \     'method': 'textDocument/diagnostic',
+  \     'registerOptions': {
+  \       'interFileDependencies': v:true,
+  \     },
+  \   },
+  \   {
+  \     'id': 'workspace-symbols',
+  \     'method': 'workspace/symbol',
+  \     'registerOptions': {},
+  \   },
+  \ ])
+
+  call ale#lsp#UnregisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'workspace-symbols',
+  \     'method': 'workspace/symbol',
+  \   },
+  \ ])
+
+  AssertEqual 1, b:conn.capabilities.pull_model
+
+Execute(Static pull model diagnostics should not be disabled by unregistering dynamic registrations):
+  call ale#lsp#UpdateCapabilities(b:conn.id, {
+  \ 'diagnosticProvider': {
+  \   'interFileDependencies': v:false,
+  \ },
+  \})
+
+  call ale#lsp#RegisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'pyright-diagnostics',
+  \     'method': 'textDocument/diagnostic',
+  \     'registerOptions': {
+  \       'interFileDependencies': v:true,
+  \     },
+  \   },
+  \ ])
+
+  call ale#lsp#UnregisterCapabilities(b:conn.id, [
+  \   {
+  \     'id': 'pyright-diagnostics',
+  \     'method': 'textDocument/diagnostic',
+  \   },
+  \ ])
+
+  AssertEqual 1, b:conn.capabilities.pull_model
+
 Execute(Results that are not dictionaries should be handled correctly):
   call ale#lsp#HandleInitResponse(b:conn, {
   \ 'jsonrpc': '2.0',