erk 0.4.5__py3-none-any.whl

erk/cli/commands/exec/scripts/preprocess_session.py

@@ -0,0 +1,777 @@
+#!/usr/bin/env python3
+"""
+Session Log Preprocessor
+
+Compresses JSONL session logs to XML format by removing metadata and deduplicating messages.
+This command is invoked via erk exec preprocess-session <log-path>.
+"""
+
+import json
+import tempfile
+from pathlib import Path
+
+import click
+
+
+def escape_xml(text: str) -> str:
+    """Minimal XML escaping for special characters."""
+    return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
+
+
+def is_empty_session(entries: list[dict]) -> bool:
+    """Check if session contains only metadata with no meaningful content.
+
+    Empty sessions are characterized by:
+    - Fewer than 3 entries (too small to be meaningful)
+    - Only metadata/system entries without substantive interaction
+
+    Args:
+        entries: List of session entries to check
+
+    Returns:
+        True if session is empty/meaningless, False otherwise
+    """
+    if len(entries) < 3:
+        return True
+
+    # Check if there's any meaningful content
+    has_user_message = False
+    has_assistant_response = False
+
+    for entry in entries:
+        entry_type = entry.get("type")
+        if entry_type == "user":
+            content = entry.get("message", {}).get("content", "")
+            if isinstance(content, list):
+                text_parts = []
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        text_parts.append(block.get("text", ""))
+                content = " ".join(text_parts)
+            if content and len(str(content).strip()) > 0:
+                has_user_message = True
+
+        elif entry_type == "assistant":
+            content_blocks = entry.get("message", {}).get("content", [])
+            for block in content_blocks:
+                if block.get("type") == "text" and block.get("text", "").strip():
+                    has_assistant_response = True
+                    break
+
+    # Session is empty if it lacks meaningful interaction
+    return not (has_user_message and has_assistant_response)
+
+
+def is_warmup_session(entries: list[dict]) -> bool:
+    """Check if session is a warmup containing only boilerplate acknowledgment.
+
+    Warmup sessions contain predictable patterns like:
+    - "I've reviewed"
+    - "I'm ready"
+    - "loaded the instructions"
+
+    Args:
+        entries: List of session entries to check
+
+    Returns:
+        True if session is a warmup, False otherwise
+    """
+    if not entries:
+        return False
+
+    # Look for warmup keyword in first user message
+    for entry in entries:
+        if entry.get("type") == "user":
+            content = entry.get("message", {}).get("content", "")
+            if isinstance(content, list):
+                text_parts = []
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        text_parts.append(block.get("text", ""))
+                content = " ".join(text_parts)
+
+            content_lower = str(content).lower()
+            if "warmup" in content_lower:
+                return True
+            break
+
+    return False
+
+
+def deduplicate_documentation_blocks(entries: list[dict]) -> list[dict]:
+    """Replace duplicate command documentation blocks with marker text.
+
+    Command documentation can appear verbatim multiple times, consuming
+    significant tokens. This function detects duplicate blocks by content hash
+    and replaces them with a reference marker.
+
+    Args:
+        entries: List of session entries
+
+    Returns:
+        Modified entries with duplicate documentation replaced by markers
+    """
+    import hashlib
+
+    seen_docs: dict[str, int] = {}  # hash -> first occurrence count
+    occurrence_counter: dict[str, int] = {}  # hash -> current occurrence
+    deduplicated = []
+
+    for entry in entries:
+        if entry.get("type") == "user":
+            content = entry.get("message", {}).get("content", "")
+            if isinstance(content, list):
+                text_parts = []
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        text_parts.append(block.get("text", ""))
+                content = " ".join(text_parts)
+
+            content_str = str(content)
+
+            # Detect command documentation by markers
+            is_doc = any(
+                marker in content_str
+                for marker in [
+                    "/erk:plan-save-issue",
+                    "/erk:plan-implement",
+                    "/gt:submit-branch",
+                    "/gt:pr-update",
+                    "command-message>",
+                    "command-name>",
+                ]
+            )
+
+            if is_doc and len(content_str) > 500:
+                # Hash the content
+                content_hash = hashlib.sha256(content_str.encode()).hexdigest()[:16]
+
+                if content_hash not in seen_docs:
+                    # First occurrence - keep it
+                    seen_docs[content_hash] = 1
+                    occurrence_counter[content_hash] = 1
+                    deduplicated.append(entry)
+                else:
+                    # Duplicate - replace with marker
+                    occurrence_counter[content_hash] += 1
+                    occurrence_num = occurrence_counter[content_hash]
+
+                    # Create marker entry
+                    marker_entry = entry.copy()
+                    marker_content = (
+                        f"[Duplicate command documentation block omitted - "
+                        f"hash {content_hash}, occurrence #{occurrence_num}]"
+                    )
+
+                    # Preserve structure
+                    if isinstance(entry.get("message", {}).get("content"), list):
+                        marker_entry["message"] = {
+                            "content": [{"type": "text", "text": marker_content}]
+                        }
+                    else:
+                        marker_entry["message"] = {"content": marker_content}
+
+                    deduplicated.append(marker_entry)
+            else:
+                deduplicated.append(entry)
+        else:
+            deduplicated.append(entry)
+
+    return deduplicated
+
+
+def truncate_parameter_value(value: str, max_length: int = 200) -> str:
+    """Truncate long parameter values while preserving identifiability.
+
+    Special handling for file paths to preserve structure.
+
+    Args:
+        value: Parameter value to truncate
+        max_length: Maximum length (default 200)
+
+    Returns:
+        Truncated value with context markers
+    """
+    if len(value) <= max_length:
+        return value
+
+    # Detect file paths - check for path separators and no spaces
+    has_slash = "/" in value
+    has_no_spaces_early = " " not in value[: min(100, len(value))]
+
+    if has_slash and has_no_spaces_early:
+        # Likely a file path - preserve start and end structure
+        parts = value.split("/")
+        if len(parts) > 3:
+            # Build path keeping first 2 parts and last 2 parts
+            first_parts = "/".join(parts[:2])
+            last_parts = "/".join(parts[-2:])
+            return f"{first_parts}/.../{last_parts}"
+
+    # General text - keep beginning and end with marker
+    keep_chars = (max_length - 20) // 2
+    truncated_count = len(value) - max_length
+    return f"{value[:keep_chars]}...[truncated {truncated_count} chars]...{value[-keep_chars:]}"
+
+
+def truncate_tool_parameters(entries: list[dict]) -> list[dict]:
+    """Truncate verbose tool parameters to reduce token usage.
+
+    Tool parameters can be extremely long (20+ lines), especially prompts.
+    This function truncates them while preserving identifiability.
+
+    Args:
+        entries: List of session entries
+
+    Returns:
+        Modified entries with truncated parameters
+    """
+    truncated = []
+
+    for entry in entries:
+        if entry.get("type") == "assistant":
+            message = entry.get("message", {})
+            content_blocks = message.get("content", [])
+
+            modified_blocks = []
+            for block in content_blocks:
+                if block.get("type") == "tool_use":
+                    # Truncate input parameters
+                    input_params = block.get("input", {})
+                    truncated_params = {}
+                    for key, value in input_params.items():
+                        value_str = str(value)
+                        if len(value_str) > 200:
+                            truncated_params[key] = truncate_parameter_value(value_str)
+                        else:
+                            truncated_params[key] = value
+
+                    # Create modified block
+                    modified_block = block.copy()
+                    modified_block["input"] = truncated_params
+                    modified_blocks.append(modified_block)
+                else:
+                    modified_blocks.append(block)
+
+            # Update entry
+            modified_entry = entry.copy()
+            modified_entry["message"] = message.copy()
+            modified_entry["message"]["content"] = modified_blocks
+            truncated.append(modified_entry)
+        else:
+            truncated.append(entry)
+
+    return truncated
+
+
+def prune_tool_result_content(result_text: str) -> str:
+    """Prune verbose tool results to first 30 lines, preserving errors.
+
+    Tool results can be extremely long. This function keeps the first 30 lines
+    (which usually contain the most relevant context) and preserves any lines
+    containing error keywords.
+
+    Args:
+        result_text: Tool result text to prune
+
+    Returns:
+        Pruned result text with error preservation
+    """
+    lines = result_text.split("\n")
+
+    if len(lines) <= 30:
+        return result_text
+
+    # Keep first 30 lines
+    kept_lines = lines[:30]
+
+    # Scan remaining lines for errors
+    error_keywords = ["error", "exception", "failed", "failure", "fatal", "warning"]
+    error_lines = []
+
+    for line in lines[30:]:
+        line_lower = line.lower()
+        if any(keyword in line_lower for keyword in error_keywords):
+            error_lines.append(line)
+
+    # Combine
+    if error_lines:
+        result_lines = kept_lines + [f"\n... [{len(lines) - 30} lines omitted] ...\n"] + error_lines
+    else:
+        result_lines = kept_lines + [f"\n... [{len(lines) - 30} lines omitted] ..."]
+
+    return "\n".join(result_lines)
+
+
+def is_log_discovery_operation(entry: dict) -> bool:
+    """Check if entry is a log discovery bash command (pwd, ls, etc.).
+
+    These are implementation mechanics that don't provide semantic value
+    for plan enhancement.
+
+    Args:
+        entry: Session entry to check
+
+    Returns:
+        True if entry is a log discovery operation, False otherwise
+    """
+    if entry.get("type") != "assistant":
+        return False
+
+    content_blocks = entry.get("message", {}).get("content", [])
+
+    for block in content_blocks:
+        if block.get("type") == "tool_use":
+            tool_name = block.get("name", "")
+            if tool_name != "Bash":
+                continue
+
+            # Check command parameter
+            input_params = block.get("input", {})
+            command = input_params.get("command", "")
+
+            # Log discovery patterns
+            log_discovery_patterns = [
+                "pwd",
+                "ls ~/.claude/projects/",
+                "ls ~/.claude",
+                "find ~/.claude",
+                "echo $SESSION_ID",
+            ]
+
+            for pattern in log_discovery_patterns:
+                if pattern in command:
+                    return True
+
+    return False
+
+
+def deduplicate_assistant_messages(entries: list[dict]) -> list[dict]:
+    """Remove duplicate assistant text when tool_use present."""
+    deduplicated = []
+    prev_assistant_text = None
+
+    for entry in entries:
+        if entry["type"] == "assistant":
+            message_content = entry["message"].get("content", [])
+
+            # Extract text and tool uses separately
+            text_blocks = [c for c in message_content if c.get("type") == "text"]
+            tool_uses = [c for c in message_content if c.get("type") == "tool_use"]
+
+            current_text = text_blocks[0]["text"] if text_blocks else None
+
+            # If text same as previous AND there's a tool_use, drop the duplicate text
+            if current_text == prev_assistant_text and tool_uses:
+                # Keep only tool_use content
+                entry["message"]["content"] = tool_uses
+
+            prev_assistant_text = current_text
+
+        deduplicated.append(entry)
+
+    return deduplicated
+
+
+def generate_compressed_xml(
+    entries: list[dict], source_label: str | None = None, enable_pruning: bool = True
+) -> str:
+    """Generate coarse-grained XML from filtered entries.
+
+    Args:
+        entries: List of session entries to convert to XML
+        source_label: Optional label for agent logs
+        enable_pruning: Whether to prune tool results (default: True)
+
+    Returns:
+        XML string representation of the session
+    """
+    xml_lines = ["<session>"]
+
+    # Add source label if provided (for agent logs)
+    if source_label:
+        xml_lines.append(f'  <meta source="{escape_xml(source_label)}" />')
+
+    # Extract session metadata once (from first entry with gitBranch)
+    for entry in entries:
+        # Check in the original entry structure (before filtering)
+        if "gitBranch" in entry:
+            branch = entry["gitBranch"]
+            xml_lines.append(f'  <meta branch="{escape_xml(branch)}" />')
+            break
+
+    for entry in entries:
+        entry_type = entry["type"]
+        message = entry.get("message", {})
+
+        if entry_type == "user":
+            # Extract user content
+            content = message.get("content", "")
+            if isinstance(content, list):
+                # Handle list of content blocks
+                text_parts = []
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        text_parts.append(block.get("text", ""))
+                    elif isinstance(block, str):
+                        text_parts.append(block)
+                content = "\n".join(text_parts)
+            xml_lines.append(f"  <user>{escape_xml(content)}</user>")
+
+        elif entry_type == "assistant":
+            # Extract text and tool uses
+            content_blocks = message.get("content", [])
+            for content in content_blocks:
+                if content.get("type") == "text":
+                    text = content.get("text", "")
+                    if text.strip():  # Only include non-empty text
+                        xml_lines.append(f"  <assistant>{escape_xml(text)}</assistant>")
+                elif content.get("type") == "tool_use":
+                    tool_name = content.get("name", "")
+                    tool_id = content.get("id", "")
+                    escaped_name = escape_xml(tool_name)
+                    escaped_id = escape_xml(tool_id)
+                    xml_lines.append(f'  <tool_use name="{escaped_name}" id="{escaped_id}">')
+                    input_params = content.get("input", {})
+                    for key, value in input_params.items():
+                        escaped_key = escape_xml(key)
+                        escaped_value = escape_xml(str(value))
+                        xml_lines.append(f'    <param name="{escaped_key}">{escaped_value}</param>')
+                    xml_lines.append("  </tool_use>")
+
+        elif entry_type == "tool_result":
+            # Handle tool results - apply pruning if enabled
+            content_blocks = message.get("content", [])
+            tool_use_id = message.get("tool_use_id", "")
+
+            # Extract result content
+            result_parts = []
+            for block in content_blocks:
+                if isinstance(block, dict):
+                    if block.get("type") == "text":
+                        result_parts.append(block.get("text", ""))
+                    elif "text" in block:
+                        result_parts.append(block["text"])
+                elif isinstance(block, str):
+                    result_parts.append(block)
+
+            result_text = "\n".join(result_parts)
+
+            # Apply pruning if enabled
+            if enable_pruning:
+                result_text = prune_tool_result_content(result_text)
+
+            xml_lines.append(f'  <tool_result tool="{escape_xml(tool_use_id)}">')
+            xml_lines.append(escape_xml(result_text))
+            xml_lines.append("  </tool_result>")
+
+    xml_lines.append("</session>")
+    return "\n".join(xml_lines)
+
+
+def process_log_file(
+    log_path: Path,
+    session_id: str | None = None,
+    source_label: str | None = None,
+    enable_filtering: bool = True,
+) -> tuple[list[dict], int, int]:
+    """Process a single JSONL log file and return filtered entries.
+
+    Args:
+        log_path: Path to the JSONL log file
+        session_id: Optional session ID to filter entries by
+        source_label: Optional label for agent logs
+        enable_filtering: Whether to apply optimization filters (default: True)
+
+    Returns:
+        Tuple of (filtered entries, total entries count, skipped entries count)
+    """
+    entries = []
+    total_entries = 0
+    skipped_entries = 0
+
+    for line in log_path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+
+        entry = json.loads(line)
+        total_entries += 1
+
+        # Filter by session ID if provided
+        if session_id is not None:
+            entry_session = entry.get("sessionId")
+            # Include if sessionId matches OR if sessionId field missing (backward compat)
+            if entry_session is not None and entry_session != session_id:
+                skipped_entries += 1
+                continue
+
+        # Filter out noise entries
+        if entry.get("type") == "file-history-snapshot":
+            continue
+
+        # Filter log discovery operations if filtering enabled
+        if enable_filtering and is_log_discovery_operation(entry):
+            continue
+
+        # Keep minimal fields but preserve gitBranch for metadata extraction
+        filtered = {
+            "type": entry["type"],
+            "message": entry.get("message", {}),
+        }
+
+        # Preserve gitBranch for metadata (will be extracted in XML generation)
+        if "gitBranch" in entry:
+            filtered["gitBranch"] = entry["gitBranch"]
+
+        # Drop usage metadata from assistant messages
+        if "usage" in filtered["message"]:
+            del filtered["message"]["usage"]
+
+        entries.append(filtered)
+
+    return entries, total_entries, skipped_entries
+
+
+def discover_agent_logs(session_log_path: Path) -> list[Path]:
+    """Discover agent logs in the same directory as the session log."""
+    log_dir = session_log_path.parent
+    agent_logs = sorted(log_dir.glob("agent-*.jsonl"))
+    return agent_logs
+
+
+def discover_planning_agent_logs(session_log_path: Path, parent_session_id: str) -> list[Path]:
+    """
+    Discover agent logs from Plan subagents only.
+
+    Algorithm:
+    1. Parse parent session JSONL to find Task tool invocations
+    2. Filter for entries where input.subagent_type == "Plan"
+    3. Extract agent IDs via temporal correlation with agent logs
+    4. Return only agent logs matching Plan subagents
+
+    Args:
+        session_log_path: Path to the main session log file
+        parent_session_id: Session ID of the parent session
+
+    Returns:
+        List of agent log paths from Plan subagents only.
+        Empty list if no Plan subagents found.
+    """
+    log_dir = session_log_path.parent
+
+    # Step 1: Find all Task tool invocations with subagent_type="Plan"
+    plan_task_timestamps: list[float] = []
+
+    for line in session_log_path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+
+        entry = json.loads(line)
+
+        # Look for assistant messages with tool_use content
+        if entry.get("type") == "assistant":
+            message = entry.get("message", {})
+            content = message.get("content", [])
+
+            if isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "tool_use":
+                        # Check if this is a Task tool with subagent_type="Plan"
+                        if block.get("name") == "Task":
+                            tool_input = block.get("input", {})
+                            if tool_input.get("subagent_type") == "Plan":
+                                # Record timestamp for correlation
+                                timestamp = message.get("timestamp")
+                                if timestamp is not None:
+                                    plan_task_timestamps.append(timestamp)
+
+    # If no Plan tasks found, return empty list (fallback to main session only)
+    if not plan_task_timestamps:
+        return []
+
+    # Step 2: Discover all agent logs
+    all_agent_logs = sorted(log_dir.glob("agent-*.jsonl"))
+
+    # Step 3: Filter agent logs by temporal correlation
+    planning_agent_logs: list[Path] = []
+
+    for agent_log in all_agent_logs:
+        # Read first entry to check sessionId and timestamp
+        if not agent_log.exists():
+            continue
+        first_line = agent_log.read_text(encoding="utf-8").splitlines()[0]
+        if not first_line.strip():
+            continue
+
+        first_entry = json.loads(first_line)
+
+        # Check if this agent log belongs to our parent session
+        if first_entry.get("sessionId") != parent_session_id:
+            continue
+
+        # Check if this agent log's timestamp correlates with a Plan Task
+        agent_timestamp = first_entry.get("message", {}).get("timestamp")
+        if agent_timestamp is None:
+            continue
+
+        # Match if within 1 second of any Plan Task timestamp
+        for plan_timestamp in plan_task_timestamps:
+            if abs(agent_timestamp - plan_timestamp) <= 1.0:
+                planning_agent_logs.append(agent_log)
+                break
+
+    return planning_agent_logs
+
+
+@click.command(name="preprocess-session")
+@click.argument("log_path", type=click.Path(exists=True, path_type=Path))
+@click.option(
+    "--session-id",
+    type=str,
+    default=None,
+    help="Filter JSONL entries by session ID before preprocessing",
+)
+@click.option(
+    "--include-agents/--no-include-agents",
+    default=True,
+    help="Include agent logs from same directory (default: True)",
+)
+@click.option(
+    "--no-filtering",
+    is_flag=True,
+    help="Disable all filtering optimizations (raw output)",
+)
+@click.option(
+    "--stdout",
+    is_flag=True,
+    help="Output XML to stdout instead of temp file",
+)
+def preprocess_session(
+    log_path: Path, session_id: str | None, include_agents: bool, no_filtering: bool, stdout: bool
+) -> None:
+    """Preprocess session log JSONL to compressed XML format.
+
+    By default, automatically discovers and includes agent logs (agent-*.jsonl)
+    from the same directory as the main session log.
+
+    All optimization filters are enabled by default for maximum token reduction:
+    - Empty session filtering
+    - Warmup session filtering
+    - Documentation deduplication
+    - Parameter truncation
+    - Tool result pruning
+    - Log discovery operation filtering
+
+    Use --no-filtering to disable all optimizations and get raw output.
+
+    Args:
+        log_path: Path to the main session JSONL file
+        session_id: Optional session ID to filter entries by
+        include_agents: Whether to include agent logs
+        no_filtering: Disable all filtering optimizations
+    """
+    enable_filtering = not no_filtering
+
+    # Process main session log
+    entries, total_entries, skipped_entries = process_log_file(
+        log_path, session_id=session_id, enable_filtering=enable_filtering
+    )
+
+    # Apply filtering operations if enabled
+    if enable_filtering:
+        # Check for empty/warmup sessions
+        if is_empty_session(entries):
+            click.echo("⚠️  Empty session detected - skipping output", err=True)
+            return
+
+        if is_warmup_session(entries):
+            click.echo("⚠️  Warmup session detected - skipping output", err=True)
+            return
+
+        # Apply documentation deduplication
+        entries = deduplicate_documentation_blocks(entries)
+
+        # Apply parameter truncation
+        entries = truncate_tool_parameters(entries)
+
+    # Apply standard deduplication (always enabled)
+    entries = deduplicate_assistant_messages(entries)
+
+    # Show diagnostic output if filtering by session ID
+    if session_id is not None:
+        click.echo(f"✅ Filtered JSONL by session ID: {session_id[:8]}...", err=True)
+        click.echo(
+            f"📊 Included {total_entries - skipped_entries} entries, "
+            f"skipped {skipped_entries} entries",
+            err=True,
+        )
+
+    # Generate main session XML
+    xml_sections = [generate_compressed_xml(entries, enable_pruning=enable_filtering)]
+
+    # Discover and process agent logs if requested
+    if include_agents:
+        agent_logs = discover_agent_logs(log_path)
+        for agent_log in agent_logs:
+            agent_entries, agent_total, agent_skipped = process_log_file(
+                agent_log, session_id=session_id, enable_filtering=enable_filtering
+            )
+
+            # Apply filtering for agent logs
+            if enable_filtering:
+                if is_empty_session(agent_entries):
+                    continue
+                if is_warmup_session(agent_entries):
+                    continue
+                agent_entries = deduplicate_documentation_blocks(agent_entries)
+                agent_entries = truncate_tool_parameters(agent_entries)
+
+            agent_entries = deduplicate_assistant_messages(agent_entries)
+
+            # Generate XML with source label
+            source_label = f"agent-{agent_log.stem.replace('agent-', '')}"
+            agent_xml = generate_compressed_xml(
+                agent_entries, source_label=source_label, enable_pruning=enable_filtering
+            )
+            xml_sections.append(agent_xml)
+
+    # Combine all XML sections
+    xml_content = "\n\n".join(xml_sections)
+
+    # Calculate compression metrics (only when filtering is enabled)
+    if enable_filtering:
+        original_size = sum(len(log_path.read_text(encoding="utf-8")) for _ in [log_path])
+        compressed_size = len(xml_content)
+        if original_size > 0:
+            reduction_pct = ((original_size - compressed_size) / original_size) * 100
+            stats_msg = (
+                f"📉 Token reduction: {reduction_pct:.1f}% "
+                f"({original_size:,} → {compressed_size:,} chars)"
+            )
+            # Route stats to stderr when stdout contains XML
+            click.echo(stats_msg, err=True)
+
+    if stdout:
+        # Output XML directly to stdout
+        click.echo(xml_content)
+    else:
+        # Write to temp file and print path (backward compatible)
+        # Use NamedTemporaryFile to avoid conflicts when multiple tests use same filename
+        filename_session_id = log_path.stem  # Extract session ID from filename
+        with tempfile.NamedTemporaryFile(
+            mode="w",
+            encoding="utf-8",
+            prefix=f"session-{filename_session_id}-",
+            suffix="-compressed.xml",
+            delete=False,
+            dir=tempfile.gettempdir(),
+        ) as f:
+            f.write(xml_content)
+            temp_file = Path(f.name)
+
+        # Print path to stdout for command capture
+        click.echo(str(temp_file))
+
+
+if __name__ == "__main__":
+    preprocess_session()