vignette/scripts/backfill-claude-token-usage.py

#!/usr/bin/env python
"""과거 Claude CLI 턴을 로컬 Claude JSONL의 실제 usage로 보수적으로 백필한다.

기본은 dry-run이다. 본문은 출력하지 않으며, DB text hash + 생성 시각 창에 후보가
정확히 하나인 기존 0/0 Claude 턴만 --apply 대상으로 삼는다.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
from pathlib import Path
import sys


def _repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def parse_args() -> argparse.Namespace:
    default_root = (
        Path(os.environ.get("USERPROFILE", ""))
        / ".claude"
        / "projects"
        / "D--workspace-vignette-apps-api"
    )
    parser = argparse.ArgumentParser(
        description="Backfill exact Claude token usage from local Claude JSONL metadata.",
    )
    parser.add_argument("--claude-project-dir", type=Path, default=default_root)
    parser.add_argument("--before-seconds", type=float, default=30.0)
    parser.add_argument("--after-seconds", type=float, default=180.0)
    parser.add_argument("--apply", action="store_true")
    parser.add_argument(
        "--expected-matches",
        type=int,
        help="Required with --apply; aborts if the current exact-match count differs.",
    )
    return parser.parse_args()


async def run(args: argparse.Namespace) -> dict[str, object]:
    api_root = _repo_root() / "apps" / "api"
    sys.path.insert(0, str(api_root))
    from app import db
    from app.services.claude_usage_backfill import (
        load_claude_usage_candidates,
        match_claude_usage,
    )

    source_root = args.claude_project_dir.resolve()
    if not source_root.is_dir():
        raise RuntimeError(f"Claude project directory not found: {source_root}")
    candidates = load_claude_usage_candidates(source_root)

    await db.init_pool()
    try:
        async with db.get_pool().acquire() as conn:
            rows = await conn.fetch(
                """
                SELECT id::text, created_at, text_masked
                FROM app.turns
                WHERE speaker = 'client'
                  AND llm_provider = 'claude_cli'
                  AND COALESCE(tokens_in, 0) <= 0
                  AND COALESCE(tokens_out, 0) <= 0
                ORDER BY created_at, id
                """
            )
            report = match_claude_usage(
                rows,
                candidates,
                before_seconds=args.before_seconds,
                after_seconds=args.after_seconds,
            )
            matches = list(report.matches)
            if args.apply:
                if args.expected_matches is None:
                    raise RuntimeError("--expected-matches is required with --apply")
                if len(matches) != args.expected_matches:
                    raise RuntimeError(
                        "exact match count changed: "
                        f"expected={args.expected_matches} actual={len(matches)}"
                    )
                async with conn.transaction():
                    for match in matches:
                        status = await conn.execute(
                            """
                            UPDATE app.turns
                            SET tokens_in = $2,
                                tokens_out = $3
                            WHERE id = $1::uuid
                              AND speaker = 'client'
                              AND llm_provider = 'claude_cli'
                              AND COALESCE(tokens_in, 0) <= 0
                              AND COALESCE(tokens_out, 0) <= 0
                            """,
                            match.turn_id,
                            match.tokens_in,
                            match.tokens_out,
                        )
                        if status != "UPDATE 1":
                            raise RuntimeError(
                                f"guarded update failed for one matched turn: {status}"
                            )
    finally:
        await db.close_pool()

    deltas = [abs(match.delta_seconds) for match in matches]
    return {
        "schema": "vignette.claude_token_backfill.v1",
        "mode": "apply" if args.apply else "dry_run",
        "source_files": len(list(source_root.glob("*.jsonl"))),
        "usage_candidates": len(candidates),
        "eligible_zero_token_turns": len(rows),
        "exact_matches": len(matches),
        "unmatched_turns": report.unmatched_turns,
        "ambiguous_turns": report.ambiguous_turns,
        "tokens_in": sum(match.tokens_in for match in matches),
        "tokens_out": sum(match.tokens_out for match in matches),
        "max_abs_delta_seconds": round(max(deltas), 3) if deltas else None,
        "applied": bool(args.apply),
    }


def main() -> int:
    args = parse_args()
    result = asyncio.run(run(args))
    print(json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())