From 3dd72e4703076ce9aea975504fff2940284b6729 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20Ch=C5=82odnicki?= <rchl2k@gmail.com>
Date: Sun, 22 Jan 2023 22:37:54 +0100
Subject: [PATCH] trigger "on_server_response_async" also for "initialize"
 response (#2172)

---
 .github/workflows/main.yml | 2 +-
 plugin/core/protocol.py    | 9 ++++++---
 plugin/core/sessions.py    | 9 +++++++--
 tox.ini                    | 2 +-
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 81428a2af..108338e0f 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -38,7 +38,7 @@ jobs:
           python-version: '3.8'
       - run: sudo apt update
       - run: sudo apt install --no-install-recommends -y x11-xserver-utils
-      - run: pip3 install mypy==0.971 flake8==5.0.4 pyright==1.1.285 yapf==0.31.0 --user
+      - run: pip3 install mypy==0.971 flake8==5.0.4 pyright==1.1.289 yapf==0.31.0 --user
       - run: echo "$HOME/.local/bin" >> $GITHUB_PATH
       # - run: mypy -p plugin
       - run: flake8 plugin tests
diff --git a/plugin/core/protocol.py b/plugin/core/protocol.py
index 0a8c39f8d..08c76d75e 100644
--- a/plugin/core/protocol.py
+++ b/plugin/core/protocol.py
@@ -1,5 +1,5 @@
 from .typing import Enum, IntEnum, IntFlag, StrEnum
-from .typing import Any, Dict, Iterable, List, Literal, Mapping, NotRequired, Optional, TypedDict, Union
+from .typing import Any, Dict, Generic, Iterable, List, Literal, Mapping, NotRequired, Optional, TypedDict, TypeVar, Union  # noqa: E501
 import sublime
 
 INT_MAX = 2**31 - 1
@@ -5997,11 +5997,14 @@ def from_exception(cls, ex: Exception) -> 'Error':
         return Error(ErrorCodes.InternalError, str(ex))
 
 
-class Response:
+T = TypeVar('T', bound=Union[None, bool, int, Uint, float, str, Mapping[str, Any], Iterable[Any]])
+
+
+class Response(Generic[T]):
 
     __slots__ = ('request_id', 'result')
 
-    def __init__(self, request_id: Any, result: Union[None, Mapping[str, Any], Iterable[Any]]) -> None:
+    def __init__(self, request_id: Any, result: T) -> None:
         self.request_id = request_id
         self.result = result
 
diff --git a/plugin/core/sessions.py b/plugin/core/sessions.py
index d95414f0e..0f8486f49 100644
--- a/plugin/core/sessions.py
+++ b/plugin/core/sessions.py
@@ -36,7 +36,9 @@
 from .protocol import FailureHandlingKind
 from .protocol import FileEvent
 from .protocol import GeneralClientCapabilities
+from .protocol import InitializeError
 from .protocol import InitializeParams
+from .protocol import InitializeResult
 from .protocol import InsertTextMode
 from .protocol import Location
 from .protocol import LocationLink
@@ -1387,13 +1389,16 @@ def initialize_async(
         self.send_request_async(
             Request.initialize(params), self._handle_initialize_success, self._handle_initialize_error)
 
-    def _handle_initialize_success(self, result: Any) -> None:
+    def _handle_initialize_success(self, result: InitializeResult) -> None:
         self.capabilities.assign(result.get('capabilities', dict()))
         if self._workspace_folders and not self._supports_workspace_folders():
             self._workspace_folders = self._workspace_folders[:1]
         self.state = ClientStates.READY
         if self._plugin_class is not None:
             self._plugin = self._plugin_class(weakref.ref(self))
+            # We've missed calling the "on_server_response_async" API as plugin was not created yet.
+            # Handle it now and use fake request ID since it shouldn't matter.
+            self._plugin.on_server_response_async('initialize', Response(-1, result))
         self.send_notification(Notification.initialized())
         self._maybe_send_did_change_configuration()
         execute_commands = self.get_capability('executeCommandProvider.commands')
@@ -1421,7 +1426,7 @@ def _handle_initialize_success(self, result: Any) -> None:
             self._init_callback(self, False)
             self._init_callback = None
 
-    def _handle_initialize_error(self, result: Any) -> None:
+    def _handle_initialize_error(self, result: InitializeError) -> None:
         self._initialize_error = (result.get('code', -1), Exception(result.get('message', 'Error initializing server')))
         # Init callback called after transport is closed to avoid pre-mature GC of Session.
         self.end_async()
diff --git a/tox.ini b/tox.ini
index 1f119e9d8..cd782f96a 100644
--- a/tox.ini
+++ b/tox.ini
@@ -22,7 +22,7 @@ per-file-ignores =
 deps =
     ; mypy==0.971
     flake8==5.0.4
-    pyright==1.1.285
+    pyright==1.1.289
 commands =
     # mypy disabled as it doesn't currently support cyclic definitions - https://github.com/python/mypy/issues/731
     ; mypy plugin