github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 00:06:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid cloning thread history for token usage replay (#34361)

Browse Source 
## What changed

- Resolve the turn ID for restored token usage before sending resume and fork responses, avoiding a clone of the full thread history.
- Add `ThreadHistoryBuilder::active_turn_id` so attribution can inspect the active turn without materializing its items.
- Stop history reconstruction at the latest persisted `TokenCount` and verify that a later tail turn does not receive the replayed usage.

GitOrigin-RevId: 69c4a9fe04770d1a938470afc933431e61670094
This commit is contained in:
Charlie Marsh
2026-07-20 14:52:42 +00:00
committed by copyberry
parent ef6b597f41
commit 1e20272fa5
5 changed files with 53 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/src/protocol/thread_history.rs
+8
View File
@@ -274,6 +274,14 @@ impl ThreadHistoryBuilder {
.or_else(|| self.turns.last().cloned())
}
/// Returns the id of the active turn without materializing its items.
pub fn active_turn_id(&self) -> Option<&str> {
self.current_turn
.as_ref()
.map(|turn| turn.id.as_str())
.or_else(|| self.turns.last().map(|turn| turn.id.as_str()))
}
pub fn turn_snapshot(&self, turn_id: &str) -> Option<Turn> {
self.current_turn
.as_ref()
codex-rs/app-server/src/request_processors.rs
+1 -1
View File
@@ -570,7 +570,7 @@ use crate::thread_state::ConnectionCapabilities;
use crate::thread_state::ThreadListenerCommand;
use crate::thread_state::ThreadState;
use crate::thread_state::ThreadStateManager;
use token_usage_replay::latest_token_usage_turn_id_from_rollout_items;
use token_usage_replay::restored_token_usage_turn_id;
use token_usage_replay::send_thread_token_usage_update_to_connection;
fn resolve_request_cwd(cwd: Option<PathBuf>) -> Result<Option<AbsolutePathBuf>, JSONRPCErrorError> {
codex-rs/app-server/src/request_processors/thread_lifecycle.rs
+4 -7
View File
@@ -577,7 +577,9 @@ pub(super) async fn handle_pending_thread_resume_request(
thread_status.clone(),
has_live_in_progress_turn,
);
let token_usage_thread = pending.include_turns.then(|| thread.clone());
let token_usage_turn_id = pending
.include_turns
.then(|| restored_token_usage_turn_id(&pending.history_items, &thread));
let mut initial_turns_page = if let Some(mut page) = pending.paginated_initial_turns_page.take()
{
if let (Some(active_turn), Some(params)) =
@@ -714,18 +716,13 @@ pub(super) async fn handle_pending_thread_resume_request(
.await;
// Match cold resume: metadata-only resume should attach the listener without
// paying the cost of turn reconstruction for historical usage replay.
if let Some(token_usage_thread) = token_usage_thread {
let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items(
&pending.history_items,
token_usage_thread.turns.as_slice(),
);
if let Some(token_usage_turn_id) = token_usage_turn_id {
// Rejoining a loaded thread has the same UI contract as a cold resume, but
// uses the live conversation state instead of reconstructing a new session.
send_thread_token_usage_update_to_connection(
outgoing,
connection_id,
conversation_id,
&token_usage_thread,
conversation.as_ref(),
token_usage_turn_id,
)
codex-rs/app-server/src/request_processors/thread_processor.rs
+7 -14
View File
@@ -3316,7 +3316,9 @@ impl ThreadRequestProcessor {
let active_permission_profile = thread_response_active_permission_profile(
config_snapshot.active_permission_profile,
);
let token_usage_thread = include_turns.then(|| thread.clone());
let token_usage_turn_id = include_turns.then(|| {
restored_token_usage_turn_id(response_history.get_rollout_items(), &thread)
});
let mut initial_turns_page = if let Some(params) = initial_turns_page.as_ref() {
let initial_turns_page_result = if paginated_resume {
self.paginated_resume_initial_turns_page(thread_id, params)
@@ -3373,11 +3375,7 @@ impl ThreadRequestProcessor {
.await;
// `excludeTurns` is explicitly the cheap resume path, so avoid
// rebuilding history only to attribute a replayed usage update.
if let Some(token_usage_thread) = token_usage_thread {
let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items(
response_history.get_rollout_items(),
token_usage_thread.turns.as_slice(),
);
if let Some(token_usage_turn_id) = token_usage_turn_id {
// The client needs restored usage before it starts another turn.
// Sending after the response preserves JSON-RPC request ordering while
// still filling the status line before the next turn lifecycle begins.
@@ -3385,7 +3383,6 @@ impl ThreadRequestProcessor {
&self.outgoing,
connection_id,
thread_id,
&token_usage_thread,
codex_thread.as_ref(),
token_usage_turn_id,
)
@@ -4263,24 +4260,20 @@ impl ThreadRequestProcessor {
let notif = thread_started_notification(thread);
let connection_id = request_id.connection_id;
let token_usage_thread = include_turns.then(|| response.thread.clone());
let token_usage_turn_id =
include_turns.then(|| restored_token_usage_turn_id(&history_items, &response.thread));
self.outgoing
.send_response_with_thread_originator(request_id, response, thread_originator)
.await;
// `excludeTurns` is the cheap fork path, so skip restored usage replay
// instead of rebuilding history only to attribute a historical update.
if let Some(token_usage_thread) = token_usage_thread {
let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items(
&history_items,
token_usage_thread.turns.as_slice(),
);
if let Some(token_usage_turn_id) = token_usage_turn_id {
// Mirror the resume contract for forks: the new thread is usable as soon
// as the response arrives, so restored usage must follow immediately.
send_thread_token_usage_update_to_connection(
&self.outgoing,
connection_id,
thread_id,
&token_usage_thread,
forked_thread.as_ref(),
token_usage_turn_id,
)
codex-rs/app-server/src/request_processors/token_usage_replay.rs
+33 -27
View File
@@ -37,16 +37,15 @@ pub(super) async fn send_thread_token_usage_update_to_connection(
outgoing: &Arc<OutgoingMessageSender>,
connection_id: ConnectionId,
thread_id: ThreadId,
thread: &Thread,
conversation: &CodexThread,
token_usage_turn_id: Option<String>,
token_usage_turn_id: String,
) {
let Some(info) = conversation.token_usage_info().await else {
return;
};
let notification = ThreadTokenUsageUpdatedNotification {
thread_id: thread_id.to_string(),
turn_id: token_usage_turn_id.unwrap_or_else(|| latest_token_usage_turn_id(thread)),
turn_id: token_usage_turn_id,
token_usage: ThreadTokenUsage::from(info),
};
outgoing
@@ -57,41 +56,36 @@ pub(super) async fn send_thread_token_usage_update_to_connection(
.await;
}
/// Identifies the turn that was active when a `TokenCount` record appeared.
pub(super) fn restored_token_usage_turn_id(
rollout_items: &[RolloutItem],
thread: &Thread,
) -> String {
latest_token_usage_turn_id_from_rollout_items(rollout_items, thread.turns.as_slice())
.unwrap_or_else(|| latest_token_usage_turn_id(thread))
}
/// Identifies the turn that was active when the latest `TokenCount` record appeared.
///
/// The id is preferred when it still appears in the rebuilt thread. The position is a
/// fallback for histories whose implicit turn ids are regenerated during reconstruction.
struct TokenUsageTurnOwner {
id: String,
position: Option<usize>,
}
pub(super) fn latest_token_usage_turn_id_from_rollout_items(
fn latest_token_usage_turn_id_from_rollout_items(
rollout_items: &[RolloutItem],
turns: &[Turn],
) -> Option<String> {
let token_count_index = rollout_items
.iter()
.rposition(|item| matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_))))?;
let mut builder = ThreadHistoryBuilder::new();
let mut token_usage_turn_owner = None;
for item in rollout_items {
if matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_))) {
token_usage_turn_owner =
builder
.active_turn_snapshot()
.map(|turn| TokenUsageTurnOwner {
id: turn.id,
position: builder.active_turn_position(),
});
}
for item in &rollout_items[..token_count_index] {
builder.handle_rollout_item(item);
}
let owner = token_usage_turn_owner?;
if turns.iter().any(|turn| turn.id == owner.id) {
Some(owner.id)
let active_turn_id = builder.active_turn_id()?;
if turns.iter().any(|turn| turn.id == active_turn_id) {
Some(active_turn_id.to_string())
} else {
owner
.position
builder
.active_turn_position()
.and_then(|position| turns.get(position))
.map(|turn| turn.id.clone())
}
@@ -145,6 +139,18 @@ mod tests {
);
}
#[test]
fn replay_attribution_uses_latest_token_count_and_ignores_tail_turn() {
let mut rollout_items = token_usage_history();
rollout_items.extend(token_usage_history());
let turns = build_turns_from_rollout_items(&rollout_items);
assert_eq!(
latest_token_usage_turn_id_from_rollout_items(&rollout_items, turns.as_slice()),
Some(turns[2].id.clone())
);
}
fn token_usage_history() -> Vec<RolloutItem> {
vec![
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {