Refactor EditingToolWithDiagnostics

* Add context manager to simplify usage, improving output handling
* Add util module ls_diagnostics, reifying DiagnosticsDiff

scripts/demo_diagnostics.py

@@ -16,7 +16,13 @@ from serena.agent import SerenaAgent
 from serena.config.serena_config import LanguageBackend, ProjectConfig, RegisteredProject, SerenaConfig
 from serena.constants import REPO_ROOT
 from serena.project import Project
-from serena.tools import CreateTextFileTool, GetDiagnosticsForFileTool, GetDiagnosticsForSymbolTool, ReplaceContentTool, Tool
+from serena.tools import (
+    CreateTextFileTool,
+    EditingToolWithDiagnostics,
+    GetDiagnosticsForFileTool,
+    GetDiagnosticsForSymbolTool,
+    ReplaceContentTool,
+)
 from solidlsp.ls_config import Language
 
 SEPARATOR = "=" * 80
@@ -69,7 +75,7 @@ def parse_edit_diagnostics_result(result: str) -> dict:
 
 
 if __name__ == "__main__":
-    Tool._ENABLE_DIAGNOSTICS = True
+    EditingToolWithDiagnostics.ENABLE_DIAGNOSTICS = True
 
     temp_dir = Path(tempfile.mkdtemp(prefix="serena_demo_", dir=REPO_PATH))
     temp_file = temp_dir / "demo_temp_diagnostics.py"
@@ -166,5 +172,5 @@ if __name__ == "__main__":
         assert "missing_one" in diagnostics_after_second_edit_json, diagnostics_after_second_edit
         assert "missing_two" in diagnostics_after_second_edit_json, diagnostics_after_second_edit
     finally:
-        agent.shutdown()
         shutil.rmtree(temp_dir, ignore_errors=True)
+        agent.shutdown()

src/serena/code_editor.py

@@ -4,7 +4,6 @@ import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Iterator, Reversible
 from contextlib import contextmanager
-from dataclasses import dataclass
 from typing import Generic, TypeVar, cast
 
 from serena.jetbrains.jetbrains_plugin_client import JetBrainsPluginClient
@@ -14,17 +13,12 @@ from solidlsp.ls import LSPFileBuffer
 from solidlsp.ls_utils import PathUtils, TextUtils
 
 from .project import Project
+from .util.ls_diagnostics import EditedFilePath
 
 log = logging.getLogger(__name__)
 TSymbol = TypeVar("TSymbol", bound=Symbol)
 
 
-@dataclass(frozen=True)
-class EditedFilePath:
-    before_relative_path: str
-    after_relative_path: str
-
-
 class CodeEditor(Generic[TSymbol], ABC):
     def __init__(self, project: Project) -> None:
         self.project_root = project.project_root

src/serena/tools/file_tools.py

@@ -12,7 +12,6 @@ from fnmatch import fnmatch
 from pathlib import Path
 from typing import Literal
 
-from serena.code_editor import EditedFilePath
 from serena.tools import SUCCESS_RESULT, EditedFileContext, EditingToolWithDiagnostics, Tool, ToolMarkerOptional
 from serena.util.file_system import scan_directory
 from serena.util.text_utils import ContentReplacer, search_files
@@ -62,30 +61,27 @@ class CreateTextFileTool(EditingToolWithDiagnostics):
         :param content: the (appropriately encoded) content to write to the file
         :return: a message indicating success or failure
         """
-        # capturing diagnostics before the edit
-        edited_file_paths = [EditedFilePath(relative_path, relative_path)]
-        diagnostics_snapshot = self._capture_published_lsp_diagnostics_snapshot(edited_file_paths)
+        with self.DiagnosticsContext(self, relative_path) as diagnostics_context:
+            # validating the destination path
+            project_root = self.get_project_root()
+            abs_path = (Path(project_root) / relative_path).resolve()
+            will_overwrite_existing = abs_path.exists()
 
-        # validating the destination path
-        project_root = self.get_project_root()
-        abs_path = (Path(project_root) / relative_path).resolve()
-        will_overwrite_existing = abs_path.exists()
+            if will_overwrite_existing:
+                self.project.validate_relative_path(relative_path, require_not_ignored=True)
+            else:
+                assert abs_path.is_relative_to(self.get_project_root()), (
+                    f"Cannot create file outside of the project directory, got {relative_path=}"
+                )
 
-        if will_overwrite_existing:
-            self.project.validate_relative_path(relative_path, require_not_ignored=True)
-        else:
-            assert abs_path.is_relative_to(self.get_project_root()), (
-                f"Cannot create file outside of the project directory, got {relative_path=}"
-            )
+            # writing the file
+            abs_path.parent.mkdir(parents=True, exist_ok=True)
+            abs_path.write_text(content, encoding=self.project.project_config.encoding, newline=self.project.line_ending.newline_str)
+            answer = f"File created: {relative_path}."
+            if will_overwrite_existing:
+                answer += " Overwrote existing file."
 
-        # writing the file
-        abs_path.parent.mkdir(parents=True, exist_ok=True)
-        abs_path.write_text(content, encoding=self.project.project_config.encoding, newline=self.project.line_ending.newline_str)
-        answer = f"File created: {relative_path}."
-        if will_overwrite_existing:
-            answer += " Overwrote existing file."
-
-        return self._format_lsp_edit_result_with_new_diagnostics(answer, edited_file_paths, diagnostics_snapshot)
+            return diagnostics_context.format_result(answer)
 
 
 class ListDirTool(Tool):
@@ -220,19 +216,14 @@ class ReplaceContentTool(EditingToolWithDiagnostics):
         Performs the replacement, with additional options not exposed in the tool.
         This function can be used internally by other tools.
         """
-        # capturing diagnostics before the edit
-        edited_file_paths = [EditedFilePath(relative_path, relative_path)]
-        diagnostics_snapshot = self._capture_published_lsp_diagnostics_snapshot(edited_file_paths)
-
-        # applying the content replacement
-        self.project.validate_relative_path(relative_path, require_not_ignored=require_not_ignored)
-        with EditedFileContext(relative_path, self.create_code_editor()) as context:
-            original_content = context.get_original_content()
-            replacer = ContentReplacer(mode=mode, allow_multiple_occurrences=allow_multiple_occurrences)
-            updated_content = replacer.replace(original_content, needle, repl)
-            context.set_updated_content(updated_content)
-
-        return self._format_lsp_edit_result_with_new_diagnostics(SUCCESS_RESULT, edited_file_paths, diagnostics_snapshot)
+        with self.DiagnosticsContext(self, relative_path) as diagnostics_context:
+            self.project.validate_relative_path(relative_path, require_not_ignored=require_not_ignored)
+            with EditedFileContext(relative_path, self.create_code_editor()) as context:
+                original_content = context.get_original_content()
+                replacer = ContentReplacer(mode=mode, allow_multiple_occurrences=allow_multiple_occurrences)
+                updated_content = replacer.replace(original_content, needle, repl)
+                context.set_updated_content(updated_content)
+            return diagnostics_context.format_result(SUCCESS_RESULT)
 
 
 class DeleteLinesTool(EditingToolWithDiagnostics, ToolMarkerOptional):
@@ -255,15 +246,10 @@ class DeleteLinesTool(EditingToolWithDiagnostics, ToolMarkerOptional):
         :param start_line: the 0-based index of the first line to be deleted
         :param end_line: the 0-based index of the last line to be deleted
         """
-        # capturing diagnostics before the edit
-        edited_file_paths = [EditedFilePath(relative_path, relative_path)]
-        diagnostics_snapshot = self._capture_published_lsp_diagnostics_snapshot(edited_file_paths)
-
-        # applying the edit
-        code_editor = self.create_code_editor()
-        code_editor.delete_lines(relative_path, start_line, end_line)
-
-        return self._format_lsp_edit_result_with_new_diagnostics(SUCCESS_RESULT, edited_file_paths, diagnostics_snapshot)
+        with self.DiagnosticsContext(self, relative_path) as diagnostics_context:
+            code_editor = self.create_code_editor()
+            code_editor.delete_lines(relative_path, start_line, end_line)
+            return diagnostics_context.format_result(SUCCESS_RESULT)
 
 
 class ReplaceLinesTool(EditingToolWithDiagnostics, ToolMarkerOptional):
@@ -292,16 +278,12 @@ class ReplaceLinesTool(EditingToolWithDiagnostics, ToolMarkerOptional):
         if not content.endswith("\n"):
             content += "\n"
 
-        # capturing diagnostics before the edit
-        edited_file_paths = [EditedFilePath(relative_path, relative_path)]
-        diagnostics_snapshot = self._capture_published_lsp_diagnostics_snapshot(edited_file_paths)
+        with self.DiagnosticsContext(self, relative_path) as diagnostics_context:
+            code_editor = self.create_code_editor()
+            code_editor.delete_lines(relative_path, start_line, end_line)
+            code_editor.insert_at_line(relative_path, start_line, content)
 
-        # applying the replacement
-        code_editor = self.create_code_editor()
-        code_editor.delete_lines(relative_path, start_line, end_line)
-        code_editor.insert_at_line(relative_path, start_line, content)
-
-        return self._format_lsp_edit_result_with_new_diagnostics(SUCCESS_RESULT, edited_file_paths, diagnostics_snapshot)
+            return diagnostics_context.format_result(SUCCESS_RESULT)
 
 
 class InsertAtLineTool(EditingToolWithDiagnostics, ToolMarkerOptional):
@@ -329,15 +311,11 @@ class InsertAtLineTool(EditingToolWithDiagnostics, ToolMarkerOptional):
         if not content.endswith("\n"):
             content += "\n"
 
-        # capturing diagnostics before the edit
-        edited_file_paths = [EditedFilePath(relative_path, relative_path)]
-        diagnostics_snapshot = self._capture_published_lsp_diagnostics_snapshot(edited_file_paths)
+        with self.DiagnosticsContext(self, relative_path) as diagnostics_context:
+            code_editor = self.create_code_editor()
+            code_editor.insert_at_line(relative_path, line, content)
 
-        # applying the insertion
-        code_editor = self.create_code_editor()
-        code_editor.insert_at_line(relative_path, line, content)
-
-        return self._format_lsp_edit_result_with_new_diagnostics(SUCCESS_RESULT, edited_file_paths, diagnostics_snapshot)
+            return diagnostics_context.format_result(SUCCESS_RESULT)
 
 
 class SearchForPatternTool(Tool):

src/serena/tools/symbol_tools.py

@@ -9,7 +9,6 @@ from collections import Counter, defaultdict
 from collections.abc import Sequence
 from typing import Any
 
-from serena.code_editor import EditedFilePath
 from serena.symbol import LanguageServerSymbol, LanguageServerSymbolDictGrouper
 from serena.tools import (
     SUCCESS_RESULT,
@@ -796,19 +795,14 @@ class ReplaceSymbolBodyTool(EditingToolWithDiagnostics):
             in the programming language, including e.g. the signature line for functions.
             IMPORTANT: The body does NOT include any preceding docstrings/comments or imports, in particular.
         """
-        # capturing diagnostics before the edit
-        edited_file_paths = [EditedFilePath(relative_path, relative_path)]
-        diagnostics_snapshot = self._capture_published_lsp_diagnostics_snapshot(edited_file_paths)
-
-        # applying the symbol replacement
-        code_editor = self.create_code_editor()
-        code_editor.replace_body(
-            name_path,
-            relative_file_path=relative_path,
-            body=body,
-        )
-
-        return self._format_lsp_edit_result_with_new_diagnostics(SUCCESS_RESULT, edited_file_paths, diagnostics_snapshot)
+        with self.DiagnosticsContext(self, relative_path) as diagnostics_context:
+            code_editor = self.create_code_editor()
+            code_editor.replace_body(
+                name_path,
+                relative_file_path=relative_path,
+                body=body,
+            )
+            return diagnostics_context.format_result(SUCCESS_RESULT)
 
 
 class InsertAfterSymbolTool(EditingToolWithDiagnostics):
@@ -831,15 +825,10 @@ class InsertAfterSymbolTool(EditingToolWithDiagnostics):
         :param body: the body/content to be inserted. The inserted code shall begin with the next line after
             the symbol.
         """
-        # capturing diagnostics before the edit
-        edited_file_paths = [EditedFilePath(relative_path, relative_path)]
-        diagnostics_snapshot = self._capture_published_lsp_diagnostics_snapshot(edited_file_paths)
-
-        # applying the insertion
-        code_editor = self.create_code_editor()
-        code_editor.insert_after_symbol(name_path, relative_file_path=relative_path, body=body)
-
-        return self._format_lsp_edit_result_with_new_diagnostics(SUCCESS_RESULT, edited_file_paths, diagnostics_snapshot)
+        with self.DiagnosticsContext(self, relative_path) as diagnostics_context:
+            code_editor = self.create_code_editor()
+            code_editor.insert_after_symbol(name_path, relative_file_path=relative_path, body=body)
+            return diagnostics_context.format_result(SUCCESS_RESULT)
 
 
 class InsertBeforeSymbolTool(EditingToolWithDiagnostics):
@@ -862,15 +851,10 @@ class InsertBeforeSymbolTool(EditingToolWithDiagnostics):
         :param relative_path: the relative path to the file containing the symbol
         :param body: the body/content to be inserted before the line in which the referenced symbol is defined
         """
-        # capturing diagnostics before the edit
-        edited_file_paths = [EditedFilePath(relative_path, relative_path)]
-        diagnostics_snapshot = self._capture_published_lsp_diagnostics_snapshot(edited_file_paths)
-
-        # applying the insertion
-        code_editor = self.create_code_editor()
-        code_editor.insert_before_symbol(name_path, relative_file_path=relative_path, body=body)
-
-        return self._format_lsp_edit_result_with_new_diagnostics(SUCCESS_RESULT, edited_file_paths, diagnostics_snapshot)
+        with self.DiagnosticsContext(self, relative_path) as diagnostics_context:
+            code_editor = self.create_code_editor()
+            code_editor.insert_before_symbol(name_path, relative_file_path=relative_path, body=body)
+            return diagnostics_context.format_result(SUCCESS_RESULT)
 
 
 class RenameSymbolTool(Tool, ToolMarkerSymbolicEdit):

src/serena/tools/tools_base.py

@@ -5,7 +5,7 @@ from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from functools import cached_property
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Protocol, Self, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Optional, Protocol, Self, TypeVar, cast
 
 from mcp import Implementation
 from mcp.server.fastmcp import Context
@@ -18,13 +18,12 @@ from serena.project import MemoriesManager, Project
 from serena.prompt_factory import PromptFactory
 from serena.util.class_decorators import singleton
 from serena.util.inspection import iter_subclasses
-from solidlsp import ls_types
+from serena.util.ls_diagnostics import DiagnosticsDiff, EditedFilePath, PublishedDiagnosticsSnapshot
 from solidlsp.ls_exceptions import SolidLSPException
-from solidlsp.lsp_protocol_handler.lsp_types import DiagnosticSeverity
 
 if TYPE_CHECKING:
     from serena.agent import SerenaAgent
-    from serena.code_editor import CodeEditor, EditedFilePath, LanguageServerCodeEditor
+    from serena.code_editor import CodeEditor, LanguageServerCodeEditor
     from serena.symbol import LanguageServerSymbolRetriever
 
 log = logging.getLogger(__name__)
@@ -32,24 +31,6 @@ T = TypeVar("T")
 SUCCESS_RESULT = "OK"
 
 
-@dataclass(frozen=True)
-class DiagnosticIdentity:
-    message: str
-    start_line: int
-    start_character: int
-    end_line: int
-    end_character: int
-    severity: int | None
-    code_repr: str | None
-    source: str | None
-
-
-@dataclass
-class PublishedDiagnosticsSnapshot:
-    generation_by_after_path: dict[str, int]
-    warning_identities_by_before_path: dict[str, set[DiagnosticIdentity]]
-
-
 class Component(ABC):
     def __init__(self, agent: "SerenaAgent"):
         self.agent = agent
@@ -428,7 +409,7 @@ class EditingToolWithDiagnostics(Tool, ToolMarkerCanEdit):
     Base class for editing tools that want to capture and report changes in LSP diagnostics before and after the edit.
     """
 
-    _ENABLE_DIAGNOSTICS: bool = False
+    ENABLE_DIAGNOSTICS: bool = False
     """
     Global flag to enable/disable diagnostics for LSP-based editing tools derived from this class.
     The feature is currently disabled, because per-edit diagnostics are a questionable feature, since individual
@@ -436,179 +417,47 @@ class EditingToolWithDiagnostics(Tool, ToolMarkerCanEdit):
     are then resolved in subsequent edits.
     """
 
-    @staticmethod
-    def _diagnostic_code_repr(code: Any) -> str | None:
-        if code is None:
-            return None
-        try:
-            return json.dumps(code, sort_keys=True, ensure_ascii=False)
-        except TypeError:
-            return repr(code)
+    DIAGNOSTICS_KEY = "diagnostics[warning-or-higher]"
 
-    @classmethod
-    def _diagnostic_identity(cls, diagnostic: ls_types.Diagnostic) -> DiagnosticIdentity:
-        diagnostic_range = diagnostic["range"]
-        start = diagnostic_range["start"]
-        end = diagnostic_range["end"]
-        return DiagnosticIdentity(
-            message=diagnostic["message"],
-            start_line=start["line"],
-            start_character=start["character"],
-            end_line=end["line"],
-            end_character=end["character"],
-            severity=diagnostic.get("severity"),
-            code_repr=cls._diagnostic_code_repr(diagnostic.get("code")),
-            source=diagnostic.get("source"),
-        )
+    class DiagnosticsContext:
+        def __init__(self, tool: "EditingToolWithDiagnostics", *edited_relative_paths: str) -> None:
+            self._tool = tool
+            self._is_diagnostics_enabled = tool.ENABLE_DIAGNOSTICS and tool.agent.is_using_language_server()
+            self._edited_files = [EditedFilePath(path, path) for path in edited_relative_paths]
+            self._before_edit_diagnostics_snapshot: PublishedDiagnosticsSnapshot | None = None
+            self._symbol_retriever: Optional["LanguageServerSymbolRetriever"] | None = None
+            if self._is_diagnostics_enabled:
+                symbol_retriever = tool.create_language_server_symbol_retriever()
+                self._before_edit_diagnostics_snapshot = PublishedDiagnosticsSnapshot(self._edited_files, symbol_retriever)
 
-    @staticmethod
-    def _diagnostic_severity_name(severity: int | None) -> str:
-        if severity is None:
-            return "Unknown"
-        try:
-            return DiagnosticSeverity(severity).name
-        except ValueError:
-            return f"Severity_{severity}"
+        def __enter__(self) -> Self:
+            return self
 
-    @staticmethod
-    def _diagnostic_output_dict(diagnostic: ls_types.Diagnostic) -> dict[str, Any]:
-        result: dict[str, Any] = {
-            "message": diagnostic["message"],
-            "range": diagnostic["range"],
-        }
-        if "code" in diagnostic:
-            result["code"] = diagnostic["code"]
-        if "source" in diagnostic:
-            result["source"] = diagnostic["source"]
-        return result
+        def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore
+            pass
 
-    @classmethod
-    def _add_grouped_diagnostic(
-        cls,
-        grouped_result: dict[str, dict[str, dict[str, list[dict[str, Any]]]]],
-        relative_path: str,
-        name_path: str,
-        diagnostic: ls_types.Diagnostic,
-    ) -> None:
-        severity_name = cls._diagnostic_severity_name(diagnostic.get("severity"))
-        grouped_result.setdefault(relative_path, {}).setdefault(severity_name, {}).setdefault(name_path, []).append(
-            cls._diagnostic_output_dict(diagnostic)
-        )
+        def format_result(
+            self,
+            base_result: str,
+        ) -> str:
+            if not self._is_diagnostics_enabled:
+                return base_result
 
-    def _capture_published_lsp_diagnostics_snapshot(
-        self,
-        edited_file_paths: Iterable["EditedFilePath"],
-    ) -> PublishedDiagnosticsSnapshot | None:
-        if not self._ENABLE_DIAGNOSTICS:
-            return None
-        if self.agent.get_language_backend() != LanguageBackend.LSP:
-            return None
+            if self._before_edit_diagnostics_snapshot is None:
+                return base_result
 
-        # collecting diagnostics state before the edit
-        symbol_retriever = self.create_language_server_symbol_retriever()
-        generation_by_after_path: dict[str, int] = {}
-        warning_identities_by_before_path: dict[str, set[DiagnosticIdentity]] = {}
+            assert self._symbol_retriever is not None
+            diagnostics_diff = DiagnosticsDiff(self._before_edit_diagnostics_snapshot, self._edited_files, self._symbol_retriever)
+            grouped_diagnostics = diagnostics_diff.get_grouped_diagnostics()
 
-        for edited_file_path in edited_file_paths:
-            try:
-                language_server = symbol_retriever.get_language_server(edited_file_path.after_relative_path)
-            except Exception:
-                return None
-
-            generation_by_after_path[edited_file_path.after_relative_path] = language_server.get_published_diagnostics_generation(
-                edited_file_path.after_relative_path
-            )
-
-            cached_diagnostics = language_server.get_cached_published_text_document_diagnostics(
-                edited_file_path.before_relative_path,
-                min_severity=2,
-            )
-            if cached_diagnostics is None:
-                try:
-                    cached_diagnostics = language_server.request_text_document_diagnostics(
-                        edited_file_path.before_relative_path,
-                        min_severity=2,
-                    )
-                except Exception:
-                    cached_diagnostics = []
-            warning_identities_by_before_path[edited_file_path.before_relative_path] = {
-                self._diagnostic_identity(diagnostic) for diagnostic in cached_diagnostics or []
-            }
-
-        return PublishedDiagnosticsSnapshot(
-            generation_by_after_path=generation_by_after_path,
-            warning_identities_by_before_path=warning_identities_by_before_path,
-        )
-
-    def _format_lsp_edit_result_with_new_diagnostics(
-        self,
-        default_result: str,
-        edited_file_paths: Iterable["EditedFilePath"],
-        before_edit_diagnostics_snapshot: PublishedDiagnosticsSnapshot | None,
-    ) -> str:
-        if not self._ENABLE_DIAGNOSTICS:
-            return default_result
-
-        # TODO: this is weird, it works because before_edit_diagnostics_snapshot is only None when diagnostics
-        #   are disabled, not when they are empty. A part of the general design flaws introduced by agent-generated code
-        if before_edit_diagnostics_snapshot is None or self.agent.get_language_backend() != LanguageBackend.LSP:
-            return default_result
-
-        # collecting diagnostics state after the edit
-        symbol_retriever = self.create_language_server_symbol_retriever()
-        grouped_result: dict[str, dict[str, dict[str, list[dict[str, Any]]]]] = {}
-        saw_diagnostics_result = False
-
-        for edited_file_path in edited_file_paths:
-            try:
-                language_server = symbol_retriever.get_language_server(edited_file_path.after_relative_path)
-            except Exception:
-                return default_result
-
-            published_diagnostics = language_server.request_published_text_document_diagnostics(
-                relative_file_path=edited_file_path.after_relative_path,
-                after_generation=before_edit_diagnostics_snapshot.generation_by_after_path.get(edited_file_path.after_relative_path, -1),
-                timeout=2.5,
-                min_severity=2,
-                allow_cached=True,
-            )
-            if not published_diagnostics:
-                try:
-                    published_diagnostics = language_server.request_text_document_diagnostics(
-                        edited_file_path.after_relative_path,
-                        min_severity=2,
-                    )
-                except Exception:
-                    published_diagnostics = None
-            if published_diagnostics is None:
-                continue
-
-            saw_diagnostics_result = True
-            existing_warning_identities = before_edit_diagnostics_snapshot.warning_identities_by_before_path.get(
-                edited_file_path.before_relative_path, set()
-            )
-            new_warning_identities: set[DiagnosticIdentity] = set()
-
-            for diagnostic in published_diagnostics:
-                diagnostic_identity = self._diagnostic_identity(diagnostic)
-                if diagnostic_identity in existing_warning_identities or diagnostic_identity in new_warning_identities:
-                    continue
-                new_warning_identities.add(diagnostic_identity)
-
-                diagnostic_start = diagnostic["range"]["start"]
-                owner_symbol = symbol_retriever.find_diagnostic_owner_symbol(
-                    relative_file_path=edited_file_path.after_relative_path,
-                    line=diagnostic_start["line"],
-                    column=diagnostic_start["character"],
-                )
-                name_path = owner_symbol.get_name_path() if owner_symbol is not None else "<file>"
-                self._add_grouped_diagnostic(grouped_result, edited_file_path.after_relative_path, name_path, diagnostic)
-
-        # preserving the normal success result when no new diagnostics were introduced
-        if not saw_diagnostics_result or not grouped_result:
-            return default_result
-
-        return f"Edit introduced new warning-or-higher diagnostics: {self._to_json(grouped_result)}"
+            if not grouped_diagnostics:
+                return base_result
+            else:
+                result_dict = {
+                    "result": base_result,
+                    EditingToolWithDiagnostics.DIAGNOSTICS_KEY: grouped_diagnostics,
+                }
+                return self._tool._to_json(result_dict)
 
 
 class EditedFileContext:

src/serena/util/ls_diagnostics.py

@@ -0,0 +1,202 @@
+import json
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from solidlsp import ls_types
+from solidlsp.lsp_protocol_handler.lsp_types import DiagnosticSeverity
+
+if TYPE_CHECKING:
+    from serena.symbol import LanguageServerSymbolRetriever
+
+
+@dataclass(frozen=True)
+class EditedFilePath:
+    before_relative_path: str
+    after_relative_path: str
+
+
+@dataclass(frozen=True)
+class DiagnosticIdentity:
+    message: str
+    start_line: int
+    start_character: int
+    end_line: int
+    end_character: int
+    severity: int | None
+    code_repr: str | None
+    source: str | None
+
+    @classmethod
+    def from_diagnostic(cls, diagnostic: ls_types.Diagnostic) -> "DiagnosticIdentity":
+        diagnostic_range = diagnostic["range"]
+        start = diagnostic_range["start"]
+        end = diagnostic_range["end"]
+        return DiagnosticIdentity(
+            message=diagnostic["message"],
+            start_line=start["line"],
+            start_character=start["character"],
+            end_line=end["line"],
+            end_character=end["character"],
+            severity=diagnostic.get("severity"),
+            code_repr=cls._diagnostic_code_repr(diagnostic.get("code")),
+            source=diagnostic.get("source"),
+        )
+
+    @staticmethod
+    def _diagnostic_code_repr(code: Any) -> str | None:
+        if code is None:
+            return None
+        try:
+            return json.dumps(code, sort_keys=True, ensure_ascii=False)
+        except TypeError:
+            return repr(code)
+
+
+class PublishedDiagnosticsSnapshot:
+    def __init__(self, edited_file_paths: Iterable[EditedFilePath], symbol_retriever: "LanguageServerSymbolRetriever"):
+        generation_by_after_path: dict[str, int] = {}
+        warning_identities_by_before_path: dict[str, set[DiagnosticIdentity]] = {}
+
+        for edited_file_path in edited_file_paths:
+            try:
+                language_server = symbol_retriever.get_language_server(edited_file_path.after_relative_path)
+            except Exception:
+                continue
+
+            generation_by_after_path[edited_file_path.after_relative_path] = language_server.get_published_diagnostics_generation(
+                edited_file_path.after_relative_path
+            )
+
+            cached_diagnostics = language_server.get_cached_published_text_document_diagnostics(
+                edited_file_path.before_relative_path,
+                min_severity=2,
+            )
+            if cached_diagnostics is None:
+                try:
+                    cached_diagnostics = language_server.request_text_document_diagnostics(
+                        edited_file_path.before_relative_path,
+                        min_severity=2,
+                    )
+                except Exception:
+                    cached_diagnostics = []
+            warning_identities_by_before_path[edited_file_path.before_relative_path] = {
+                DiagnosticIdentity.from_diagnostic(diagnostic) for diagnostic in cached_diagnostics or []
+            }
+
+        self.generation_by_after_path = generation_by_after_path
+        self.warning_identities_by_before_path = warning_identities_by_before_path
+
+
+class DiagnosticsDiff:
+    def __init__(
+        self,
+        before_snapshot: PublishedDiagnosticsSnapshot,
+        edited_files: Iterable[EditedFilePath],
+        symbol_retriever: "LanguageServerSymbolRetriever",
+    ):
+        grouped_result: dict[str, dict[str, dict[str, list[dict[str, Any]]]]] = {}
+
+        for edited_file_path in edited_files:
+            try:
+                language_server = symbol_retriever.get_language_server(edited_file_path.after_relative_path)
+            except:
+                continue
+
+            published_diagnostics = language_server.request_published_text_document_diagnostics(
+                relative_file_path=edited_file_path.after_relative_path,
+                after_generation=before_snapshot.generation_by_after_path.get(edited_file_path.after_relative_path, -1),
+                timeout=2.5,
+                min_severity=2,
+                allow_cached=True,
+            )
+            if not published_diagnostics:
+                try:
+                    published_diagnostics = language_server.request_text_document_diagnostics(
+                        edited_file_path.after_relative_path,
+                        min_severity=2,
+                    )
+                except:
+                    published_diagnostics = None
+            if published_diagnostics is None:
+                continue
+
+            existing_warning_identities = before_snapshot.warning_identities_by_before_path.get(
+                edited_file_path.before_relative_path, set()
+            )
+            new_warning_identities: set[DiagnosticIdentity] = set()
+
+            for diagnostic in published_diagnostics:
+                diagnostic_identity = DiagnosticIdentity.from_diagnostic(diagnostic)
+                if diagnostic_identity in existing_warning_identities or diagnostic_identity in new_warning_identities:
+                    continue
+                new_warning_identities.add(diagnostic_identity)
+
+                diagnostic_start = diagnostic["range"]["start"]
+                owner_symbol = symbol_retriever.find_diagnostic_owner_symbol(
+                    relative_file_path=edited_file_path.after_relative_path,
+                    line=diagnostic_start["line"],
+                    column=diagnostic_start["character"],
+                )
+                name_path = owner_symbol.get_name_path() if owner_symbol is not None else "<file>"
+                self._add_grouped_diagnostic(grouped_result, edited_file_path.after_relative_path, name_path, diagnostic)
+
+        self._grouped_result = grouped_result
+
+    def get_grouped_diagnostics(self) -> dict[str, dict[str, dict[str, list[dict[str, Any]]]]]:
+        """
+        Returns a nested dictionary of the form:
+        {
+            relative_file_path: {
+                severity_name: {
+                    name_path: [
+                        diagnostic_dict,
+                        ...
+                    ],
+                    ...
+                },
+                ...
+            },
+            ...
+        }
+        where:
+        - relative_file_path is the relative path of the file containing the diagnostic
+        - severity_name is the name of the diagnostic severity (e.g. "Warning", "Error")
+        - name_path is the name path of the symbol that owns the diagnostic (or "<file>" if no owner symbol was found)
+        - diagnostic_dict is a dictionary containing the diagnostic's message, range, and optionally code and source
+        """
+        return self._grouped_result
+
+    @classmethod
+    def _add_grouped_diagnostic(
+        cls,
+        grouped_result: dict[str, dict[str, dict[str, list[dict[str, Any]]]]],
+        relative_path: str,
+        name_path: str,
+        diagnostic: ls_types.Diagnostic,
+    ) -> None:
+        severity_name = cls._diagnostic_severity_name(diagnostic.get("severity"))
+        grouped_result.setdefault(relative_path, {}).setdefault(severity_name, {}).setdefault(name_path, []).append(
+            cls._diagnostic_output_dict(diagnostic)
+        )
+
+    @staticmethod
+    def _diagnostic_severity_name(severity: int | None) -> str:
+        if severity is None:
+            return "Unknown"
+        try:
+            return DiagnosticSeverity(severity).name
+        except ValueError:
+            return f"Severity_{severity}"
+
+    @staticmethod
+    def _diagnostic_output_dict(diagnostic: ls_types.Diagnostic) -> dict[str, Any]:
+        result: dict[str, Any] = {
+            "message": diagnostic["message"],
+            "range": diagnostic["range"],
+        }
+        if "code" in diagnostic:
+            result["code"] = diagnostic["code"]
+        if "source" in diagnostic:
+            result["source"] = diagnostic["source"]
+        return result

test/serena/test_serena_agent.py

@@ -18,6 +18,7 @@ from serena.project import Project
 from serena.tools import (
     SUCCESS_RESULT,
     ActivateProjectTool,
+    EditingToolWithDiagnostics,
     FindDefiningSymbolAtLocationTool,
     FindDefiningSymbolTool,
     FindImplementationsTool,
@@ -812,9 +813,9 @@ def read_project_file(project: Project, relative_path: str) -> str:
 
 def parse_edit_diagnostics_result(result: str) -> dict:
     """Utility function to parse the diagnostic payload returned by edit tools."""
-    prefix = "Edit introduced new warning-or-higher diagnostics: "
-    assert result.startswith(prefix), result
-    return json.loads(result[len(prefix) :])
+    assert EditingToolWithDiagnostics.DIAGNOSTICS_KEY in result
+    d = json.loads(result)
+    return d[EditingToolWithDiagnostics.DIAGNOSTICS_KEY]
 
 
 @contextmanager
@@ -1201,21 +1202,24 @@ class TestSerenaAgent:
         """Tests that file-level edits report newly introduced diagnostics."""
         relative_path = os.path.join("test_repo", "services.py")
         replace_content_tool = serena_agent.get_tool(ReplaceContentTool)
-        replace_content_tool._ENABLE_DIAGNOSTICS = True
+        try:
+            replace_content_tool.ENABLE_DIAGNOSTICS = True
 
-        with project_file_modification_context(serena_agent, relative_path):
-            result = replace_content_tool.apply(
-                relative_path=relative_path,
-                needle="return container",
-                repl="return missing_container",
-                mode="literal",
-            )
+            with project_file_modification_context(serena_agent, relative_path):
+                result = replace_content_tool.apply(
+                    relative_path=relative_path,
+                    needle="return container",
+                    repl="return missing_container",
+                    mode="literal",
+                )
 
-        diagnostics = parse_edit_diagnostics_result(result)
-        relative_path_result = diagnostics[relative_path]
-        diagnostic_messages = json.dumps(relative_path_result)
-        assert "missing_container" in diagnostic_messages
-        assert "create_service_container" in diagnostic_messages
+            diagnostics = parse_edit_diagnostics_result(result)
+            relative_path_result = diagnostics[relative_path]
+            diagnostic_messages = json.dumps(relative_path_result)
+            assert "missing_container" in diagnostic_messages
+            assert "create_service_container" in diagnostic_messages
+        finally:
+            replace_content_tool.ENABLE_DIAGNOSTICS = False
 
     @pytest.mark.parametrize(
         "serena_agent",
@@ -1229,23 +1233,26 @@ class TestSerenaAgent:
         """Tests that symbol-level edits report newly introduced diagnostics."""
         relative_path = os.path.join("test_repo", "services.py")
         replace_symbol_body_tool = serena_agent.get_tool(ReplaceSymbolBodyTool)
-        replace_symbol_body_tool._ENABLE_DIAGNOSTICS = True
+        try:
+            replace_symbol_body_tool.ENABLE_DIAGNOSTICS = True
 
-        with project_file_modification_context(serena_agent, relative_path):
-            result = replace_symbol_body_tool.apply(
-                name_path="create_service_container",
-                relative_path=relative_path,
-                body="""
-def create_service_container() -> dict[str, Any]:
-    return missing_container
-""",
-            )
+            with project_file_modification_context(serena_agent, relative_path):
+                result = replace_symbol_body_tool.apply(
+                    name_path="create_service_container",
+                    relative_path=relative_path,
+                    body="""
+    def create_service_container() -> dict[str, Any]:
+        return missing_container
+    """,
+                )
 
-        diagnostics = parse_edit_diagnostics_result(result)
-        relative_path_result = diagnostics[relative_path]
-        diagnostic_messages = json.dumps(relative_path_result)
-        assert "missing_container" in diagnostic_messages
-        assert "create_service_container" in diagnostic_messages
+            diagnostics = parse_edit_diagnostics_result(result)
+            relative_path_result = diagnostics[relative_path]
+            diagnostic_messages = json.dumps(relative_path_result)
+            assert "missing_container" in diagnostic_messages
+            assert "create_service_container" in diagnostic_messages
+        finally:
+            replace_symbol_body_tool.ENABLE_DIAGNOSTICS = False
 
     @pytest.mark.parametrize(
         "serena_agent",