From 061020ce226e7f0790ee3baffe97c6f2b4bdef82 Mon Sep 17 00:00:00 2001 From: Dominik Jain Date: Wed, 29 Apr 2026 23:18:47 +0200 Subject: [PATCH] Refactor EditingToolWithDiagnostics * Add context manager to simplify usage, improving output handling * Add util module ls_diagnostics, reifying DiagnosticsDiff --- scripts/demo_diagnostics.py | 12 +- src/serena/code_editor.py | 8 +- src/serena/tools/file_tools.py | 100 +++++-------- src/serena/tools/symbol_tools.py | 48 +++---- src/serena/tools/tools_base.py | 227 +++++------------------------- src/serena/util/ls_diagnostics.py | 202 ++++++++++++++++++++++++++ test/serena/test_serena_agent.py | 69 +++++---- 7 files changed, 343 insertions(+), 323 deletions(-) create mode 100644 src/serena/util/ls_diagnostics.py diff --git a/scripts/demo_diagnostics.py b/scripts/demo_diagnostics.py index 7aa85361..2e5cd30e 100644 --- a/scripts/demo_diagnostics.py +++ b/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() diff --git a/src/serena/code_editor.py b/src/serena/code_editor.py index 7d52f2b8..ec834273 100644 --- a/src/serena/code_editor.py +++ b/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 diff --git a/src/serena/tools/file_tools.py b/src/serena/tools/file_tools.py index 4572f2fc..4b4313eb 100644 --- a/src/serena/tools/file_tools.py +++ b/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): diff --git a/src/serena/tools/symbol_tools.py b/src/serena/tools/symbol_tools.py index 036c1d71..7f9c17f5 100644 --- a/src/serena/tools/symbol_tools.py +++ b/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): diff --git a/src/serena/tools/tools_base.py b/src/serena/tools/tools_base.py index ca8e5a49..11d92dce 100644 --- a/src/serena/tools/tools_base.py +++ b/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 "" - 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: diff --git a/src/serena/util/ls_diagnostics.py b/src/serena/util/ls_diagnostics.py new file mode 100644 index 00000000..6df5ce34 --- /dev/null +++ b/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 "" + 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 "" 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 diff --git a/test/serena/test_serena_agent.py b/test/serena/test_serena_agent.py index 974cdc4b..b3e2427b 100644 --- a/test/serena/test_serena_agent.py +++ b/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",