Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions src/autoskillit/core/types/_type_token.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,12 @@ def to_dict(self) -> dict[str, Any]:
return {
"input_tokens": self.input_tokens,
"output_tokens": self.output_tokens,
"cache_read_tokens": self.cache_read_tokens,
"cache_write_tokens": self.cache_write_tokens,
"cache_read_tokens": self.cache_read_tokens
if self.cache_read_tokens is not None
else 0,
"cache_write_tokens": self.cache_write_tokens
if self.cache_write_tokens is not None
else 0,
"provider": self.provider,
}

Expand Down
24 changes: 14 additions & 10 deletions src/autoskillit/execution/session_log.py
Original file line number Diff line number Diff line change
Expand Up @@ -449,16 +449,20 @@ def flush_session_log(
# Write per-session telemetry files; gate on data presence, not session identity
label = _resolve_session_label(step_name, dispatch_id)
if token_usage is not None:
_cache_write = token_usage.get(
"cache_write_tokens", token_usage.get("cache_creation_input_tokens", 0)
_cw_raw = token_usage.get("cache_write_tokens")
_cache_write = (
_cw_raw
if _cw_raw is not None
else (token_usage.get("cache_creation_input_tokens") or 0)
)
_cache_read = token_usage.get(
"cache_read_tokens", token_usage.get("cache_read_input_tokens", 0)
_cr_raw = token_usage.get("cache_read_tokens")
_cache_read = (
_cr_raw if _cr_raw is not None else (token_usage.get("cache_read_input_tokens") or 0)
)
tu_data = {
"session_label": label,
"input_tokens": token_usage.get("input_tokens", 0),
"output_tokens": token_usage.get("output_tokens", 0),
"input_tokens": token_usage.get("input_tokens") or 0,
"output_tokens": token_usage.get("output_tokens") or 0,
"cache_write_tokens": _cache_write,
"cache_read_tokens": _cache_read,
"timing_seconds": timing_seconds if timing_seconds is not None else 0.0,
Expand Down Expand Up @@ -513,10 +517,10 @@ def flush_session_log(
"peak_rss_kb": peak_rss_kb,
"peak_oom_score": peak_oom_score,
"step_name": step_name,
"input_tokens": token_usage.get("input_tokens", 0) if token_usage else 0,
"output_tokens": token_usage.get("output_tokens", 0) if token_usage else 0,
"cache_write_tokens": token_usage.get("cache_write_tokens", 0) if token_usage else 0,
"cache_read_tokens": token_usage.get("cache_read_tokens", 0) if token_usage else 0,
"input_tokens": (token_usage.get("input_tokens") or 0) if token_usage else 0,
"output_tokens": (token_usage.get("output_tokens") or 0) if token_usage else 0,
"cache_write_tokens": (token_usage.get("cache_write_tokens") or 0) if token_usage else 0,
"cache_read_tokens": (token_usage.get("cache_read_tokens") or 0) if token_usage else 0,
"write_call_count": write_call_count,
"fs_writes_detected": fs_writes_detected,
"git_writes_detected": git_writes_detected,
Expand Down
28 changes: 16 additions & 12 deletions src/autoskillit/hooks/token_summary_hook.py
Original file line number Diff line number Diff line change
Expand Up @@ -239,19 +239,23 @@ def _load_sessions(
_model = data.get("model_identifier", "")
if _model and not entry["model"]:
entry["model"] = _model
entry["input_tokens"] += data.get("input_tokens", 0)
entry["output_tokens"] += data.get("output_tokens", 0)
entry["cache_write_tokens"] += data.get(
"cache_write_tokens", data.get(_V1_CACHE_WRITE_KEY, 0)
entry["input_tokens"] += data.get("input_tokens") or 0
entry["output_tokens"] += data.get("output_tokens") or 0
_cw = data.get("cache_write_tokens")
entry["cache_write_tokens"] += (
_cw if _cw is not None else (data.get(_V1_CACHE_WRITE_KEY) or 0)
)
entry["cache_read_tokens"] += data.get(
"cache_read_tokens", data.get(_V1_CACHE_READ_KEY, 0)
_cr = data.get("cache_read_tokens")
entry["cache_read_tokens"] += (
_cr if _cr is not None else (data.get(_V1_CACHE_READ_KEY) or 0)
)
_raw_timing = data.get("timing_seconds")
entry["elapsed_seconds"] += float(_raw_timing) if _raw_timing is not None else 0.0
entry["invocation_count"] += 1
entry["loc_insertions"] = entry.get("loc_insertions", 0) + data.get("loc_insertions", 0)
entry["loc_deletions"] = entry.get("loc_deletions", 0) + data.get("loc_deletions", 0)
entry["loc_insertions"] = entry.get("loc_insertions", 0) + (
data.get("loc_insertions") or 0
)
entry["loc_deletions"] = entry.get("loc_deletions", 0) + (data.get("loc_deletions") or 0)
_raw_peak = data.get("peak_context", 0)
if isinstance(_raw_peak, int) and _raw_peak > entry["peak_context"]:
entry["peak_context"] = _raw_peak
Expand Down Expand Up @@ -389,10 +393,10 @@ def _format_model_table(aggregated: dict[str, dict[str, Any]]) -> str:
}
md = model_data[model]
md["_steps"].add(entry.get("step_name", ""))
md["input_tokens"] += entry.get("input_tokens", 0)
md["output_tokens"] += entry.get("output_tokens", 0)
md["cache_write_tokens"] += entry.get("cache_write_tokens", 0)
md["cache_read_tokens"] += entry.get("cache_read_tokens", 0)
md["input_tokens"] += entry.get("input_tokens") or 0
md["output_tokens"] += entry.get("output_tokens") or 0
md["cache_write_tokens"] += entry.get("cache_write_tokens") or 0
md["cache_read_tokens"] += entry.get("cache_read_tokens") or 0
md["elapsed_seconds"] += entry.get("elapsed_seconds", 0.0)

if not model_data:
Expand Down
26 changes: 14 additions & 12 deletions src/autoskillit/pipeline/tokens.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,10 +140,10 @@ def record(
_model = model or _primary_model(token_usage)
if _model and not e.model:
e.model = _model
e.input_tokens += token_usage.get("input_tokens", 0)
e.output_tokens += token_usage.get("output_tokens", 0)
e.cache_write_tokens += token_usage.get("cache_write_tokens", 0)
e.cache_read_tokens += token_usage.get("cache_read_tokens", 0)
e.input_tokens += token_usage.get("input_tokens") or 0
e.output_tokens += token_usage.get("output_tokens") or 0
e.cache_write_tokens += token_usage.get("cache_write_tokens") or 0
e.cache_read_tokens += token_usage.get("cache_read_tokens") or 0
e.invocation_count += 1
e.loc_insertions += loc_insertions
e.loc_deletions += loc_deletions
Expand Down Expand Up @@ -335,20 +335,22 @@ def load_from_log_dir(
_model = data.get("configured_model") or data.get("model_identifier", "")
if _model and not e.model:
e.model = _model
e.input_tokens += data.get("input_tokens", 0)
e.output_tokens += data.get("output_tokens", 0)
e.cache_write_tokens += data.get(
"cache_write_tokens", data.get("cache_creation_input_tokens", 0)
e.input_tokens += data.get("input_tokens") or 0
e.output_tokens += data.get("output_tokens") or 0
_cw = data.get("cache_write_tokens")
e.cache_write_tokens += (
_cw if _cw is not None else (data.get("cache_creation_input_tokens") or 0)
)
e.cache_read_tokens += data.get(
"cache_read_tokens", data.get("cache_read_input_tokens", 0)
_cr = data.get("cache_read_tokens")
e.cache_read_tokens += (
_cr if _cr is not None else (data.get("cache_read_input_tokens") or 0)
)
# timing_seconds is the on-disk key name written by session_log.py;
# elapsed_seconds is the in-memory field name on TokenEntry.
_raw_timing = data.get("timing_seconds")
e.elapsed_seconds += float(_raw_timing) if _raw_timing is not None else 0.0
e.loc_insertions += data.get("loc_insertions", 0)
e.loc_deletions += data.get("loc_deletions", 0)
e.loc_insertions += data.get("loc_insertions") or 0
e.loc_deletions += data.get("loc_deletions") or 0
_raw_peak = data.get("peak_context", 0)
if isinstance(_raw_peak, int) and _raw_peak > e.peak_context:
e.peak_context = _raw_peak
Expand Down
22 changes: 22 additions & 0 deletions tests/core/test_canonical_token_usage.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,28 @@ def test_merge_raises_on_mismatched_providers(self):
CanonicalTokenUsage.merge(a, b)


class TestToDictNoneCoercion:
def test_to_dict_coerces_none_to_zero_for_int_fields(self):
ctu = CanonicalTokenUsage(
input_tokens=100,
output_tokens=50,
cache_read_tokens=None,
cache_write_tokens=None,
provider="codex",
)
d = ctu.to_dict()
for key in ("input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens"):
assert isinstance(d[key], int), f"{key} should be int, got {type(d[key])}"

def test_codex_to_dict_coerces_none_cache_fields(self):
raw = {"input_tokens": 200, "output_tokens": 80, "cached_input_tokens": 50}
ctu = CanonicalTokenUsage.from_codex_dict(raw)
assert ctu.cache_write_tokens is None
d = ctu.to_dict()
assert d["cache_write_tokens"] == 0
assert d["cache_read_tokens"] == 50


class TestFrozen:
def test_cannot_mutate_fields(self):
usage = CanonicalTokenUsage(
Expand Down
2 changes: 1 addition & 1 deletion tests/execution/backends/test_codex_result_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -404,7 +404,7 @@ def test_cache_read_mapped_and_cache_write_none(self) -> None:
result = CodexResultParser().parse_stdout(ndjson)
canonical = result.raw["canonical_token_usage"]
assert canonical["cache_read_tokens"] == 30
assert canonical["cache_write_tokens"] is None
assert canonical["cache_write_tokens"] == 0

def test_no_turn_completed_yields_none_canonical(self) -> None:
"""Scenario 4: no turn.completed event — canonical_token_usage is None."""
Expand Down
20 changes: 20 additions & 0 deletions tests/execution/test_session_log_flush.py
Original file line number Diff line number Diff line change
Expand Up @@ -641,6 +641,26 @@ def test_token_usage_json_includes_peak_context_and_turn_count(tmp_path):
assert tu["turn_count"] == 14


def test_token_usage_json_coerces_none_cache_fields_to_zero(tmp_path):
"""Codex-originated None for cache_write_tokens must persist as 0 on disk."""
_flush(
tmp_path,
step_name="implement",
token_usage={
"input_tokens": 200,
"output_tokens": 80,
"cache_write_tokens": None,
"cache_read_tokens": 50,
},
timing_seconds=0.0,
proc_snapshots=None,
success=True,
)
tu = json.loads((tmp_path / "sessions" / "test-session-001" / "token_usage.json").read_text())
assert tu["cache_write_tokens"] == 0
assert tu["cache_read_tokens"] == 50


def test_step_timing_json_schema(tmp_path):
"""step_timing.json contains step_name and total_seconds."""
_flush(tmp_path, step_name="plan", timing_seconds=20.0, proc_snapshots=None, success=False)
Expand Down
40 changes: 40 additions & 0 deletions tests/infra/test_token_summary_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1206,3 +1206,43 @@ def test_reader_agreement_contract(
if scenario == "non-anthropic":
assert not disk_model.startswith("claude-")
assert disk_model not in ("sonnet", "opus", "haiku")


def test_load_sessions_handles_null_cache_write(tmp_path: Path) -> None:
"""_load_sessions must handle null cache_write_tokens in token_usage.json without crash."""
from autoskillit.hooks.token_summary_hook import _load_sessions

log_root = tmp_path / "logs"
log_root.mkdir()
pipeline_id = "test-null-cache"

_write_sessions(
log_root,
[
{
"dir_name": "s1",
"cwd": "/some/worktree",
"kitchen_id": pipeline_id,
"step_name": "plan",
},
],
)
sess_dir = log_root / "sessions" / "s1"
sess_dir.mkdir(parents=True, exist_ok=True)
(sess_dir / "token_usage.json").write_text(
json.dumps(
{
"session_label": "plan",
"input_tokens": 100,
"output_tokens": 50,
"cache_write_tokens": None,
"cache_read_tokens": None,
"timing_seconds": 0.0,
}
)
)

result = _load_sessions(log_root, pipeline_id)
assert "plan" in result
assert result["plan"]["cache_write_tokens"] == 0
assert result["plan"]["cache_read_tokens"] == 0
73 changes: 73 additions & 0 deletions tests/pipeline/test_tokens_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -438,6 +438,79 @@ def test_load_null_timing_seconds(self, tmp_path):
report = log.get_report()
assert report[0]["elapsed_seconds"] == 0.0

def test_load_null_cache_fields(self, tmp_path):
"""TokenLog must handle null cache fields without TypeError."""
_write_session(
tmp_path,
"s001",
{
"session_label": "implement",
"input_tokens": 100,
"output_tokens": 50,
"cache_write_tokens": None,
"cache_read_tokens": None,
"timing_seconds": 0.0,
"schema_version": 2,
},
)
log = DefaultTokenLog()
n = log.load_from_log_dir(tmp_path)
assert n == 1
report = log.get_report()
assert report[0]["cache_write_tokens"] == 0
assert report[0]["cache_read_tokens"] == 0

@pytest.mark.parametrize(
"null_field",
[
"input_tokens",
"output_tokens",
"cache_write_tokens",
"cache_read_tokens",
"loc_insertions",
"loc_deletions",
],
)
def test_load_null_int_field(self, tmp_path, null_field):
"""Each int field can be null on disk; load_from_log_dir must handle it."""
payload = {
"session_label": "implement",
"input_tokens": 100,
"output_tokens": 50,
"cache_write_tokens": 0,
"cache_read_tokens": 0,
"loc_insertions": 5,
"loc_deletions": 3,
"timing_seconds": 0.0,
"schema_version": 2,
}
payload[null_field] = None
_write_session(tmp_path, "s001", payload)
log = DefaultTokenLog()
n = log.load_from_log_dir(tmp_path)
assert n == 1

def test_load_preserves_zero_cache_write_with_v1_fallback(self, tmp_path):
"""cache_write_tokens: 0 must not consult v1 fallback key (cache_creation_input_tokens)."""
_write_session(
tmp_path,
"s001",
{
"session_label": "implement",
"input_tokens": 100,
"output_tokens": 50,
"cache_write_tokens": 0,
"cache_creation_input_tokens": 500,
"cache_read_tokens": 0,
"timing_seconds": 0.0,
"schema_version": 2,
},
)
log = DefaultTokenLog()
log.load_from_log_dir(tmp_path)
report = log.get_report()
assert report[0]["cache_write_tokens"] == 0


class TestLoadFromLogDirSchemaVersionCompat:
"""Backward compatibility: load_from_log_dir reads both v1 and v2 token_usage.json."""
Expand Down
Loading