fix(lsp): _get_workspace_folders does not handle root_dir() function #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: Justin M. Keyes <justinkz@gmail.com>

runtime/lua/vim/lsp.lua

@@ -56,7 +56,7 @@ function lsp._unsupported_method(method)
   return msg
 end
 
----@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
@@ -68,6 +68,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 all_clients = {}
 --- @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 all_clients = {}
 --- @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()|.
@@ -1367,6 +1367,27 @@ 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
+
 -- Export for internal use only.
 Client._all = all_clients
 

runtime/lua/vim/lsp/health.lua

@@ -89,10 +89,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',
-          client.root_dir and vim.fn.fnamemodify(client.root_dir, ':~')
-        ) or nil
+          '- Root %s:\n    %s',
+          vim.tbl_count(root_dirs) > 1 and 'directories' or 'directory',
+          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),