From 4730ed9abb792cf2570b4feb7163c47efa79c9ef Mon Sep 17 00:00:00 2001
From: Ben Younes <2910651+ousamabenyounes@users.noreply.github.com>
Date: Wed, 12 Aug 2026 14:20:21 +0100
Subject: [PATCH] fix(watch): refuse a shrink caused by an extractor failure
 during update (#2663)

---
 graphify/watch.py   | 11 +++++++++
 tests/test_watch.py | 56 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/graphify/watch.py b/graphify/watch.py
index 4f2b143f..4aadf749 100644
--- a/graphify/watch.py
+++ b/graphify/watch.py
@@ -847,6 +847,7 @@ def _check_shrink(
     *,
     had_explicit_deletions: bool = False,
     rebuilt_sources: "set[str] | None" = None,
+    failed_sources: "set[str] | None" = None,
 ) -> bool:
     """Return True (ok to proceed) or False (shrink refused).
 
@@ -866,6 +867,8 @@ def _check_shrink(
     NOT touch — e.g. a dropped semantic/doc node) refuses the write. This lets a
     plain ``graphify update`` after deleting a function refresh the graph without
     ``--force`` (#1116 left stale nodes write-blocked even though build dropped them).
+    Files in ``failed_sources`` never account for lost nodes: extraction did not
+    complete, so their disappearance is the silent shrink this guard protects.
     """
     if force or not existing_data:
         return True
@@ -890,6 +893,8 @@ def _check_shrink(
 
         def _accounted(n: dict) -> bool:
             sf = n.get("source_file")
+            if sf and failed_sources and _norm_source_file(sf) in failed_sources:
+                return False
             return (not sf
                     or sf in rebuilt_sources
                     or _norm_source_file(sf) in rebuilt_sources)
@@ -1407,6 +1412,10 @@ def _rebuild_code(
         else:
             rebuilt_sources = {(_nsf(str(p), _rebuilt_root) or str(p)) for p in extract_targets}
         rebuilt_sources |= set(deleted_paths)
+        failed_sources = {
+            _nsf(source, _rebuilt_root) or source
+            for source in _failed_ast_sources
+        }
         out.mkdir(exist_ok=True)
 
         if no_cluster:
@@ -1454,6 +1463,7 @@ def _rebuild_code(
                     force, existing_graph_data, candidate_graph_data,
                     had_explicit_deletions=bool(deleted_paths),
                     rebuilt_sources=rebuilt_sources,
+                    failed_sources=failed_sources,
                 ):
                     return False
                 from graphify.export import backup_if_protected as _backup
@@ -1654,6 +1664,7 @@ def _rebuild_code(
                 tmp=graph_tmp,
                 had_explicit_deletions=bool(deleted_paths),
                 rebuilt_sources=rebuilt_sources,
+                failed_sources=failed_sources,
             ):
                 return False
             from graphify.export import backup_if_protected as _backup
diff --git a/tests/test_watch.py b/tests/test_watch.py
index 401ddde9..00572dde 100644
--- a/tests/test_watch.py
+++ b/tests/test_watch.py
@@ -9,6 +9,7 @@ import pytest
 
 from graphify.watch import (
     _notify_only,
+    _rebuild_code,
     _WATCHED_EXTENSIONS,
     _rebuild_lock,
     _check_shrink,
@@ -1256,6 +1257,61 @@ def test_check_shrink_blocks_shrink_outside_rebuilt_sources(capsys):
     assert "Refusing to overwrite" in capsys.readouterr().err
 
 
+def test_check_shrink_blocks_loss_from_failed_rebuilt_source(capsys):
+    """A failed extractor must not account for nodes it dropped during rebuild."""
+    existing = {"nodes": [
+        {"id": "app", "source_file": "app.py"},
+        {"id": "table", "source_file": "schema.sql"},
+    ], "links": []}
+    new = {"nodes": [{"id": "app", "source_file": "app.py"}], "links": []}
+
+    ok = _check_shrink(
+        False,
+        existing,
+        new,
+        rebuilt_sources={"app.py", "schema.sql"},
+        failed_sources={"schema.sql"},
+    )
+
+    assert ok is False
+    assert "Refusing to overwrite" in capsys.readouterr().err
+
+
+@pytest.mark.parametrize("no_cluster", [False, True])
+def test_rebuild_refuses_loss_from_failed_source(tmp_path, monkeypatch, no_cluster):
+    """A failed AST extractor must not overwrite its last good graph."""
+    previous_sql_node_count = 10
+    (tmp_path / "app.py").write_text(
+        "def alpha(): return beta()\ndef beta(): return 1\n",
+        encoding="utf-8",
+    )
+    (tmp_path / "schema.sql").write_text(
+        "create table cliente (id int primary key);\n",
+        encoding="utf-8",
+    )
+    out = tmp_path / "graphify-out"
+    out.mkdir()
+    graph_path = out / "graph.json"
+    existing_nodes = [
+        {"id": "app.py", "label": "app.py", "source_file": "app.py", "type": "file", "_origin": "ast"},
+        {"id": "app.py:alpha", "label": "alpha", "source_file": "app.py", "type": "function", "_origin": "ast"},
+        {"id": "app.py:beta", "label": "beta", "source_file": "app.py", "type": "function", "_origin": "ast"},
+        {"id": "schema.sql", "label": "schema.sql", "source_file": "schema.sql", "type": "file", "_origin": "ast"},
+    ]
+    existing_nodes.extend(
+        {"id": f"sql:table_{index}", "label": f"table_{index}", "source_file": "schema.sql", "type": "table", "_origin": "ast"}
+        for index in range(previous_sql_node_count)
+    )
+    existing = {"nodes": existing_nodes, "links": []}
+    graph_path.write_text(json.dumps(existing), encoding="utf-8")
+    monkeypatch.setitem(sys.modules, "tree_sitter_sql", None)
+
+    ok = _rebuild_code(tmp_path, force=False, no_cluster=no_cluster)
+
+    assert ok is False
+    assert json.loads(graph_path.read_text(encoding="utf-8")) == existing
+
+
 def test_check_shrink_allows_growth():
     """new > existing is always fine."""
     ok = _check_shrink(