Persist response token usage in rollout history (#41912)

## Why

Resumed threads need to continue per-turn and per-thread token totals without scanning arbitrarily far beyond the latest compaction checkpoint.

## What changed

- Add durable `TokenUsageRecord` rollout items with response, turn, thread, session, and root-turn attribution.
- Restore accumulated usage on resume and snapshot the latest record plus the compaction response ID in compaction checkpoints.
- Preserve root-turn lineage in persisted turn context, while ensuring forked child threads start with their own usage totals.

## Testing

- Cover usage accumulation across multiple responses and resumed turns.
- Cover local and remote compaction checkpoints, invalid remote compaction output, rollout reconstruction, and fork isolation.

GitOrigin-RevId: ef9e0c4a9102a08a2c382be4cdac68c84353c90a
This commit is contained in:
Owen Lin authored and copyberry committed 2026-08-31 21:15:38 +00:00
1 parent a7913390f7
commit 5f79a92e39
67 files changed
+932 -162

No files matched your search

+46
View File
@@ -1,11 +1,57 @@
use super::*;
use crate::session::tests::make_session_configuration_for_tests;
use crate::state::AutoCompactWindowSnapshot;
use codex_protocol::SessionId;
use codex_protocol::ThreadId;
use codex_protocol::protocol::CreditsSnapshot;
use codex_protocol::protocol::RateLimitWindow;
use codex_protocol::protocol::SpendControlLimitSnapshot;
use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TokenUsageRecord;
use pretty_assertions::assert_eq;
#[tokio::test]
async fn record_token_usage_continues_restored_totals() {
let thread_id = ThreadId::new();
let session_id = SessionId::from(ThreadId::new());
let usage = |total_tokens| TokenUsage {
total_tokens,
..TokenUsage::default()
};
let mut restored = SessionState::new(make_session_configuration_for_tests().await);
restored.latest_token_usage_record = Some(TokenUsageRecord {
thread_id,
turn_id: "turn-b".to_string(),
session_id,
root_turn_id: "root-turn".to_string(),
response_id: "response-c".to_string(),
usage: usage(30),
turn_token_usage: usage(30),
thread_token_usage: usage(230),
});
let after_resume = restored.record_token_usage(
thread_id,
"turn-b",
session_id,
"root-turn".to_string(),
"response-d".to_string(),
&usage(20),
);
assert_eq!(
after_resume,
TokenUsageRecord {
thread_id,
turn_id: "turn-b".to_string(),
session_id,
root_turn_id: "root-turn".to_string(),
response_id: "response-d".to_string(),
usage: usage(20),
turn_token_usage: usage(50),
thread_token_usage: usage(250),
}
);
}
#[tokio::test]
// Verifies connector merging deduplicates repeated IDs.
async fn merge_connector_selection_deduplicates_entries() {