261 lines
8.7 KiB
Python
261 lines
8.7 KiB
Python
"""Compactor — runs out-of-band, mutates JSONL, signals wrapper to resume.
|
|
|
|
Invocation (by on_tool_use.py or /aircontext-now):
|
|
python compactor.py --project <root> --transcript <jsonl> --session-id <sid>
|
|
|
|
Lifecycle:
|
|
1. Acquire single-instance lock (compactor.lock).
|
|
2. Mark state.compaction_in_progress=true (already set by hook, idempotent).
|
|
3. Load active chain. If too short, abort.
|
|
4. Backup JSONL.
|
|
5. Slice tail (preserved) vs head (to compress).
|
|
6. Build conversation text + system prompt (rules.md).
|
|
7. Call LLM backend.
|
|
8. Append [summary, continuation] messages with parentUuid=null chain head.
|
|
9. Write state.compaction_ready=true so wrapper restarts claude.
|
|
"""
|
|
from __future__ import annotations
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
_HERE = Path(__file__).resolve().parent.parent
|
|
if str(_HERE) not in sys.path:
|
|
sys.path.insert(0, str(_HERE))
|
|
|
|
from core.backends.base import build_backend # noqa: E402
|
|
from core.config_loader import load_config # noqa: E402
|
|
from core.jsonl_ops import ( # noqa: E402
|
|
append_isolated_summary,
|
|
backup_jsonl,
|
|
find_active_chain,
|
|
load_messages,
|
|
rotate_snapshots,
|
|
)
|
|
from core.state import StateFile # noqa: E402
|
|
|
|
|
|
MIN_CHAIN_TO_COMPACT = 6
|
|
|
|
|
|
def _serialise_message_for_llm(msg: dict[str, Any]) -> str | None:
|
|
"""Reduce a JSONL message to plain text the LLM can read.
|
|
|
|
Returns None for messages that should be skipped entirely
|
|
(file-history-snapshot, queue-operation, etc.).
|
|
"""
|
|
t = msg.get("type")
|
|
inner = msg.get("message") or {}
|
|
if t == "user":
|
|
content = inner.get("content")
|
|
if isinstance(content, str):
|
|
return f"USER: {content}"
|
|
if isinstance(content, list):
|
|
parts = [p.get("text", "") if isinstance(p, dict) else str(p) for p in content]
|
|
return "USER: " + " ".join(p for p in parts if p)
|
|
if t == "assistant":
|
|
content = inner.get("content")
|
|
if isinstance(content, str):
|
|
return f"ASSISTANT: {content}"
|
|
if isinstance(content, list):
|
|
parts: list[str] = []
|
|
for p in content:
|
|
if not isinstance(p, dict):
|
|
continue
|
|
if p.get("type") == "text":
|
|
parts.append(p.get("text", ""))
|
|
elif p.get("type") == "tool_use":
|
|
parts.append(
|
|
f"[tool_use {p.get('name')}({json.dumps(p.get('input', {}), ensure_ascii=False)[:300]})]"
|
|
)
|
|
elif p.get("type") == "thinking":
|
|
pass # drop
|
|
return "ASSISTANT: " + " ".join(s for s in parts if s)
|
|
if t == "tool_result":
|
|
# `message.content` for tool_result is the result text
|
|
content = inner.get("content") if isinstance(inner, dict) else msg.get("content")
|
|
if isinstance(content, list):
|
|
content = " ".join(
|
|
(p.get("text", "") if isinstance(p, dict) else str(p)) for p in content
|
|
)
|
|
if not isinstance(content, str):
|
|
return None
|
|
return f"TOOL_RESULT: {content}"
|
|
if t in ("system",):
|
|
content = inner.get("content")
|
|
if isinstance(content, str):
|
|
return f"SYSTEM: {content}"
|
|
return None
|
|
|
|
|
|
def _truncate_long_tool_results(rendered: list[str], max_lines: int) -> list[str]:
|
|
out: list[str] = []
|
|
for r in rendered:
|
|
if not r.startswith("TOOL_RESULT: "):
|
|
out.append(r)
|
|
continue
|
|
body = r[len("TOOL_RESULT: "):]
|
|
lines = body.splitlines()
|
|
if len(lines) <= max_lines:
|
|
out.append(r)
|
|
continue
|
|
head = "\n".join(lines[:50])
|
|
tail = "\n".join(lines[-50:])
|
|
out.append(
|
|
"TOOL_RESULT: "
|
|
+ head
|
|
+ f"\n... [{len(lines) - 100} lines omitted by AirContext pre-truncation] ...\n"
|
|
+ tail
|
|
)
|
|
return out
|
|
|
|
|
|
def run(project: Path, transcript: Path, session_id: str) -> int:
|
|
air = project / "AirContext"
|
|
cfg_path = air / "config.yaml"
|
|
state = StateFile(air / "state.json")
|
|
|
|
try:
|
|
cfg = load_config(cfg_path)
|
|
except Exception as e:
|
|
print(f"[compactor] config load failed: {e}", file=sys.stderr)
|
|
state.update(compaction_in_progress=False)
|
|
return 2
|
|
|
|
if state.read().get("paused"):
|
|
state.update(compaction_in_progress=False)
|
|
return 0
|
|
|
|
# Single-instance lock (best-effort; sufficient for single-user dev environment).
|
|
lock = air / "compactor.lock"
|
|
try:
|
|
fd = os.open(str(lock), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
|
os.write(fd, str(os.getpid()).encode())
|
|
os.close(fd)
|
|
except FileExistsError:
|
|
print("[compactor] another compactor is already running; exiting", file=sys.stderr)
|
|
return 0
|
|
|
|
try:
|
|
return _run_locked(project, transcript, session_id, cfg, state, air)
|
|
finally:
|
|
try:
|
|
os.unlink(lock)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _run_locked(
|
|
project: Path,
|
|
transcript: Path,
|
|
session_id: str,
|
|
cfg: dict,
|
|
state: StateFile,
|
|
air: Path,
|
|
) -> int:
|
|
state.update(compaction_in_progress=True, last_session_id=session_id)
|
|
msgs = load_messages(transcript)
|
|
chain = find_active_chain(msgs)
|
|
if len(chain) < MIN_CHAIN_TO_COMPACT:
|
|
print(
|
|
f"[compactor] chain too short ({len(chain)}); skipping",
|
|
file=sys.stderr,
|
|
)
|
|
state.update(compaction_in_progress=False)
|
|
return 0
|
|
|
|
comp_cfg = cfg.get("compaction") or {}
|
|
safety = cfg.get("safety") or {}
|
|
preserve_tail = int(comp_cfg.get("preserve_tail_messages", 10))
|
|
drop_lines = int(comp_cfg.get("drop_tool_results_over_lines", 1000))
|
|
rules_file = comp_cfg.get("rules_file", "rules.md")
|
|
continuation_prompt = comp_cfg.get("continuation_prompt", "") or ""
|
|
|
|
rules_path = air / rules_file
|
|
rules_text = rules_path.read_text(encoding="utf-8") if rules_path.exists() else ""
|
|
|
|
head = chain[:-preserve_tail] if preserve_tail > 0 else chain
|
|
tail = chain[-preserve_tail:] if preserve_tail > 0 else []
|
|
|
|
rendered = [r for r in (_serialise_message_for_llm(m) for m in head) if r]
|
|
rendered = _truncate_long_tool_results(rendered, drop_lines)
|
|
conversation_text = "\n\n".join(rendered)
|
|
if not conversation_text.strip():
|
|
print("[compactor] nothing to summarise", file=sys.stderr)
|
|
state.update(compaction_in_progress=False)
|
|
return 0
|
|
|
|
system_prompt = (
|
|
"You are a context-compression engine for a Claude Code session. "
|
|
"Apply the user-supplied compression rules below to the conversation "
|
|
"transcript that follows. Output ONLY the compressed summary text. "
|
|
"Do not add preamble, headers, or apologies.\n\n"
|
|
"=== USER COMPRESSION RULES ===\n"
|
|
f"{rules_text}\n"
|
|
"=== END RULES ==="
|
|
)
|
|
|
|
backend = build_backend(cfg)
|
|
print(
|
|
f"[compactor] calling {cfg['backend']['endpoint']} model={cfg['backend']['model']} "
|
|
f"head_msgs={len(head)} tail_preserved={len(tail)}",
|
|
file=sys.stderr,
|
|
)
|
|
try:
|
|
summary_text = backend.summarise(system_prompt, conversation_text)
|
|
except Exception as e:
|
|
print(f"[compactor] LLM call failed: {e}", file=sys.stderr)
|
|
state.update(compaction_in_progress=False)
|
|
return 3
|
|
|
|
if safety.get("dry_run"):
|
|
print(
|
|
f"[compactor] dry_run=true; summary length={len(summary_text)}; "
|
|
"JSONL not modified",
|
|
file=sys.stderr,
|
|
)
|
|
state.update(compaction_in_progress=False, last_compaction_unix=int(time.time()))
|
|
return 0
|
|
|
|
if safety.get("backup", True):
|
|
backup_jsonl(transcript, air / "snapshots", session_id)
|
|
|
|
template = chain[-1] # use the latest message's metadata as template
|
|
summary_uuid, cont_uuid = append_isolated_summary(
|
|
transcript,
|
|
summary_text=summary_text,
|
|
continuation_text=continuation_prompt or None,
|
|
template_message=template,
|
|
)
|
|
print(
|
|
f"[compactor] appended summary={summary_uuid} continuation={cont_uuid}",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
rotate_snapshots(air / "snapshots", int(safety.get("max_snapshots", 50)))
|
|
|
|
state.update(
|
|
compaction_in_progress=False,
|
|
compaction_ready=True,
|
|
pending_resume_session_id=session_id,
|
|
last_compaction_unix=int(time.time()),
|
|
)
|
|
return 0
|
|
|
|
|
|
def main() -> int:
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--project", required=True)
|
|
p.add_argument("--transcript", required=True)
|
|
p.add_argument("--session-id", required=True)
|
|
args = p.parse_args()
|
|
return run(Path(args.project), Path(args.transcript), args.session_id)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|