fix: repair Mermaid labels and report validation accurately

.github/workflows/ci.yml

@@ -10,6 +10,20 @@ permissions:
   contents: read
 
 jobs:
+  mermaid:
+    name: Mermaid syntax
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v4
+        with:
+          node-version: "24"
+          cache: npm
+          cache-dependency-path: tests/mermaid/package-lock.json
+      - run: npm ci --prefix tests/mermaid --ignore-scripts --no-audit --no-fund
+      - run: npm test --prefix tests/mermaid
+
   test:
     name: Test
     runs-on: ubuntu-latest

.gitignore

@@ -91,3 +91,10 @@ Thumbs.db
 !tests/test_updater_write_guard.py
 !tests/test_updater_orchestrator.py
 !tests/test_usage_ledger.py
+
+# Mermaid regression tests and development-only parser harness
+!tests/test_mermaid_validation.py
+!tests/test_mermaid_writers.py
+!tests/mermaid/
+!tests/mermaid/**
+tests/mermaid/node_modules/

codewiki/mcp/tools/doc_writer.py

@@ -45,14 +45,18 @@ def _ensure_parent_dirs(path: Path) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
 
 
-async def _validate_mermaid(file_path: str, relative_path: str) -> str:
+async def _validate_mermaid(file_path: str, relative_path: str, *, repair: bool = True) -> str:
     """Run Mermaid validation and return the result string."""
     try:
-        from codewiki.src.be.utils import validate_mermaid_diagrams
+        from codewiki.src.be.utils import (
+            validate_and_repair_mermaid_diagrams,
+            validate_mermaid_diagrams,
+        )
 
-        return await validate_mermaid_diagrams(file_path, relative_path)
+        validator = validate_and_repair_mermaid_diagrams if repair else validate_mermaid_diagrams
+        return await validator(file_path, relative_path)
     except Exception as e:
-        return f"Mermaid validation skipped: {e}"
+        return f"Warning: Mermaid validation unavailable: {e}"
 
 
 def _save_history(session: SessionState, doc_path: Path, content: str) -> None:
@@ -142,7 +146,7 @@ async def handle_edit_doc_file(
         await asyncio.to_thread(doc_path.write_text, old_content, "utf-8")
 
         # Validate Mermaid after undo
-        mermaid_result = await _validate_mermaid(str(doc_path), filename)
+        mermaid_result = await _validate_mermaid(str(doc_path), filename, repair=False)
         return json.dumps(
             {
                 "status": "undone",
@@ -157,7 +161,7 @@ async def handle_edit_doc_file(
             {"error": f"File not found: {filename}. Use write_doc_file to create it."}
         )
 
-    current_content = await asyncio.to_thread(doc_path.read_text, "utf-8")
+    current_content = (await asyncio.to_thread(doc_path.read_bytes)).decode("utf-8")
 
     if command == "str_replace":
         old_str = arguments.get("old_str")
@@ -214,6 +218,12 @@ async def handle_edit_doc_file(
     # Mermaid validation
     mermaid_result = await _validate_mermaid(str(doc_path), filename)
 
+    # Validation may have repaired labels; show the persisted text to the agent.
+    saved_lines = (await asyncio.to_thread(doc_path.read_text, "utf-8")).split("\n")
+    snippet = "\n".join(
+        f"{i + 1:6}\t{saved_lines[i]}" for i in range(start, min(end, len(saved_lines)))
+    )
+
     result = {
         "status": "edited",
         "command": command,

codewiki/src/be/agent_tools/str_replace_editor.py

@@ -17,7 +17,7 @@
 from pydantic import BeforeValidator
 from pydantic_ai import RunContext, Tool
 
-from ..utils import validate_mermaid_diagrams
+from ..utils import validate_and_repair_mermaid_diagrams, validate_mermaid_diagrams
 from .deps import CodeWikiDeps
 
 # Configure logging and monitoring
@@ -422,6 +422,7 @@ def __init__(self, REGISTRY, absolute_docs_path=None):
         self._encoding = None
         self.REGISTRY = REGISTRY
         self.logs = []
+        self.did_write = False
         self.absolute_docs_path = Path(absolute_docs_path) if absolute_docs_path else None
 
     def _get_display_path(self, path: Path) -> str:
@@ -436,11 +437,20 @@ def _get_display_path(self, path: Path) -> str:
 
     @property
     def _file_history(self):
-        return defaultdict(list, json.loads(self.REGISTRY.get("file_history", "{}")))
+        stored = self.REGISTRY.get("file_history", "{}")
+        stored = json.loads(stored) if isinstance(stored, str) else stored
+        return defaultdict(list, {Path(key): value for key, value in stored.items()})
 
     @_file_history.setter
     def _file_history(self, value: dict):
-        self.REGISTRY["file_history"] = json.dumps(value)
+        self.REGISTRY["file_history"] = json.dumps(
+            {str(key): items for key, items in value.items()}
+        )
+
+    def _remember(self, path: Path, content: str) -> None:
+        history = self._file_history
+        history[path].append(content)
+        self._file_history = history
 
     def __call__(
         self,
@@ -525,8 +535,9 @@ def create_file(self, path: Path, file_text: str):
                 f"The parent directory {self._get_display_path(path.parent)} does not exist. Please create it first."
             )
             return
-        self.write_file(path, file_text)
-        self._file_history[path].append(file_text)
+        if not self.write_file(path, file_text):
+            return
+        self._remember(path, file_text)
         self.logs.append(f"File created successfully at: {self._get_display_path(path)}")
 
     def _inside_docs(self, path: Path) -> bool:
@@ -630,7 +641,8 @@ def view(self, path: Path, view_range: list[int] | None = None):
     def str_replace(self, path: Path, old_str: str, new_str: str | None):
         """Implement the str_replace command, which replaces old_str with new_str in the file content"""
         # Read the file content
-        file_content = self.read_file(path).expandtabs()
+        history_content = self.read_file(path)
+        file_content = history_content.expandtabs()
         old_str = old_str.expandtabs()
         new_str = new_str.expandtabs() if new_str is not None else ""
 
@@ -666,7 +678,8 @@ def str_replace(self, path: Path, old_str: str, new_str: str | None):
         new_file_content = file_content.replace(old_str, new_str)
 
         # Write the new content to the file
-        self.write_file(path, new_file_content)
+        if not self.write_file(path, new_file_content):
+            return
 
         post_edit_lint = ""
         if USE_LINTER:
@@ -691,7 +704,7 @@ def str_replace(self, path: Path, old_str: str, new_str: str | None):
                 epilogue = LINT_WARNING_TEMPLATE.format(errors=errors)
 
         # Save the content to history
-        self._file_history[path].append(file_content)
+        self._remember(path, history_content)
 
         # Create a snippet of the edited section
         replacement_line = file_content.split(old_str)[0].count("\n")
@@ -720,7 +733,8 @@ def str_replace(self, path: Path, old_str: str, new_str: str | None):
 
     def insert(self, path: Path, insert_line: int, new_str: str):
         """Implement the insert command, which inserts new_str at the specified line in the file content."""
-        file_text = self.read_file(path).expandtabs()
+        history_content = self.read_file(path)
+        file_text = history_content.expandtabs()
         new_str = new_str.expandtabs()
         file_text_lines = file_text.split("\n")
         n_lines_file = len(file_text_lines)
@@ -744,8 +758,9 @@ def insert(self, path: Path, insert_line: int, new_str: str):
         new_file_text = "\n".join(new_file_text_lines)
         snippet = "\n".join(snippet_lines)
 
-        self.write_file(path, new_file_text)
-        self._file_history[path].append(file_text)
+        if not self.write_file(path, new_file_text):
+            return
+        self._remember(path, history_content)
 
         # todo: Also expand these windows
 
@@ -760,12 +775,16 @@ def insert(self, path: Path, insert_line: int, new_str: str):
 
     def undo_edit(self, path: Path):
         """Implement the undo_edit command."""
-        if not self._file_history[path]:
+        history = self._file_history
+        if not history[path]:
             self.logs.append(f"No edit history found for {self._get_display_path(path)}.")
             return
 
-        old_text = self._file_history[path].pop()
-        self.write_file(path, old_text)
+        old_text = history[path][-1]
+        if not self.write_file(path, old_text):
+            return
+        history[path].pop()
+        self._file_history = history
 
         self.logs.append(
             f"Last edit to {self._get_display_path(path)} undone successfully. {self._make_output(old_text, self._get_display_path(path))}"
@@ -782,7 +801,8 @@ def read_file(self, path: Path):
         exception = None
         for self._encoding, errors in encodings:
             try:
-                text = path.read_text(encoding=self._encoding, errors=errors)
+                with path.open(encoding=self._encoding, errors=errors, newline="") as stream:
+                    text = stream.read()
             except UnicodeDecodeError as e:
                 exception = e
             else:
@@ -797,12 +817,15 @@ def read_file(self, path: Path):
     def write_file(self, path: Path, file: str):
         """Write the content of a file to a given path; raise a ToolError if an error occurs."""
         try:
-            path.write_text(file, encoding=self._encoding or "utf-8")
+            with path.open("w", encoding=self._encoding or "utf-8", newline="") as stream:
+                stream.write(file)
+            self.did_write = True
+            return True
         except Exception as e:  # noqa: BLE001 — reported to the agent as a tool message
             self.logs.append(
                 f"Ran into {e} while trying to write to {self._get_display_path(path)}"
             )
-            return
+            return False
 
     def _make_output(
         self,
@@ -928,8 +951,18 @@ async def str_replace_editor(
 
     result = "\n".join(tool.logs)
 
-    if command != "view" and path.endswith(".md"):
-        mermaid_validation = await validate_mermaid_diagrams(absolute_path, path)
+    if tool.did_write and path.endswith(".md"):
+        before_validation = Path(absolute_path).read_bytes()
+        validator = (
+            validate_mermaid_diagrams
+            if command == "undo_edit"
+            else validate_and_repair_mermaid_diagrams
+        )
+        mermaid_validation = await validator(absolute_path, path)
+        if Path(absolute_path).read_bytes() != before_validation:
+            result = "File saved with repaired Mermaid labels.\n" + tool._make_output(
+                Path(absolute_path).read_text(encoding="utf-8"), path
+            )
         result = result + "\n---------- Mermaid validation ----------\n" + mermaid_validation
 
     return result

codewiki/src/be/caw_toolkit.py

@@ -143,7 +143,10 @@ async def str_replace_editor(
         insert_line: int | str | None = None,
     ) -> str:
         from codewiki.src.be.agent_tools.str_replace_editor import EditTool
-        from codewiki.src.be.utils import validate_mermaid_diagrams
+        from codewiki.src.be.utils import (
+            validate_and_repair_mermaid_diagrams,
+            validate_mermaid_diagrams,
+        )
 
         # ``Literal`` annotations would be the cleanest way to constrain these,
         # but ``from __future__ import annotations`` turns them into forward refs
@@ -226,8 +229,18 @@ async def str_replace_editor(
 
         result = "\n".join(edit_tool.logs)
 
-        if command != "view" and path.endswith(".md"):
-            mermaid_validation = await validate_mermaid_diagrams(absolute_path, path)
+        if edit_tool.did_write and path.endswith(".md"):
+            before_validation = Path(absolute_path).read_bytes()
+            validator = (
+                validate_mermaid_diagrams
+                if command == "undo_edit"
+                else validate_and_repair_mermaid_diagrams
+            )
+            mermaid_validation = await validator(absolute_path, path)
+            if Path(absolute_path).read_bytes() != before_validation:
+                result = "File saved with repaired Mermaid labels.\n" + edit_tool._make_output(
+                    Path(absolute_path).read_text(encoding="utf-8"), path
+                )
             result = result + "\n---------- Mermaid validation ----------\n" + mermaid_validation
 
         return result

codewiki/src/be/documentation_generator.py

@@ -22,6 +22,7 @@
     organize_docs,
     relative_link,
 )
+from codewiki.src.be.utils import validate_mermaid_content
 from codewiki.src.be.module_naming import (
     dedupe_module_tree_names,
     find_missing_module_docs,
@@ -403,31 +404,50 @@ async def generate_parent_module_docs(
 
         phase = "overview" if not module_path else "parent_modules"
         try:
-            with self.usage_ledger.in_phase(phase, target=module_name):
-                parent_docs = self.backend.complete(
-                    prompt,
-                    system_prompt=format_overview_system_prompt(self.config.get_prompt_addition()),
-                )
-            if not parent_docs:
-                raise RuntimeError(
-                    f"LLM returned empty content for {module_name} overview "
-                    f"(possible output truncation at max_tokens)"
-                )
-
-            # Parse and save parent documentation. Subscription-CLI backends
-            # (claude-code / codex) sometimes ignore the <OVERVIEW> wrapper and
-            # return raw markdown; fall back to the response as-is in that case
-            # rather than crashing with an index error.
-            if "<OVERVIEW>" in parent_docs and "</OVERVIEW>" in parent_docs:
-                parent_content = parent_docs.split("<OVERVIEW>")[1].split("</OVERVIEW>")[0].strip()
-            else:
-                logger.warning(
-                    f"Overview response for {module_name} missing <OVERVIEW> wrapper; "
-                    f"using raw response as markdown."
+            request = prompt
+            for attempt in range(2):
+                with self.usage_ledger.in_phase(phase, target=module_name):
+                    parent_docs = self.backend.complete(
+                        request,
+                        system_prompt=format_overview_system_prompt(
+                            self.config.get_prompt_addition()
+                        ),
+                    )
+                if not parent_docs:
+                    raise RuntimeError(
+                        f"LLM returned empty content for {module_name} overview "
+                        f"(possible output truncation at max_tokens)"
+                    )
+                # Subscription backends may return raw Markdown instead of
+                # the requested wrapper. Validate either form before saving.
+                if "<OVERVIEW>" in parent_docs and "</OVERVIEW>" in parent_docs:
+                    parent_content = (
+                        parent_docs.split("<OVERVIEW>")[1].split("</OVERVIEW>")[0].strip()
+                    )
+                else:
+                    parent_content = parent_docs.strip()
+                report = await validate_mermaid_content(parent_content, repair=True)
+                if not report.errors:
+                    if report.unavailable:
+                        logger.warning(report.format(parent_docs_path))
+                    elif report.repairs:
+                        logger.info(report.format(parent_docs_path))
+                    break
+                if attempt == 1:
+                    raise RuntimeError(
+                        f"Invalid Mermaid in {module_name} overview after one correction attempt:\n"
+                        + report.format(parent_docs_path)
+                    )
+                request = (
+                    prompt
+                    + "\n\nCorrect the Mermaid syntax in this previous draft. Return the complete "
+                    "overview, preserving the other content.\n<PREVIOUS_DRAFT>\n"
+                    + report.content
+                    + "\n</PREVIOUS_DRAFT>\n"
+                    + report.format(parent_docs_path)
                 )
-                parent_content = parent_docs.strip()
             os.makedirs(os.path.dirname(parent_docs_path), exist_ok=True)
-            file_manager.save_text(parent_content, parent_docs_path)
+            file_manager.save_text(report.content, parent_docs_path)
 
             logger.debug(f"Successfully generated parent documentation for: {module_name}")
             return module_tree