mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
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:
1 parent
a7913390f7
commit
5f79a92e39
67 files changed
+932
-162
No files matched your search
@@ -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() {
|
||||
|
||||
Reference in new issue
Block a user