From cceff23eecaf4801e9bf50f237e046df1dc040dc Mon Sep 17 00:00:00 2001
From: Janos Wortmann <jwortmann@outlook.com>
Date: Mon, 25 Nov 2024 13:15:36 +0100
Subject: [PATCH 1/5] Patch dynamic registration

---
 plugin/core/types.py     | 22 ++++++++++++++++------
 plugin/session_buffer.py |  5 +++++
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/plugin/core/types.py b/plugin/core/types.py
index b5d5b932f..105b3c58c 100644
--- a/plugin/core/types.py
+++ b/plugin/core/types.py
@@ -528,6 +528,13 @@ class Capabilities(DottedDict):
     (from Server -> Client).
     """
 
+    __slots__ = ('_d', '_registrations', '_registration_options')
+
+    def __init__(self, d: dict[str, Any] | None = None) -> None:
+        super().__init__(d)
+        self._registrations: dict[str, set[str]] = {}
+        self._registration_options: dict[str, Any] = {}
+
     def register(
         self,
         registration_id: str,
@@ -535,10 +542,8 @@ def register(
         registration_path: str,
         options: dict[str, Any]
     ) -> None:
-        stored_registration_id = self.get(registration_path)
-        if isinstance(stored_registration_id, str):
-            msg = "{} is already registered at {} with ID {}, overwriting"
-            debug(msg.format(capability_path, registration_path, stored_registration_id))
+        self._registrations.setdefault(capability_path, set()).add(registration_id)
+        self._registration_options[registration_id] = options
         self.set(capability_path, options)
         self.set(registration_path, registration_id)
 
@@ -548,13 +553,15 @@ def unregister(
         capability_path: str,
         registration_path: str
     ) -> dict[str, Any] | None:
+        self._registrations.get(capability_path, set()).discard(registration_id)
+        self._registration_options.pop(registration_id, None)
         stored_registration_id = self.get(registration_path)
         if not isinstance(stored_registration_id, str):
             debug("stored registration ID at", registration_path, "is not a string")
             return None
         elif stored_registration_id != registration_id:
-            msg = "stored registration ID ({}) is not the same as the provided registration ID ({})"
-            debug(msg.format(stored_registration_id, registration_id))
+            # msg = "stored registration ID ({}) is not the same as the provided registration ID ({})"
+            # debug(msg.format(stored_registration_id, registration_id))
             return None
         else:
             discarded = self.get(capability_path)
@@ -562,6 +569,9 @@ def unregister(
             self.remove(registration_path)
             return discarded
 
+    def get_all(self, path: str) -> list[Any]:
+        return [self._registration_options[registration_id] for registration_id in self._registrations.get(path, [])]
+
     def assign(self, d: ServerCapabilities) -> None:
         textsync = normalize_text_sync(d.pop("textDocumentSync", None))
         super().assign(cast(dict, d))
diff --git a/plugin/session_buffer.py b/plugin/session_buffer.py
index e0d4cb49f..003ad9622 100644
--- a/plugin/session_buffer.py
+++ b/plugin/session_buffer.py
@@ -271,6 +271,11 @@ def get_capability(self, capability_path: str) -> Any | None:
         value = self.capabilities.get(capability_path)
         return value if value is not None else self.session.capabilities.get(capability_path)
 
+    def get_capability_2(self, capability_path: str) -> list[Any]:
+        if self.session.config.is_disabled_capability(capability_path):
+            return []
+        return self.capabilities.get_all(capability_path) + self.session.capabilities.get_all(capability_path)
+
     def has_capability(self, capability_path: str) -> bool:
         value = self.get_capability(capability_path)
         return value is not False and value is not None

From 5951f1338e67ef98b8df88a2abdfe14a7bd84b3f Mon Sep 17 00:00:00 2001
From: Janos Wortmann <jwortmann@outlook.com>
Date: Fri, 29 Nov 2024 14:46:10 +0100
Subject: [PATCH 2/5] Add support for diagnostic streams

---
 plugin/core/constants.py           |  2 +
 plugin/core/diagnostics_storage.py | 99 ++++++++++++++++--------------
 plugin/core/registry.py            |  2 +-
 plugin/core/sessions.py            |  4 +-
 plugin/core/types.py               |  8 ++-
 plugin/core/url.py                 | 11 +++-
 plugin/core/windows.py             |  3 +-
 plugin/goto_diagnostic.py          | 12 ++--
 plugin/session_buffer.py           | 82 +++++++++++++++----------
 9 files changed, 132 insertions(+), 91 deletions(-)

diff --git a/plugin/core/constants.py b/plugin/core/constants.py
index beae8df1c..66aff195b 100644
--- a/plugin/core/constants.py
+++ b/plugin/core/constants.py
@@ -12,6 +12,8 @@
 SublimeKind = Tuple[int, str, str]
 
 
+INSTALLED_PACKAGES_PATH = sublime.installed_packages_path()
+PACKAGES_PATH = sublime.packages_path()
 ST_VERSION = int(sublime.version())
 
 
diff --git a/plugin/core/diagnostics_storage.py b/plugin/core/diagnostics_storage.py
index 1e1ebcf17..8f8d89a99 100644
--- a/plugin/core/diagnostics_storage.py
+++ b/plugin/core/diagnostics_storage.py
@@ -1,49 +1,71 @@
 from __future__ import annotations
-from .protocol import Diagnostic, DiagnosticSeverity, DocumentUri
-from .url import parse_uri
+from .protocol import Diagnostic
+from .protocol import DiagnosticSeverity
+from .protocol import DocumentUri
+from .protocol import Point
+from .url import normalize_uri
 from .views import diagnostic_severity
-from collections import OrderedDict
+from collections.abc import MutableMapping
 from typing import Callable, Iterator, Tuple, TypeVar
+import itertools
 import functools
 
 ParsedUri = Tuple[str, str]
 T = TypeVar('T')
 
 
-# NOTE: OrderedDict can only be properly typed in Python >=3.8.
-class DiagnosticsStorage(OrderedDict):
-    # From the specs:
-    #
-    #   When a file changes it is the server’s responsibility to re-compute
-    #   diagnostics and push them to the client. If the computed set is empty
-    #   it has to push the empty array to clear former diagnostics. Newly
-    #   pushed diagnostics always replace previously pushed diagnostics. There
-    #   is no merging that happens on the client side.
-    #
-    # https://microsoft.github.io/language-server-protocol/specification#textDocument_publishDiagnostics
+class DiagnosticsStorage(MutableMapping):
 
-    def add_diagnostics_async(self, document_uri: DocumentUri, diagnostics: list[Diagnostic]) -> None:
-        """
-        Add `diagnostics` for `document_uri` to the store, replacing previously received `diagnoscis`
-        for this `document_uri`. If `diagnostics` is the empty list, `document_uri` is removed from
-        the store. The item received is moved to the end of the store.
-        """
-        uri = parse_uri(document_uri)
-        if not diagnostics:
-            # received "clear diagnostics" message for this uri
-            self.pop(uri, None)
-            return
-        self[uri] = diagnostics
-        self.move_to_end(uri)  # maintain incoming order
+    def __init__(self) -> None:
+        super().__init__()
+        self._d: dict[tuple[DocumentUri, str], list[Diagnostic]] = dict()
+        self._identifiers = {''}
+        self._uris: set[DocumentUri] = set()
+
+    def __getitem__(self, key: DocumentUri, /) -> list[Diagnostic]:
+        uri = normalize_uri(key)
+        return sorted(
+            itertools.chain.from_iterable(self._d.get((uri, identifier), []) for identifier in self._identifiers),
+            key=lambda diagnostic: Point.from_lsp(diagnostic['range']['start'])
+        )
+
+    def __setitem__(self, key: DocumentUri | tuple[DocumentUri, str], value: list[Diagnostic], /) -> None:
+        uri, identifier = (normalize_uri(key), '') if isinstance(key, DocumentUri) else (normalize_uri(key[0]), key[1])
+        if identifier not in self._identifiers:
+            raise ValueError(f'identifier {identifier} must be registered first')
+        if value:
+            self._uris.add(uri)
+            self._d[(uri, identifier)] = value
+        else:
+            self._uris.discard(uri)
+            self._d.pop((uri, identifier), None)
+
+    def __delitem__(self, key: DocumentUri, /) -> None:
+        uri = normalize_uri(key)
+        self._uris.discard(uri)
+        for identifier in self._identifiers:
+            self._d.pop((uri, identifier), None)
+
+    def __iter__(self) -> Iterator[DocumentUri]:
+        return iter(self._uris)
+
+    def __len__(self) -> int:
+        return len(self._uris)
+
+    def register(self, identifier: str) -> None:
+        self._identifiers.add(identifier)
+
+    def unregister(self, identifier: str) -> None:
+        self._identifiers.discard(identifier)
 
     def filter_map_diagnostics_async(
-        self, pred: Callable[[Diagnostic], bool], f: Callable[[ParsedUri, Diagnostic], T]
-    ) -> Iterator[tuple[ParsedUri, list[T]]]:
+        self, pred: Callable[[Diagnostic], bool], f: Callable[[DocumentUri, Diagnostic], T]
+    ) -> Iterator[tuple[DocumentUri, list[T]]]:
         """
         Yields `(uri, results)` items with `results` being a list of `f(diagnostic)` for each
         diagnostic for this `uri` with `pred(diagnostic) == True`, filtered by `bool(f(diagnostic))`.
         Only `uri`s with non-empty `results` are returned. Each `uri` is guaranteed to be yielded
-        not more than once. Items and results are ordered as they came in from the server.
+        not more than once.
         """
         for uri, diagnostics in self.items():
             results: list[T] = list(filter(None, map(functools.partial(f, uri), filter(pred, diagnostics))))
@@ -51,12 +73,11 @@ def filter_map_diagnostics_async(
                 yield uri, results
 
     def filter_map_diagnostics_flat_async(self, pred: Callable[[Diagnostic], bool],
-                                          f: Callable[[ParsedUri, Diagnostic], T]) -> Iterator[tuple[ParsedUri, T]]:
+                                          f: Callable[[DocumentUri, Diagnostic], T]) -> Iterator[tuple[DocumentUri, T]]:
         """
         Flattened variant of `filter_map_diagnostics_async()`. Yields `(uri, result)` items for each
         of the `result`s per `uri` instead. Each `uri` can be yielded more than once. Items are
-        grouped by `uri` and each `uri` group is guaranteed to appear not more than once. Items are
-        ordered as they came in from the server.
+        grouped by `uri` and each `uri` group is guaranteed to appear not more than once.
         """
         for uri, results in self.filter_map_diagnostics_async(pred, f):
             for result in results:
@@ -71,18 +92,6 @@ def sum_total_errors_and_warnings_async(self) -> tuple[int, int]:
             sum(map(severity_count(DiagnosticSeverity.Warning), self.values())),
         )
 
-    def diagnostics_by_document_uri(self, document_uri: DocumentUri) -> list[Diagnostic]:
-        """
-        Returns possibly empty list of diagnostic for `document_uri`.
-        """
-        return self.get(parse_uri(document_uri), [])
-
-    def diagnostics_by_parsed_uri(self, uri: ParsedUri) -> list[Diagnostic]:
-        """
-        Returns possibly empty list of diagnostic for `uri`.
-        """
-        return self.get(uri, [])
-
 
 def severity_count(severity: int) -> Callable[[list[Diagnostic]], int]:
     def severity_count(diagnostics: list[Diagnostic]) -> int:
diff --git a/plugin/core/registry.py b/plugin/core/registry.py
index 4893920fd..2a4142eb3 100644
--- a/plugin/core/registry.py
+++ b/plugin/core/registry.py
@@ -244,7 +244,7 @@ def navigate_diagnostics(view: sublime.View, point: int | None, forward: bool =
         return
     diagnostics: list[Diagnostic] = []
     for session in wm.get_sessions():
-        diagnostics.extend(session.diagnostics.diagnostics_by_document_uri(uri))
+        diagnostics.extend(session.diagnostics[uri])
     if not diagnostics:
         return
     # Sort diagnostics by location
diff --git a/plugin/core/sessions.py b/plugin/core/sessions.py
index b4392348d..93bc84c2a 100644
--- a/plugin/core/sessions.py
+++ b/plugin/core/sessions.py
@@ -1374,7 +1374,7 @@ def register_session_buffer_async(self, sb: SessionBufferProtocol) -> None:
             data.check_applicable(sb)
         uri = sb.get_uri()
         if uri:
-            diagnostics = self.diagnostics.diagnostics_by_document_uri(uri)
+            diagnostics = self.diagnostics[uri]
             if diagnostics:
                 self._publish_diagnostics_to_session_buffer_async(sb, diagnostics, version=None)
 
@@ -2086,7 +2086,7 @@ def m_textDocument_publishDiagnostics(self, params: PublishDiagnosticsParams) ->
         if isinstance(reason, str):
             return debug("ignoring unsuitable diagnostics for", uri, "reason:", reason)
         diagnostics = params["diagnostics"]
-        self.diagnostics.add_diagnostics_async(uri, diagnostics)
+        self.diagnostics[uri] = diagnostics
         mgr.on_diagnostics_updated()
         sb = self.get_session_buffer_for_uri_async(uri)
         if sb:
diff --git a/plugin/core/types.py b/plugin/core/types.py
index 105b3c58c..a78eb2d70 100644
--- a/plugin/core/types.py
+++ b/plugin/core/types.py
@@ -532,8 +532,8 @@ class Capabilities(DottedDict):
 
     def __init__(self, d: dict[str, Any] | None = None) -> None:
         super().__init__(d)
-        self._registrations: dict[str, set[str]] = {}
-        self._registration_options: dict[str, Any] = {}
+        self._registrations: dict[str, set[str | None]] = {}
+        self._registration_options: dict[str | None, Any] = {}
 
     def register(
         self,
@@ -577,6 +577,10 @@ def assign(self, d: ServerCapabilities) -> None:
         super().assign(cast(dict, d))
         if textsync:
             self.update(textsync)
+        diagnostic_provider_options = d.get('diagnosticProvider')
+        if diagnostic_provider_options:
+            self._registrations.setdefault('diagnosticProvider', set()).add(None)
+            self._registration_options[None] = diagnostic_provider_options
 
     def should_notify_did_open(self) -> bool:
         return "textDocumentSync.didOpen" in self
diff --git a/plugin/core/url.py b/plugin/core/url.py
index 1e9f09418..916027a5c 100644
--- a/plugin/core/url.py
+++ b/plugin/core/url.py
@@ -1,4 +1,7 @@
 from __future__ import annotations
+from .constants import INSTALLED_PACKAGES_PATH
+from .constants import PACKAGES_PATH
+from .protocol import DocumentUri
 from typing import Any
 from typing_extensions import deprecated
 from urllib.parse import urljoin
@@ -11,14 +14,18 @@
 import sublime
 
 
+def normalize_uri(uri: DocumentUri) -> DocumentUri:
+    return unparse_uri(parse_uri(uri))
+
+
 def filename_to_uri(file_name: str) -> str:
     """
     Convert a file name obtained from view.file_name() into an URI
     """
-    prefix = sublime.installed_packages_path()
+    prefix = INSTALLED_PACKAGES_PATH
     if file_name.startswith(prefix):
         return _to_resource_uri(file_name, prefix)
-    prefix = sublime.packages_path()
+    prefix = PACKAGES_PATH
     if file_name.startswith(prefix) and not os.path.exists(file_name):
         return _to_resource_uri(file_name, prefix)
     path = pathname2url(file_name)
diff --git a/plugin/core/windows.py b/plugin/core/windows.py
index 9234724ef..941e309aa 100644
--- a/plugin/core/windows.py
+++ b/plugin/core/windows.py
@@ -478,8 +478,9 @@ def update_diagnostics_panel_async(self) -> None:
         max_severity = userprefs().diagnostics_panel_include_severity_level
         contributions: OrderedDict[str, list[tuple[str, int | None, str | None, str | None]]] = OrderedDict()
         for session in self._sessions:
-            for (_, path), contribution in session.diagnostics.filter_map_diagnostics_async(
+            for uri, contribution in session.diagnostics.filter_map_diagnostics_async(
                     is_severity_included(max_severity), lambda _, diagnostic: format_diagnostic_for_panel(diagnostic)):
+                path = parse_uri(uri)[1]
                 seen = path in contributions
                 contributions.setdefault(path, []).extend(contribution)
                 if not seen:
diff --git a/plugin/goto_diagnostic.py b/plugin/goto_diagnostic.py
index 5f3018481..7dc9d08e5 100644
--- a/plugin/goto_diagnostic.py
+++ b/plugin/goto_diagnostic.py
@@ -64,9 +64,8 @@ def is_enabled(self, uri: DocumentUri | None = None, diagnostic: dict | None = N
                 return False
         max_severity = userprefs().diagnostics_panel_include_severity_level
         if uri:
-            parsed_uri = parse_uri(uri)
             return any(diagnostic for session in get_sessions(self.window)
-                       for diagnostic in session.diagnostics.diagnostics_by_parsed_uri(parsed_uri)
+                       for diagnostic in session.diagnostics[uri]
                        if is_severity_included(max_severity)(diagnostic))
         return any(diagnostic for session in get_sessions(self.window)
                    for diagnostics in session.diagnostics.values()
@@ -112,12 +111,13 @@ def get_list_items(self) -> tuple[list[sublime.ListInputItem], int]:
         severities_per_path: OrderedDict[ParsedUri, list[DiagnosticSeverity]] = OrderedDict()
         self.first_locations: dict[ParsedUri, tuple[Session, Location]] = dict()
         for session in get_sessions(self.window):
-            for parsed_uri, severity in session.diagnostics.filter_map_diagnostics_flat_async(
+            for uri, severity in session.diagnostics.filter_map_diagnostics_flat_async(
                     is_severity_included(max_severity), lambda _, diagnostic: diagnostic_severity(diagnostic)):
+                parsed_uri = parse_uri(uri)
                 severities_per_path.setdefault(parsed_uri, []).append(severity)
                 if parsed_uri not in self.first_locations:
                     severities_per_path.move_to_end(parsed_uri)
-                    diagnostics = session.diagnostics.diagnostics_by_parsed_uri(parsed_uri)
+                    diagnostics = session.diagnostics[uri]
                     if diagnostics:
                         self.first_locations[parsed_uri] = session, diagnostic_location(parsed_uri, diagnostics[0])
         # build items
@@ -190,6 +190,7 @@ def __init__(self, window: sublime.Window, view: sublime.View, uri: DocumentUri)
         self.window = window
         self.view = view
         self.sessions = list(get_sessions(window))
+        self.uri = uri
         self.parsed_uri = parse_uri(uri)
 
     def name(self) -> str:
@@ -199,8 +200,7 @@ def list_items(self) -> list[sublime.ListInputItem]:
         list_items: list[sublime.ListInputItem] = []
         max_severity = userprefs().diagnostics_panel_include_severity_level
         for i, session in enumerate(self.sessions):
-            for diagnostic in filter(is_severity_included(max_severity),
-                                     session.diagnostics.diagnostics_by_parsed_uri(self.parsed_uri)):
+            for diagnostic in filter(is_severity_included(max_severity), session.diagnostics[self.uri]):
                 lines = diagnostic["message"].splitlines()
                 first_line = lines[0] if lines else ""
                 if len(lines) > 1:
diff --git a/plugin/session_buffer.py b/plugin/session_buffer.py
index 003ad9622..7438dd863 100644
--- a/plugin/session_buffer.py
+++ b/plugin/session_buffer.py
@@ -132,7 +132,7 @@ def __init__(self, session_view: SessionViewProtocol, buffer_id: int, uri: Docum
         self.diagnostics_flags = 0
         self._diagnostics_are_visible = False
         self.document_diagnostic_needs_refresh = False
-        self._document_diagnostic_pending_request: PendingDocumentDiagnosticRequest | None = None
+        self._document_diagnostic_pending_requests: dict[str, PendingDocumentDiagnosticRequest | None] = {}
         self._last_synced_version = 0
         self._last_text_change_time = 0.0
         self._diagnostics_debouncer_async = DebouncerNonThreadSafe(async_thread=True)
@@ -322,9 +322,10 @@ def on_text_changed_async(self, view: sublime.View, change_count: int,
                           lambda: view.is_valid() and change_count == view.change_count(), async_thread=True)
 
     def _cancel_pending_requests_async(self) -> None:
-        if self._document_diagnostic_pending_request:
-            self.session.cancel_request(self._document_diagnostic_pending_request.request_id)
-            self._document_diagnostic_pending_request = None
+        for identifier, pending_request in self._document_diagnostic_pending_requests.items():
+            if pending_request:
+                self.session.cancel_request(pending_request.request_id)
+                self._document_diagnostic_pending_requests[identifier] = None
         if self.semantic_tokens.pending_response:
             self.session.cancel_request(self.semantic_tokens.pending_response)
             self.semantic_tokens.pending_response = None
@@ -503,50 +504,67 @@ def do_document_diagnostic_async(
             return
         if version is None:
             version = view.change_count()
-        if self._document_diagnostic_pending_request:
-            if self._document_diagnostic_pending_request.version == version and not forced_update:
-                return
-            self.session.cancel_request(self._document_diagnostic_pending_request.request_id)
-        params: DocumentDiagnosticParams = {'textDocument': text_document_identifier(view)}
-        identifier = self.get_capability("diagnosticProvider.identifier")
-        if identifier:
-            params['identifier'] = identifier
-        result_id = self.session.diagnostics_result_ids.get(self._last_known_uri)
-        if result_id is not None:
-            params['previousResultId'] = result_id
-        request_id = self.session.send_request_async(
-            Request.documentDiagnostic(params, view),
-            partial(self._on_document_diagnostic_async, version),
-            partial(self._on_document_diagnostic_error_async, version)
-        )
-        self._document_diagnostic_pending_request = PendingDocumentDiagnosticRequest(version, request_id)
+        # if self._document_diagnostic_pending_request:
+        #     if self._document_diagnostic_pending_request.version == version and not forced_update:
+        #         return
+        #     self.session.cancel_request(self._document_diagnostic_pending_request.request_id)
+        for identifier, pending_request in self._document_diagnostic_pending_requests.items():
+            if pending_request:
+                self.session.cancel_request(pending_request.request_id)
+                self._document_diagnostic_pending_requests[identifier] = None
+        _params: DocumentDiagnosticParams = {'textDocument': text_document_identifier(view)}
+        identifiers = set()
+        for registration in self.get_capability_2('diagnosticProvider'):
+            identifiers.add(registration.get('identifier', ''))
+        for identifier in identifiers:
+            params = _params.copy()
+            if identifier:
+                params['identifier'] = identifier
+            result_id = self.session.diagnostics_result_ids.get(self._last_known_uri)
+            if result_id is not None:
+                params['previousResultId'] = result_id
+            request_id = self.session.send_request_async(
+                Request.documentDiagnostic(params, view),
+                partial(self._on_document_diagnostic_async, identifier, version),
+                partial(self._on_document_diagnostic_error_async, identifier, version)
+            )
+            self._document_diagnostic_pending_requests[identifier] = \
+                PendingDocumentDiagnosticRequest(version, request_id)
 
-    def _on_document_diagnostic_async(self, version: int, response: DocumentDiagnosticReport) -> None:
-        self._document_diagnostic_pending_request = None
-        self._if_view_unchanged(self._apply_document_diagnostic_async, version)(response)
+    def _on_document_diagnostic_async(self, identifier: str, version: int, response: DocumentDiagnosticReport) -> None:
+        self._document_diagnostic_pending_requests[identifier] = None
+        view = self.some_view()
+        if view and view.change_count() == version:
+            self._apply_document_diagnostic_async(identifier, version, response)
+            mgr = self.session.manager()
+            if mgr:
+                mgr.on_diagnostics_updated()
 
     def _apply_document_diagnostic_async(
-        self, view: sublime.View | None, response: DocumentDiagnosticReport
+        self, identifier: str, version: int, response: DocumentDiagnosticReport
     ) -> None:
         self.session.diagnostics_result_ids[self._last_known_uri] = response.get('resultId')
         if is_full_document_diagnostic_report(response):
-            self.session.m_textDocument_publishDiagnostics(
-                {'uri': self._last_known_uri, 'diagnostics': response['items']})
+            self.session.diagnostics[(self._last_known_uri, identifier)] = response['items']
+            self.on_diagnostics_async(
+                self.session.diagnostics[self._last_known_uri], version, self.session.session_views_by_visibility()[0])
         for uri, diagnostic_report in response.get('relatedDocuments', {}):
             sb = self.session.get_session_buffer_for_uri_async(uri)
             if sb:
                 cast(SessionBuffer, sb)._apply_document_diagnostic_async(
-                    None, cast(DocumentDiagnosticReport, diagnostic_report))
+                    identifier, version, cast(DocumentDiagnosticReport, diagnostic_report))
 
-    def _on_document_diagnostic_error_async(self, version: int, error: ResponseError) -> None:
-        self._document_diagnostic_pending_request = None
+    def _on_document_diagnostic_error_async(self, identifier: str, version: int, error: ResponseError) -> None:
+        self._document_diagnostic_pending_requests[identifier] = None
         if error['code'] == LSPErrorCodes.ServerCancelled:
             data = error.get('data')
             if is_diagnostic_server_cancellation_data(data) and data['retriggerRequest']:
                 # Retrigger the request after a short delay, but only if there were no additional changes to the buffer
                 # (in that case the request will be retriggered automatically anyway)
-                sublime.set_timeout_async(
-                    lambda: self._if_view_unchanged(self.do_document_diagnostic_async, version)(version), 500)
+                pass
+                # TODO retrigger for this identifier
+                # sublime.set_timeout_async(
+                #     lambda: self._if_view_unchanged(self.do_document_diagnostic_async, version)(version), 500)
 
     def set_document_diagnostic_pending_refresh(self, needs_refresh: bool = True) -> None:
         self.document_diagnostic_needs_refresh = needs_refresh

From 0b6669b33eb064dc25fdfddce5d2c60a44ada949 Mon Sep 17 00:00:00 2001
From: Janos Wortmann <jwortmann@outlook.com>
Date: Mon, 2 Dec 2024 13:52:33 +0100
Subject: [PATCH 3/5] Fix some bugs

---
 plugin/core/diagnostics_storage.py |  2 ++
 plugin/core/sessions.py            | 26 +++++++++++++++++++-------
 plugin/session_buffer.py           |  4 ++--
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/plugin/core/diagnostics_storage.py b/plugin/core/diagnostics_storage.py
index 8f8d89a99..941181b7b 100644
--- a/plugin/core/diagnostics_storage.py
+++ b/plugin/core/diagnostics_storage.py
@@ -53,9 +53,11 @@ def __len__(self) -> int:
         return len(self._uris)
 
     def register(self, identifier: str) -> None:
+        """ Register an identifier for a diagnostics stream. """
         self._identifiers.add(identifier)
 
     def unregister(self, identifier: str) -> None:
+        """ Unregister an identifier for a diagnostics stream. """
         self._identifiers.discard(identifier)
 
     def filter_map_diagnostics_async(
diff --git a/plugin/core/sessions.py b/plugin/core/sessions.py
index 93bc84c2a..4c5aee1b2 100644
--- a/plugin/core/sessions.py
+++ b/plugin/core/sessions.py
@@ -1281,7 +1281,7 @@ def __init__(self, manager: Manager, logger: Logger, workspace_folders: list[Wor
         self.state = ClientStates.STARTING
         self.capabilities = Capabilities()
         self.diagnostics = DiagnosticsStorage()
-        self.diagnostics_result_ids: dict[DocumentUri, str | None] = {}
+        self.diagnostics_result_ids: dict[tuple[DocumentUri, str], str | None] = {}
         self.workspace_diagnostics_pending_response: int | None = None
         self.exiting = False
         self._registrations: dict[str, _RegistrationData] = {}
@@ -1544,7 +1544,11 @@ def initialize_async(
             Request.initialize(params), self._handle_initialize_success, self._handle_initialize_error)
 
     def _handle_initialize_success(self, result: InitializeResult) -> None:
-        self.capabilities.assign(result.get('capabilities', dict()))
+        capabilities = result.get('capabilities', dict())
+        self.capabilities.assign(capabilities)
+        if diagnostic_provider := capabilities.get('diagnosticProvider'):
+            if identifier := diagnostic_provider.get('identifier'):
+                self.diagnostics.register(identifier)
         if self._workspace_folders and not self._supports_workspace_folders():
             self._workspace_folders = self._workspace_folders[:1]
         self.state = ClientStates.READY
@@ -1937,12 +1941,13 @@ def session_views_by_visibility(self) -> tuple[set[SessionViewProtocol], set[Ses
     # --- Workspace Pull Diagnostics -----------------------------------------------------------------------------------
 
     def do_workspace_diagnostics_async(self) -> None:
+        # TODO consider separate diagnostic streams (identifiers)
         if self.workspace_diagnostics_pending_response:
             # The server is probably leaving the request open intentionally, in order to continuously stream updates via
             # $/progress notifications.
             return
         previous_result_ids: list[PreviousResultId] = [
-            {'uri': uri, 'value': result_id} for uri, result_id in self.diagnostics_result_ids.items()
+            {'uri': uri, 'value': result_id} for (uri, _), result_id in self.diagnostics_result_ids.items()
             if result_id is not None
         ]
         params: WorkspaceDiagnosticParams = {'previousResultIds': previous_result_ids}
@@ -1951,11 +1956,11 @@ def do_workspace_diagnostics_async(self) -> None:
             params['identifier'] = identifier
         self.workspace_diagnostics_pending_response = self.send_request_async(
             Request.workspaceDiagnostic(params),
-            self._on_workspace_diagnostics_async,
+            functools.partial(self._on_workspace_diagnostics_async, identifier or ''),
             self._on_workspace_diagnostics_error_async)
 
     def _on_workspace_diagnostics_async(
-        self, response: WorkspaceDiagnosticReport, reset_pending_response: bool = True
+        self, identifier: str, response: WorkspaceDiagnosticReport, reset_pending_response: bool = True
     ) -> None:
         if reset_pending_response:
             self.workspace_diagnostics_pending_response = None
@@ -1984,7 +1989,7 @@ def _on_workspace_diagnostics_async(
                 sb = self.get_session_buffer_for_uri_async(uri)
                 if sb and sb.version != version:
                     continue
-            self.diagnostics_result_ids[uri] = diagnostic_report.get('resultId')
+            self.diagnostics_result_ids[(uri, identifier)] = diagnostic_report.get('resultId')
             if is_workspace_full_document_diagnostic_report(diagnostic_report):
                 self.m_textDocument_publishDiagnostics({'uri': uri, 'diagnostics': diagnostic_report['items']})
 
@@ -2135,6 +2140,9 @@ def m_client_registerCapability(self, params: RegistrationParams, request_id: An
                         watcher = self._watcher_impl.create(folder.path, [pattern], kind, ignores, self)
                         file_watchers.append(watcher)
                 self._dynamic_file_watchers[registration_id] = file_watchers
+            elif capability_path == 'diagnosticProvider':
+                if (identifier := options.get('identifier')) is not None:
+                    self.diagnostics.register(identifier)
         self.send_response(Response(request_id, None))
 
     def m_client_unregisterCapability(self, params: UnregistrationParams, request_id: Any) -> None:
@@ -2157,6 +2165,9 @@ def m_client_unregisterCapability(self, params: UnregistrationParams, request_id
                 if isinstance(discarded, dict):
                     for sv in self.session_views_async():
                         sv.on_capability_removed_async(registration_id, discarded)
+            if capability_path == 'diagnosticProvider':
+                if data and (identifier := data.options.get('identifier')) is not None:
+                    self.diagnostics.unregister(identifier)
         self.send_response(Response(request_id, None))
 
     def m_window_showDocument(self, params: Any, request_id: Any) -> None:
@@ -2210,8 +2221,9 @@ def m___progress(self, params: ProgressParams) -> None:
             request_id = int(token[len(_PARTIAL_RESULT_PROGRESS_PREFIX):])
             request = self._response_handlers[request_id][0]
             if request.method == "workspace/diagnostic":
+                # TODO somehow get the identifier (probably needs to be stored based on progress token)
                 self._on_workspace_diagnostics_async(
-                    cast(WorkspaceDiagnosticReport, value), reset_pending_response=False)
+                    '', cast(WorkspaceDiagnosticReport, value), reset_pending_response=False)
             return
         # Work Done Progress
         # https://microsoft.github.io/language-server-protocol/specifications/specification-current/#workDoneProgress
diff --git a/plugin/session_buffer.py b/plugin/session_buffer.py
index 7438dd863..5197a3068 100644
--- a/plugin/session_buffer.py
+++ b/plugin/session_buffer.py
@@ -520,7 +520,7 @@ def do_document_diagnostic_async(
             params = _params.copy()
             if identifier:
                 params['identifier'] = identifier
-            result_id = self.session.diagnostics_result_ids.get(self._last_known_uri)
+            result_id = self.session.diagnostics_result_ids.get((self._last_known_uri, identifier))
             if result_id is not None:
                 params['previousResultId'] = result_id
             request_id = self.session.send_request_async(
@@ -543,7 +543,7 @@ def _on_document_diagnostic_async(self, identifier: str, version: int, response:
     def _apply_document_diagnostic_async(
         self, identifier: str, version: int, response: DocumentDiagnosticReport
     ) -> None:
-        self.session.diagnostics_result_ids[self._last_known_uri] = response.get('resultId')
+        self.session.diagnostics_result_ids[(self._last_known_uri, identifier)] = response.get('resultId')
         if is_full_document_diagnostic_report(response):
             self.session.diagnostics[(self._last_known_uri, identifier)] = response['items']
             self.on_diagnostics_async(

From ca6c0149d53e44ba5e2a9d0018ade3035f8400c5 Mon Sep 17 00:00:00 2001
From: Janos Wortmann <jwortmann@outlook.com>
Date: Thu, 12 Dec 2024 11:22:33 +0100
Subject: [PATCH 4/5] Attempt to fix workspace pull diagnostics

---
 plugin/core/diagnostics_storage.py |  9 ++--
 plugin/core/sessions.py            | 83 ++++++++++++++++++------------
 plugin/session_buffer.py           | 27 ++++++----
 3 files changed, 72 insertions(+), 47 deletions(-)

diff --git a/plugin/core/diagnostics_storage.py b/plugin/core/diagnostics_storage.py
index 941181b7b..986692960 100644
--- a/plugin/core/diagnostics_storage.py
+++ b/plugin/core/diagnostics_storage.py
@@ -18,8 +18,8 @@ class DiagnosticsStorage(MutableMapping):
 
     def __init__(self) -> None:
         super().__init__()
-        self._d: dict[tuple[DocumentUri, str], list[Diagnostic]] = dict()
-        self._identifiers = {''}
+        self._d: dict[tuple[DocumentUri, str | None], list[Diagnostic]] = dict()
+        self._identifiers: set[str | None] = {None}
         self._uris: set[DocumentUri] = set()
 
     def __getitem__(self, key: DocumentUri, /) -> list[Diagnostic]:
@@ -29,8 +29,9 @@ def __getitem__(self, key: DocumentUri, /) -> list[Diagnostic]:
             key=lambda diagnostic: Point.from_lsp(diagnostic['range']['start'])
         )
 
-    def __setitem__(self, key: DocumentUri | tuple[DocumentUri, str], value: list[Diagnostic], /) -> None:
-        uri, identifier = (normalize_uri(key), '') if isinstance(key, DocumentUri) else (normalize_uri(key[0]), key[1])
+    def __setitem__(self, key: DocumentUri | tuple[DocumentUri, str | None], value: list[Diagnostic], /) -> None:
+        uri, identifier = (normalize_uri(key), None) if isinstance(key, DocumentUri) else \
+            (normalize_uri(key[0]), key[1])
         if identifier not in self._identifiers:
             raise ValueError(f'identifier {identifier} must be registered first')
         if value:
diff --git a/plugin/core/sessions.py b/plugin/core/sessions.py
index 665d74088..c37fdec88 100644
--- a/plugin/core/sessions.py
+++ b/plugin/core/sessions.py
@@ -31,6 +31,7 @@
 from .protocol import DiagnosticTag
 from .protocol import DidChangeWatchedFilesRegistrationOptions
 from .protocol import DidChangeWorkspaceFoldersParams
+from .protocol import DocumentDiagnosticReport
 from .protocol import DocumentDiagnosticReportKind
 from .protocol import DocumentLink
 from .protocol import DocumentUri
@@ -663,6 +664,11 @@ def on_diagnostics_async(
     ) -> None:
         ...
 
+    def on_document_diagnostic_async(
+        self, identifier: str | None, version: int, response: DocumentDiagnosticReport
+    ) -> None:
+        ...
+
     def get_document_link_at_point(self, view: sublime.View, point: int) -> DocumentLink | None:
         ...
 
@@ -1282,8 +1288,8 @@ def __init__(self, manager: Manager, logger: Logger, workspace_folders: list[Wor
         self.state = ClientStates.STARTING
         self.capabilities = Capabilities()
         self.diagnostics = DiagnosticsStorage()
-        self.diagnostics_result_ids: dict[tuple[DocumentUri, str], str | None] = {}
-        self.workspace_diagnostics_pending_response: int | None = None
+        self.diagnostics_result_ids: dict[tuple[DocumentUri, str | None], str | None] = {}
+        self.workspace_diagnostics_pending_responses: dict[str | None, int | None] = {}
         self.exiting = False
         self._registrations: dict[str, _RegistrationData] = {}
         self._init_callback: InitCallback | None = None
@@ -1451,15 +1457,23 @@ def can_handle(self, view: sublime.View, scheme: str, capability: str | None, in
                 return self.has_capability(capability)
         return False
 
+    @deprecated("Use has_provider instead")
     def has_capability(self, capability: str) -> bool:
         value = self.get_capability(capability)
         return value is not False and value is not None
 
+    @deprecated("Use get_providers instead")
     def get_capability(self, capability: str) -> Any | None:
         if self.config.is_disabled_capability(capability):
             return None
         return self.capabilities.get(capability)
 
+    def get_providers(self, capability_name: str) -> list[Any]:
+        return self.capabilities.get_all(capability_name)
+
+    def has_provider(self, capability_name: str) -> bool:
+        return bool(self.get_providers(capability_name))
+
     def should_notify_did_open(self) -> bool:
         return self.capabilities.should_notify_did_open()
 
@@ -1944,29 +1958,29 @@ def session_views_by_visibility(self) -> tuple[set[SessionViewProtocol], set[Ses
     # --- Workspace Pull Diagnostics -----------------------------------------------------------------------------------
 
     def do_workspace_diagnostics_async(self) -> None:
-        # TODO consider separate diagnostic streams (identifiers)
-        if self.workspace_diagnostics_pending_response:
-            # The server is probably leaving the request open intentionally, in order to continuously stream updates via
-            # $/progress notifications.
-            return
-        previous_result_ids: list[PreviousResultId] = [
-            {'uri': uri, 'value': result_id} for (uri, _), result_id in self.diagnostics_result_ids.items()
-            if result_id is not None
-        ]
-        params: WorkspaceDiagnosticParams = {'previousResultIds': previous_result_ids}
-        identifier = self.get_capability("diagnosticProvider.identifier")
-        if identifier:
-            params['identifier'] = identifier
-        self.workspace_diagnostics_pending_response = self.send_request_async(
-            Request.workspaceDiagnostic(params),
-            functools.partial(self._on_workspace_diagnostics_async, identifier or ''),
-            self._on_workspace_diagnostics_error_async)
+        for provider in self.get_providers('diagnosticProvider'):
+            identifier = provider.get('identifier')
+            if self.workspace_diagnostics_pending_responses[identifier]:
+                # The server is probably leaving the request open intentionally, in order to continuously stream updates
+                # via $/progress notifications.
+                return
+            previous_result_ids: list[PreviousResultId] = [
+                {'uri': uri, 'value': result_id} for (uri, id_), result_id in self.diagnostics_result_ids.items()
+                if id_ == identifier and result_id is not None
+            ]
+            params: WorkspaceDiagnosticParams = {'previousResultIds': previous_result_ids}
+            if identifier:
+                params['identifier'] = identifier
+            self.workspace_diagnostics_pending_responses[identifier] = self.send_request_async(
+                Request.workspaceDiagnostic(params),
+                functools.partial(self._on_workspace_diagnostics_async, identifier),
+                functools.partial(self._on_workspace_diagnostics_error_async, identifier))
 
     def _on_workspace_diagnostics_async(
-        self, identifier: str, response: WorkspaceDiagnosticReport, reset_pending_response: bool = True
+        self, identifier: str | None, response: WorkspaceDiagnosticReport, reset_pending_response: bool = True
     ) -> None:
         if reset_pending_response:
-            self.workspace_diagnostics_pending_response = None
+            self.workspace_diagnostics_pending_responses[identifier] = None
         if not response['items']:
             return
         window = sublime.active_window()
@@ -1983,20 +1997,23 @@ def _on_workspace_diagnostics_async(
                 uri = unparse_uri((scheme, path))
             # Note: 'version' is a mandatory field, but some language servers have serialization bugs with null values.
             version = diagnostic_report.get('version')
-            # Skip if outdated
-            # Note: this is just a necessary, but not a sufficient condition to decide whether the diagnostics for this
-            # file are likely not accurate anymore, because changes in another file in the meanwhile could have affected
-            # the diagnostics in this file. If this is the case, a new request is already queued, or updated partial
-            # results are expected to be streamed by the server.
-            if isinstance(version, int):
+            if version is not None:
                 sb = self.get_session_buffer_for_uri_async(uri)
-                if sb and sb.version != version:
+                if not sb:
+                    # There should always be a SessionBuffer if version != None
                     continue
+                if sb.version != version:
+                    # Skip if outdated
+                    continue
+                if is_workspace_full_document_diagnostic_report(diagnostic_report):
+                    diagnostic_report = cast(DocumentDiagnosticReport, diagnostic_report)
+                    sb.on_document_diagnostic_async(identifier, version, diagnostic_report)
+            else:
+                # TODO support diagnostics for unopened docuements (version == None)
+                pass
             self.diagnostics_result_ids[(uri, identifier)] = diagnostic_report.get('resultId')
-            if is_workspace_full_document_diagnostic_report(diagnostic_report):
-                self.m_textDocument_publishDiagnostics({'uri': uri, 'diagnostics': diagnostic_report['items']})
 
-    def _on_workspace_diagnostics_error_async(self, error: ResponseError) -> None:
+    def _on_workspace_diagnostics_error_async(self, identifier: str | None, error: ResponseError) -> None:
         if error['code'] == LSPErrorCodes.ServerCancelled:
             data = error.get('data')
             if is_diagnostic_server_cancellation_data(data) and data['retriggerRequest']:
@@ -2005,12 +2022,12 @@ def _on_workspace_diagnostics_error_async(self, error: ResponseError) -> None:
                 # infinite cycles of cancel -> retrigger, in case the server is busy.
 
                 def _retrigger_request() -> None:
-                    self.workspace_diagnostics_pending_response = None
+                    self.workspace_diagnostics_pending_responses[identifier] = None
                     self.do_workspace_diagnostics_async()
 
                 sublime.set_timeout_async(_retrigger_request, WORKSPACE_DIAGNOSTICS_TIMEOUT)
                 return
-        self.workspace_diagnostics_pending_response = None
+        self.workspace_diagnostics_pending_responses[identifier] = None
 
     # --- server request handlers --------------------------------------------------------------------------------------
 
diff --git a/plugin/session_buffer.py b/plugin/session_buffer.py
index 5197a3068..d982642c1 100644
--- a/plugin/session_buffer.py
+++ b/plugin/session_buffer.py
@@ -132,7 +132,7 @@ def __init__(self, session_view: SessionViewProtocol, buffer_id: int, uri: Docum
         self.diagnostics_flags = 0
         self._diagnostics_are_visible = False
         self.document_diagnostic_needs_refresh = False
-        self._document_diagnostic_pending_requests: dict[str, PendingDocumentDiagnosticRequest | None] = {}
+        self._document_diagnostic_pending_requests: dict[str | None, PendingDocumentDiagnosticRequest | None] = {}
         self._last_synced_version = 0
         self._last_text_change_time = 0.0
         self._diagnostics_debouncer_async = DebouncerNonThreadSafe(async_thread=True)
@@ -265,21 +265,26 @@ def unregister_capability_async(
         for sv in self.session_views:
             sv.on_capability_removed_async(registration_id, discarded)
 
+    @deprecated("Use get_providers instead")
     def get_capability(self, capability_path: str) -> Any | None:
         if self.session.config.is_disabled_capability(capability_path):
             return None
         value = self.capabilities.get(capability_path)
         return value if value is not None else self.session.capabilities.get(capability_path)
 
-    def get_capability_2(self, capability_path: str) -> list[Any]:
-        if self.session.config.is_disabled_capability(capability_path):
+    def get_providers(self, capability_name: str) -> list[Any]:
+        if self.session.config.is_disabled_capability(capability_name):
             return []
-        return self.capabilities.get_all(capability_path) + self.session.capabilities.get_all(capability_path)
+        return self.capabilities.get_all(capability_name) + self.session.capabilities.get_all(capability_name)
 
+    @deprecated("Use has_provider instead")
     def has_capability(self, capability_path: str) -> bool:
         value = self.get_capability(capability_path)
         return value is not False and value is not None
 
+    def has_provider(self, capability_name: str) -> bool:
+        return bool(self.get_providers(capability_name))
+
     def text_sync_kind(self) -> TextDocumentSyncKind:
         value = self.capabilities.text_sync_kind()
         return value if value != TextDocumentSyncKind.None_ else self.session.text_sync_kind()
@@ -514,8 +519,8 @@ def do_document_diagnostic_async(
                 self._document_diagnostic_pending_requests[identifier] = None
         _params: DocumentDiagnosticParams = {'textDocument': text_document_identifier(view)}
         identifiers = set()
-        for registration in self.get_capability_2('diagnosticProvider'):
-            identifiers.add(registration.get('identifier', ''))
+        for provider in self.get_providers('diagnosticProvider'):
+            identifiers.add(provider.get('identifier', ''))
         for identifier in identifiers:
             params = _params.copy()
             if identifier:
@@ -525,13 +530,15 @@ def do_document_diagnostic_async(
                 params['previousResultId'] = result_id
             request_id = self.session.send_request_async(
                 Request.documentDiagnostic(params, view),
-                partial(self._on_document_diagnostic_async, identifier, version),
+                partial(self.on_document_diagnostic_async, identifier, version),
                 partial(self._on_document_diagnostic_error_async, identifier, version)
             )
             self._document_diagnostic_pending_requests[identifier] = \
                 PendingDocumentDiagnosticRequest(version, request_id)
 
-    def _on_document_diagnostic_async(self, identifier: str, version: int, response: DocumentDiagnosticReport) -> None:
+    def on_document_diagnostic_async(
+        self, identifier: str | None, version: int, response: DocumentDiagnosticReport
+    ) -> None:
         self._document_diagnostic_pending_requests[identifier] = None
         view = self.some_view()
         if view and view.change_count() == version:
@@ -541,7 +548,7 @@ def _on_document_diagnostic_async(self, identifier: str, version: int, response:
                 mgr.on_diagnostics_updated()
 
     def _apply_document_diagnostic_async(
-        self, identifier: str, version: int, response: DocumentDiagnosticReport
+        self, identifier: str | None, version: int, response: DocumentDiagnosticReport
     ) -> None:
         self.session.diagnostics_result_ids[(self._last_known_uri, identifier)] = response.get('resultId')
         if is_full_document_diagnostic_report(response):
@@ -554,7 +561,7 @@ def _apply_document_diagnostic_async(
                 cast(SessionBuffer, sb)._apply_document_diagnostic_async(
                     identifier, version, cast(DocumentDiagnosticReport, diagnostic_report))
 
-    def _on_document_diagnostic_error_async(self, identifier: str, version: int, error: ResponseError) -> None:
+    def _on_document_diagnostic_error_async(self, identifier: str | None, version: int, error: ResponseError) -> None:
         self._document_diagnostic_pending_requests[identifier] = None
         if error['code'] == LSPErrorCodes.ServerCancelled:
             data = error.get('data')

From ef3f6980f289e111b2b7fdbf1b8b80c723f593f0 Mon Sep 17 00:00:00 2001
From: Janos Wortmann <jwortmann@outlook.com>
Date: Mon, 16 Dec 2024 12:53:57 +0100
Subject: [PATCH 5/5] Fix missing identifier case

---
 plugin/session_buffer.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/plugin/session_buffer.py b/plugin/session_buffer.py
index d982642c1..21a09af51 100644
--- a/plugin/session_buffer.py
+++ b/plugin/session_buffer.py
@@ -518,9 +518,11 @@ def do_document_diagnostic_async(
                 self.session.cancel_request(pending_request.request_id)
                 self._document_diagnostic_pending_requests[identifier] = None
         _params: DocumentDiagnosticParams = {'textDocument': text_document_identifier(view)}
+        # Not all diagnostic streams (identifiers) which are stored in the Session's DiagnosticStorage must necessarily
+        # be applicable to this SessionBuffer in case only a subset of them was registered for this DocumentUri.
         identifiers = set()
         for provider in self.get_providers('diagnosticProvider'):
-            identifiers.add(provider.get('identifier', ''))
+            identifiers.add(provider.get('identifier'))
         for identifier in identifiers:
             params = _params.copy()
             if identifier: