diff --git a/Sensor/README.md b/Sensor/README.md index bd81ecf..69a731a 100644 --- a/Sensor/README.md +++ b/Sensor/README.md @@ -233,6 +233,34 @@ agent mode populates the richest version: } ``` +### `token_usage` + +Parsers may attach normalized per-turn and cumulative token counters to an +`AgentEvent`. Only counters reported as nonnegative integers are included: + +```json +{ + "token_usage": { + "last_turn": { + "input_tokens": 120, + "cached_input_tokens": 80, + "output_tokens": 30, + "reasoning_output_tokens": 10, + "total_tokens": 150 + }, + "cumulative": { + "input_tokens": 420, + "cached_input_tokens": 200, + "cache_write_input_tokens": 40, + "output_tokens": 90, + "reasoning_output_tokens": 25, + "total_tokens": 510 + }, + "model_context_window": 128000 + } +} +``` + ## Adding a New Parser ADR Sensor is designed to be extensible. To add support for a new AI agent: diff --git a/Sensor/adr_sensor/parsers/codex_parser.py b/Sensor/adr_sensor/parsers/codex_parser.py index 4987c95..8901e52 100644 --- a/Sensor/adr_sensor/parsers/codex_parser.py +++ b/Sensor/adr_sensor/parsers/codex_parser.py @@ -14,6 +14,15 @@ from ..utils.timestamp_utils import normalize_timestamp from .base_parser import BaseParser +_TOKEN_USAGE_FIELDS = ( + "input_tokens", + "cached_input_tokens", + "cache_write_input_tokens", + "output_tokens", + "reasoning_output_tokens", + "total_tokens", +) + class CodexParser(BaseParser): """Parser for OpenAI Codex CLI JSONL log files.""" @@ -50,6 +59,7 @@ def parse_jsonl_file(self, file_path: Path) -> Optional[AgentEvent]: "timestamp": None, "cwd": None, "model": None, + "token_usage": None, "messages": [], "pending_tool_calls": {}, } @@ -99,6 +109,7 @@ def parse_jsonl_file(self, file_path: Path) -> Optional[AgentEvent]: session_id=f"codex_{session_data['id']}", project_path=session_data["cwd"], model=session_data["model"], + token_usage=session_data["token_usage"], chat_history=chat_history, raw_log_path=str(file_path), ) @@ -169,6 +180,34 @@ def _normalize_tool_output(self, output: Any) -> Optional[str]: return truncate_middle(text, max_length=1000, edge_chars=400) + def _normalize_token_usage(self, info: Any) -> Optional[Dict[str, Any]]: + """Normalize a Codex token_count snapshot into the public session shape.""" + if not isinstance(info, dict): + return None + + token_usage: Dict[str, Any] = {} + for source_key, target_key in ( + ("last_token_usage", "last_turn"), + ("total_token_usage", "cumulative"), + ): + raw_counts = info.get(source_key) + if not isinstance(raw_counts, dict): + continue + + counts = {} + for field in _TOKEN_USAGE_FIELDS: + value = raw_counts.get(field) + if isinstance(value, int) and not isinstance(value, bool) and value >= 0: + counts[field] = value + if counts: + token_usage[target_key] = counts + + context_window = info.get("model_context_window") + if isinstance(context_window, int) and not isinstance(context_window, bool) and context_window >= 0: + token_usage["model_context_window"] = context_window + + return token_usage or None + def _process_event(self, event: Dict[str, Any], session_data: Dict[str, Any]): """Process a single event.""" evt_type = event.get("type") @@ -184,6 +223,12 @@ def _process_event(self, event: Dict[str, Any], session_data: Dict[str, Any]): if payload.get("model"): session_data["model"] = payload.get("model") + elif evt_type == "event_msg": + if isinstance(payload, dict) and payload.get("type") == "token_count": + token_usage = self._normalize_token_usage(payload.get("info")) + if token_usage is not None: + session_data["token_usage"] = token_usage + elif evt_type == "response_item": item_type = payload.get("type") diff --git a/Sensor/adr_sensor/schemas/agent_event_schema.py b/Sensor/adr_sensor/schemas/agent_event_schema.py index 19b18bb..b535a55 100644 --- a/Sensor/adr_sensor/schemas/agent_event_schema.py +++ b/Sensor/adr_sensor/schemas/agent_event_schema.py @@ -69,6 +69,9 @@ class AgentEvent: chunk_sequence: Optional[int] = None is_truncated: bool = False + # Normalized token accounting (populated by some parsers) + token_usage: Optional[Dict[str, Any]] = None + # UUID for this log entry uuid: str = field(init=False) diff --git a/Sensor/tests/test_parsers.py b/Sensor/tests/test_parsers.py index 4112a14..f801b83 100644 --- a/Sensor/tests/test_parsers.py +++ b/Sensor/tests/test_parsers.py @@ -402,20 +402,36 @@ def test_output_list_with_unexpected_items(self, tmp_path): tool = [t for m in CodexParser().parse_jsonl_file(jsonl_file).chat_history for t in m.tools][0] assert tool.result == "bare string\nkept" - def test_event_msg_records_are_ignored(self, tmp_path): - """event_msg records (token_count, web_search_end, ...) must not break parsing. - - They are currently unparsed; this pins that they are skipped cleanly rather - than raising or polluting chat_history. - """ + def test_event_msg_token_count_is_normalized(self, tmp_path): + """The latest valid token_count snapshot is attached to the session.""" jsonl_file = tmp_path / "rollout-eventmsg.jsonl" events = [ {"type": "session_meta", "payload": {"id": "s-em", "timestamp": "2026-08-08T17:00:00.000Z"}}, - {"type": "event_msg", "payload": {"type": "token_count", - "info": {"total_token_usage": {"input_tokens": 10}}}}, + { + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, + "total_token_usage": {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}, + "model_context_window": 64000, + }, + }, + }, + { + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": {"input_tokens": 8, "output_tokens": 3, "total_tokens": 11}, + "total_token_usage": {"input_tokens": 18, "output_tokens": 7, "total_tokens": 25}, + "model_context_window": 128000, + }, + }, + }, {"type": "event_msg", "payload": {"type": "web_search_end", "call_id": "w1", "query": "anything"}}, {"type": "response_item", "payload": {"type": "message", "role": "user", - "content": [{"type": "input_text", "text": "hello there"}]}}, + "content": [{"type": "input_text", "text": "hello there"}]}}, {"type": "response_item", "payload": {"type": "custom_tool_call", "call_id": "c1", "name": "exec", "input": "true"}}, ] @@ -424,9 +440,94 @@ def test_event_msg_records_are_ignored(self, tmp_path): f.write(json.dumps(e) + "\n") entry = CodexParser().parse_jsonl_file(jsonl_file) + assert entry.token_usage == { + "last_turn": {"input_tokens": 8, "output_tokens": 3, "total_tokens": 11}, + "cumulative": {"input_tokens": 18, "output_tokens": 7, "total_tokens": 25}, + "model_context_window": 128000, + } assert len(entry.chat_history) == 2 # user message + assistant tool turn assert len([t for m in entry.chat_history for t in m.tools]) == 1 + def test_token_count_filters_invalid_values_and_ignores_malformed_snapshots(self, tmp_path): + jsonl_file = tmp_path / "rollout-token-validation.jsonl" + events = [ + {"type": "session_meta", "payload": {"id": "s-tokens", "timestamp": "2026-08-08T17:00:00Z"}}, + { + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": { + "input_tokens": 0, + "cached_input_tokens": True, + "cache_write_input_tokens": -1, + "output_tokens": 2.5, + "reasoning_output_tokens": "3", + "total_tokens": 4, + "unknown_tokens": 99, + }, + "total_token_usage": { + "input_tokens": 20, + "cached_input_tokens": 5, + "cache_write_input_tokens": 2, + "output_tokens": 6, + "reasoning_output_tokens": 1, + "total_tokens": 27, + }, + "model_context_window": 128000, + }, + }, + }, + {"type": "event_msg", "payload": {"type": "token_count", "info": None}}, + { + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": [], + "total_token_usage": {"input_tokens": False, "output_tokens": -2}, + "model_context_window": True, + }, + }, + }, + ] + with open(jsonl_file, "w") as f: + for event in events: + f.write(json.dumps(event) + "\n") + + entry = CodexParser().parse_jsonl_file(jsonl_file) + + assert entry.token_usage == { + "last_turn": {"input_tokens": 0, "total_tokens": 4}, + "cumulative": { + "input_tokens": 20, + "cached_input_tokens": 5, + "cache_write_input_tokens": 2, + "output_tokens": 6, + "reasoning_output_tokens": 1, + "total_tokens": 27, + }, + "model_context_window": 128000, + } + + def test_token_count_accepts_context_window_only_snapshot(self, tmp_path): + jsonl_file = tmp_path / "rollout-context-window.jsonl" + events = [ + {"type": "session_meta", "payload": {"id": "s-context-window"}}, + { + "type": "event_msg", + "payload": { + "type": "token_count", + "info": {"model_context_window": 128000}, + }, + }, + ] + jsonl_file.write_text("\n".join(json.dumps(event) for event in events), encoding="utf-8") + + entry = CodexParser().parse_jsonl_file(jsonl_file) + + assert entry.token_usage == {"model_context_window": 128000} + def test_parse_no_directory(self): """Test parse_all when directory doesn't exist.""" parser = CodexParser() diff --git a/Sensor/tests/test_schemas.py b/Sensor/tests/test_schemas.py index c671d2a..a147716 100644 --- a/Sensor/tests/test_schemas.py +++ b/Sensor/tests/test_schemas.py @@ -1,8 +1,10 @@ """Tests for ADR Sensor schemas.""" -import pytest +import json from datetime import datetime, timezone +import pytest + from adr_sensor.schemas.agent_event_schema import AgentEvent, ChatMessage, ToolUsage @@ -58,6 +60,7 @@ def test_create_basic(self): ) assert event.source == "claude" assert event.session_id == "test_session_1" + assert event.token_usage is None assert event.uuid is not None assert len(event.uuid) == 64 # SHA-256 hex @@ -74,6 +77,32 @@ def test_uuid_deterministic(self): event2 = AgentEvent(**kwargs) assert event1.uuid == event2.uuid + def test_new_fields_preserve_existing_positional_arguments(self): + event = AgentEvent( + datetime(2025, 1, 1, tzinfo=timezone.utc), + "codex", + "test-session", + [], + "user-id", + "/workspace", + "model-id", + "test-host", + "test-user", + "/logs/session.jsonl", + {"originator": "cli"}, + True, + 3, + 2, + True, + ) + + assert event.session_context == {"originator": "cli"} + assert event.is_chunked is True + assert event.total_chunks == 3 + assert event.chunk_sequence == 2 + assert event.is_truncated is True + assert event.token_usage is None + def test_uuid_unique_for_different_sessions(self): event1 = AgentEvent( timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc), @@ -152,6 +181,24 @@ def test_to_dict(self): assert d["source"] == "claude" assert d["timestamp"] == "2025-01-01T00:00:00+00:00" + def test_token_usage_serialization(self): + token_usage = { + "last_turn": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15}, + "cumulative": {"input_tokens": 30, "output_tokens": 8, "total_tokens": 38}, + "model_context_window": 128000, + } + event = AgentEvent( + timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc), + source="codex", + session_id="test", + hostname="test-host", + username="test-user", + token_usage=token_usage, + ) + + assert event.to_dict()["token_usage"] == token_usage + assert json.loads(event.to_json())["token_usage"] == token_usage + def test_to_json(self): event = AgentEvent( timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc),