diff --git a/tests/test_trace.py b/tests/test_trace.py index 57592d1..cdb7c3e 100644 --- a/tests/test_trace.py +++ b/tests/test_trace.py @@ -232,3 +232,82 @@ def msg(eid, parent, role): msg("x", "b", "user"), msg("y", "x", "assistant"), # abandoned msg("c", "b", "user"), msg("d", "c", "toolResult"), msg("e", "d", "assistant")]) assert trace.pi_prompt_count(p) == 2 + + +# Valid JSON that is not an object. A reader that takes every line for a record calls `.get` on +# it. The array holds each marker a reader looks for before it parses a line. +NON_OBJECT_LINES = [ + "null", "42", '"text"', + '["usage", "session_meta", "turn_context", "token_count", "response_item", ' + '"session.start", "type", "system"]'] + + +def _around_non_objects(records): + import json + return NON_OBJECT_LINES + [json.dumps(r) for r in records] + NON_OBJECT_LINES + + +def test_a_non_object_line_does_not_discard_the_calls_around_it(): + claude = {"type": "assistant", "message": {"content": [ + {"type": "tool_use", "id": "t", "name": "Bash", "input": {"command": "ls"}}]}} + codex = {"type": "item.completed", + "item": {"type": "command_execution", "command": "ls", "exit_code": 0}} + start = {"type": "tool.execution_start", + "data": {"toolCallId": "c", "toolName": "bash", "arguments": {"command": "ls"}}} + done = {"type": "tool.execution_complete", "data": {"toolCallId": "c", "success": True}} + assert [c[:2] for c in trace.claude_calls(_around_non_objects([claude]))] == [ + ("Bash", {"command": "ls"})] + assert [c[:2] for c in trace.codex_exec_calls(_around_non_objects([codex]))] == [ + ("shell", {"command": "ls"})] + assert [c[:2] for c in trace.copilot_event_calls(_around_non_objects([start, done]))] == [ + ("bash", {"command": "ls"})] + + +def test_a_non_object_line_does_not_stop_a_session_from_being_read(tmp_path): + from tokeneconomizer.trace.copilot import copilot_session_cwd + claude = {"type": "assistant", "message": {"id": "m1", "usage": {"input_tokens": 5}}} + meta = {"type": "session_meta", "payload": {"cwd": "/w", "cli_version": "1.0"}} + count = {"type": "event_msg", "payload": {"type": "token_count", "info": { + "total_token_usage": {"input_tokens": 10, "cached_input_tokens": 4, + "output_tokens": 3}}}} + start = {"type": "session.start", "data": {"context": {"cwd": "/repo"}}} + p = tmp_path / "claude.jsonl" + p.write_text("\n".join(_around_non_objects([claude]))) + assert trace.claude_session_usage(str(p))["input_tokens"] == 5 + r = tmp_path / "rollout.jsonl" + r.write_text("\n".join(_around_non_objects([meta, count]))) + assert trace.codex_session_meta(str(r)) == {"cwd": "/w", "cli_version": "1.0"} + usage = trace.codex_session_usage(str(r)) + assert (usage["input_tokens"], usage["output_tokens"], usage["requests"]) == (6, 3, 1) + assert copilot_session_cwd(_around_non_objects([start])) == "/repo" + + +def test_a_non_object_line_does_not_stop_a_rollout_from_being_converted(tmp_path): + done = {"type": "event_msg", "payload": {"type": "item_completed", "item": { + "type": "CommandExecution", "command": ["ls"], "exit_code": 0}}} + call = {"type": "response_item", "payload": { + "type": "function_call", "name": "x", "call_id": "c", "arguments": "{}"}} + r, out = tmp_path / "rollout.jsonl", tmp_path / "exec.jsonl" + r.write_text("\n".join(_around_non_objects([done, call]))) + assert trace.codex_rollout_to_exec(str(r), str(out)) == 1 + assert trace.codex_capability_trace(str(r), str(out)) == 2 + + +def test_a_non_object_line_does_not_empty_a_tool_catalog(): + system = {"type": "system", "tools": ["Bash"]} + checkpoint = {"type": "session.usage_checkpoint", "data": {"promptCacheBreakState": [ + {"models": {"m": {"tools": [{"name": "view"}]}}}]}} + text = "\n".join(_around_non_objects([system])) + assert trace.parse_tool_catalog_claude(text, snapshot={}) == {"Bash": ""} + text = "\n".join(_around_non_objects([checkpoint])) + assert trace.parse_tool_names_copilot(text) == {"view": ""} + + +def test_the_loop_guard_and_the_capability_audit_skip_a_non_object_line(tmp_path): + from tokeneconomizer.loopguard.state import completed_calls + from tokeneconomizer.toolprune.codex import inspect_usage + assert [completed_calls(line, {}) for line in NON_OBJECT_LINES] == [[]] * 4 + reconnect = {"type": "error", "message": "Reconnecting... 1/5"} + p = tmp_path / "exec.jsonl" + p.write_text("\n".join(_around_non_objects([reconnect]))) + assert inspect_usage(str(p))["reconnects"] == 1 diff --git a/tokeneconomizer/loopguard/state.py b/tokeneconomizer/loopguard/state.py index 722f8d2..9763be3 100644 --- a/tokeneconomizer/loopguard/state.py +++ b/tokeneconomizer/loopguard/state.py @@ -63,6 +63,8 @@ def completed_calls(line, pend): ev = json.loads(line) except Exception: return [] + if not isinstance(ev, dict): + return [] t = ev.get("type") out = [] # ---- Codex exec --json schema ---- diff --git a/tokeneconomizer/toolprune/codex.py b/tokeneconomizer/toolprune/codex.py index 4cab69c..78b1e14 100644 --- a/tokeneconomizer/toolprune/codex.py +++ b/tokeneconomizer/toolprune/codex.py @@ -208,6 +208,8 @@ def inspect_usage(transcript_path): event = json.loads(line) except Exception: continue + if not isinstance(event, dict): + continue phase = str(event.get("type") or "") if phase == "error" and "reconnecting" in str( event.get("message") or event.get("error") or "").lower(): diff --git a/tokeneconomizer/trace/catalog.py b/tokeneconomizer/trace/catalog.py index 7dfcd87..fbcf357 100644 --- a/tokeneconomizer/trace/catalog.py +++ b/tokeneconomizer/trace/catalog.py @@ -127,7 +127,8 @@ def parse_tool_catalog_claude(transcript_text, max_desc=200, snapshot=None): ev = json.loads(line) except Exception: continue - if ev.get("type") == "system" and isinstance(ev.get("tools"), list): + if isinstance(ev, dict) and ev.get("type") == "system" and isinstance( + ev.get("tools"), list): names = ev["tools"] break return {n: (snap.get(n, "") or "")[:max_desc] for n in names} @@ -204,6 +205,8 @@ def parse_tool_names_copilot(transcript_text): ev = json.loads(line) except Exception: continue + if not isinstance(ev, dict): + continue for pcb in ((ev.get("data") or {}).get("promptCacheBreakState") or []): for model in (pcb.get("models") or {}).values(): for tool in (model.get("tools") or []): diff --git a/tokeneconomizer/trace/claude.py b/tokeneconomizer/trace/claude.py index 3fe97ce..d6c6dcd 100644 --- a/tokeneconomizer/trace/claude.py +++ b/tokeneconomizer/trace/claude.py @@ -33,6 +33,8 @@ def claude_calls(lines): ev = json.loads(line) except Exception: continue + if not isinstance(ev, dict): + continue if ev.get("type", "assistant") not in CONVERSATION: continue msg = ev.get("message") or {} @@ -293,7 +295,7 @@ def claude_session_usage(transcript_path, include_subagents=True): ev = json.loads(line) except ValueError: continue - if ev.get("type") != "assistant": + if not isinstance(ev, dict) or ev.get("type") != "assistant": continue msg = ev.get("message") or {} usage = msg.get("usage") diff --git a/tokeneconomizer/trace/codex.py b/tokeneconomizer/trace/codex.py index 5754091..5e34060 100644 --- a/tokeneconomizer/trace/codex.py +++ b/tokeneconomizer/trace/codex.py @@ -46,7 +46,7 @@ def codex_exec_calls(lines): ev = json.loads(line) except Exception: continue - if ev.get("type") != "item.completed": + if not isinstance(ev, dict) or ev.get("type") != "item.completed": continue item = ev.get("item") or {} kind = item.get("type") @@ -148,6 +148,8 @@ def codex_rollout_to_exec(rollout_path, out_path): ev = json.loads(line) except ValueError: continue + if not isinstance(ev, dict): + continue payload = ev.get("payload") if isinstance(ev.get("payload"), dict) else {} if ev.get("type") == "session_meta" and thread is None: thread = payload.get("id") or payload.get("session_id") @@ -184,6 +186,8 @@ def codex_capability_trace(rollout_path, out_path): ev = json.loads(line) except ValueError: continue + if not isinstance(ev, dict): + continue p = ev.get("payload") if isinstance(ev.get("payload"), dict) else {} if ev.get("type") != "response_item" or p.get("type") not in ( "custom_tool_call", "function_call"): @@ -226,6 +230,8 @@ def codex_session_meta(rollout_path): ev = json.loads(line) except ValueError: continue + if not isinstance(ev, dict): + continue p = ev.get("payload") if isinstance(ev.get("payload"), dict) else {} if ev.get("type") == "session_meta": meta.setdefault("cwd", p.get("cwd")) @@ -249,6 +255,8 @@ def codex_session_usage(rollout_path): ev = json.loads(line) except ValueError: continue + if not isinstance(ev, dict): + continue p = ev.get("payload") if isinstance(ev.get("payload"), dict) else {} if p.get("type") != "token_count": continue diff --git a/tokeneconomizer/trace/copilot.py b/tokeneconomizer/trace/copilot.py index bd305db..909ed52 100644 --- a/tokeneconomizer/trace/copilot.py +++ b/tokeneconomizer/trace/copilot.py @@ -32,6 +32,8 @@ def copilot_event_calls(lines): ev = json.loads(line) except Exception: continue + if not isinstance(ev, dict): + continue t = ev.get("type") d = ev.get("data", {}) if isinstance(ev.get("data"), dict) else {} if t == "tool.execution_start": @@ -53,9 +55,12 @@ def copilot_session_cwd(lines): if '"session.start"' not in line: continue try: - d = json.loads(line).get("data") or {} + ev = json.loads(line) except ValueError: continue + if not isinstance(ev, dict): + continue + d = ev.get("data") or {} cwd = (d.get("context") or {}).get("cwd") return cwd if isinstance(cwd, str) else None return None