Files
aircontext/scripts/core/jsonl_ops.py
2026-05-18 11:45:08 +08:00

159 lines
5.0 KiB
Python

"""JSONL transcript parsing + isolated-summary append.
The Claude Code transcript is a JSON-Lines file where each line is a message.
Messages form a DAG via `parentUuid`; the "active chain" is the path from the
latest leaf back to the first message with parentUuid==null.
We append a NEW chain by writing two `user` messages at the tail of the file:
1. summary (parentUuid: null) — becomes a fresh chain root
2. continuation (parentUuid: summary.uuid) — becomes the latest leaf
On `claude --resume <id>` the leaf-selection algorithm picks the continuation
message (newest leaf) and walks back to the summary, so the loaded context is
exactly those two messages plus the system/CLAUDE.md/etc. fixed cost.
"""
from __future__ import annotations
import datetime as _dt
import json
import shutil
import uuid as _uuid
from pathlib import Path
from typing import Any, Iterator
def iter_messages(path: Path) -> Iterator[dict[str, Any]]:
with path.open("r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
yield json.loads(line)
except json.JSONDecodeError:
continue
def load_messages(path: Path) -> list[dict[str, Any]]:
return list(iter_messages(path))
def find_active_chain(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Return messages on the active chain, ordered from root to leaf.
Mirrors `claude --resume` behaviour:
1. Build uuid -> message map.
2. Find leaves: uuids not referenced as a parentUuid by any other message.
3. Pick leaf with latest timestamp.
4. Walk parentUuid pointers back until parentUuid is null.
"""
by_uuid = {m["uuid"]: m for m in messages if "uuid" in m}
referenced = {m.get("parentUuid") for m in messages if m.get("parentUuid")}
leaves = [m for m in messages if m.get("uuid") and m["uuid"] not in referenced]
if not leaves:
return []
def _ts(m: dict[str, Any]) -> str:
return m.get("timestamp", "")
leaf = max(leaves, key=_ts)
chain: list[dict[str, Any]] = []
cur: dict[str, Any] | None = leaf
seen: set[str] = set()
while cur is not None and cur.get("uuid") not in seen:
chain.append(cur)
seen.add(cur["uuid"])
parent_uuid = cur.get("parentUuid")
if not parent_uuid:
break
cur = by_uuid.get(parent_uuid)
chain.reverse()
return chain
def _meta_from(reference: dict[str, Any]) -> dict[str, Any]:
"""Copy non-content metadata fields from a reference message.
We deliberately only copy metadata fields, never the content/message field,
so callers control what payload the new message carries.
"""
keep = ("sessionId", "cwd", "version", "gitBranch", "userType")
return {k: reference[k] for k in keep if k in reference}
def _now_iso() -> str:
return _dt.datetime.now(_dt.timezone.utc).isoformat(timespec="milliseconds").replace(
"+00:00", "Z"
)
def build_user_message(
*,
parent_uuid: str | None,
content: str,
template: dict[str, Any],
) -> dict[str, Any]:
msg: dict[str, Any] = {
"type": "user",
"uuid": str(_uuid.uuid4()),
"parentUuid": parent_uuid,
"timestamp": _now_iso(),
"isSidechain": False,
**_meta_from(template),
"message": {"role": "user", "content": content},
}
msg.setdefault("userType", "external")
return msg
def backup_jsonl(path: Path, snapshots_dir: Path, session_id: str) -> Path:
snapshots_dir.mkdir(parents=True, exist_ok=True)
ts = _dt.datetime.now().strftime("%Y%m%d-%H%M%S")
dst = snapshots_dir / f"{ts}-{session_id}.jsonl"
shutil.copyfile(path, dst)
return dst
def append_isolated_summary(
path: Path,
*,
summary_text: str,
continuation_text: str | None,
template_message: dict[str, Any],
) -> tuple[str, str | None]:
"""Append summary (parentUuid=null) and optional continuation.
Returns (summary_uuid, continuation_uuid_or_None).
"""
summary = build_user_message(
parent_uuid=None, content=summary_text, template=template_message
)
lines: list[str] = [json.dumps(summary, ensure_ascii=False) + "\n"]
cont_uuid: str | None = None
if continuation_text:
cont = build_user_message(
parent_uuid=summary["uuid"],
content=continuation_text,
template=template_message,
)
cont_uuid = cont["uuid"]
lines.append(json.dumps(cont, ensure_ascii=False) + "\n")
with path.open("a", encoding="utf-8") as f:
f.writelines(lines)
return summary["uuid"], cont_uuid
def rotate_snapshots(snapshots_dir: Path, max_keep: int) -> None:
if not snapshots_dir.exists():
return
files = sorted(
(p for p in snapshots_dir.iterdir() if p.is_file() and p.suffix == ".jsonl"),
key=lambda p: p.stat().st_mtime,
)
excess = len(files) - max_keep
for p in files[:max(0, excess)]:
try:
p.unlink()
except OSError:
pass