505 lines
16 KiB
Python
505 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Context Compaction Manager
|
|
|
|
Handles pre-compact state saving and post-compact resume injection.
|
|
Monitors context usage and triggers appropriate hooks.
|
|
|
|
Usage:
|
|
python3 context_compact.py save [--workflow-dir .workflow/versions/v001]
|
|
python3 context_compact.py resume [--workflow-dir .workflow/versions/v001]
|
|
python3 context_compact.py status [--workflow-dir .workflow/versions/v001]
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
|
|
# ============================================================================
|
|
# Configuration
|
|
# ============================================================================
|
|
|
|
DEFAULT_WORKFLOW_DIR = ".workflow/versions/v001"
|
|
STATE_FILE = "context_state.json"
|
|
RESUME_PROMPT_FILE = "resume_prompt.md"
|
|
MODIFIED_FILES_FILE = "modified_files.json"
|
|
|
|
# Context thresholds (percentage)
|
|
THRESHOLDS = {
|
|
"warning": 0.70,
|
|
"save": 0.80,
|
|
"compact": 0.90,
|
|
"critical": 0.95
|
|
}
|
|
|
|
|
|
# ============================================================================
|
|
# File Operations
|
|
# ============================================================================
|
|
|
|
def load_json(filepath: str) -> dict:
|
|
"""Load JSON file."""
|
|
if not os.path.exists(filepath):
|
|
return {}
|
|
try:
|
|
with open(filepath, 'r') as f:
|
|
return json.load(f)
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
|
|
|
|
def save_json(filepath: str, data: dict):
|
|
"""Save data to JSON file."""
|
|
os.makedirs(os.path.dirname(filepath), exist_ok=True)
|
|
with open(filepath, 'w') as f:
|
|
json.dump(data, f, indent=2, default=str)
|
|
|
|
|
|
def save_text(filepath: str, content: str):
|
|
"""Save text to file."""
|
|
os.makedirs(os.path.dirname(filepath), exist_ok=True)
|
|
with open(filepath, 'w') as f:
|
|
f.write(content)
|
|
|
|
|
|
# ============================================================================
|
|
# Git Operations
|
|
# ============================================================================
|
|
|
|
def get_git_status() -> List[Dict[str, str]]:
|
|
"""Get list of modified files from git."""
|
|
try:
|
|
result = subprocess.run(
|
|
['git', 'status', '--porcelain'],
|
|
capture_output=True, text=True, check=True
|
|
)
|
|
files = []
|
|
for line in result.stdout.strip().split('\n'):
|
|
if line:
|
|
status = line[:2].strip()
|
|
path = line[3:]
|
|
action = {
|
|
'M': 'modified',
|
|
'A': 'added',
|
|
'D': 'deleted',
|
|
'?': 'untracked',
|
|
'R': 'renamed'
|
|
}.get(status[0] if status else '?', 'unknown')
|
|
files.append({'path': path, 'action': action, 'summary': ''})
|
|
return files
|
|
except subprocess.CalledProcessError:
|
|
return []
|
|
|
|
|
|
def get_recent_commits(count: int = 5) -> List[Dict[str, str]]:
|
|
"""Get recent commit messages."""
|
|
try:
|
|
result = subprocess.run(
|
|
['git', 'log', f'-{count}', '--oneline'],
|
|
capture_output=True, text=True, check=True
|
|
)
|
|
commits = []
|
|
for line in result.stdout.strip().split('\n'):
|
|
if line:
|
|
parts = line.split(' ', 1)
|
|
commits.append({
|
|
'hash': parts[0],
|
|
'message': parts[1] if len(parts) > 1 else ''
|
|
})
|
|
return commits
|
|
except subprocess.CalledProcessError:
|
|
return []
|
|
|
|
|
|
def create_checkpoint(message: str = "WIP: Pre-compaction checkpoint"):
|
|
"""Create a git checkpoint with uncommitted changes."""
|
|
try:
|
|
# Check if there are changes
|
|
result = subprocess.run(
|
|
['git', 'status', '--porcelain'],
|
|
capture_output=True, text=True
|
|
)
|
|
if result.stdout.strip():
|
|
# Stage all changes
|
|
subprocess.run(['git', 'add', '-A'], check=True)
|
|
# Commit
|
|
subprocess.run(['git', 'commit', '-m', message], check=True)
|
|
print(f"Created checkpoint: {message}")
|
|
return True
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Warning: Could not create checkpoint: {e}")
|
|
return False
|
|
|
|
|
|
# ============================================================================
|
|
# Workflow State Operations
|
|
# ============================================================================
|
|
|
|
def load_workflow_state(workflow_dir: str) -> dict:
|
|
"""Load current workflow state."""
|
|
state_path = os.path.join(workflow_dir, 'workflow_state.json')
|
|
return load_json(state_path)
|
|
|
|
|
|
def load_active_tasks(workflow_dir: str) -> List[dict]:
|
|
"""Load tasks that are in progress."""
|
|
tasks_dir = os.path.join(workflow_dir, 'tasks')
|
|
active_tasks = []
|
|
|
|
if os.path.exists(tasks_dir):
|
|
for filename in os.listdir(tasks_dir):
|
|
if filename.endswith('.yml') or filename.endswith('.json'):
|
|
task_path = os.path.join(tasks_dir, filename)
|
|
task = load_json(task_path) if filename.endswith('.json') else {}
|
|
if task.get('status') == 'in_progress':
|
|
active_tasks.append(task)
|
|
|
|
return active_tasks
|
|
|
|
|
|
def get_pending_tasks(workflow_dir: str) -> List[dict]:
|
|
"""Get pending tasks in priority order."""
|
|
tasks_dir = os.path.join(workflow_dir, 'tasks')
|
|
pending = []
|
|
|
|
if os.path.exists(tasks_dir):
|
|
for filename in os.listdir(tasks_dir):
|
|
if filename.endswith('.json'):
|
|
task_path = os.path.join(tasks_dir, filename)
|
|
task = load_json(task_path)
|
|
if task.get('status') == 'pending':
|
|
pending.append(task)
|
|
|
|
# Sort by layer, then by ID
|
|
pending.sort(key=lambda t: (t.get('layer', 999), t.get('id', '')))
|
|
return pending
|
|
|
|
|
|
# ============================================================================
|
|
# Context State Management
|
|
# ============================================================================
|
|
|
|
def capture_context_state(
|
|
workflow_dir: str,
|
|
context_percentage: float = 0.0,
|
|
active_work: Optional[dict] = None,
|
|
decisions: Optional[List[dict]] = None,
|
|
blockers: Optional[List[dict]] = None
|
|
) -> dict:
|
|
"""Capture current context state for later resume."""
|
|
|
|
workflow_state = load_workflow_state(workflow_dir)
|
|
active_tasks = load_active_tasks(workflow_dir)
|
|
pending_tasks = get_pending_tasks(workflow_dir)
|
|
modified_files = get_git_status()
|
|
|
|
# Determine active task
|
|
active_task = active_tasks[0] if active_tasks else None
|
|
|
|
# Build next actions from pending tasks
|
|
next_actions = []
|
|
for task in pending_tasks[:5]: # Top 5 pending
|
|
next_actions.append({
|
|
'action': task.get('type', 'implement'),
|
|
'target': task.get('title', task.get('id', 'unknown')),
|
|
'priority': len(next_actions) + 1,
|
|
'context_needed': [task.get('context', {}).get('context_snapshot_path', '')]
|
|
})
|
|
|
|
# Build context state
|
|
state = {
|
|
'session_id': f"compact_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
|
|
'captured_at': datetime.now().isoformat(),
|
|
|
|
'context_usage': {
|
|
'tokens_used': 0, # Would need to be passed in
|
|
'tokens_max': 0,
|
|
'percentage': context_percentage,
|
|
'threshold_triggered': THRESHOLDS['save']
|
|
},
|
|
|
|
'workflow_position': {
|
|
'workflow_id': workflow_state.get('id', 'unknown'),
|
|
'current_phase': workflow_state.get('current_phase', 'UNKNOWN'),
|
|
'active_task_id': active_task.get('id') if active_task else None,
|
|
'layer': active_task.get('layer', 1) if active_task else 1
|
|
},
|
|
|
|
'active_work': active_work or {
|
|
'entity_id': active_task.get('entity_id', '') if active_task else '',
|
|
'entity_type': active_task.get('type', '') if active_task else '',
|
|
'action': 'implementing' if active_task else 'pending',
|
|
'file_path': None,
|
|
'progress_notes': ''
|
|
},
|
|
|
|
'next_actions': next_actions,
|
|
'modified_files': modified_files,
|
|
'decisions': decisions or [],
|
|
'blockers': blockers or []
|
|
}
|
|
|
|
return state
|
|
|
|
|
|
def save_context_state(workflow_dir: str, state: dict):
|
|
"""Save context state to file."""
|
|
state_path = os.path.join(workflow_dir, STATE_FILE)
|
|
save_json(state_path, state)
|
|
print(f"Saved context state to: {state_path}")
|
|
|
|
# Also save modified files separately for quick access
|
|
modified_path = os.path.join(workflow_dir, MODIFIED_FILES_FILE)
|
|
save_json(modified_path, state.get('modified_files', []))
|
|
|
|
|
|
def generate_resume_prompt(state: dict) -> str:
|
|
"""Generate human-readable resume prompt."""
|
|
|
|
lines = [
|
|
"## Context Recovery - Resuming Previous Session",
|
|
"",
|
|
"### Session Info",
|
|
f"- **Original Session**: {state.get('session_id', 'unknown')}",
|
|
f"- **Captured At**: {state.get('captured_at', 'unknown')}",
|
|
f"- **Context Usage**: {state.get('context_usage', {}).get('percentage', 0) * 100:.1f}%",
|
|
"",
|
|
"### Workflow Position",
|
|
f"- **Phase**: {state.get('workflow_position', {}).get('current_phase', 'UNKNOWN')}",
|
|
f"- **Active Task**: {state.get('workflow_position', {}).get('active_task_id', 'None')}",
|
|
f"- **Layer**: {state.get('workflow_position', {}).get('layer', 1)}",
|
|
"",
|
|
"### What Was Being Worked On",
|
|
]
|
|
|
|
active = state.get('active_work', {})
|
|
lines.extend([
|
|
f"- **Entity**: {active.get('entity_id', 'None')} ({active.get('entity_type', '')})",
|
|
f"- **Action**: {active.get('action', 'unknown')}",
|
|
f"- **File**: {active.get('file_path', 'None')}",
|
|
f"- **Progress**: {active.get('progress_notes', 'No notes')}",
|
|
"",
|
|
"### Next Actions (Priority Order)",
|
|
])
|
|
|
|
for action in state.get('next_actions', []):
|
|
context = ', '.join(action.get('context_needed', [])) or 'None'
|
|
lines.append(f"{action.get('priority', '?')}. **{action.get('action', '')}** {action.get('target', '')}")
|
|
lines.append(f" - Context needed: {context}")
|
|
|
|
if state.get('modified_files'):
|
|
lines.extend(["", "### Recent Changes"])
|
|
for f in state.get('modified_files', []):
|
|
lines.append(f"- `{f.get('path', '')}` - {f.get('action', '')}: {f.get('summary', '')}")
|
|
|
|
if state.get('decisions'):
|
|
lines.extend(["", "### Key Decisions Made"])
|
|
for d in state.get('decisions', []):
|
|
lines.append(f"- **{d.get('topic', '')}**: {d.get('decision', '')}")
|
|
|
|
if state.get('blockers'):
|
|
lines.extend(["", "### Current Blockers"])
|
|
for b in state.get('blockers', []):
|
|
lines.append(f"- {b.get('issue', '')} ({b.get('status', '')}): {b.get('notes', '')}")
|
|
|
|
lines.extend([
|
|
"",
|
|
"---",
|
|
"**Action Required**: Continue from the next action listed above.",
|
|
"",
|
|
"To load full context, read the following files:",
|
|
])
|
|
|
|
# Add context files to read
|
|
for action in state.get('next_actions', [])[:3]:
|
|
for ctx in action.get('context_needed', []):
|
|
if ctx:
|
|
lines.append(f"- `{ctx}`")
|
|
|
|
return '\n'.join(lines)
|
|
|
|
|
|
# ============================================================================
|
|
# Commands
|
|
# ============================================================================
|
|
|
|
def cmd_save(args):
|
|
"""Save current context state (pre-compact hook)."""
|
|
workflow_dir = args.workflow_dir
|
|
|
|
print("=" * 60)
|
|
print("PRE-COMPACT: Saving Context State")
|
|
print("=" * 60)
|
|
|
|
# Capture state
|
|
state = capture_context_state(
|
|
workflow_dir,
|
|
context_percentage=args.percentage / 100 if args.percentage else 0.80
|
|
)
|
|
|
|
# Save state
|
|
save_context_state(workflow_dir, state)
|
|
|
|
# Generate resume prompt
|
|
resume_prompt = generate_resume_prompt(state)
|
|
resume_path = os.path.join(workflow_dir, RESUME_PROMPT_FILE)
|
|
save_text(resume_path, resume_prompt)
|
|
print(f"Generated resume prompt: {resume_path}")
|
|
|
|
# Create git checkpoint if requested
|
|
if args.checkpoint:
|
|
create_checkpoint("WIP: Pre-compaction checkpoint")
|
|
|
|
print()
|
|
print("=" * 60)
|
|
print("State saved successfully!")
|
|
print(f" Session ID: {state['session_id']}")
|
|
print(f" Active Task: {state['workflow_position']['active_task_id']}")
|
|
print(f" Next Actions: {len(state['next_actions'])}")
|
|
print("=" * 60)
|
|
|
|
return 0
|
|
|
|
|
|
def cmd_resume(args):
|
|
"""Display resume prompt (post-compact hook)."""
|
|
workflow_dir = args.workflow_dir
|
|
|
|
# Load state
|
|
state_path = os.path.join(workflow_dir, STATE_FILE)
|
|
state = load_json(state_path)
|
|
|
|
if not state:
|
|
print("No saved context state found.")
|
|
print(f"Expected: {state_path}")
|
|
return 1
|
|
|
|
# Generate and display resume prompt
|
|
if args.json:
|
|
print(json.dumps(state, indent=2))
|
|
else:
|
|
resume_prompt = generate_resume_prompt(state)
|
|
print(resume_prompt)
|
|
|
|
return 0
|
|
|
|
|
|
def cmd_status(args):
|
|
"""Show current context state status."""
|
|
workflow_dir = args.workflow_dir
|
|
|
|
print("=" * 60)
|
|
print("CONTEXT STATE STATUS")
|
|
print("=" * 60)
|
|
|
|
# Check for saved state
|
|
state_path = os.path.join(workflow_dir, STATE_FILE)
|
|
if os.path.exists(state_path):
|
|
state = load_json(state_path)
|
|
print(f"\nSaved State Found:")
|
|
print(f" Session ID: {state.get('session_id', 'unknown')}")
|
|
print(f" Captured At: {state.get('captured_at', 'unknown')}")
|
|
print(f" Phase: {state.get('workflow_position', {}).get('current_phase', 'UNKNOWN')}")
|
|
print(f" Active Task: {state.get('workflow_position', {}).get('active_task_id', 'None')}")
|
|
else:
|
|
print(f"\nNo saved state at: {state_path}")
|
|
|
|
# Check workflow state
|
|
workflow_state = load_workflow_state(workflow_dir)
|
|
if workflow_state:
|
|
print(f"\nCurrent Workflow:")
|
|
print(f" ID: {workflow_state.get('id', 'unknown')}")
|
|
print(f" Phase: {workflow_state.get('current_phase', 'UNKNOWN')}")
|
|
|
|
# Check git status
|
|
modified = get_git_status()
|
|
if modified:
|
|
print(f"\nModified Files: {len(modified)}")
|
|
for f in modified[:5]:
|
|
print(f" - {f['action']}: {f['path']}")
|
|
if len(modified) > 5:
|
|
print(f" ... and {len(modified) - 5} more")
|
|
|
|
return 0
|
|
|
|
|
|
def cmd_clear(args):
|
|
"""Clear saved context state."""
|
|
workflow_dir = args.workflow_dir
|
|
|
|
files_to_remove = [
|
|
os.path.join(workflow_dir, STATE_FILE),
|
|
os.path.join(workflow_dir, RESUME_PROMPT_FILE),
|
|
os.path.join(workflow_dir, MODIFIED_FILES_FILE)
|
|
]
|
|
|
|
removed = 0
|
|
for filepath in files_to_remove:
|
|
if os.path.exists(filepath):
|
|
os.remove(filepath)
|
|
print(f"Removed: {filepath}")
|
|
removed += 1
|
|
|
|
if removed == 0:
|
|
print("No state files found to remove.")
|
|
else:
|
|
print(f"\nCleared {removed} state file(s).")
|
|
|
|
return 0
|
|
|
|
|
|
# ============================================================================
|
|
# Main CLI
|
|
# ============================================================================
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Context Compaction Manager for Guardrail Orchestrator"
|
|
)
|
|
parser.add_argument(
|
|
'--workflow-dir', '-w',
|
|
default=DEFAULT_WORKFLOW_DIR,
|
|
help='Workflow directory path'
|
|
)
|
|
|
|
subparsers = parser.add_subparsers(dest='command', help='Commands')
|
|
|
|
# Save command
|
|
save_parser = subparsers.add_parser('save', help='Save context state (pre-compact)')
|
|
save_parser.add_argument('--percentage', '-p', type=float, help='Current context percentage (0-100)')
|
|
save_parser.add_argument('--checkpoint', '-c', action='store_true', help='Create git checkpoint')
|
|
|
|
# Resume command
|
|
resume_parser = subparsers.add_parser('resume', help='Display resume prompt (post-compact)')
|
|
resume_parser.add_argument('--json', '-j', action='store_true', help='Output as JSON')
|
|
|
|
# Status command
|
|
subparsers.add_parser('status', help='Show context state status')
|
|
|
|
# Clear command
|
|
subparsers.add_parser('clear', help='Clear saved context state')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.command == 'save':
|
|
return cmd_save(args)
|
|
elif args.command == 'resume':
|
|
return cmd_resume(args)
|
|
elif args.command == 'status':
|
|
return cmd_status(args)
|
|
elif args.command == 'clear':
|
|
return cmd_clear(args)
|
|
else:
|
|
parser.print_help()
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|