feat(lsp): implement textDocument/diagnostic (#24128)

10 files changed, 422 insertions, 74 deletions

diff --git a/runtime/doc/lsp.txt b/runtime/doc/lsp.txt
index edc9f50c8d..899066ff00 100644
--- a/runtime/doc/lsp.txt
+++ b/runtime/doc/lsp.txt
@@ -175,6 +175,7 @@ won't run if your server doesn't support them.
 - textDocument/completion
 - textDocument/declaration*
 - textDocument/definition
+- textDocument/diagnostic
 - textDocument/documentHighlight
 - textDocument/documentSymbol
 - textDocument/formatting
@@ -553,6 +554,27 @@ to the callback in the "data" table. The token fields are documented in
 Note: doing anything other than calling
 |vim.lsp.semantic_tokens.highlight_token()| is considered experimental.
 
+LspNotify                                                          *LspNotify*
+
+This event is triggered after each successful notification sent to an LSP server.
+
+When used from Lua, the client_id, LSP method, and parameters are sent in the
+"data" table. Example: >lua
+
+    vim.api.nvim_create_autocmd('LspNotify', {
+      callback = function(args)
+        local bufnr = args.buf
+        local client_id = args.data.client_id
+        local method = args.data.method
+        local params = args.data.params
+
+        -- do something with the notification
+        if method == 'textDocument/...' then
+          update_buffer(bufnr)
+        end
+      end,
+    })
+<
 
 LspRequest                                                        *LspRequest*
 
@@ -1328,12 +1350,44 @@ workspace_symbol({query}, {options})          *vim.lsp.buf.workspace_symbol()*
 ==============================================================================
 Lua module: vim.lsp.diagnostic                                *lsp-diagnostic*
 
-get_namespace({client_id})                *vim.lsp.diagnostic.get_namespace()*
+                                          *vim.lsp.diagnostic.get_namespace()*
+get_namespace({client_id}, {is_pull})
     Get the diagnostic namespace associated with an LSP client
-    |vim.diagnostic|.
+    |vim.diagnostic| for diagnostics
 
     Parameters: ~
       • {client_id}  (integer) The id of the LSP client
+      • {is_pull}    (boolean) Whether the namespace is for a pull or push
+                     client
+
+                                          *vim.lsp.diagnostic.on_diagnostic()*
+on_diagnostic({_}, {result}, {ctx}, {config})
+    |lsp-handler| for the method "textDocument/diagnostic"
+
+    See |vim.diagnostic.config()| for configuration options. Handler-specific
+    configuration can be set using |vim.lsp.with()|: >lua
+
+     vim.lsp.handlers["textDocument/diagnostic"] = vim.lsp.with(
+       vim.lsp.diagnostic.on_diagnostic, {
+         -- Enable underline, use default values
+         underline = true,
+         -- Enable virtual text, override spacing to 4
+         virtual_text = {
+           spacing = 4,
+         },
+         -- Use a function to dynamically turn signs off
+         -- and on, using buffer local variables
+         signs = function(namespace, bufnr)
+           return vim.b[bufnr].show_signs == true
+         end,
+         -- Disable a feature
+         update_in_insert = false,
+       }
+     )
+<
+
+    Parameters: ~
+      • {config}  (table) Configuration table (see |vim.diagnostic.config()|).
 
                                  *vim.lsp.diagnostic.on_publish_diagnostics()*
 on_publish_diagnostics({_}, {result}, {ctx}, {config})
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 6294a8b505..4e24bb6dac 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -117,8 +117,11 @@ The following new APIs and features were added.
 • Builtin TUI can now recognize "super" (|<D-|) and "meta" (|<T-|) modifiers in a
   terminal emulator that supports |tui-csiu|.
 
-• Implemented LSP inlay hints: |vim.lsp.inlay_hint()|
-  https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_inlayHint
+• LSP
+  • Implemented LSP inlay hints: |vim.lsp.inlay_hint()|
+    https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_inlayHint
+  • Implemented pull diagnostic textDocument/diagnostic: |vim.lsp.diagnostic.on_diagnostic()|
+    https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_diagnostic
 
 • Bundled treesitter parser and queries (highlight, folds) for Markdown,
   Python, and Bash.
diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 65cce6af47..fa2a888a54 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -61,6 +61,7 @@ lsp._request_name_to_capability = {
   ['textDocument/semanticTokens/full'] = { 'semanticTokensProvider' },
   ['textDocument/semanticTokens/full/delta'] = { 'semanticTokensProvider' },
   ['textDocument/inlayHint'] = { 'inlayHintProvider' },
+  ['textDocument/diagnostic'] = { 'diagnosticProvider' },
   ['inlayHint/resolve'] = { 'inlayHintProvider', 'resolveProvider' },
 }
 
@@ -954,6 +955,9 @@ function lsp._set_defaults(client, bufnr)
       vim.keymap.set('n', 'K', vim.lsp.buf.hover, { buffer = bufnr })
     end
   end)
+  if client.supports_method('textDocument/diagnostic') then
+    lsp.diagnostic._enable(bufnr)
+  end
 end
 
 --- @class lsp.ClientConfig
@@ -1567,7 +1571,23 @@ function lsp.start_client(config)
     if method ~= 'textDocument/didChange' then
       changetracking.flush(client)
     end
-    return rpc.notify(method, params)
+
+    local result = rpc.notify(method, params)
+
+    if result then
+      vim.schedule(function()
+        nvim_exec_autocmds('LspNotify', {
+          modeline = false,
+          data = {
+            client_id = client.id,
+            method = method,
+            params = params,
+          },
+        })
+      end)
+    end
+
+    return result
   end
 
   ---@private
diff --git a/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index c2cf7c6ba5..34be13096d 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -1,9 +1,14 @@
 ---@brief lsp-diagnostic
 
+local util = require('vim.lsp.util')
 local protocol = require('vim.lsp.protocol')
 
+local api = vim.api
+
 local M = {}
 
+local augroup = api.nvim_create_augroup('vim_lsp_diagnostic', {})
+
 local DEFAULT_CLIENT_ID = -1
 
 local function get_client_id(client_id)
@@ -154,19 +159,43 @@ local function diagnostic_vim_to_lsp(diagnostics)
 end
 
 ---@type table<integer,integer>
-local _client_namespaces = {}
+local _client_push_namespaces = {}
+---@type table<integer,integer>
+local _client_pull_namespaces = {}
 
---- Get the diagnostic namespace associated with an LSP client |vim.diagnostic|.
+--- Get the diagnostic namespace associated with an LSP client |vim.diagnostic| for diagnostics
 ---
 ---@param client_id integer The id of the LSP client
-function M.get_namespace(client_id)
+---@param is_pull boolean Whether the namespace is for a pull or push client
+function M.get_namespace(client_id, is_pull)
   vim.validate({ client_id = { client_id, 'n' } })
-  if not _client_namespaces[client_id] then
-    local client = vim.lsp.get_client_by_id(client_id)
-    local name = string.format('vim.lsp.%s.%d', client and client.name or 'unknown', client_id)
-    _client_namespaces[client_id] = vim.api.nvim_create_namespace(name)
+
+  local namespace_table
+  local key
+  local name
+  local client = vim.lsp.get_client_by_id(client_id)
+
+  if is_pull then
+    namespace_table = _client_pull_namespaces
+    local server_id = vim.tbl_get(client.server_capabilities, 'diagnosticProvider', 'identifier')
+    key = string.format('%d:%s', client_id, server_id or 'nil')
+    name = string.format(
+      'vim.lsp.%s.%d.%s',
+      client and client.name or 'unknown',
+      client_id,
+      server_id or 'nil'
+    )
+  else
+    namespace_table = _client_push_namespaces
+    key = client_id
+    name = string.format('vim.lsp.%s.%d', client and client.name or 'unknown', client_id)
   end
-  return _client_namespaces[client_id]
+
+  if not namespace_table[key] then
+    namespace_table[key] = api.nvim_create_namespace(name)
+  end
+
+  return namespace_table[key]
 end
 
 --- |lsp-handler| for the method "textDocument/publishDiagnostics"
@@ -209,7 +238,7 @@ function M.on_publish_diagnostics(_, result, ctx, config)
   end
 
   client_id = get_client_id(client_id)
-  local namespace = M.get_namespace(client_id)
+  local namespace = M.get_namespace(client_id, false)
 
   if config then
     for _, opt in pairs(config) do
@@ -229,7 +258,75 @@ function M.on_publish_diagnostics(_, result, ctx, config)
   vim.diagnostic.set(namespace, bufnr, diagnostic_lsp_to_vim(diagnostics, bufnr, client_id))
 end
 
---- Clear diagnostics and diagnostic cache.
+--- |lsp-handler| for the method "textDocument/diagnostic"
+---
+--- See |vim.diagnostic.config()| for configuration options. Handler-specific
+--- configuration can be set using |vim.lsp.with()|:
+--- <pre>lua
+--- vim.lsp.handlers["textDocument/diagnostic"] = vim.lsp.with(
+---   vim.lsp.diagnostic.on_diagnostic, {
+---     -- Enable underline, use default values
+---     underline = true,
+---     -- Enable virtual text, override spacing to 4
+---     virtual_text = {
+---       spacing = 4,
+---     },
+---     -- Use a function to dynamically turn signs off
+---     -- and on, using buffer local variables
+---     signs = function(namespace, bufnr)
+---       return vim.b[bufnr].show_signs == true
+---     end,
+---     -- Disable a feature
+---     update_in_insert = false,
+---   }
+--- )
+--- </pre>
+---
+---@param config table Configuration table (see |vim.diagnostic.config()|).
+function M.on_diagnostic(_, result, ctx, config)
+  local client_id = ctx.client_id
+  local uri = ctx.params.textDocument.uri
+  local fname = vim.uri_to_fname(uri)
+
+  if result == nil then
+    return
+  end
+
+  if result.kind == 'unchanged' then
+    return
+  end
+
+  local diagnostics = result.items
+  if #diagnostics == 0 and vim.fn.bufexists(fname) == 0 then
+    return
+  end
+  local bufnr = vim.fn.bufadd(fname)
+
+  if not bufnr then
+    return
+  end
+
+  client_id = get_client_id(client_id)
+
+  local namespace = M.get_namespace(client_id, true)
+
+  if config then
+    for _, opt in pairs(config) do
+      if type(opt) == 'table' and not opt.severity and opt.severity_limit then
+        opt.severity = { min = severity_lsp_to_vim(opt.severity_limit) }
+      end
+    end
+
+    -- Persist configuration to ensure buffer reloads use the same
+    -- configuration. To make lsp.with configuration work (See :help
+    -- lsp-handler-configuration)
+    vim.diagnostic.config(config, namespace)
+  end
+
+  vim.diagnostic.set(namespace, bufnr, diagnostic_lsp_to_vim(diagnostics, bufnr, client_id))
+end
+
+--- Clear push diagnostics and diagnostic cache.
 ---
 --- Diagnostic producers should prefer |vim.diagnostic.reset()|. However,
 --- this method signature is still used internally in some parts of the LSP
@@ -243,7 +340,7 @@ function M.reset(client_id, buffer_client_map)
   vim.schedule(function()
     for bufnr, client_ids in pairs(buffer_client_map) do
       if client_ids[client_id] then
-        local namespace = M.get_namespace(client_id)
+        local namespace = M.get_namespace(client_id, false)
         vim.diagnostic.reset(namespace, bufnr)
       end
     end
@@ -275,7 +372,7 @@ function M.get_line_diagnostics(bufnr, line_nr, opts, client_id)
   end
 
   if client_id then
-    opts.namespace = M.get_namespace(client_id)
+    opts.namespace = M.get_namespace(client_id, false)
   end
 
   if not line_nr then
@@ -287,4 +384,70 @@ function M.get_line_diagnostics(bufnr, line_nr, opts, client_id)
   return diagnostic_vim_to_lsp(vim.diagnostic.get(bufnr, opts))
 end
 
+--- Clear diagnostics from pull based clients
+--- @private
+local function clear(bufnr)
+  for _, namespace in pairs(_client_pull_namespaces) do
+    vim.diagnostic.reset(namespace, bufnr)
+  end
+end
+
+--- autocmd ids for LspNotify handlers per buffer
+--- @private
+--- @type table<integer,integer>
+local _autocmd_ids = {}
+
+--- Disable pull diagnostics for a buffer
+--- @private
+local function disable(bufnr)
+  if not _autocmd_ids[bufnr] then
+    return
+  end
+  api.nvim_del_autocmd(_autocmd_ids[bufnr])
+  _autocmd_ids[bufnr] = nil
+  clear(bufnr)
+end
+
+--- Enable pull diagnostics for a buffer
+---@param bufnr (integer) Buffer handle, or 0 for current
+---@private
+function M._enable(bufnr)
+  if bufnr == nil or bufnr == 0 then
+    bufnr = api.nvim_get_current_buf()
+  end
+
+  if _autocmd_ids[bufnr] then
+    return
+  end
+
+  _autocmd_ids[bufnr] = api.nvim_create_autocmd('LspNotify', {
+    buffer = bufnr,
+    callback = function(opts)
+      if opts.data.method ~= 'textDocument/didChange' then
+        return
+      end
+      util._refresh('textDocument/diagnostic', { bufnr = bufnr, only_visible = true })
+    end,
+    group = augroup,
+  })
+
+  api.nvim_buf_attach(bufnr, false, {
+    on_reload = function()
+      util._refresh('textDocument/diagnostic', { bufnr = bufnr })
+    end,
+    on_detach = function()
+      disable(bufnr)
+    end,
+  })
+
+  api.nvim_create_autocmd('LspDetach', {
+    buffer = bufnr,
+    callback = function()
+      disable(bufnr)
+    end,
+    once = true,
+    group = augroup,
+  })
+end
+
 return M
diff --git a/runtime/lua/vim/lsp/handlers.lua b/runtime/lua/vim/lsp/handlers.lua
index ce3db68618..d887183972 100644
--- a/runtime/lua/vim/lsp/handlers.lua
+++ b/runtime/lua/vim/lsp/handlers.lua
@@ -214,6 +214,10 @@ M['textDocument/publishDiagnostics'] = function(...)
   return require('vim.lsp.diagnostic').on_publish_diagnostics(...)
 end
 
+M['textDocument/diagnostic'] = function(...)
+  return require('vim.lsp.diagnostic').on_diagnostic(...)
+end
+
 M['textDocument/codeLens'] = function(...)
   return require('vim.lsp.codelens').on_codelens(...)
 end
diff --git a/runtime/lua/vim/lsp/inlay_hint.lua b/runtime/lua/vim/lsp/inlay_hint.lua
index bec6f33e93..1c29e8a866 100644
--- a/runtime/lua/vim/lsp/inlay_hint.lua
+++ b/runtime/lua/vim/lsp/inlay_hint.lua
@@ -87,54 +87,6 @@ function M.on_inlayhint(err, result, ctx, _)
   api.nvim__buf_redraw_range(bufnr, 0, -1)
 end
 
-local function resolve_bufnr(bufnr)
-  return bufnr > 0 and bufnr or api.nvim_get_current_buf()
-end
-
---- Refresh inlay hints for a buffer
----
----@param opts (nil|table) Optional arguments
----  - bufnr (integer, default: 0): Buffer whose hints to refresh
----  - only_visible (boolean, default: false): Whether to only refresh hints for the visible regions of the buffer
----
-local function refresh(opts)
-  opts = opts or {}
-  local bufnr = resolve_bufnr(opts.bufnr or 0)
-  local bufstate = bufstates[bufnr]
-  if not (bufstate and bufstate.enabled) then
-    return
-  end
-  local only_visible = opts.only_visible or false
-  local buffer_windows = {}
-  for _, winid in ipairs(api.nvim_list_wins()) do
-    if api.nvim_win_get_buf(winid) == bufnr then
-      table.insert(buffer_windows, winid)
-    end
-  end
-  for _, window in ipairs(buffer_windows) do
-    local first = vim.fn.line('w0', window)
-    local last = vim.fn.line('w$', window)
-    local params = {
-      textDocument = util.make_text_document_params(bufnr),
-      range = {
-        start = { line = first - 1, character = 0 },
-        ['end'] = { line = last, character = 0 },
-      },
-    }
-    vim.lsp.buf_request(bufnr, 'textDocument/inlayHint', params)
-  end
-  if not only_visible then
-    local params = {
-      textDocument = util.make_text_document_params(bufnr),
-      range = {
-        start = { line = 0, character = 0 },
-        ['end'] = { line = api.nvim_buf_line_count(bufnr), character = 0 },
-      },
-    }
-    vim.lsp.buf_request(bufnr, 'textDocument/inlayHint', params)
-  end
-end
-
 --- |lsp-handler| for the method `textDocument/inlayHint/refresh`
 ---@private
 function M.on_refresh(err, _, ctx, _)
@@ -144,8 +96,11 @@ function M.on_refresh(err, _, ctx, _)
   for _, bufnr in ipairs(vim.lsp.get_buffers_by_client_id(ctx.client_id)) do
     for _, winid in ipairs(api.nvim_list_wins()) do
       if api.nvim_win_get_buf(winid) == bufnr then
-        refresh({ bufnr = bufnr })
-        break
+        local bufstate = bufstates[bufnr]
+        if bufstate and bufstate.enabled then
+          util._refresh('textDocument/inlayHint', { bufnr = bufnr })
+          break
+        end
       end
     end
   end
@@ -156,7 +111,9 @@ end
 --- Clear inlay hints
 ---@param bufnr (integer) Buffer handle, or 0 for current
 local function clear(bufnr)
-  bufnr = resolve_bufnr(bufnr)
+  if bufnr == nil or bufnr == 0 then
+    bufnr = api.nvim_get_current_buf()
+  end
   if not bufstates[bufnr] then
     return
   end
@@ -175,17 +132,19 @@ end
 
 local function make_request(request_bufnr)
   reset_timer(request_bufnr)
-  refresh({ bufnr = request_bufnr })
+  util._refresh('textDocument/inlayHint', { bufnr = request_bufnr })
 end
 
 --- Enable inlay hints for a buffer
 ---@param bufnr (integer) Buffer handle, or 0 for current
 local function enable(bufnr)
-  bufnr = resolve_bufnr(bufnr)
+  if bufnr == nil or bufnr == 0 then
+    bufnr = api.nvim_get_current_buf()
+  end
   local bufstate = bufstates[bufnr]
   if not (bufstate and bufstate.enabled) then
     bufstates[bufnr] = { enabled = true, timer = nil, applied = {} }
-    refresh({ bufnr = bufnr })
+    util._refresh('textDocument/inlayHint', { bufnr = bufnr })
     api.nvim_buf_attach(bufnr, true, {
       on_lines = function(_, cb_bufnr)
         if not bufstates[cb_bufnr].enabled then
@@ -201,7 +160,7 @@ local function enable(bufnr)
         if bufstates[cb_bufnr] and bufstates[cb_bufnr].enabled then
           bufstates[cb_bufnr] = { enabled = true, applied = {} }
         end
-        refresh({ bufnr = cb_bufnr })
+        util._refresh('textDocument/inlayHint', { bufnr = cb_bufnr })
       end,
       on_detach = function(_, cb_bufnr)
         clear(cb_bufnr)
@@ -222,7 +181,9 @@ end
 --- Disable inlay hints for a buffer
 ---@param bufnr (integer) Buffer handle, or 0 for current
 local function disable(bufnr)
-  bufnr = resolve_bufnr(bufnr)
+  if bufnr == nil or bufnr == 0 then
+    bufnr = api.nvim_get_current_buf()
+  end
   if bufstates[bufnr] and bufstates[bufnr].enabled then
     clear(bufnr)
     bufstates[bufnr].enabled = nil
@@ -233,7 +194,9 @@ end
 --- Toggle inlay hints for a buffer
 ---@param bufnr (integer) Buffer handle, or 0 for current
 local function toggle(bufnr)
-  bufnr = resolve_bufnr(bufnr)
+  if bufnr == nil or bufnr == 0 then
+    bufnr = api.nvim_get_current_buf()
+  end
   local bufstate = bufstates[bufnr]
   if bufstate and bufstate.enabled then
     disable(bufnr)
diff --git a/runtime/lua/vim/lsp/protocol.lua b/runtime/lua/vim/lsp/protocol.lua
index 5bc0baf241..537a5eda39 100644
--- a/runtime/lua/vim/lsp/protocol.lua
+++ b/runtime/lua/vim/lsp/protocol.lua
@@ -641,6 +641,9 @@ function protocol.make_client_capabilities()
       },
     },
     textDocument = {
+      diagnostic = {
+        dynamicRegistration = false,
+      },
       inlayHint = {
         dynamicRegistration = true,
         resolveSupport = {
diff --git a/runtime/lua/vim/lsp/util.lua b/runtime/lua/vim/lsp/util.lua
index 9a6114c35b..0b06d2bbb5 100644
--- a/runtime/lua/vim/lsp/util.lua
+++ b/runtime/lua/vim/lsp/util.lua
@@ -2183,6 +2183,46 @@ function M.lookup_section(settings, section)
   return settings
 end
 
+---@private
+--- Request updated LSP information for a buffer.
+---
+---@param method string LSP method to call
+---@param opts (nil|table) Optional arguments
+---  - bufnr (integer, default: 0): Buffer to refresh
+---  - only_visible (boolean, default: false): Whether to only refresh for the visible regions of the buffer
+function M._refresh(method, opts)
+  opts = opts or {}
+  local bufnr = opts.bufnr
+  if bufnr == nil or bufnr == 0 then
+    bufnr = api.nvim_get_current_buf()
+  end
+  local only_visible = opts.only_visible or false
+  for _, window in ipairs(api.nvim_list_wins()) do
+    if api.nvim_win_get_buf(window) == bufnr then
+      local first = vim.fn.line('w0', window)
+      local last = vim.fn.line('w$', window)
+      local params = {
+        textDocument = M.make_text_document_params(bufnr),
+        range = {
+          start = { line = first - 1, character = 0 },
+          ['end'] = { line = last, character = 0 },
+        },
+      }
+      vim.lsp.buf_request(bufnr, method, params)
+    end
+  end
+  if not only_visible then
+    local params = {
+      textDocument = M.make_text_document_params(bufnr),
+      range = {
+        start = { line = 0, character = 0 },
+        ['end'] = { line = api.nvim_buf_line_count(bufnr), character = 0 },
+      },
+    }
+    vim.lsp.buf_request(bufnr, method, params)
+  end
+end
+
 M._get_line_byte_from_position = get_line_byte_from_position
 
 ---@nodoc
diff --git a/src/nvim/auevents.lua b/src/nvim/auevents.lua
index 41d7ee9b47..f023ee1340 100644
--- a/src/nvim/auevents.lua
+++ b/src/nvim/auevents.lua
@@ -73,6 +73,7 @@ return {
     'LspAttach',              -- after an LSP client attaches to a buffer
     'LspDetach',              -- after an LSP client detaches from a buffer
     'LspRequest',             -- after an LSP request is started, canceled, or completed
+    'LspNotify',              -- after an LSP notice has been sent to the server
     'LspTokenUpdate',         -- after a visible LSP token is updated
     'LspProgress',            -- after a LSP progress update
     'MenuPopup',              -- just before popup menu is displayed
@@ -154,6 +155,7 @@ return {
     DiagnosticChanged=true,
     LspAttach=true,
     LspDetach=true,
+    LspNotify=true,
     LspRequest=true,
     LspProgress=true,
     LspTokenUpdate=true,
diff --git a/test/functional/plugin/lsp/diagnostic_spec.lua b/test/functional/plugin/lsp/diagnostic_spec.lua
index f58016bf01..d1c3fd6b1e 100644
--- a/test/functional/plugin/lsp/diagnostic_spec.lua
+++ b/test/functional/plugin/lsp/diagnostic_spec.lua
@@ -1,10 +1,13 @@
 local helpers = require('test.functional.helpers')(after_each)
+local lsp_helpers = require('test.functional.plugin.lsp.helpers')
 
 local clear = helpers.clear
 local exec_lua = helpers.exec_lua
 local eq = helpers.eq
 local neq = require('test.helpers').neq
 
+local create_server_definition = lsp_helpers.create_server_definition
+
 describe('vim.lsp.diagnostic', function()
   local fake_uri
 
@@ -265,4 +268,97 @@ describe('vim.lsp.diagnostic', function()
       eq(exec_lua([[return #vim.diagnostic.get(...)]], bufnr), 0)
     end)
   end)
+
+  describe('vim.lsp.diagnostic.on_diagnostic', function()
+    before_each(function()
+      exec_lua(create_server_definition)
+      exec_lua([[
+        server = _create_server({
+          capabilities = {
+            diagnosticProvider = {
+            }
+          }
+        })
+
+        function get_extmarks(bufnr, client_id)
+          local namespace = vim.lsp.diagnostic.get_namespace(client_id, true)
+          local ns = vim.diagnostic.get_namespace(namespace)
+          local extmarks = {}
+          if ns.user_data.virt_text_ns then
+            for _, e in pairs(vim.api.nvim_buf_get_extmarks(bufnr, ns.user_data.virt_text_ns, 0, -1, {details=true})) do
+              table.insert(extmarks, e)
+            end
+          end
+          if ns.user_data.underline_ns then
+            for _, e in pairs(vim.api.nvim_buf_get_extmarks(bufnr, ns.user_data.underline_ns, 0, -1, {details=true})) do
+              table.insert(extmarks, e)
+            end
+          end
+          return extmarks
+        end
+
+        client_id = vim.lsp.start({ name = 'dummy', cmd = server.cmd })
+      ]])
+    end)
+
+    it('adds diagnostics to vim.diagnostics', function()
+      local diags = exec_lua([[
+        vim.lsp.diagnostic.on_diagnostic(nil,
+          {
+            kind = 'full',
+            items = {
+              make_error('Pull Diagnostic', 4, 4, 4, 4),
+            }
+          },
+          {
+            params = {
+              textDocument = { uri = fake_uri },
+            },
+            uri = fake_uri,
+            client_id = client_id,
+          },
+          {}
+        )
+
+        return vim.diagnostic.get(diagnostic_bufnr)
+      ]])
+      eq(1, #diags)
+      eq('Pull Diagnostic', diags[1].message)
+    end)
+
+    it('allows configuring the virtual text via vim.lsp.with', function()
+      local expected_spacing = 10
+      local extmarks = exec_lua(
+        [[
+        Diagnostic = vim.lsp.with(vim.lsp.diagnostic.on_diagnostic, {
+          virtual_text = {
+            spacing = ...,
+          },
+        })
+
+        Diagnostic(nil,
+          {
+            kind = 'full',
+            items = {
+              make_error('Pull Diagnostic', 4, 4, 4, 4),
+            }
+          },
+          {
+            params = {
+              textDocument = { uri = fake_uri },
+            },
+            uri = fake_uri,
+            client_id = client_id,
+          },
+          {}
+        )
+
+        return get_extmarks(diagnostic_bufnr, client_id)
+      ]],
+        expected_spacing
+      )
+      eq(2, #extmarks)
+      eq(expected_spacing, #extmarks[1][4].virt_text[1][1])
+    end)
+  end)
 end)