Skip to content

Account for last token count on resume #8677

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aibrahim-oai merged 6 commits into from
Jan 2, 2026
Merged

Account for last token count on resume #8677

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
be68e48
account for last token count on resume
aibrahim-oai Jan 2, 2026
b73bfb7
account for last token count on resume
aibrahim-oai Jan 2, 2026
16aebc4
integration
aibrahim-oai Jan 2, 2026
064b423
tests
aibrahim-oai Jan 2, 2026
f1e423e
event
aibrahim-oai Jan 2, 2026
2fd86d9
Merge branch 'main' into resume-token-count
aibrahim-oai Jan 2, 2026
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
94 changes: 94 additions & 0 deletions codex-rs/core/src/codex.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -810,6 +810,13 @@ impl Session {
.await;
}

// Seed usage info from the recorded rollout so UIs can show token counts
// immediately on resume/fork.
if let Some(info) = Self::last_token_info_from_rollout(&rollout_items) {
let mut state = self.state.lock().await;
state.set_token_info(Some(info));
}

// If persisting, persist all rollout items as-is (recorder filters)
if persist && !rollout_items.is_empty() {
self.persist_rollout_items(&rollout_items).await;
Expand All @@ -820,6 +827,13 @@ impl Session {
}
}

fn last_token_info_from_rollout(rollout_items: &[RolloutItem]) -> Option<TokenUsageInfo> {
rollout_items.iter().rev().find_map(|item| match item {
RolloutItem::EventMsg(EventMsg::TokenCount(ev)) => ev.info.clone(),
_ => None,
Comment on lines +830 to +833
Copy link
Contributor

@chatgpt-codex-connector chatgpt-codex-connector bot Jan 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Filter sub-agent TokenCount events when seeding resume usage

This helper blindly picks the last EventMsg::TokenCount in the rollout, but review sub-agent events are forwarded into the parent session (see process_review_events in core/src/tasks/review.rs, which passes through all “other” events) and are persisted to the rollout. If the last TokenCount came from a /review run, this will seed token_info with the reviewer model’s usage instead of the main conversation’s context, so the resumed session’s auto-compaction decisions and UI token indicators can be wrong. Consider scoping to the main session’s turn context or excluding sub-agent sources when choosing the last token info.

Useful? React with 👍 / 👎.

})
}

pub(crate) async fn update_settings(
&self,
updates: SessionSettingsUpdate,
Expand Down Expand Up @@ -2759,6 +2773,9 @@ mod tests {
use crate::protocol::RateLimitSnapshot;
use crate::protocol::RateLimitWindow;
use crate::protocol::ResumedHistory;
use crate::protocol::TokenCountEvent;
use crate::protocol::TokenUsage;
use crate::protocol::TokenUsageInfo;
use crate::state::TaskKind;
use crate::tasks::SessionTask;
use crate::tasks::SessionTaskContext;
Expand Down Expand Up @@ -2813,6 +2830,83 @@ mod tests {
assert_eq!(expected, actual);
}

#[tokio::test]
async fn record_initial_history_seeds_token_info_from_rollout() {
let (session, turn_context) = make_session_and_context().await;
let (mut rollout_items, _expected) = sample_rollout(&session, &turn_context);

let info1 = TokenUsageInfo {
total_token_usage: TokenUsage {
input_tokens: 10,
cached_input_tokens: 0,
output_tokens: 20,
reasoning_output_tokens: 0,
total_tokens: 30,
},
last_token_usage: TokenUsage {
input_tokens: 3,
cached_input_tokens: 0,
output_tokens: 4,
reasoning_output_tokens: 0,
total_tokens: 7,
},
model_context_window: Some(1_000),
};
let info2 = TokenUsageInfo {
total_token_usage: TokenUsage {
input_tokens: 100,
cached_input_tokens: 50,
output_tokens: 200,
reasoning_output_tokens: 25,
total_tokens: 375,
},
last_token_usage: TokenUsage {
input_tokens: 10,
cached_input_tokens: 0,
output_tokens: 20,
reasoning_output_tokens: 5,
total_tokens: 35,
},
model_context_window: Some(2_000),
};

rollout_items.push(RolloutItem::EventMsg(EventMsg::TokenCount(
TokenCountEvent {
info: Some(info1),
rate_limits: None,
},
)));
rollout_items.push(RolloutItem::EventMsg(EventMsg::TokenCount(
TokenCountEvent {
info: None,
rate_limits: None,
},
)));
rollout_items.push(RolloutItem::EventMsg(EventMsg::TokenCount(
TokenCountEvent {
info: Some(info2.clone()),
rate_limits: None,
},
)));
rollout_items.push(RolloutItem::EventMsg(EventMsg::TokenCount(
TokenCountEvent {
info: None,
rate_limits: None,
},
)));

session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ConversationId::default(),
history: rollout_items,
rollout_path: PathBuf::from("/tmp/resume.jsonl"),
}))
.await;

let actual = session.state.lock().await.token_info();
assert_eq!(actual, Some(info2));
}

#[tokio::test]
async fn record_initial_history_reconstructs_forked_transcript() {
let (session, turn_context) = make_session_and_context().await;
Expand Down
4 changes: 2 additions & 2 deletions codex-rs/protocol/src/protocol.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -844,7 +844,7 @@ pub struct TaskStartedEvent {
pub model_context_window: Option<i64>,
}

#[derive(Debug, Clone, Deserialize, Serialize, Default, JsonSchema, TS)]
#[derive(Debug, Clone, Deserialize, Serialize, Default, PartialEq, Eq, JsonSchema, TS)]
pub struct TokenUsage {
#[ts(type = "number")]
pub input_tokens: i64,
Expand All @@ -858,7 +858,7 @@ pub struct TokenUsage {
pub total_tokens: i64,
}

#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, TS)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)]
pub struct TokenUsageInfo {
pub total_token_usage: TokenUsage,
pub last_token_usage: TokenUsage,
Expand Down
Loading