vignette/apps/api/app/services/claude_usage_backfill.py

"""Claude Code JSONL에서 과거 턴의 실제 토큰 사용량을 안전하게 복구한다."""

from __future__ import annotations

import hashlib
import json
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterable, Mapping, Sequence


@dataclass(frozen=True, slots=True)
class ClaudeUsageCandidate:
    text_digest: bytes
    occurred_at: datetime
    tokens_in: int
    tokens_out: int
    model: str
    source_key: str


@dataclass(frozen=True, slots=True)
class ClaudeUsageMatch:
    turn_id: str
    tokens_in: int
    tokens_out: int
    model: str
    source_key: str
    delta_seconds: float


@dataclass(frozen=True, slots=True)
class ClaudeUsageMatchReport:
    matches: tuple[ClaudeUsageMatch, ...]
    unmatched_turns: int
    ambiguous_turns: int


def normalize_text(value: object) -> str:
    return str(value or "").replace("\r\n", "\n").strip()


def text_digest(value: object) -> bytes:
    return hashlib.sha256(normalize_text(value).encode("utf-8")).digest()


def _safe_usage_int(usage: Mapping[str, object], key: str) -> int:
    try:
        return max(0, int(usage.get(key) or 0))
    except (TypeError, ValueError):
        return 0


def _parse_timestamp(value: object) -> datetime | None:
    if not value:
        return None
    try:
        return datetime.fromisoformat(str(value).replace("Z", "+00:00"))
    except ValueError:
        return None


def _assistant_text(message: Mapping[str, object]) -> str:
    content = message.get("content")
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return ""
    return "".join(
        str(block.get("text") or "")
        for block in content
        if isinstance(block, dict) and block.get("type") == "text"
    )


def load_claude_usage_candidates(root: Path) -> list[ClaudeUsageCandidate]:
    """본문을 외부로 노출하지 않고 assistant text hash와 usage만 읽는다."""

    candidates: list[ClaudeUsageCandidate] = []
    for path in root.glob("*.jsonl"):
        try:
            lines = path.open("r", encoding="utf-8", errors="replace")
        except OSError:
            continue
        with lines:
            for line_number, raw in enumerate(lines, start=1):
                try:
                    item = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    continue
                if item.get("type") != "assistant":
                    continue
                message = item.get("message")
                if not isinstance(message, dict):
                    continue
                usage = message.get("usage")
                if not isinstance(usage, dict):
                    continue
                occurred_at = _parse_timestamp(item.get("timestamp"))
                text = normalize_text(_assistant_text(message))
                if occurred_at is None or not text:
                    continue
                tokens_in = sum(
                    _safe_usage_int(usage, key)
                    for key in (
                        "input_tokens",
                        "cache_read_input_tokens",
                        "cache_creation_input_tokens",
                    )
                )
                tokens_out = _safe_usage_int(usage, "output_tokens")
                if tokens_in <= 0 and tokens_out <= 0:
                    continue
                candidates.append(
                    ClaudeUsageCandidate(
                        text_digest=text_digest(text),
                        occurred_at=occurred_at,
                        tokens_in=tokens_in,
                        tokens_out=tokens_out,
                        model=str(message.get("model") or ""),
                        source_key=f"{path.name}:{item.get('uuid') or line_number}",
                    )
                )
    return candidates


def match_claude_usage(
    rows: Iterable[Mapping[str, object]],
    candidates: Sequence[ClaudeUsageCandidate],
    *,
    before_seconds: float = 30.0,
    after_seconds: float = 180.0,
) -> ClaudeUsageMatchReport:
    """동일 본문 해시와 제한 시간창에 후보가 정확히 하나인 턴만 복구 대상으로 삼는다."""

    by_digest: dict[bytes, list[ClaudeUsageCandidate]] = defaultdict(list)
    for candidate in candidates:
        by_digest[candidate.text_digest].append(candidate)

    matches: list[ClaudeUsageMatch] = []
    unmatched = 0
    ambiguous = 0
    used_sources: set[str] = set()
    for row in rows:
        created_at = row.get("created_at")
        if not isinstance(created_at, datetime):
            unmatched += 1
            continue
        options: list[tuple[float, ClaudeUsageCandidate]] = []
        for candidate in by_digest.get(text_digest(row.get("text_masked")), []):
            delta = (created_at - candidate.occurred_at).total_seconds()
            if -before_seconds <= delta <= after_seconds:
                options.append((delta, candidate))
        if not options:
            unmatched += 1
            continue
        if len(options) != 1:
            ambiguous += 1
            continue
        delta, candidate = options[0]
        if candidate.source_key in used_sources:
            ambiguous += 1
            continue
        used_sources.add(candidate.source_key)
        matches.append(
            ClaudeUsageMatch(
                turn_id=str(row.get("id") or ""),
                tokens_in=candidate.tokens_in,
                tokens_out=candidate.tokens_out,
                model=candidate.model,
                source_key=candidate.source_key,
                delta_seconds=delta,
            )
        )
    return ClaudeUsageMatchReport(
        matches=tuple(matches),
        unmatched_turns=unmatched,
        ambiguous_turns=ambiguous,
    )