diff --git a/.gitignore b/.gitignore index 6343d01..9cf1595 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ __pycache__/ .venv/ eval_cache/ REVIEW.md +.worktrees diff --git a/README.md b/README.md index 26371a8..8924840 100644 --- a/README.md +++ b/README.md @@ -224,6 +224,18 @@ PLUGIN_SETTING_BUDGET=80000 # token budget PLUGIN_SETTING_FALLBACK_TO_BUILTIN=true # fall back to LLM on error ``` +### Default token budget + +| Entry point | Default | Notes | +|---|---|---| +| `compact.py --budget` | 80,000 | CLI flag | +| `compact_codex.py --budget` | 80,000 | CLI flag | +| `compact-session.sh` (slash cmd) | 80,000 | `PLUGIN_SETTING_BUDGET` env | +| `supercompact-precompact.sh` (PreCompact hook) | 80,000 | `PLUGIN_SETTING_BUDGET` env | +| Patched `cli.js` (LLM fallback path) | `max(40% × preCompactTokenCount, 40000)` | Intentionally percentage-based: scales with context size when Claude's auto-compact fires | + +The patched `cli.js` path is the odd one out on purpose — when Claude's built-in auto-compact fires, the context size is known at call time (`preCompactTokenCount`), so scaling the target by a percentage gives consistent behavior regardless of how large the conversation grew. All the other entry points are explicitly invoked (manually or by the PreCompact hook) and use a fixed default. + The Claude Code plugin also provides a `/supercompact` slash command for manual compaction: ``` diff --git a/compact.py b/compact.py index 07b8b9e..ab8721a 100644 --- a/compact.py +++ b/compact.py @@ -126,8 +126,14 @@ def cmd_compact(args: argparse.Namespace) -> int: total_kept = result.user_tokens + result.short_system_tokens + result.scored_kept_tokens if total_kept > args.budget: console.print(f"\n[yellow]Primary method '{args.method}' couldn't reach budget ({total_kept:,} > {args.budget:,}). Falling back to NUCLEAR truncation.[/yellow]") + # Preserve the primary pass's dropped turns so the archive captures them. + # nuclear_compact returns a fresh SelectionResult with no dropped_turns + # (it truncates in place instead of discarding turns). + primary_dropped = result.dropped_turns from lib.nuclear import nuclear_compact result = nuclear_compact(result.kept_turns, args.budget) + result.dropped_turns = primary_dropped + result.scored_dropped_tokens = sum(st.tokens for st in primary_dropped) t_elapsed = time.monotonic() - t_start diff --git a/lib/nuclear.py b/lib/nuclear.py index 37b8cf9..6ffb1a2 100644 --- a/lib/nuclear.py +++ b/lib/nuclear.py @@ -15,23 +15,17 @@ def nuclear_compact(turns: list[Turn], budget: int) -> SelectionResult: Maintains the structural integrity of the JSONL (UUIDs, turn types) but severely truncates long tool outputs and text blocks. """ - console.print(f"Running NUCLEAR compaction with budget {budget:,} tokens...") - kept_turns = copy.deepcopy(turns) - total_tokens_before = sum(turn_tokens(t) for t in turns) - - # Simple strategy: iterative truncation - # First, truncate anything over 10,000 chars. Then 5000, 2000, 1000, 500 - # Stop when we fit the budget. - + + # Iterative truncation: shrink long blocks until we fit, or exhaust steps. current_tokens = total_tokens_before - + final_max_len: int | None = None + for max_len in [10000, 5000, 2000, 1000, 500, 250, 100]: if current_tokens <= budget: break - - console.print(f" Over budget ({current_tokens:,} > {budget:,}). Squashing to max {max_len} chars per block...") + final_max_len = max_len for i, turn in enumerate(kept_turns): for record in turn.lines: @@ -66,6 +60,13 @@ def nuclear_compact(turns: list[Turn], budget: int) -> SelectionResult: current_tokens = sum(turn_tokens(t) for t in kept_turns) + if final_max_len is not None: + status = "fit budget" if current_tokens <= budget else "gave up" + console.print( + f"Nuclear: {total_tokens_before:,} -> {current_tokens:,} tokens " + f"(budget {budget:,}, min block {final_max_len} chars, {status})" + ) + result = SelectionResult( kept_turns=kept_turns, budget=budget, diff --git a/lib/parser.py b/lib/parser.py index 4383083..9447e18 100644 --- a/lib/parser.py +++ b/lib/parser.py @@ -120,15 +120,66 @@ def parse_jsonl(path: Path) -> list[Turn]: return turns -def extract_text(turn: Turn) -> str: +_ATTACHMENT_TEXT_FIELDS = ( + "content", + "stdout", + "stderr", + "prompt", + "snippet", + "command", +) + +_ATTACHMENT_LIST_FIELDS = ( + "addedLines", + "addedNames", + "removedNames", +) + + +def _extract_attachment_text(att: dict) -> list[str]: + """Pull Claude-visible text out of an attachment record. + + Attachments carry hook output (stdout/stderr/command), file reads + (content/displayPath), slash command listings (content), prompt + previews, system reminders, etc. All of this gets injected into the + model's context — so it has to count toward the token budget. + """ + out: list[str] = [] + subtype = att.get("type", "") + if subtype: + out.append(f"[attachment:{subtype}]") + for field in _ATTACHMENT_TEXT_FIELDS: + v = att.get(field) + if isinstance(v, str) and v: + out.append(v) + for field in _ATTACHMENT_LIST_FIELDS: + v = att.get(field) + if isinstance(v, list): + for item in v: + out.append(str(item)) + return out + + +def extract_text(turn: Turn, *, truncate: bool = True) -> str: """Extract human-readable text from a turn for scoring/display. Concatenates message content strings, thinking text, tool_use names/inputs, - and tool_result content into a single string. + tool_result content, and attachment payloads into a single string. + + With ``truncate=True`` (default), tool_use inputs are capped so relevance + scorers aren't dominated by large file contents. With ``truncate=False``, + the full payload is preserved — use this for accurate token counting, + since the API receives the untruncated tool input. """ parts: list[str] = [] for record in turn.lines: + # Attachment records carry hook output / file reads / slash command + # output — injected into context but not wrapped in message.content. + attachment = record.get("attachment") + if isinstance(attachment, dict): + parts.extend(_extract_attachment_text(attachment)) + msg = record.get("message", {}) content = msg.get("content") @@ -150,11 +201,11 @@ def extract_text(turn: Turn) -> str: if isinstance(inp, dict): for k, v in inp.items(): v_str = str(v) - if len(v_str) > 500: + if truncate and len(v_str) > 500: v_str = v_str[:500] + "..." parts.append(f" {k}: {v_str}") elif isinstance(inp, str): - parts.append(inp[:1000]) + parts.append(inp[:1000] if truncate else inp) elif btype == "tool_result": result_content = block.get("content", "") if isinstance(result_content, str): diff --git a/lib/selector.py b/lib/selector.py index 1a36120..45f77e9 100644 --- a/lib/selector.py +++ b/lib/selector.py @@ -86,7 +86,10 @@ def select_turns( tc = token_counts.get(last_system.index, 0) kept_indices.add(last_system.index) used_tokens += tc - # Track it in scored_kept if it was scored + # The tail turn is long (short system turns are already in kept_indices). + # Its tokens must flow into scored_kept_tokens so total_kept accounting + # reflects reality and callers can detect budget overshoot. + result.scored_kept_tokens += tc if last_system.index in scored_map: result.kept_scored.append(scored_map[last_system.index]) diff --git a/lib/tokenizer.py b/lib/tokenizer.py index 2b880af..6eb1237 100644 --- a/lib/tokenizer.py +++ b/lib/tokenizer.py @@ -1,26 +1,50 @@ -"""Token counting using Qwen3 tokenizer.""" +"""Token counting using Qwen3 tokenizer with graceful fallback.""" from __future__ import annotations -from transformers import AutoTokenizer +import os +import sys from .parser import Turn, extract_text _tokenizer = None - - -def _get_tokenizer() -> AutoTokenizer: - global _tokenizer - if _tokenizer is None: +_fallback = False + + +def _get_tokenizer(): + """Load Qwen3 tokenizer lazily. On any failure (missing model, perm + error, offline, transformers absent) switch to a byte-heuristic + fallback so EITF/dedup/setcover methods still work without a model.""" + global _tokenizer, _fallback + if _tokenizer is not None or _fallback: + return _tokenizer + try: + from transformers import AutoTokenizer _tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-Embedding-0.6B") + except Exception as e: + _fallback = True + print( + f"[tokenizer] WARN: Qwen tokenizer unavailable ({type(e).__name__}: {e}); " + f"using byte-heuristic fallback (len/4).", + file=sys.stderr, + ) return _tokenizer def estimate_tokens(text: str) -> int: - """Count tokens using the Qwen3 tokenizer.""" - return len(_get_tokenizer().encode(text, add_special_tokens=False)) + """Count tokens using the Qwen3 tokenizer, or byte-heuristic fallback.""" + tok = _get_tokenizer() + if tok is None: + # Rough heuristic: ~4 bytes/token for English; slightly under-counts + # CJK but safe for budget enforcement. + return max(1, len(text) // 4) if text else 0 + return len(tok.encode(text, add_special_tokens=False)) def turn_tokens(turn: Turn) -> int: - """Count the tokens of an entire turn.""" - return estimate_tokens(extract_text(turn)) + """Count the tokens of an entire turn. + + Uses the untruncated extraction so tool_use payloads (which the API + receives in full) are counted at their real size, not the scoring stub. + """ + return estimate_tokens(extract_text(turn, truncate=False)) diff --git a/plugins/claude-code/hooks-handlers/supercompact-precompact.sh b/plugins/claude-code/hooks-handlers/supercompact-precompact.sh index 0e70ec3..f20da99 100755 --- a/plugins/claude-code/hooks-handlers/supercompact-precompact.sh +++ b/plugins/claude-code/hooks-handlers/supercompact-precompact.sh @@ -67,7 +67,9 @@ JSONL_SIZE=$(wc -l < "${JSONL_FILE}") echo "$(date -Iseconds) Transcript: ${JSONL_FILE} (${JSONL_SIZE} lines)" >> "${LOG_DIR}/hook.log" # 1. Back up the full transcript before Claude's compaction destroys detail -BACKUP_FILE="${JSONL_FILE}.pre-compact-full" +# Use a timestamped filename so consecutive compactions don't overwrite each other. +BACKUP_STAMP=$(date +%Y%m%d-%H%M%S) +BACKUP_FILE="${JSONL_FILE}.pre-compact-full.${BACKUP_STAMP}" cp "${JSONL_FILE}" "${BACKUP_FILE}" echo "$(date -Iseconds) Full backup saved: ${BACKUP_FILE}" >> "${LOG_DIR}/hook.log" @@ -90,16 +92,20 @@ if uv run python compact.py "${JSONL_FILE}" \ SC_SIZE=$(wc -l < "${SC_OUTPUT}") echo "$(date -Iseconds) Supercompact (${METHOD}): ${JSONL_SIZE} -> ${SC_SIZE} lines" >> "${LOG_DIR}/hook.log" - # Back up the pre-compaction JSONL (in addition to the .pre-compact-full backup above) - cp "${JSONL_FILE}" "${JSONL_FILE}.pre-supercompact" + # Back up the pre-compaction JSONL (in addition to the .pre-compact-full backup above). + # Timestamped so repeated compactions don't clobber prior backups. + cp "${JSONL_FILE}" "${JSONL_FILE}.pre-supercompact.${BACKUP_STAMP}" # Swap in the compacted version mv "${SC_OUTPUT}" "${JSONL_FILE}" echo "$(date -Iseconds) SUCCESS: Swapped compacted JSONL in-place" >> "${LOG_DIR}/hook.log" - # Clean up old backups (keep last 3 of each type) - ls -t "${JSONL_FILE}.pre-compact-full"* 2>/dev/null | tail -n +4 | xargs rm -f 2>/dev/null || true - ls -t "${JSONL_FILE}.pre-supercompact"* 2>/dev/null | tail -n +4 | xargs rm -f 2>/dev/null || true + # Clean up old backups (keep last 3 of each type). Now that names are + # timestamped, ls -t reliably orders them by creation time. + # shellcheck disable=SC2012 + ls -t "${JSONL_FILE}".pre-compact-full.* 2>/dev/null | tail -n +4 | xargs -r rm -f 2>/dev/null || true + # shellcheck disable=SC2012 + ls -t "${JSONL_FILE}".pre-supercompact.* 2>/dev/null | tail -n +4 | xargs -r rm -f 2>/dev/null || true # Restart Claude to load the compacted context — this kills the process # before the Anthropic API compaction call fires. unleash-refresh adds diff --git a/plugins/claude-code/scripts/patcher.py b/plugins/claude-code/scripts/patcher.py index b361d03..bda7b13 100644 --- a/plugins/claude-code/scripts/patcher.py +++ b/plugins/claude-code/scripts/patcher.py @@ -113,8 +113,15 @@ def build_replacement(m: re.Match, supercompact_dir: str) -> str: f'_fs=require("fs");' f'const _home=process.env.HOME||"/root";' f'const _rawMethod=process.env.PLUGIN_SETTING_METHOD||"eitf";' + # Only local, fast methods fit inside Claude Code's 30s execSync timeout. + # Server-backed methods (embed, llama-*) would either time out or spin + # up a model server mid-compact. For those, we skip supercompact and + # let the outer try/catch route to the built-in LLM fallback (which is + # precisely what the user would get if the patch weren't installed). f'const _allowedMethods={{"eitf":1,"setcover":1,"dedup":1}};' - f'const _method=_allowedMethods[_rawMethod]?_rawMethod:"eitf";' + f'if(!_allowedMethods[_rawMethod])' + f'throw new Error("SUPERCOMPACT_EITF:unsupported_method:"+_rawMethod);' + f'const _method=_rawMethod;' f'const _cfgBudget=process.env.PLUGIN_SETTING_BUDGET;' f'const _parsedBudget=_cfgBudget?parseInt(_cfgBudget,10):0;' f'const _budget=(_parsedBudget>0&&_parsedBudget===(_parsedBudget|0))?_parsedBudget' diff --git a/plugins/codex-cli/codex_parser.py b/plugins/codex-cli/codex_parser.py index 52ccce1..802afae 100644 --- a/plugins/codex-cli/codex_parser.py +++ b/plugins/codex-cli/codex_parser.py @@ -183,7 +183,7 @@ def parse_codex_jsonl(path: Path) -> list: return turns -def extract_codex_text(turn) -> str: +def extract_codex_text(turn, *, truncate: bool = True) -> str: """Extract human-readable text from a Codex turn for scoring/display. Handles Codex's payload-wrapped format: @@ -193,6 +193,10 @@ def extract_codex_text(turn) -> str: - Reasoning: payload.content[].text (reasoning_text type) - Turn context: payload.user_instructions - Compacted: payload.message + + With ``truncate=False``, function call arguments and outputs are kept at + full length — use this for accurate token counting (the API sees them + untruncated). Default ``truncate=True`` keeps the short form for scoring. """ parts: list[str] = [] @@ -244,7 +248,7 @@ def extract_codex_text(turn) -> str: arguments = payload.get("arguments", "") parts.append(f"[function_call: {name}]") if isinstance(arguments, str) and arguments: - if len(arguments) > 500: + if truncate and len(arguments) > 500: arguments = arguments[:500] + "..." parts.append(arguments) @@ -252,7 +256,7 @@ def extract_codex_text(turn) -> str: elif payload_type == "function_call_output": output = payload.get("output", "") if isinstance(output, str): - if len(output) > 1000: + if truncate and len(output) > 1000: output = output[:1000] + "..." parts.append(output) diff --git a/plugins/codex-cli/compact_codex.py b/plugins/codex-cli/compact_codex.py index a2c84d6..11edb91 100644 --- a/plugins/codex-cli/compact_codex.py +++ b/plugins/codex-cli/compact_codex.py @@ -139,9 +139,10 @@ def _run_compaction( total_tokens = sum(token_counts.values()) console.print(f" {total_tokens:,} tokens total") - if total_tokens <= budget: + target_with_padding = int(budget * 1.10) + if total_tokens <= target_with_padding: console.print( - f"[green]Already within budget ({total_tokens:,} <= {budget:,}), " + f"[green]Already within budget padding ({total_tokens:,} <= {target_with_padding:,}), " f"nothing to compact.[/green]" ) return 0 @@ -184,6 +185,18 @@ def _run_compaction( short_threshold=short_threshold, ) + total_kept = result.user_tokens + result.short_system_tokens + result.scored_kept_tokens + if total_kept > budget: + console.print( + f"\n[yellow]Primary method '{method}' couldn't reach budget " + f"({total_kept:,} > {budget:,}). Falling back to NUCLEAR truncation.[/yellow]" + ) + primary_dropped = result.dropped_turns + from lib.nuclear import nuclear_compact + result = nuclear_compact(result.kept_turns, budget) + result.dropped_turns = primary_dropped + result.scored_dropped_tokens = sum(st.tokens for st in primary_dropped) + t_elapsed = time.monotonic() - t_start # Display stats