From 371aa1c5663ccb73b12b105ed09e85813d1c985b Mon Sep 17 00:00:00 2001
From: Riley Bruins <ribru17@hotmail.com>
Date: Tue, 8 Jul 2025 18:41:50 -0700
Subject: [PATCH] feat(lsp): diagnostic related documents support

---
 runtime/doc/news.txt                          |  2 +
 runtime/lua/vim/lsp/diagnostic.lua            | 16 ++++++
 runtime/lua/vim/lsp/protocol.lua              |  1 +
 .../functional/plugin/lsp/diagnostic_spec.lua | 53 +++++++++++++++++++
 4 files changed, 72 insertions(+)

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 220c5fcb38..a73b02fefa 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -214,6 +214,8 @@ LSP
   jump to the problematic location.
 • Support for `textDocument/linkedEditingRange`: |lsp-linked_editing_range|
   https://microsoft.github.io/language-server-protocol/specification/#textDocument_linkedEditingRange
+• Support for related documents in pull diagnostics:
+  https://microsoft.github.io/language-server-protocol/specifications/specification-current/#relatedFullDocumentDiagnosticReport
 
 LUA
 
diff --git a/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index fcb3662adb..4eae02c195 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -283,6 +283,22 @@ function M.on_diagnostic(error, result, ctx)
   end
 
   handle_diagnostics(ctx.params.textDocument.uri, client_id, result.items, true)
+
+  for uri, related_result in pairs(result.relatedDocuments or {}) do
+    if related_result.kind == 'full' then
+      handle_diagnostics(uri, client_id, related_result.items, true)
+    end
+
+    local related_bufnr = vim.uri_to_bufnr(uri)
+    local related_bufstate = bufstates[related_bufnr]
+      -- Create a new bufstate if it doesn't exist for the related document. This will not enable
+      -- diagnostic pulling by itself, but will allow previous result IDs to be passed correctly the
+      -- next time this buffer's diagnostics are pulled.
+      or { pull_kind = 'document', client_result_id = {} }
+    bufstates[related_bufnr] = related_bufstate
+
+    related_bufstate.client_result_id[client_id] = related_result.resultId
+  end
 end
 
 --- Clear push diagnostics and diagnostic cache.
diff --git a/runtime/lua/vim/lsp/protocol.lua b/runtime/lua/vim/lsp/protocol.lua
index 314717037a..372affecc4 100644
--- a/runtime/lua/vim/lsp/protocol.lua
+++ b/runtime/lua/vim/lsp/protocol.lua
@@ -350,6 +350,7 @@ function protocol.make_client_capabilities()
         },
         dataSupport = true,
         relatedInformation = true,
+        relatedDocumentSupport = true,
       },
       inlayHint = {
         dynamicRegistration = true,
diff --git a/test/functional/plugin/lsp/diagnostic_spec.lua b/test/functional/plugin/lsp/diagnostic_spec.lua
index 146575d0ca..72d1dfebab 100644
--- a/test/functional/plugin/lsp/diagnostic_spec.lua
+++ b/test/functional/plugin/lsp/diagnostic_spec.lua
@@ -525,5 +525,58 @@ describe('vim.lsp.diagnostic', function()
         end)
       )
     end)
+
+    it('handles relatedDocuments diagnostics', function()
+      local fake_uri_2 = 'file:///fake/uri2'
+      ---@type vim.Diagnostic[], vim.Diagnostic[], string?
+      local diagnostics, related_diagnostics, relatedPreviousResultId = exec_lua(function()
+        local second_buf = vim.uri_to_bufnr(fake_uri_2)
+        vim.fn.bufload(second_buf)
+
+        -- Attach the client to both buffers.
+        vim.api.nvim_win_set_buf(0, second_buf)
+        vim.lsp.start({ name = 'dummy', cmd = _G.server.cmd })
+
+        vim.lsp.diagnostic.on_diagnostic(nil, {
+          kind = 'full',
+          relatedDocuments = {
+            [fake_uri_2] = {
+              kind = 'full',
+              resultId = 'spongebob',
+              items = {
+                {
+                  range = _G.make_range(4, 4, 4, 4),
+                  message = 'related bad!',
+                },
+              },
+            },
+          },
+          items = {},
+        }, {
+          params = {
+            textDocument = { uri = fake_uri },
+          },
+          uri = fake_uri,
+          client_id = client_id,
+          bufnr = diagnostic_bufnr,
+        }, {})
+
+        vim.api.nvim_exec_autocmds('LspNotify', {
+          buffer = second_buf,
+          data = {
+            method = vim.lsp.protocol.Methods.textDocument_didChange,
+            client_id = client_id,
+          },
+        })
+
+        return vim.diagnostic.get(diagnostic_bufnr),
+          vim.diagnostic.get(second_buf),
+          _G.params.previousResultId
+      end)
+      eq(0, #diagnostics)
+      eq(1, #related_diagnostics)
+      eq('related bad!', related_diagnostics[1].message)
+      eq('spongebob', relatedPreviousResultId)
+    end)
   end)
 end)