modules/lsp/onAttach: handle dynamic registrations

Some LSP servers register capabilities after LspAttach, so the global onAttach body must also run from client/registerCapability for buffers already attached to that client.

Keep the initial LspAttach autocmd for servers that advertise capabilities up front, and preserve the wrapped handler's return value.

Closes #3323

modules/lsp/on-attach.nix

@@ -47,13 +47,38 @@ in
               -- client and bufnr are supplied to the builtin `on_attach` callback,
               -- so make them available in scope for our global `onAttach` impl
               local client = vim.lsp.get_client_by_id(event.data.client_id)
-              local bufnr = event.buf
-              ${cfg.onAttach}
+              if client == nil then
+                return
+              end
+
+              __nixvim_lsp_on_attach(client, event.buf)
             end
           end
         '';
         desc = "Run LSP onAttach";
       }
     ];
+
+    extraConfigLua = ''
+      local function __nixvim_lsp_on_attach(client, bufnr)
+        ${cfg.onAttach}
+      end
+
+      vim.lsp.handlers["client/registerCapability"] = (function(overridden)
+        return function(err, res, ctx, ...)
+          local result = overridden(err, res, ctx, ...)
+          local client = vim.lsp.get_client_by_id(ctx.client_id)
+          if client == nil then
+            return result
+          end
+
+          for bufnr, _ in pairs(client.attached_buffers) do
+            __nixvim_lsp_on_attach(client, bufnr)
+          end
+
+          return result
+        end
+      end)(vim.lsp.handlers["client/registerCapability"])
+    '';
   };
 }

tests/test-sources/modules/lsp.nix

@@ -173,6 +173,83 @@
       '';
     };
 
+  on-attach-dynamic-registration = {
+    lsp.onAttach = ''
+      table.insert(_G.__nixvim_lsp_on_attach_calls, {
+        client = client.id,
+        bufnr = bufnr,
+      })
+    '';
+
+    extraConfigLuaPre = ''
+      _G.__nixvim_lsp_on_attach_calls = {}
+      vim.lsp.handlers["client/registerCapability"] = function()
+        return "registered"
+      end
+    '';
+
+    extraConfigLuaPost = ''
+      local get_client_by_id = vim.lsp.get_client_by_id
+      vim.lsp.get_client_by_id = function(client_id)
+        if client_id == "stub_id" then
+          return {
+            id = "stub_id",
+            attached_buffers = {
+              [1] = true,
+            },
+          }
+        end
+      end
+
+      local attach_buf = vim.api.nvim_create_buf(false, true)
+      vim.api.nvim_exec_autocmds("LspAttach", {
+        group = "nixvim_lsp_on_attach",
+        buffer = attach_buf,
+        modeline = false,
+        data = {
+          client_id = "stub_id",
+        },
+      })
+
+      local result = vim.lsp.handlers["client/registerCapability"](nil, nil, { client_id = "stub_id" })
+      vim.lsp.get_client_by_id = get_client_by_id
+
+      if result ~= "registered" then
+        print("Expected client/registerCapability handler to return original result, got", result)
+      end
+
+      local expected_calls = {
+        { client = "stub_id", bufnr = attach_buf },
+        { client = "stub_id", bufnr = 1 },
+      }
+
+      for index, expected in ipairs(expected_calls) do
+        local actual = _G.__nixvim_lsp_on_attach_calls[index]
+        if actual == nil then
+          print("Missing lsp.onAttach call", index)
+        elseif actual.client ~= expected.client or actual.bufnr ~= expected.bufnr then
+          print(
+            "Unexpected lsp.onAttach call",
+            index,
+            "expected",
+            vim.inspect(expected),
+            "got",
+            vim.inspect(actual)
+          )
+        end
+      end
+
+      if #_G.__nixvim_lsp_on_attach_calls ~= #expected_calls then
+        print(
+          "Expected",
+          #expected_calls,
+          "lsp.onAttach calls, got",
+          #_G.__nixvim_lsp_on_attach_calls
+        )
+      end
+    '';
+  };
+
   package-fallback =
     { lib, config, ... }:
     {