185 lines
6.7 KiB
JavaScript
185 lines
6.7 KiB
JavaScript
#!/usr/bin/env node
|
|
/**
|
|
* Patch: Aggressive Text Pruning for Claude Code
|
|
*
|
|
* Extends Claude Code's micro-compaction (Vd function) to also prune old
|
|
* user/assistant TEXT content, not just tool results. This keeps context
|
|
* lean and prevents full compaction from ever being needed.
|
|
*
|
|
* What it does:
|
|
* After Vd() runs (pruning tool results), this patch adds a second pass
|
|
* that truncates old conversation text. It keeps the last N turns intact
|
|
* and replaces older text with brief summaries.
|
|
*
|
|
* How it works:
|
|
* Patches cli.js to insert a textPrune() function call after Vd().
|
|
* The function:
|
|
* 1. Counts total text tokens in the message array
|
|
* 2. If above threshold (configurable via CLAUDE_TEXT_PRUNE_THRESHOLD)
|
|
* 3. Keeps the last CLAUDE_TEXT_PRUNE_KEEP turns intact
|
|
* 4. Truncates older text blocks to first line + "[earlier context pruned]"
|
|
* 5. Preserves tool_use/tool_result structure (never breaks the API contract)
|
|
*
|
|
* Safety:
|
|
* - Only modifies text content blocks, never tool_use or tool_result
|
|
* - Always keeps last N turns fully intact
|
|
* - Preserves message structure (role, type, ids)
|
|
* - Falls back gracefully if anything fails
|
|
* - Can be reverted by running: node patch-aggressive-prune.mjs --revert
|
|
*
|
|
* Usage:
|
|
* node patch-aggressive-prune.mjs # Apply patch
|
|
* node patch-aggressive-prune.mjs --revert # Revert patch
|
|
* node patch-aggressive-prune.mjs --check # Check if patched
|
|
*/
|
|
|
|
import { readFileSync, writeFileSync, copyFileSync, existsSync } from 'fs';
|
|
import { join, dirname } from 'path';
|
|
import { fileURLToPath } from 'url';
|
|
|
|
const __filename = fileURLToPath(import.meta.url);
|
|
const __dirname = dirname(__filename);
|
|
const PROJECT_ROOT = join(__dirname, '../..');
|
|
const CLI_PATH = join(PROJECT_ROOT, 'node_modules/@anthropic-ai/claude-agent-sdk/cli.js');
|
|
const BACKUP_PATH = CLI_PATH + '.backup';
|
|
|
|
const PATCH_MARKER = '/*AGGRESSIVE_TEXT_PRUNE_PATCH*/';
|
|
|
|
// The text pruning function to inject
|
|
const TEXT_PRUNE_FUNCTION = `
|
|
/*AGGRESSIVE_TEXT_PRUNE_PATCH*/
|
|
function _aggressiveTextPrune(messages) {
|
|
try {
|
|
var KEEP = parseInt(process.env.CLAUDE_TEXT_PRUNE_KEEP || '10', 10);
|
|
var THRESHOLD = parseInt(process.env.CLAUDE_TEXT_PRUNE_THRESHOLD || '60000', 10);
|
|
var MAX_OLD_TEXT = parseInt(process.env.CLAUDE_TEXT_PRUNE_MAX_CHARS || '150', 10);
|
|
|
|
// Count text tokens roughly (4 chars per token)
|
|
var totalChars = 0;
|
|
for (var i = 0; i < messages.length; i++) {
|
|
var m = messages[i];
|
|
if ((m.type === 'user' || m.type === 'assistant') && Array.isArray(m.message?.content)) {
|
|
for (var j = 0; j < m.message.content.length; j++) {
|
|
var c = m.message.content[j];
|
|
if (c.type === 'text') totalChars += (c.text || '').length;
|
|
}
|
|
}
|
|
}
|
|
|
|
var totalTokensEst = Math.ceil(totalChars / 4);
|
|
if (totalTokensEst < THRESHOLD) return messages;
|
|
|
|
// Find turn boundaries (user message = new turn)
|
|
var turnStarts = [];
|
|
for (var i = 0; i < messages.length; i++) {
|
|
if (messages[i].type === 'user') turnStarts.push(i);
|
|
}
|
|
|
|
// Keep last KEEP turns intact
|
|
var cutoffIdx = turnStarts.length > KEEP ? turnStarts[turnStarts.length - KEEP] : 0;
|
|
if (cutoffIdx === 0) return messages;
|
|
|
|
var pruned = [];
|
|
var prunedChars = 0;
|
|
for (var i = 0; i < messages.length; i++) {
|
|
var m = messages[i];
|
|
if (i >= cutoffIdx) {
|
|
pruned.push(m);
|
|
continue;
|
|
}
|
|
if ((m.type === 'user' || m.type === 'assistant') && Array.isArray(m.message?.content)) {
|
|
var newContent = [];
|
|
for (var j = 0; j < m.message.content.length; j++) {
|
|
var c = m.message.content[j];
|
|
if (c.type === 'text' && c.text && c.text.length > MAX_OLD_TEXT) {
|
|
var firstLine = c.text.split('\\n')[0].slice(0, MAX_OLD_TEXT);
|
|
prunedChars += c.text.length - firstLine.length - 30;
|
|
newContent.push({ ...c, text: firstLine + '\\n[earlier context pruned]' });
|
|
} else {
|
|
newContent.push(c);
|
|
}
|
|
}
|
|
pruned.push({ ...m, message: { ...m.message, content: newContent } });
|
|
} else {
|
|
pruned.push(m);
|
|
}
|
|
}
|
|
|
|
if (prunedChars > 1000) {
|
|
process.stderr?.write?.('[TextPrune] Pruned ~' + Math.round(prunedChars/4) + ' tokens of old text (kept last ' + KEEP + ' turns)\\n');
|
|
}
|
|
return pruned;
|
|
} catch(e) {
|
|
return messages;
|
|
}
|
|
}
|
|
/*END_AGGRESSIVE_TEXT_PRUNE_PATCH*/`;
|
|
|
|
// The injection point: after Vd() call, before CT2() call
|
|
const VD_CALL_PATTERN = 'z=await Vd(F,void 0,Y);if(F=z.messages,';
|
|
const PATCHED_PATTERN = 'z=await Vd(F,void 0,Y);if(F=_aggressiveTextPrune(z.messages),';
|
|
|
|
function check() {
|
|
const src = readFileSync(CLI_PATH, 'utf8');
|
|
const isPatched = src.includes(PATCH_MARKER);
|
|
console.log(isPatched ? 'PATCHED' : 'NOT PATCHED');
|
|
return isPatched;
|
|
}
|
|
|
|
function apply() {
|
|
if (check()) {
|
|
console.log('Already patched. Use --revert first to re-apply.');
|
|
return;
|
|
}
|
|
|
|
const src = readFileSync(CLI_PATH, 'utf8');
|
|
|
|
// Verify the injection point exists
|
|
if (!src.includes(VD_CALL_PATTERN)) {
|
|
console.error('ERROR: Could not find Vd() call pattern in cli.js.');
|
|
console.error('Claude Code may have been updated. Pattern expected:');
|
|
console.error(' ' + VD_CALL_PATTERN);
|
|
process.exit(1);
|
|
}
|
|
|
|
// Backup
|
|
if (!existsSync(BACKUP_PATH)) {
|
|
copyFileSync(CLI_PATH, BACKUP_PATH);
|
|
console.log('Backup saved to:', BACKUP_PATH);
|
|
}
|
|
|
|
// Inject the function at the top of the file (after the first line)
|
|
let patched = src;
|
|
const firstNewline = patched.indexOf('\n');
|
|
patched = patched.slice(0, firstNewline + 1) + TEXT_PRUNE_FUNCTION + '\n' + patched.slice(firstNewline + 1);
|
|
|
|
// Patch the Vd() call site to also run our text pruner
|
|
patched = patched.replace(VD_CALL_PATTERN, PATCHED_PATTERN);
|
|
|
|
writeFileSync(CLI_PATH, patched);
|
|
console.log('PATCH APPLIED successfully.');
|
|
console.log('');
|
|
console.log('Configuration (via env vars in settings.json):');
|
|
console.log(' CLAUDE_TEXT_PRUNE_KEEP=10 # Keep last N turns intact');
|
|
console.log(' CLAUDE_TEXT_PRUNE_THRESHOLD=60000 # Start pruning above this token count');
|
|
console.log(' CLAUDE_TEXT_PRUNE_MAX_CHARS=150 # Truncate old text to this many chars');
|
|
console.log('');
|
|
console.log('Restart Claude Code for the patch to take effect.');
|
|
}
|
|
|
|
function revert() {
|
|
if (!existsSync(BACKUP_PATH)) {
|
|
console.error('No backup found at:', BACKUP_PATH);
|
|
console.error('Cannot revert. Reinstall with: npm install @anthropic-ai/claude-agent-sdk');
|
|
process.exit(1);
|
|
}
|
|
|
|
copyFileSync(BACKUP_PATH, CLI_PATH);
|
|
console.log('REVERTED to original cli.js from backup.');
|
|
}
|
|
|
|
const arg = process.argv[2];
|
|
if (arg === '--revert') revert();
|
|
else if (arg === '--check') check();
|
|
else apply();
|