okstra 0.200.0 → 0.201.0

package/runtime/python/okstra_token_usage/claude.py

@@ -4,12 +4,12 @@ from __future__ import annotations
 import json
 import sys
 import re
-from datetime import datetime
 from pathlib import Path
 
 from .cursor import MAX_NEEDLES, fresh_cache, load_cache, save_cache
 from .paths import claude_project_dir, ts_in_window
 from okstra_ctl.wrapper_status import read_wrapper_status
+from okstra_ctl.usage_cells import duration_ms_from_bounds
 
 # lead 의 단계 체크포인트 라인(prompts/lead/okstra-lead-contract.md "Progress reporting") — phase id 가
 # `phase-<digit>` 로 시작하는 라인만 마커로 인정해 일반 대화의 오탐을 줄인다.
@@ -45,19 +45,13 @@ def claude_code_status_total(status_path: Path) -> dict:
     }
 
 
-def _event_from_record(rec: dict) -> dict | None:
-    """jsonl 레코드 1개 → 압축 이벤트. 집계에 기여하지 않으면 None.
-
-    키: t=timestamp, i/o=input/output, c=cache_creation 합, c5/c1=ephemeral
-    5m/1h, r=cache_read, u=tool_use 수. 0/부재 필드는 생략(캐시 크기 절약).
-    ts-only 레코드도 보존한다 — 임의 윈도우의 first/last ts 산출에 필요.
-    """
-    msg = rec.get("message")
-    if not isinstance(msg, dict):
-        msg = {}
+def _usage_from_message(msg: dict) -> dict:
+    """응답 식별자를 남겨 여러 출력 조각의 사용량을 한 번만 집계한다."""
     ev: dict = {}
     usage = msg.get("usage")
-    if usage:
+    if isinstance(usage, dict):
+        if isinstance(msg.get("id"), str) and msg["id"]:
+            ev["m"] = msg["id"]
         for src, key in (("input_tokens", "i"), ("output_tokens", "o"),
                          ("cache_read_input_tokens", "r")):
             v = usage.get(src, 0) or 0
@@ -79,6 +73,15 @@ def _event_from_record(rec: dict) -> dict | None:
         elif cc_total:
             # API 분해가 없으면 전부 5m 티어로(1.25x — 더 싼 가정, 기존 동작).
             ev["c5"] = cc_total
+    return ev
+
+
+def _event_from_record(rec: dict) -> dict | None:
+    """사용량과 별개로 각 조각의 도구·진행·시각 이벤트를 보존한다."""
+    msg = rec.get("message")
+    if not isinstance(msg, dict):
+        msg = {}
+    ev = _usage_from_message(msg)
     if rec.get("type") == "assistant":
         tools = sum(1 for b in (msg.get("content") or [])
                     if isinstance(b, dict) and b.get("type") == "tool_use")
@@ -102,6 +105,31 @@ def _event_from_record(rec: dict) -> dict | None:
     return ev or None
 
 
+def _response_usage_events(events: list[dict], since: str | None, until: str | None):
+    """응답별 누적 스냅샷의 창 안 증가량만 세며 도구 이벤트는 그대로 둔다."""
+    latest: dict[str, int] = {}
+    baseline: dict[str, dict] = {}
+    for index, event in enumerate(events):
+        message_id = event.get("m")
+        timestamp = event.get("t")
+        if message_id and (not timestamp or ts_in_window(timestamp, None, until)):
+            latest[message_id] = index
+            if timestamp and since and not ts_in_window(timestamp, since, None):
+                baseline[message_id] = event
+    usage_keys = {"i", "o", "c", "c5", "c1", "r"}
+    for index, event in enumerate(events):
+        message_id = event.get("m")
+        if not message_id:
+            yield event
+            continue
+        projected = {key: value for key, value in event.items() if key not in usage_keys}
+        if latest.get(message_id) == index:
+            previous = baseline.get(message_id) or {}
+            projected.update({key: max(0, event.get(key, 0) - previous.get(key, 0))
+                              for key in usage_keys})
+        yield projected
+
+
 def _session_meta_from_record(rec: dict) -> tuple[str | None, str | None]:
     """레코드에서 (agentName, model) 후보 추출 — 둘 다 first-non-null 정책."""
     agent = rec.get("agentName") or None
@@ -213,7 +241,7 @@ def _totals_from_events(events: list[dict], agent_name: str | None,
     progress_markers: list[dict] = []
     first_ts: str | None = None
     last_ts: str | None = None
-    for ev in events:
+    for ev in _response_usage_events(events, since, until):
         ts = ev.get("t")
         if ts and not ts_in_window(ts, since, until):
             continue
@@ -231,14 +259,7 @@ def _totals_from_events(events: list[dict], agent_name: str | None,
                 first_ts = ts
             if last_ts is None or ts > last_ts:
                 last_ts = ts
-    duration_ms = 0
-    if first_ts and last_ts:
-        try:
-            a = datetime.fromisoformat(first_ts.replace("Z", "+00:00"))
-            b = datetime.fromisoformat(last_ts.replace("Z", "+00:00"))
-            duration_ms = max(0, int((b - a).total_seconds() * 1000))
-        except ValueError:
-            duration_ms = 0
+    duration_ms = duration_ms_from_bounds(first_ts, last_ts) or 0
     # '처리 토큰' total 에서 cache_read 는 제외한다. claude 는 매 턴 직전까지의
     # 컨텍스트 전체를 캐시에서 재읽기(cache_read)하므로, 단순 합산하면 같은 토큰을
     # 턴 수만큼 중복 카운트해 처리량이 비현실적으로 부풀려진다(예: in-session

package/runtime/python/okstra_token_usage/codex.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 
 import json
 import os
+import re
 from datetime import datetime, timezone
 from pathlib import Path
 from .jsonl_io import iter_jsonl
@@ -12,6 +13,22 @@ from .paths import CODEX_SESSIONS, codex_session_roots, ts_in_window
 _DEFAULT_CODEX_SESSIONS = CODEX_SESSIONS
 
 
+def codex_wrapper_session_ids(log_path: Path) -> set[str]:
+    """사용자 프롬프트 앞 실행기 헤더에서만 세션 ID를 읽는다."""
+    try:
+        with log_path.open(encoding="utf-8", errors="replace") as stream:
+            for line in stream:
+                clean = re.sub(r"\x1b\[[0-9;]*m", "", line).strip()
+                if clean == "user" or clean.endswith("[worker]user"):
+                    break
+                match = re.search(r"(?:^|\[worker\])session id:\s*([\w-]+)\s*$", clean)
+                if match:
+                    return {match.group(1)}
+    except OSError:
+        return set()
+    return set()
+
+
 def codex_session_total(jsonl_path: Path) -> dict:
     """Return last token_count snapshot from a codex rollout jsonl."""
     last: dict | None = None