fix(schema-drift): atomic drift merge, unlimited sample mode, baseline root check (#108)

Author: clean6378-max-it

tests/test_schema_drift.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import threading
 from pathlib import Path
 
 import pytest
@@ -168,3 +169,46 @@ def test_schema_drift_disabled_skips_fingerprint(self, monkeypatch):
         monkeypatch.setenv("CLAUDE_CODE_CHAT_BROWSER_SCHEMA_DRIFT", "0")
         parse_session(str(UNKNOWN_FIELD_FIXTURE))
         assert get_schema_report()["has_drift"] is False
+
+    def test_sample_limit_one_skips_later_records(self, monkeypatch):
+        monkeypatch.setenv("CLAUDE_CODE_CHAT_BROWSER_SCHEMA_DRIFT_SAMPLE", "1")
+        parse_session(str(UNKNOWN_FIELD_FIXTURE))
+        report = get_schema_report()
+        assert "tool" not in report["new_fields"]
+
+    def test_sample_limit_zero_fingerprints_all_records(self, monkeypatch):
+        monkeypatch.setenv("CLAUDE_CODE_CHAT_BROWSER_SCHEMA_DRIFT_SAMPLE", "0")
+        parse_session(str(UNKNOWN_FIELD_FIXTURE))
+        report = get_schema_report()
+        assert "tool" in report["new_fields"]
+
+    def test_non_object_baseline_root_is_non_fatal(self, tmp_path, monkeypatch):
+        bad_baseline = tmp_path / "schema_baseline.json"
+        bad_baseline.write_text('["not", "an", "object"]', encoding="utf-8")
+        monkeypatch.setattr("utils.schema_drift.BASELINE_PATH", bad_baseline)
+        clear_baseline_cache()
+        assert record_parse_drift({"type"}) is None
+
+    def test_concurrent_record_parse_drift_merges_all_fields(self):
+        start = threading.Barrier(10)
+        errors: list[BaseException] = []
+
+        def worker(field: str) -> None:
+            try:
+                start.wait(timeout=5)
+                record_parse_drift({"type", field})
+            except BaseException as exc:
+                errors.append(exc)
+
+        threads = [
+            threading.Thread(target=worker, args=(f"concurrent_field_{i}",)) for i in range(10)
+        ]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+
+        assert errors == []
+        report = get_schema_report()
+        for i in range(10):
+            assert f"concurrent_field_{i}" in report["new_fields"]

utils/jsonl_parser.py

@@ -175,7 +175,11 @@ def parse_session(filepath: str) -> SessionDict:
     messages: list[MessageDict] = []
     metadata = _new_session_metadata_builder(session_id)
     observed_field_paths: set[str] = set()
-    schema_samples_remaining = schema_drift_sample_limit() if is_schema_drift_enabled() else 0
+    if is_schema_drift_enabled():
+        sample_limit = schema_drift_sample_limit()
+        schema_samples_remaining: int | None = None if sample_limit == 0 else sample_limit
+    else:
+        schema_samples_remaining = -1
 
     with open(filepath, "r", encoding="utf-8", errors="replace") as f:
         for line in f:
@@ -190,9 +194,12 @@ def parse_session(filepath: str) -> SessionDict:
             if not isinstance(entry, dict):
                 continue
 
-            if schema_samples_remaining > 0:
+            if schema_samples_remaining != -1 and (
+                schema_samples_remaining is None or schema_samples_remaining > 0
+            ):
                 observed_field_paths |= collect_field_paths(entry)
-                schema_samples_remaining -= 1
+                if schema_samples_remaining is not None:
+                    schema_samples_remaining -= 1
 
             entry_type = entry.get("type")
             ts = _entry_timestamp(entry)

utils/schema_drift.py

@@ -41,7 +41,7 @@ def is_schema_drift_enabled() -> bool:
 
 
 def schema_drift_sample_limit() -> int:
-    """Max JSONL records per session to fingerprint (default 3). Set 0 to disable sampling cap."""
+    """Max JSONL records per session to fingerprint (default 3). 0 means no cap (all records)."""
     raw = os.environ.get("CLAUDE_CODE_CHAT_BROWSER_SCHEMA_DRIFT_SAMPLE", "3").strip()
     try:
         return max(0, int(raw))
@@ -75,6 +75,8 @@ def _load_baseline() -> tuple[frozenset[str], frozenset[str]]:
         return _baseline_cache
     try:
         raw = json.loads(BASELINE_PATH.read_text(encoding="utf-8"))
+        if not isinstance(raw, dict):
+            raise ValueError("schema_baseline.json: root must be an object")
         fields = raw.get("fields", {})
         if not isinstance(fields, dict):
             raise ValueError("schema_baseline.json: 'fields' must be an object")
@@ -122,30 +124,31 @@ def record_parse_drift(observed_paths: set[str]) -> SchemaDriftReport | None:
     except (OSError, json.JSONDecodeError, ValueError, TypeError):
         return None
 
+    genuinely_new: list[str] = []
+    missing_fields: list[str] = []
     with _lock:
         global _last_report
         prior_new = set(_last_report["new_fields"])
         genuinely_new = sorted(set(report["new_fields"]) - prior_new)
         merged_new = sorted(prior_new | set(report["new_fields"]))
+        missing_fields = list(report["missing_fields"])
+        _last_report = {
+            "known_fields": report["known_fields"],
+            "new_fields": merged_new,
+            "missing_fields": missing_fields,
+            "has_drift": bool(merged_new or missing_fields),
+        }
 
     if genuinely_new:
         _log.warning(
             "schema drift: new JSONL field paths not in baseline: %s",
             genuinely_new,
         )
-    if report["missing_fields"]:
+    if missing_fields:
         _log.warning(
             "schema drift: missing required JSONL field paths in sampled records: %s",
-            report["missing_fields"],
+            missing_fields,
         )
-
-    with _lock:
-        _last_report = {
-            "known_fields": report["known_fields"],
-            "new_fields": merged_new,
-            "missing_fields": list(report["missing_fields"]),
-            "has_drift": bool(merged_new or report["missing_fields"]),
-        }
     return report