ai-agent/scripts/schema-drift.py

#!/usr/bin/env python3
"""Detect Claude Code transcript-schema drift before it breaks the parser.

The ai-agent viewer parses Claude Code's on-disk transcripts
(~/.claude/projects/<proj>/<session>.jsonl) with backend/conversations.py.
That format is internal to Claude Code and undocumented, so a CLI update can
change it silently. This script is the early-warning system, two checks in one:

1. **Structural footprint** — walk recent transcripts and collect every key
   path / content-block type the records use (values are never recorded, only
   shapes, so the baseline is safe to commit). Compare against the committed
   baseline (scripts/schema-baseline.json): a NEW path means Claude Code
   started emitting something the parser has never seen; a MISSING path that
   the parser *relies on* would show up in check 2.

2. **Parser canary** — run the real backend/conversations.py over the same
   files and cross-check its summary against raw record counts (assistant
   turns with a usage block, token totals, model/timestamps extracted). This
   catches semantic breakage a structure diff can't (e.g. `usage` renamed:
   the new key is reported by check 1, and the zeroed token counts fail here).

A previously unseen Claude Code `version` stamp is reported as INFO only —
updates land ~daily and almost never drift; failing on every release would
just train you to ignore the check.

Usage:
  scripts/schema-drift.py                 # check newest transcripts vs baseline
  scripts/schema-drift.py --update        # re-baseline (after reviewing drift)
  scripts/schema-drift.py --days 30 --max-files 300   # wider sweep

Exit codes: 0 = no drift, 1 = drift or canary failure, 2 = usage/setup error.
"""

import argparse
import json
import os
import pathlib
import sys
import time

HERE = pathlib.Path(__file__).resolve().parent
DEFAULT_BASELINE = HERE / "schema-baseline.json"
DEFAULT_DIR = pathlib.Path.home() / ".claude" / "projects"

sys.path.insert(0, str(HERE.parent / "backend"))
try:
    import conversations  # noqa: E402  (the real parser — the thing we protect)
except Exception as e:  # pragma: no cover
    print(f"error: cannot import backend/conversations.py: {e}", file=sys.stderr)
    sys.exit(2)

# Subtrees whose *children* vary legitimately per tool / per feature — recording
# them would make every new tool look like schema drift. We record the node
# itself (so a rename of the container is still caught) but stop descending.
STOP_DESCENT_KEYS = {"input", "toolUseResult", "attachment", "usage"}
# …except usage: its keys are exactly what norm_usage() bills from, so those we
# DO want, one level deep (input_tokens, cache_creation.*, …).
USAGE_DEPTH = 2
MAX_DEPTH = 6


def _walk(node, prefix: str, depth: int, out: set[str]) -> None:
    if depth > MAX_DEPTH:
        return
    if isinstance(node, dict):
        for k, v in node.items():
            path = f"{prefix}.{k}"
            out.add(path)
            if k == "usage":
                _walk_usage(v, path, out)
            elif k in STOP_DESCENT_KEYS:
                continue
            else:
                _walk(v, path, depth + 1, out)
    elif isinstance(node, list):
        for item in node:
            if isinstance(item, dict) and isinstance(item.get("type"), str):
                # Content blocks: fold the block type into the path segment so
                # the block-type vocabulary (text/tool_use/thinking/…) — the
                # thing the parser branches on — is part of the footprint.
                bpath = f"{prefix}[{item['type']}]"
                out.add(bpath)
                _walk({k: v for k, v in item.items() if k != "type"},
                      bpath, depth + 1, out)
            else:
                _walk(item, f"{prefix}[]", depth + 1, out)


def _walk_usage(node, prefix: str, out: set[str]) -> None:
    if not isinstance(node, dict):
        return
    for k, v in node.items():
        out.add(f"{prefix}.{k}")
        if isinstance(v, dict):
            for k2 in v:
                out.add(f"{prefix}.{k}.{k2}")


def footprint(records) -> tuple[set[str], set[str]]:
    """(paths, cc_versions) for one transcript's records."""
    paths: set[str] = set()
    versions: set[str] = set()
    for o in records:
        t = o.get("type") or "?"
        v = o.get("version")
        if isinstance(v, str):
            versions.add(v)
        _walk(o, t, 0, paths)
    return paths, versions


# ── parser canary ────────────────────────────────────────────────────────────

def canary(path: pathlib.Path, records: list[dict]) -> list[str]:
    """Invariant violations between raw records and the parser's summary."""
    problems: list[str] = []
    s = conversations.parse_conversation(path)

    raw_usage_turns = sum(
        1 for o in records
        if o.get("type") == "assistant"
        and isinstance(o.get("message"), dict)
        and isinstance(o["message"].get("usage"), dict))
    if s["assistantTurns"] != raw_usage_turns:
        problems.append(
            f"assistantTurns={s['assistantTurns']} but {raw_usage_turns} raw "
            "assistant records carry a usage block")
    if raw_usage_turns and not s["tokens"]:
        problems.append("usage blocks present but summary tokens == 0")

    raw_models = {
        o["message"]["model"] for o in records
        if o.get("type") == "assistant"
        and isinstance(o.get("message"), dict)
        and conversations._real_model(o["message"].get("model"))}
    if raw_models and not s.get("models"):
        problems.append(f"raw models {sorted(raw_models)} but summary has none")

    if any(o.get("timestamp") for o in records) and not (
            s["startedAt"] and s["endedAt"]):
        problems.append("timestamps present but startedAt/endedAt missing")

    raw_session = next((o["sessionId"] for o in records
                        if isinstance(o.get("sessionId"), str)), None)
    if raw_session and s["sessionId"] != raw_session:
        problems.append("sessionId not extracted")
    return problems


# ── file selection ───────────────────────────────────────────────────────────

def pick_files(dirs: list[pathlib.Path], days: float,
               max_files: int) -> list[pathlib.Path]:
    """Newest .jsonl files modified within `days`, across all dirs (falls back
    to the newest few overall so the check never runs on nothing)."""
    cutoff = time.time() - days * 86400
    all_files: list[tuple[float, pathlib.Path]] = []
    for d in dirs:
        if not d.is_dir():
            continue
        for root, _dirs, names in os.walk(d):
            for n in names:
                if not n.endswith(".jsonl"):
                    continue
                p = pathlib.Path(root) / n
                try:
                    all_files.append((p.stat().st_mtime, p))
                except OSError:
                    continue
    all_files.sort(reverse=True)
    recent = [p for m, p in all_files if m >= cutoff]
    return (recent or [p for _, p in all_files[:5]])[:max_files]


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--dir", action="append", type=pathlib.Path,
                    help=f"transcript dir(s) to scan (default {DEFAULT_DIR})")
    ap.add_argument("--days", type=float, default=3.0,
                    help="scan files modified in the last N days (default 3)")
    ap.add_argument("--max-files", type=int, default=100,
                    help="cap on files scanned (default 100)")
    ap.add_argument("--baseline", type=pathlib.Path, default=DEFAULT_BASELINE)
    ap.add_argument("--update", action="store_true",
                    help="write the observed footprint as the new baseline")
    args = ap.parse_args()

    files = pick_files(args.dir or [DEFAULT_DIR], args.days, args.max_files)
    if not files:
        print("error: no .jsonl transcripts found", file=sys.stderr)
        return 2

    seen_paths: set[str] = set()
    seen_versions: set[str] = set()
    canary_failures: list[str] = []
    for p in files:
        records = list(conversations._iter_records(p))
        if not records:
            continue
        paths, versions = footprint(records)
        seen_paths |= paths
        seen_versions |= versions
        for problem in canary(p, records):
            canary_failures.append(f"{p.name}: {problem}")

    print(f"scanned {len(files)} transcript(s); "
          f"{len(seen_paths)} schema paths; "
          f"cc versions: {', '.join(sorted(seen_versions)) or '?'}")

    if args.update or not args.baseline.exists():
        if not args.baseline.exists() and not args.update:
            print("no baseline yet — writing one (first run)")
        if canary_failures:
            print(f"CANARY FAILED ({len(canary_failures)}) — fix the parser "
                  "before trusting this baseline:")
            for f in canary_failures:
                print(f"  ! {f}")
        args.baseline.write_text(json.dumps({
            "note": "structural footprint of Claude Code transcript JSONL — "
                    "key paths & block types only, no content",
            "ccVersions": sorted(seen_versions),
            "paths": sorted(seen_paths),
        }, indent=1) + "\n")
        print(f"baseline written: {args.baseline}")
        return 0

    base = json.loads(args.baseline.read_text())
    base_paths = set(base.get("paths") or [])
    base_versions = set(base.get("ccVersions") or [])

    new_paths = sorted(seen_paths - base_paths)
    gone_paths = sorted(base_paths - seen_paths)
    new_versions = sorted(seen_versions - base_versions)

    ok = True
    if new_versions:
        print(f"INFO: new Claude Code version(s) since baseline: "
              f"{', '.join(new_versions)}")
    if gone_paths:
        # Absence in a small recent sample is usually just "feature unused
        # this week" — informational unless the canary also failed.
        print(f"INFO: {len(gone_paths)} baseline path(s) unseen in this sample "
              "(likely just unused features)")
    if new_paths:
        ok = False
        print(f"DRIFT: {len(new_paths)} new schema path(s) not in baseline:")
        for p in new_paths:
            print(f"  + {p}")
        print("→ review whether backend/conversations.py should handle these, "
              "then re-run with --update to accept.")
    if canary_failures:
        ok = False
        print(f"CANARY FAILED ({len(canary_failures)}):")
        for f in canary_failures:
            print(f"  ! {f}")

    if ok:
        print("OK: no schema drift, parser invariants hold")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())