fix(lsp): _get_workspace_folders does not handle root_dir() function (#36141)

backport #36071

* fix(lsp): type of root_dir should be annotated with string|fun|nil
* feat(lsp): support root_dir as function in _get_workspace_folders
* feat(lsp): let checkhealth support root_dir() function

Examples:

    vim.lsp: Active Clients ~
    - lua_ls (id: 1)
      - Version: <Unknown>
      - Root directories:
          ~/foo/bar
          ~/dev/neovim

Co-authored-by: atusy <30277794+atusy@users.noreply.github.com>

runtime/doc/lsp.txt

@@ -1257,7 +1257,8 @@ Lua module: vim.lsp.client                                        *lsp-client*
                                 ephemerally while executing |LspRequest|
                                 autocmds when replies are received from the
                                 server.
-      • {root_dir}              (`string?`) See |vim.lsp.ClientConfig|.
+      • {root_dir}?             (`string|fun(bufnr: integer, on_dir:fun(root_dir?:string))`)
+                                See |vim.lsp.ClientConfig|.
       • {rpc}                   (`vim.lsp.rpc.PublicClient`) RPC client
                                 object, for low level interaction with the
                                 client. See |vim.lsp.rpc.start()|.
@@ -1407,9 +1408,10 @@ Lua module: vim.lsp.client                                        *lsp-client*
                                You can only modify the
                                `client.offset_encoding` here before any
                                notifications are sent.
-      • {root_dir}?            (`string`) Directory where the LSP server will
+      • {root_dir}?            (`string|fun(bufnr: integer, on_dir:fun(root_dir?:string))`)
-                               base its workspaceFolders, rootUri, and
+                               Directory where the LSP server will base its
-                               rootPath on initialization.
+                               workspaceFolders, rootUri, and rootPath on
+                               initialization.
       • {settings}?            (`lsp.LSPObject`) Map of language
                                server-specific settings, decided by the
                                client. Sent to the LS if requested via

runtime/lua/vim/lsp.lua

@@ -54,7 +54,7 @@ function lsp._unsupported_method(method)
 end
 
 ---@private
----@param workspace_folders string|lsp.WorkspaceFolder[]?
+---@param workspace_folders string|lsp.WorkspaceFolder[]|fun(bufnr: integer, on_dir:fun(root_dir?:string))?
 ---@return lsp.WorkspaceFolder[]?
 function lsp._get_workspace_folders(workspace_folders)
   if type(workspace_folders) == 'table' then
@@ -66,6 +66,15 @@ function lsp._get_workspace_folders(workspace_folders)
         name = workspace_folders,
       },
     }
+  elseif type(workspace_folders) == 'function' then
+    local name = lsp.client._resolve_root_dir(1000, 0, workspace_folders)
+    return name
+      and {
+        {
+          uri = vim.uri_from_fname(name),
+          name = name,
+        },
+      }
   end
 end
 

runtime/lua/vim/lsp/client.lua

@@ -117,7 +117,7 @@ local validate = vim.validate
 --- @field on_init? elem_or_list<fun(client: vim.lsp.Client, init_result: lsp.InitializeResult)>
 ---
 --- Directory where the LSP server will base its workspaceFolders, rootUri, and rootPath on initialization.
---- @field root_dir? string
+--- @field root_dir? string|fun(bufnr: integer, on_dir:fun(root_dir?:string))
 ---
 --- Map of language server-specific settings, decided by the client. Sent to the LS if requested via
 --- `workspace/configuration`. Keys are case-sensitive.
@@ -190,7 +190,7 @@ local validate = vim.validate
 --- @field requests table<integer,{ type: string, bufnr: integer, method: string}?>
 ---
 --- See [vim.lsp.ClientConfig].
---- @field root_dir string?
+--- @field root_dir? string|fun(bufnr: integer, on_dir:fun(root_dir?:string))
 ---
 --- RPC client object, for low level interaction with the client.
 --- See |vim.lsp.rpc.start()|.
@@ -1209,4 +1209,25 @@ function Client:_remove_workspace_folder(dir)
   end
 end
 
+--- Gets root_dir, waiting up to `ms_` for a potentially async `root_dir()` result.
+---
+--- @param ms_ integer
+--- @param buf integer
+--- @return string|nil
+function Client._resolve_root_dir(ms_, buf, root_dir)
+  if root_dir == nil or type(root_dir) == 'string' then
+    return root_dir --[[@type string|nil]]
+  end
+
+  local dir = nil --[[@type string|nil]]
+  root_dir(buf, function(d)
+    dir = d
+  end)
+  -- root_dir() may be async, wait for a result.
+  vim.wait(ms_, function()
+    return not not dir
+  end)
+  return dir
+end
+
 return Client

runtime/lua/vim/lsp/health.lua

@@ -57,10 +57,23 @@ local function check_active_clients()
         end
         dirs_info = ('- Workspace folders:\n    %s'):format(table.concat(wfolders, '\n    '))
       else
+        local root_dirs = {} ---@type table<string, boolean>
+        local timeout = 1
+        local timeoutmsg = ('root_dir() took > %ds'):format(timeout)
+        for buf, _ in pairs(client.attached_buffers) do
+          local dir = client._resolve_root_dir(1000, buf, client.root_dir)
+          root_dirs[dir or timeoutmsg] = true
+        end
         dirs_info = string.format(
-          '- Root directory: %s',
+          '- Root %s:\n    %s',
-          client.root_dir and vim.fn.fnamemodify(client.root_dir, ':~')
+          vim.tbl_count(root_dirs) > 1 and 'directories' or 'directory',
-        ) or nil
+          vim
+            .iter(root_dirs)
+            :map(function(k, _)
+              return k == timeoutmsg and timeoutmsg or vim.fn.fnamemodify(k, ':~')
+            end)
+            :join('\n    ')
+        )
       end
       report_info(table.concat({
         string.format('%s (id: %d)', client.name, client.id),

test/functional/plugin/lsp_spec.lua

@@ -1941,6 +1941,102 @@ describe('LSP', function()
         end)
       )
     end)
+
+    it('vim.lsp.start sends workspace folder when root_dir callback', function()
+      exec_lua(create_server_definition)
+
+      local expected_root_dir = tmpname(false)
+      mkdir(expected_root_dir)
+
+      local result = exec_lua(function(root_dir)
+        local server = _G._create_server()
+        local calls = {}
+
+        local client_id = vim.lsp.start({
+          name = 'cb-root',
+          cmd = server.cmd,
+          root_dir = function(bufnr, on_dir)
+            table.insert(calls, bufnr)
+            on_dir(root_dir)
+          end,
+        }, { attach = false })
+
+        vim.wait(1000, function()
+          return #server.messages > 0
+        end)
+
+        local initialize = server.messages[1]
+        if client_id then
+          vim.lsp.stop_client(client_id)
+        end
+
+        return {
+          calls = calls,
+          ---@diagnostic disable-next-line: no-unknown
+          workspace_folders = initialize and initialize.params.workspaceFolders,
+        }
+      end, expected_root_dir)
+
+      eq({ 0 }, result.calls)
+      eq({
+        {
+          uri = vim.uri_from_fname(expected_root_dir),
+          name = expected_root_dir,
+        },
+      }, result.workspace_folders)
+    end)
+
+    it('vim.lsp.start does not reuse client when root_dir callbacks differ', function()
+      exec_lua(create_server_definition)
+
+      local root1 = tmpname(false)
+      local root2 = tmpname(false)
+      mkdir(root1)
+      mkdir(root2)
+
+      local roots = exec_lua(function(root1_, root2_)
+        local server = _G._create_server()
+        local client_ids = {}
+
+        client_ids[1] = vim.lsp.start({
+          name = 'cb-root',
+          cmd = server.cmd,
+          root_dir = function(_, on_dir)
+            on_dir(root1_)
+          end,
+        }, { attach = false })
+
+        vim.wait(1000, function()
+          return #vim.lsp.get_clients() >= 1
+        end)
+
+        client_ids[2] = vim.lsp.start({
+          name = 'cb-root',
+          cmd = server.cmd,
+          root_dir = function(_, on_dir)
+            on_dir(root2_)
+          end,
+        }, { attach = false })
+
+        vim.wait(1000, function()
+          return #vim.lsp.get_clients() >= 2
+        end)
+
+        local clients = vim.lsp.get_clients()
+        local folders = {}
+        for i, client in ipairs(clients) do
+          local folder = client.workspace_folders and client.workspace_folders[1]
+          folders[i] = folder and folder.name or client.root_dir
+        end
+
+        vim.lsp.stop_client(client_ids)
+
+        return folders
+      end, root1, root2)
+
+      table.sort(roots)
+      eq({ root1, root2 }, roots)
+    end)
   end)
 
   describe('parsing tests', function()