Multilspy: sleep in TS language server before requesting references

src/multilspy/language_server.py

@@ -14,19 +14,16 @@ import os
 import pathlib
 import pickle
 import re
-from site import abs_paths
 import threading
 from collections import defaultdict
 from contextlib import asynccontextmanager, contextmanager
 from copy import copy
-from typing import Dict, List, Optional, Tuple, Union
-from fnmatch import fnmatch
 from pathlib import Path, PurePath
 from typing import AsyncIterator, Dict, Iterator, List, Optional, Tuple, Union, cast
 
 import pathspec
 
-from serena.text_utils import LineType, MatchedConsecutiveLines, TextLine, search_files, search_text
+from serena.text_utils import LineType, MatchedConsecutiveLines, TextLine, search_files
 from . import multilspy_types
 from .lsp_protocol_handler import lsp_types as LSPTypes
 from .lsp_protocol_handler.lsp_constants import LSPConstants
@@ -38,9 +35,6 @@ from .lsp_protocol_handler.server import (
 )
 from .multilspy_config import Language, MultilspyConfig
 from .multilspy_exceptions import MultilspyException
-from .multilspy_utils import PathUtils, FileUtils, TextUtils
-from pathlib import PurePath
-from typing import AsyncIterator, Iterator, List, Dict, Optional, Union, Tuple
 from .multilspy_logger import MultilspyLogger
 from .multilspy_utils import FileUtils, PathUtils, TextUtils
 from .type_helpers import ensure_all_methods_implemented
@@ -581,6 +575,16 @@ class LanguageServer:
 
         return ret
 
+    # For some reason, the LS may need longer to process this, so we just retry
+    async def _send_references_request(self, relative_file_path: str, line: int, column: int):
+        return await self.server.send.references(
+            {
+                "textDocument": {"uri": PathUtils.path_to_uri(os.path.join(self.repository_root_path, relative_file_path))},
+                "position": {"line": line, "character": column},
+                "context": {"includeDeclaration": False},
+            }
+        )
+
     async def request_references(
         self, relative_file_path: str, line: int, column: int
     ) -> List[multilspy_types.Location]:
@@ -606,13 +610,7 @@ class LanguageServer:
 
         with self.open_file(relative_file_path):
             try:
-                response = await self.server.send.references(
-                    {
-                        "textDocument": {"uri": PathUtils.path_to_uri(os.path.join(self.repository_root_path, relative_file_path))},
-                        "position": {"line": line, "character": column},
-                        "context": {"includeDeclaration": False},
-                    }
-                )
+                response = await self._send_references_request(relative_file_path, line=line, column=column)
             except Exception as e:
                 # Catch LSP internal error (-32603) and raise a more informative exception
                 if isinstance(e, Error) and getattr(e, 'code', None) == -32603:

src/multilspy/language_servers/gopls/gopls.py

@@ -156,18 +156,16 @@ class Gopls(LanguageServer):
             await self.server.stop()
 
     @override
-    async def request_references(
-        self, relative_file_path: str, line: int, column: int
-    ) -> List[multilspy_types.Location]:
+    async def _send_references_request(self, relative_file_path: str, line: int, column: int) -> List[multilspy_types.Location]:
         NUM_COLS_TO_TRY = 10
-        """selectionRange in Gopls always return the wrong column, 
-        the one at the beginnning of the declaration instead of the symbol. We do a dirty hack here.
+        """selectionRange in Gopls always contains the wrong column, 
+        the one at the beginning of the declaration instead of the symbol. We do a dirty hack here (loop over some columns).
         See https://github.com/golang/go/issues/73521
         """
         actual_column = column
         for actual_column in range(column, column+NUM_COLS_TO_TRY):
             try:
-                return await super().request_references(relative_file_path, line, actual_column)
+                return await super()._send_references_request(relative_file_path, line, actual_column)
             except (Error, RuntimeError, MultilspyException) as e:
                 self.logger.log(f"Cannot find symbol at {line=}, column={actual_column}, trying to bump column by 1", logging.INFO)
-        raise RuntimeError(f"Failed to find references for symbol at {line=}, {column=}")
+        raise MultilspyException(f"Failed to find references for symbol in\n{relative_file_path=} at {line=}, {column=}")

src/multilspy/language_servers/typescript_language_server/typescript_language_server.py

@@ -4,22 +4,23 @@ Provides TypeScript specific instantiation of the LanguageServer class. Contains
 
 import asyncio
 import json
-import shutil
 import logging
 import os
-import subprocess
 import pathlib
+import shutil
+import subprocess
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from typing import AsyncIterator, Tuple
+from time import sleep
 
 from overrides import override
 
-from multilspy.multilspy_logger import MultilspyLogger
 from multilspy.language_server import LanguageServer
-from multilspy.lsp_protocol_handler.server import ProcessLaunchInfo
 from multilspy.lsp_protocol_handler.lsp_types import InitializeParams
+from multilspy.lsp_protocol_handler.server import ProcessLaunchInfo
 from multilspy.multilspy_config import MultilspyConfig
-from multilspy.multilspy_utils import PlatformUtils, PlatformId
+from multilspy.multilspy_logger import MultilspyLogger
+from multilspy.multilspy_utils import PlatformId, PlatformUtils
 
 # Platform-specific imports
 if os.name != 'nt':  # Unix-like systems
@@ -82,7 +83,7 @@ class TypeScriptLanguageServer(LanguageServer):
         ]
         assert platform_id in valid_platforms, f"Platform {platform_id} is not supported for multilspy javascript/typescript at the moment"
 
-        with open(os.path.join(os.path.dirname(__file__), "runtime_dependencies.json"), "r") as f:
+        with open(os.path.join(os.path.dirname(__file__), "runtime_dependencies.json")) as f:
             d = json.load(f)
             del d["_description"]
 
@@ -132,7 +133,7 @@ class TypeScriptLanguageServer(LanguageServer):
         """
         Returns the initialize params for the TypeScript Language Server.
         """
-        with open(os.path.join(os.path.dirname(__file__), "initialize_params.json"), "r") as f:
+        with open(os.path.join(os.path.dirname(__file__), "initialize_params.json")) as f:
             d = json.load(f)
 
         del d["_description"]
@@ -222,3 +223,14 @@ class TypeScriptLanguageServer(LanguageServer):
 
             await self.server.shutdown()
             await self.server.stop()
+
+    @override
+    # For some reason, the LS may need longer to process this, so we just retry
+    async def _send_references_request(self, relative_file_path: str, line: int, column: int):
+        # TODO: The LS doesn't return references contained in other files if it doesn't sleep. This is
+        #   despite the LS having processed requests already. I don't know what causes this, but sleeping
+        #   one second helps. It may be that sleeping only once is enough but that's hard to reliably test.
+        #   It may be that even this 1sec is not enough in larger TS projects, at some point we should find what
+        #   causes this and solve it.
+        sleep(1)
+        return await super()._send_references_request(relative_file_path, line, column)