coze_lab 0.1.50 → 0.1.51

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "coze_lab",
-  "version": "0.1.50",
+  "version": "0.1.51",
   "description": "Configure local AI agents (Claude Code, Codex, OpenClaw) to report traces to CozeLoop",
   "keywords": [
     "cozeloop",

package/scripts/codex/cozeloop_hook.py

@@ -695,6 +695,32 @@ def _assistant_bounds(turn):
     return min(values), max(values)
 
 
+def extract_token_events(entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """扫 rollout event_msg/token_count 事件, 与 turn 完全解耦.
+
+    每条事件独立采集, 依赖 _line_number 做全局水位线幂等. 只收 last_token_usage 非空的事件
+    (空的是心跳/占位, 不代表真实模型调用). total_token_usage 一并保留供 tag 上报.
+    """
+    events: List[Dict[str, Any]] = []
+    for entry in entries:
+        if entry.get("type") != "event_msg":
+            continue
+        payload = entry.get("payload") or {}
+        if payload.get("type") != "token_count":
+            continue
+        info = payload.get("info") or {}
+        last_usage = info.get("last_token_usage")
+        if not isinstance(last_usage, dict) or not last_usage:
+            continue
+        events.append({
+            "line": entry.get("_line_number", 0),
+            "timestamp": entry.get("timestamp"),
+            "usage": last_usage,
+            "total_usage": info.get("total_token_usage") or {},
+        })
+    return events
+
+
 def group_messages_into_turns(entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """Group raw JSONL entries into conversation turns.
 
@@ -1361,7 +1387,7 @@ def send_turns_to_cozeloop(turns: List[Dict[str, Any]], session_id: str, model_n
 # tool span 挂到固定 root、当场 flush。codex 的实时单元是 tool_call（turn 内平铺，按 call_id 关联
 # result）。全局水位线 state["rt_last_tool"] 跟踪已发的 tool_call 数。
 
-def send_steps_realtime(turns, session_id, state, model_name="codex", coze_tags_override=None, is_terminal=False):
+def send_steps_realtime(turns, session_id, state, model_name="codex", coze_tags_override=None, is_terminal=False, token_events=None):
     import cozeloop
     from cozeloop.spec.tracespec import Runtime
     token = get_fresh_token()
@@ -1494,6 +1520,62 @@ def send_steps_realtime(turns, session_id, state, model_name="codex", coze_tags_
         if sent > 0:
             state["rt_last_upload_ts"] = time.time()
 
+        # --- token_count 独立增量, 与 tool span 同一 root 下, 按行号水位线幂等 ---
+        # 每条 token_count 挂一个 model_usage_<line> span:
+        #   last_token_usage → set_input_tokens/set_output_tokens (SDK 官方求和字段)
+        #   total_token_usage → tags (只观察, 不参与后端聚合, 否则会重复累加)
+        # 只在 flush 成功后推进 rt_last_token_line, 保证上报失败下次能重发.
+        token_sent = 0
+        last_token_line = state.get("rt_last_token_line", -1)
+        pending_tokens = [ev for ev in (token_events or []) if ev["line"] > last_token_line]
+        max_flushed_line = last_token_line
+        for ev in pending_tokens:
+            line = ev["line"]
+            last = ev.get("usage") or {}
+            total = ev.get("total_usage") or {}
+            in_tk = _usage_int(last, "input_tokens")
+            out_tk = _usage_int(last, "output_tokens")
+            start_dt = _parse_ts_value(ev.get("timestamp"))
+            tk_span = client.start_span(name=f"model_usage_{line}", span_type="model",
+                                        start_time=start_dt, child_of=root_ctx)
+            tk_span.set_runtime(Runtime(library="codex-cli"))
+            try: tk_span.set_model_name(model_name)
+            except Exception: pass
+            if start_dt is not None:
+                _set_finish_time_safe(tk_span, start_dt)
+            if in_tk > 0:
+                try: tk_span.set_input_tokens(in_tk)
+                except Exception: pass
+            if out_tk > 0:
+                try: tk_span.set_output_tokens(out_tk)
+                except Exception: pass
+            tk_tags = {
+                "token_event_line": line,
+                "token_source": "codex_token_count",
+                "cached_input_tokens": _usage_int(last, "cached_input_tokens"),
+                "reasoning_output_tokens": _usage_int(last, "reasoning_output_tokens"),
+                "total_tokens": _usage_int(last, "total_tokens"),
+            }
+            if total:
+                tk_tags["total_input_tokens"] = _usage_int(total, "input_tokens")
+                tk_tags["total_output_tokens"] = _usage_int(total, "output_tokens")
+                tk_tags["total_cached_input_tokens"] = _usage_int(total, "cached_input_tokens")
+                tk_tags["total_reasoning_output_tokens"] = _usage_int(total, "reasoning_output_tokens")
+                tk_tags["total_all_tokens"] = _usage_int(total, "total_tokens")
+            tk_span.set_tags(tk_tags)
+            tk_span.finish()
+            client.flush()
+            if upload_events:
+                hook_log(f"[rt] upload failed after token flush line={line} detail={upload_events[-1][:500]}")
+                if max_flushed_line > last_token_line:
+                    state["rt_last_token_line"] = max_flushed_line
+                return None
+            max_flushed_line = line
+            token_sent += 1
+        if token_sent > 0:
+            state["rt_last_token_line"] = max_flushed_line
+            state["rt_last_upload_ts"] = time.time()
+
         if is_terminal:
             # 收尾：补最终 assistant 输出
             last_out = ""
@@ -1519,7 +1601,7 @@ def send_steps_realtime(turns, session_id, state, model_name="codex", coze_tags_
             hook_log(f"[rt] upload failed state not advanced detail={upload_events[-1][:500]}")
             return None
 
-        hook_log(f"[rt] sent {sent} tool span(s), last_tool={state['rt_last_tool']}, terminal={is_terminal}")
+        hook_log(f"[rt] sent {sent} tool span(s), {token_sent} token span(s), last_tool={state['rt_last_tool']}, last_token_line={state.get('rt_last_token_line', -1)}, terminal={is_terminal}")
         return True
     except Exception as e:
         hook_log(f"[rt] error: {e}")
@@ -1637,7 +1719,9 @@ def main():
 
     # Group entries into turns
     turns = group_messages_into_turns(entries)
-    debug_log(f"Grouped into {len(turns)} turns")
+    # 独立扫 token_count 事件, 与 turn 解耦, 走独立行号水位线上报.
+    token_events = extract_token_events(entries)
+    debug_log(f"Grouped into {len(turns)} turns, {len(token_events)} token events")
 
     # If this is a subagent, save data for parent to include later
     if is_subagent:
@@ -1677,7 +1761,8 @@ def main():
         # 方案 B 实时：把新完成的 tool_call 作为 span 挂稳定 root，结束即可见。不节流。
         ok = send_steps_realtime(turns, session_id, state, model_name=model_name,
                                  coze_tags_override=coze_tags_override,
-                                 is_terminal=is_terminal_event)
+                                 is_terminal=is_terminal_event,
+                                 token_events=token_events)
         if ok is None:
             hook_log(f"realtime send failed session_id={session_id}")
             debug_log("Realtime send failed, state not advanced")