claude-session-fix-thinking

#!/usr/bin/env python3
#
# claude-session-fix-thinking: Fix "thinking blocks cannot be modified" errors
# Author: Mitesh Ashar
# Version: 1.0.0
# Last Updated: 2026-02-14
#
# Diagnoses and fixes Claude Code session JSONL files where streaming
# interleaving has corrupted thinking blocks, causing API rejection.
#
# Usage:
#   claude-session-fix-thinking diagnose <session-id>   # Show corruption points
#   claude-session-fix-thinking fix <session-id>         # Fix and back up
#   claude-session-fix-thinking nuke <session-id>        # Strip ALL thinking blocks
#
# The session JSONL is auto-discovered under ~/.claude/projects/
# Backups are created alongside the JSONL as <session-id>.jsonl.backup
#
# Dependencies: python3
#

import json
import os
import sys
import shutil
from pathlib import Path
import time

CLAUDE_DIR = Path(os.environ['CLAUDE_CONFIG_DIR']) if 'CLAUDE_CONFIG_DIR' in os.environ else Path.home() / ".claude"


def find_session_jsonl(session_id):
    """Find the JSONL file for a session ID across all projects."""
    for jsonl in CLAUDE_DIR.glob(f"projects/*/{session_id}.jsonl"):
        return jsonl
    # Also check without project nesting
    for jsonl in CLAUDE_DIR.glob(f"{session_id}.jsonl"):
        return jsonl
    return None


def load_lines(path):
    with open(path) as f:
        return f.readlines()


def find_corruption_points(lines):
    """Find consecutive assistant lines with different msg_ids involving thinking blocks.

    Returns list of dicts with keys:
        host_line, intruder_line, host_id, intruder_id,
        host_has_thinking, intruder_has_thinking, pattern
    """
    corruptions = []

    # Build list of (line_num, msg_id, has_thinking, block_types) for assistant lines
    assistant_runs = []  # groups of consecutive assistant lines (ignoring progress)
    current_run = []

    for i, line in enumerate(lines):
        obj = json.loads(line)
        t = obj.get('type', '')
        if t == 'assistant':
            msg = obj.get('message', {})
            msg_id = msg.get('id', '')
            content = msg.get('content', [])
            block_types = []
            has_thinking = False
            for b in (content if isinstance(content, list) else []):
                if isinstance(b, dict):
                    bt = b.get('type', '?')
                    block_types.append(bt)
                    if bt in ('thinking', 'redacted_thinking'):
                        has_thinking = True
            current_run.append({
                'line': i, 'msg_id': msg_id,
                'has_thinking': has_thinking, 'block_types': block_types,
            })
        elif t == 'progress':
            continue  # progress lines don't break assistant runs
        else:
            if current_run:
                assistant_runs.append(current_run)
                current_run = []

    if current_run:
        assistant_runs.append(current_run)

    # Within each run, find adjacent entries with different msg_ids + thinking
    for run in assistant_runs:
        for j in range(len(run) - 1):
            a, b = run[j], run[j + 1]
            if a['msg_id'] != b['msg_id'] and (a['has_thinking'] or b['has_thinking']):
                # Determine pattern
                if a['has_thinking'] and b['has_thinking']:
                    pattern = 'both_thinking'
                elif b['has_thinking'] and not a['has_thinking']:
                    pattern = 'intruder_thinking_after'
                else:
                    pattern = 'host_thinking_before_intruder'

                corruptions.append({
                    'host_line': a['line'],
                    'intruder_line': b['line'],
                    'host_id': a['msg_id'],
                    'intruder_id': b['msg_id'],
                    'host_blocks': a['block_types'],
                    'intruder_blocks': b['block_types'],
                    'host_has_thinking': a['has_thinking'],
                    'intruder_has_thinking': b['has_thinking'],
                    'pattern': pattern,
                })

    return corruptions


def find_trailing_junk(lines):
    """Find trailing API error messages, progress, system lines to strip."""
    remove_from = len(lines)
    for i in range(len(lines) - 1, -1, -1):
        obj = json.loads(lines[i])
        t = obj.get('type', '')
        is_error = obj.get('isApiErrorMessage', False)
        if t in ('progress', 'system', 'file-history-snapshot') or is_error:
            remove_from = i
        elif t == 'user':
            content = obj.get('message', {}).get('content', [])
            if not content or (isinstance(content, list) and len(content) == 0):
                remove_from = i
            else:
                break
        else:
            break
    return remove_from


def get_thinking_summary(obj):
    """Extract thinking block summary from a JSONL line."""
    msg = obj.get('message', {})
    content = msg.get('content', [])
    for b in (content if isinstance(content, list) else []):
        if isinstance(b, dict) and b.get('type') == 'thinking':
            text = b.get('thinking', '')
            sig = (b.get('signature', '') or 'none')[:35]
            return f"thinking({len(text)}ch, sig={sig}...)"
        elif isinstance(b, dict) and b.get('type') == 'redacted_thinking':
            data = b.get('data', '')
            return f"redacted_thinking({len(data)}ch)"
    return None


# ──────────────────────────────────────────────
# Commands
# ──────────────────────────────────────────────

def cmd_diagnose(session_id):
    """Diagnose a session for thinking block corruption."""
    path = find_session_jsonl(session_id)
    if not path:
        print(f"ERROR: No JSONL found for session {session_id}", file=sys.stderr)
        sys.exit(1)

    print(f"Session: {session_id}")
    print(f"File: {path} ({path.stat().st_size // 1024}KB)")

    lines = load_lines(path)
    print(f"Lines: {len(lines)}")

    # Count thinking blocks
    thinking_count = 0
    for line in lines:
        obj = json.loads(line)
        if obj.get('type') != 'assistant':
            continue
        content = obj.get('message', {}).get('content', [])
        for b in (content if isinstance(content, list) else []):
            if isinstance(b, dict) and b.get('type') in ('thinking', 'redacted_thinking'):
                thinking_count += 1
    print(f"Thinking blocks: {thinking_count}")

    # Find corruption
    corruptions = find_corruption_points(lines)
    if not corruptions:
        print("\nNo corruption points found. Session JSONL looks clean.")
        print("If the error persists at runtime, the corruption may be in-memory only.")
        print("Try: claude --resume (reload from JSONL may clear it)")
        return

    print(f"\nCorruption points: {len(corruptions)}")
    for i, c in enumerate(corruptions):
        print(f"\n  [{i+1}] Lines {c['host_line']}-{c['intruder_line']} ({c['pattern']})")
        print(f"      Host:     id={c['host_id'][:30]}... blocks={c['host_blocks']}")
        print(f"      Intruder: id={c['intruder_id'][:30]}... blocks={c['intruder_blocks']}")

        host_obj = json.loads(lines[c['host_line']])
        intruder_obj = json.loads(lines[c['intruder_line']])

        ht = get_thinking_summary(host_obj)
        it = get_thinking_summary(intruder_obj)
        if ht:
            print(f"      Host thinking: {ht}")
        if it:
            print(f"      Intruder thinking: {it}")

    # Trailing junk
    trim_from = find_trailing_junk(lines)
    trailing = len(lines) - trim_from
    if trailing > 0:
        print(f"\nTrailing junk: {trailing} lines (errors/progress/system)")

    print(f"\nTo fix: claude-session-fix-thinking fix {session_id}")
    if corruptions:
        print(f"Nuclear: claude-session-fix-thinking nuke {session_id}")


def cmd_fix(session_id):
    """Fix thinking block corruption with targeted repairs."""
    path = find_session_jsonl(session_id)
    if not path:
        print(f"ERROR: No JSONL found for session {session_id}", file=sys.stderr)
        sys.exit(1)

    lines = load_lines(path)
    corruptions = find_corruption_points(lines)

    if not corruptions:
        print("No corruption points found. Nothing to fix.")
        print("If the error persists, try: claude-session-fix-thinking nuke " + session_id)
        return

    print(f"Session: {session_id}")
    print(f"File: {path}")
    print(f"Lines: {len(lines)}")
    print(f"Corruption points: {len(corruptions)}")

    lines_to_remove = set()
    lines_to_modify = {}

    for i, c in enumerate(corruptions):
        host_line = c['host_line']
        intruder_line = c['intruder_line']

        print(f"\n  [{i+1}] Lines {host_line}-{intruder_line} ({c['pattern']})")

        if c['pattern'] == 'both_thinking':
            # Both have thinking — merge intruder into host
            host_obj = json.loads(lines[host_line])
            intruder_obj = json.loads(lines[intruder_line])

            host_content = host_obj['message']['content']
            intruder_content = intruder_obj['message']['content']

            # Find thinking blocks
            host_tb = next((b for b in host_content if isinstance(b, dict) and b.get('type') == 'thinking'), None)
            intruder_tb = next((b for b in intruder_content if isinstance(b, dict) and b.get('type') == 'thinking'), None)

            if host_tb and intruder_tb:
                merged_text = host_tb['thinking'] + "\n\n" + intruder_tb['thinking']
                host_tb['thinking'] = merged_text
                lines_to_modify[host_line] = json.dumps(host_obj, ensure_ascii=False) + "\n"
                lines_to_remove.add(intruder_line)
                print(f"       Merged thinking: {len(host_tb['thinking'])-len(intruder_tb['thinking'])-2}+{len(intruder_tb['thinking'])}={len(merged_text)}ch")
            else:
                # Fallback: just remove the intruder's thinking-only line
                lines_to_remove.add(intruder_line)
                print(f"       Removed intruder line {intruder_line}")

        elif c['pattern'] == 'intruder_thinking_after':
            # Intruder has thinking, host doesn't — remove the thinking line
            # Check if intruder line is thinking-only (single block)
            intruder_obj = json.loads(lines[intruder_line])
            intruder_blocks = intruder_obj.get('message', {}).get('content', [])
            if len(intruder_blocks) == 1 and isinstance(intruder_blocks[0], dict) and intruder_blocks[0].get('type') == 'thinking':
                lines_to_remove.add(intruder_line)
                print(f"       Removed thinking-only line {intruder_line}")
            else:
                # Multi-block intruder — just strip the thinking block from content
                new_content = [b for b in intruder_blocks if not (isinstance(b, dict) and b.get('type') in ('thinking', 'redacted_thinking'))]
                intruder_obj['message']['content'] = new_content
                lines_to_modify[intruder_line] = json.dumps(intruder_obj, ensure_ascii=False) + "\n"
                print(f"       Stripped thinking from line {intruder_line} (kept {len(new_content)} other blocks)")

        elif c['pattern'] == 'host_thinking_before_intruder':
            # Host has thinking, intruder doesn't — remove host's thinking
            host_obj = json.loads(lines[host_line])
            host_blocks = host_obj.get('message', {}).get('content', [])
            if len(host_blocks) == 1 and isinstance(host_blocks[0], dict) and host_blocks[0].get('type') == 'thinking':
                lines_to_remove.add(host_line)
                print(f"       Removed thinking-only line {host_line}")
            else:
                new_content = [b for b in host_blocks if not (isinstance(b, dict) and b.get('type') in ('thinking', 'redacted_thinking'))]
                host_obj['message']['content'] = new_content
                lines_to_modify[host_line] = json.dumps(host_obj, ensure_ascii=False) + "\n"
                print(f"       Stripped thinking from line {host_line} (kept {len(new_content)} other blocks)")

    # Strip trailing junk
    trim_from = find_trailing_junk(lines)
    trailing = len(lines) - trim_from
    if trailing > 0:
        for k in range(trim_from, len(lines)):
            lines_to_remove.add(k)
        print(f"\n  Stripping {trailing} trailing error/progress/system lines")

    # Build fixed output
    fixed_lines = []
    for idx, line in enumerate(lines):
        if idx in lines_to_remove:
            continue
        if idx in lines_to_modify:
            fixed_lines.append(lines_to_modify[idx])
        else:
            fixed_lines.append(line)

    print(f"\n  Summary: {len(lines)} → {len(fixed_lines)} lines ({len(lines_to_remove)} removed, {len(lines_to_modify)} modified)")

    # Backup alongside the JSONL with unix timestamp
    backup_path = path.with_suffix(f".jsonl.{int(time.time())}.backup")
    shutil.copy2(path, backup_path)
    print(f"  Backup: {backup_path}")

    # Write
    with open(path, 'w') as f:
        f.writelines(fixed_lines)
    print(f"  Written: {path}")
    print(f"\n  Resume: claude --resume {session_id}")


def cmd_nuke(session_id):
    """Nuclear option: strip ALL thinking blocks from the session."""
    path = find_session_jsonl(session_id)
    if not path:
        print(f"ERROR: No JSONL found for session {session_id}", file=sys.stderr)
        sys.exit(1)

    lines = load_lines(path)
    print(f"Session: {session_id}")
    print(f"File: {path}")
    print(f"Lines: {len(lines)}")

    fixed_lines = []
    thinking_removed = 0
    lines_removed = 0

    for line in lines:
        obj = json.loads(line)
        t = obj.get('type', '')

        # Skip API error messages
        if obj.get('isApiErrorMessage', False):
            lines_removed += 1
            continue

        if t == 'assistant':
            msg = obj.get('message', {})
            content = msg.get('content', [])
            if isinstance(content, list):
                new_content = []
                for b in content:
                    if isinstance(b, dict) and b.get('type') in ('thinking', 'redacted_thinking'):
                        thinking_removed += 1
                    else:
                        new_content.append(b)

                if not new_content:
                    # Entire line was just thinking — remove it
                    lines_removed += 1
                    continue

                if len(new_content) != len(content):
                    msg['content'] = new_content
                    obj['message'] = msg
                    line = json.dumps(obj, ensure_ascii=False) + "\n"

        fixed_lines.append(line)

    # Also strip trailing junk from fixed_lines
    while fixed_lines:
        obj = json.loads(fixed_lines[-1])
        t = obj.get('type', '')
        if t in ('progress', 'system', 'file-history-snapshot'):
            fixed_lines.pop()
            lines_removed += 1
        elif t == 'user':
            content = obj.get('message', {}).get('content', [])
            if not content:
                fixed_lines.pop()
                lines_removed += 1
            else:
                break
        else:
            break

    print(f"  Thinking blocks removed: {thinking_removed}")
    print(f"  Lines removed: {lines_removed}")
    print(f"  Result: {len(fixed_lines)} lines")

    # Backup alongside the JSONL with unix timestamp
    backup_path = path.with_suffix(f".jsonl.{int(time.time())}.backup")
    shutil.copy2(path, backup_path)
    print(f"  Backup: {backup_path}")

    # Write
    with open(path, 'w') as f:
        f.writelines(fixed_lines)
    print(f"  Written: {path}")
    print(f"\n  Resume: claude --resume {session_id}")


# ──────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────

def main():
    if len(sys.argv) < 3:
        print("Usage: claude-session-fix-thinking <command> <session-id>")
        print()
        print("Commands:")
        print("  diagnose  Show corruption points without modifying anything")
        print("  fix       Fix corruption points (targeted repair + backup)")
        print("  nuke      Strip ALL thinking blocks (nuclear option + backup)")
        sys.exit(1)

    cmd = sys.argv[1]
    session_id = sys.argv[2]

    if cmd == 'diagnose':
        cmd_diagnose(session_id)
    elif cmd == 'fix':
        cmd_fix(session_id)
    elif cmd == 'nuke':
        cmd_nuke(session_id)
    else:
        print(f"Unknown command: {cmd}", file=sys.stderr)
        print("Commands: diagnose, fix, nuke", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()