Skip to content

Commit e183f1d

Browse files
committed
fix(codex): handle spaced session types
Address the remaining CodeRabbit and cubic review threads: tolerate whitespace around Codex session type fields before fast-path classification, and stop discarding table stdout write errors. The whitespace fallback is skipped for compact non-session type rows so the common Codex hot path stays fast while still accepting valid spaced JSON session rows. Validation: cargo test --manifest-path rust/Cargo.toml --workspace codex --no-fail-fast; cargo build --manifest-path rust/Cargo.toml --release --bin ccusage; pnpm run format; pnpm typecheck; pnpm run test. Local real CODEX_HOME table benchmark is 1.02x faster than codexusage, and bunx JSON/table parity diffs are clean.
1 parent e5f30f9 commit e183f1d

2 files changed

Lines changed: 81 additions & 3 deletions

File tree

‎rust/crates/ccusage-terminal/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ impl SimpleTable {
8282
let stdout = io::stdout();
8383
let mut stdout = stdout.lock();
8484
for line in self.render_lines() {
85-
let _ = writeln!(stdout, "{line}");
85+
writeln!(stdout, "{line}").expect("failed to write table line");
8686
}
8787
}
8888

‎rust/crates/ccusage/src/codex_loader.rs‎

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ static TURN_CONTEXT_TYPE_FINDER: LazyLock<Finder<'static>> =
2323
LazyLock::new(|| Finder::new(br#""type":"turn_context""#));
2424
static TOKEN_COUNT_TYPE_FINDER: LazyLock<Finder<'static>> =
2525
LazyLock::new(|| Finder::new(br#""type":"token_count""#));
26+
static COMPACT_TYPE_FIELD_FINDER: LazyLock<Finder<'static>> =
27+
LazyLock::new(|| Finder::new(br#""type":"#));
28+
static TYPE_KEY_FINDER: LazyLock<Finder<'static>> = LazyLock::new(|| Finder::new(br#""type""#));
2629
static USAGE_FIELD_FINDER: LazyLock<Finder<'static>> =
2730
LazyLock::new(|| Finder::new(br#""usage":"#));
2831
static INPUT_TOKENS_FIELD_FINDER: LazyLock<Finder<'static>> =
@@ -336,12 +339,18 @@ fn add_codex_exec_event(
336339
}
337340

338341
fn codex_line_usage_kind(line: &[u8]) -> Option<CodexLineKind> {
342+
let has_event_msg = EVENT_MSG_TYPE_FINDER.find(line).is_some();
339343
if TURN_CONTEXT_TYPE_FINDER.find(line).is_some()
340-
|| (EVENT_MSG_TYPE_FINDER.find(line).is_some()
341-
&& TOKEN_COUNT_TYPE_FINDER.find(line).is_some())
344+
|| (has_event_msg && TOKEN_COUNT_TYPE_FINDER.find(line).is_some())
342345
{
343346
return Some(CodexLineKind::Session);
344347
}
348+
if has_event_msg || COMPACT_TYPE_FIELD_FINDER.find(line).is_none() {
349+
let (has_turn_context, has_event_msg, has_token_count) = codex_line_type_flags(line);
350+
if has_turn_context || (has_event_msg && has_token_count) {
351+
return Some(CodexLineKind::Session);
352+
}
353+
}
345354
if USAGE_FIELD_FINDER.find(line).is_some()
346355
|| INPUT_TOKENS_FIELD_FINDER.find(line).is_some()
347356
|| PROMPT_TOKENS_FIELD_FINDER.find(line).is_some()
@@ -351,6 +360,48 @@ fn codex_line_usage_kind(line: &[u8]) -> Option<CodexLineKind> {
351360
None
352361
}
353362

363+
fn codex_line_type_flags(line: &[u8]) -> (bool, bool, bool) {
364+
let mut start = 0;
365+
let mut has_turn_context = false;
366+
let mut has_event_msg = false;
367+
let mut has_token_count = false;
368+
while let Some(index) = TYPE_KEY_FINDER.find(&line[start..]) {
369+
let key_start = start + index;
370+
let mut cursor = skip_json_whitespace(line, key_start + br#""type""#.len());
371+
if line.get(cursor) != Some(&b':') {
372+
start = key_start + br#""type""#.len();
373+
continue;
374+
}
375+
cursor = skip_json_whitespace(line, cursor + 1);
376+
if line.get(cursor) != Some(&b'"') {
377+
start = cursor.saturating_add(1);
378+
continue;
379+
}
380+
cursor += 1;
381+
has_turn_context |= json_string_value_matches(line, cursor, b"turn_context");
382+
has_event_msg |= json_string_value_matches(line, cursor, b"event_msg");
383+
has_token_count |= json_string_value_matches(line, cursor, b"token_count");
384+
if has_turn_context || (has_event_msg && has_token_count) {
385+
return (has_turn_context, has_event_msg, has_token_count);
386+
}
387+
start = cursor.saturating_add(1);
388+
}
389+
(has_turn_context, has_event_msg, has_token_count)
390+
}
391+
392+
fn json_string_value_matches(line: &[u8], start: usize, value: &[u8]) -> bool {
393+
line.get(start..start + value.len())
394+
.is_some_and(|candidate| candidate == value)
395+
&& line.get(start + value.len()) == Some(&b'"')
396+
}
397+
398+
fn skip_json_whitespace(line: &[u8], mut index: usize) -> usize {
399+
while matches!(line.get(index), Some(b' ' | b'\n' | b'\r' | b'\t')) {
400+
index += 1;
401+
}
402+
index
403+
}
404+
354405
fn codex_session_timestamp(value: Option<&CodexTimestamp<'_>>) -> Option<String> {
355406
match value? {
356407
CodexTimestamp::String(text) => {
@@ -1138,6 +1189,33 @@ mod tests {
11381189
fs::remove_dir_all(dir).unwrap();
11391190
}
11401191

1192+
#[test]
1193+
fn loads_session_usage_with_spaced_type_fields() {
1194+
let dir = temp_dir("spaced-session-type");
1195+
let file = dir.join("session.jsonl");
1196+
fs::write(
1197+
&file,
1198+
[
1199+
r#"{ "timestamp": "2026-01-02T00:00:00.000Z", "type" : "turn_context", "payload": { "model": "gpt-5" } }"#,
1200+
r#"{ "timestamp": "2026-01-02T00:00:01.000Z", "type" : "event_msg", "payload": { "type" : "token_count", "info": { "total_token_usage": { "input_tokens": 100, "cached_input_tokens": 10, "output_tokens": 50, "total_tokens": 150 }, "model": "gpt-5" } } }"#,
1201+
]
1202+
.join("\n"),
1203+
)
1204+
.unwrap();
1205+
1206+
let events = load_codex_events_from_directory(&dir, true).unwrap();
1207+
1208+
assert_eq!(events.len(), 1);
1209+
assert_eq!(events[0].timestamp, "2026-01-02T00:00:01.000Z");
1210+
assert_eq!(events[0].model.as_deref(), Some("gpt-5"));
1211+
assert_eq!(events[0].input_tokens, 100);
1212+
assert_eq!(events[0].cached_input_tokens, 10);
1213+
assert_eq!(events[0].output_tokens, 50);
1214+
assert_eq!(events[0].total_tokens, 150);
1215+
1216+
fs::remove_dir_all(dir).unwrap();
1217+
}
1218+
11411219
#[test]
11421220
fn loads_headless_usage_with_token_count_text_content() {
11431221
let dir = temp_dir("headless-token-count-content");

0 commit comments

Comments
 (0)