Skip to content

Commit 0269096

Browse files
pakrym-oaipakrym-oai
authored
Move token usage/context information to session level (#3221)
Move context information into the main loop so it can be used to interrupt the loop or start auto-compaction.
1 parent 70a6d4b commit 0269096

File tree

9 files changed

+151
-117
lines changed

9 files changed

+151
-117
lines changed

‎codex-rs/core/src/client.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -398,9 +398,15 @@ impl From<ResponseCompletedUsage> for TokenUsage {
398398
fn from(val: ResponseCompletedUsage) -> Self {
399399
TokenUsage {
400400
input_tokens: val.input_tokens,
401-
cached_input_tokens: val.input_tokens_details.map(|d| d.cached_tokens),
401+
cached_input_tokens: val
402+
.input_tokens_details
403+
.map(|d| d.cached_tokens)
404+
.unwrap_or(0),
402405
output_tokens: val.output_tokens,
403-
reasoning_output_tokens: val.output_tokens_details.map(|d| d.reasoning_tokens),
406+
reasoning_output_tokens: val
407+
.output_tokens_details
408+
.map(|d| d.reasoning_tokens)
409+
.unwrap_or(0),
404410
total_tokens: val.total_tokens,
405411
}
406412
}

‎codex-rs/core/src/codex.rs‎

Lines changed: 31 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ use crate::protocol::SessionConfiguredEvent;
9999
use crate::protocol::StreamErrorEvent;
100100
use crate::protocol::Submission;
101101
use crate::protocol::TaskCompleteEvent;
102+
use crate::protocol::TokenUsageInfo;
102103
use crate::protocol::TurnDiffEvent;
103104
use crate::protocol::WebSearchBeginEvent;
104105
use crate::rollout::RolloutRecorder;
@@ -261,6 +262,7 @@ struct State {
261262
pending_approvals: HashMap<String, oneshot::Sender<ReviewDecision>>,
262263
pending_input: Vec<ResponseInputItem>,
263264
history: ConversationHistory,
265+
token_info: Option<TokenUsageInfo>,
264266
}
265267

266268
/// Context for an initialized model agent
@@ -1767,15 +1769,23 @@ async fn try_run_turn(
17671769
response_id: _,
17681770
token_usage,
17691771
} => {
1770-
if let Some(token_usage) = token_usage {
1771-
sess.tx_event
1772-
.send(Event {
1773-
id: sub_id.to_string(),
1774-
msg: EventMsg::TokenCount(token_usage),
1775-
})
1776-
.await
1777-
.ok();
1778-
}
1772+
let info = {
1773+
let mut st = sess.state.lock_unchecked();
1774+
let info = TokenUsageInfo::new_or_append(
1775+
&st.token_info,
1776+
&token_usage,
1777+
turn_context.client.get_model_context_window(),
1778+
);
1779+
st.token_info = info.clone();
1780+
info
1781+
};
1782+
sess.tx_event
1783+
.send(Event {
1784+
id: sub_id.to_string(),
1785+
msg: EventMsg::TokenCount(crate::protocol::TokenCountEvent { info }),
1786+
})
1787+
.await
1788+
.ok();
17791789

17801790
let unified_diff = turn_diff_tracker.get_unified_diff();
17811791
if let Ok(Some(unified_diff)) = unified_diff {
@@ -2841,13 +2851,21 @@ async fn drain_to_completed(
28412851
response_id: _,
28422852
token_usage,
28432853
}) => {
2844-
// some providers don't return token usage, so we default
2845-
// TODO: consider approximate token usage
2846-
let token_usage = token_usage.unwrap_or_default();
2854+
let info = {
2855+
let mut st = sess.state.lock_unchecked();
2856+
let info = TokenUsageInfo::new_or_append(
2857+
&st.token_info,
2858+
&token_usage,
2859+
turn_context.client.get_model_context_window(),
2860+
);
2861+
st.token_info = info.clone();
2862+
info
2863+
};
2864+
28472865
sess.tx_event
28482866
.send(Event {
28492867
id: sub_id.to_string(),
2850-
msg: EventMsg::TokenCount(token_usage),
2868+
msg: EventMsg::TokenCount(crate::protocol::TokenCountEvent { info }),
28512869
})
28522870
.await
28532871
.ok();

‎codex-rs/exec/src/event_processor_with_human_output.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -189,8 +189,14 @@ impl EventProcessor for EventProcessorWithHumanOutput {
189189
}
190190
return CodexStatus::InitiateShutdown;
191191
}
192-
EventMsg::TokenCount(token_usage) => {
193-
ts_println!(self, "tokens used: {}", token_usage.blended_total());
192+
EventMsg::TokenCount(ev) => {
193+
if let Some(usage_info) = ev.info {
194+
ts_println!(
195+
self,
196+
"tokens used: {}",
197+
usage_info.total_token_usage.blended_total()
198+
);
199+
}
194200
}
195201
EventMsg::AgentMessageDelta(AgentMessageDeltaEvent { delta }) => {
196202
if !self.answer_started {

‎codex-rs/protocol/src/protocol.rs‎

Lines changed: 63 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -417,9 +417,9 @@ pub enum EventMsg {
417417
/// Agent has completed all actions
418418
TaskComplete(TaskCompleteEvent),
419419

420-
/// Token count event, sent periodically to report the number of tokens
421-
/// used in the current session.
422-
TokenCount(TokenUsage),
420+
/// Usage update for the current session, including totals and last turn.
421+
/// Optional means unknown — UIs should not display when `None`.
422+
TokenCount(TokenCountEvent),
423423

424424
/// Agent text output message
425425
AgentMessage(AgentMessageEvent),
@@ -521,12 +521,54 @@ pub struct TaskStartedEvent {
521521
#[derive(Debug, Clone, Deserialize, Serialize, Default)]
522522
pub struct TokenUsage {
523523
pub input_tokens: u64,
524-
pub cached_input_tokens: Option<u64>,
524+
pub cached_input_tokens: u64,
525525
pub output_tokens: u64,
526-
pub reasoning_output_tokens: Option<u64>,
526+
pub reasoning_output_tokens: u64,
527527
pub total_tokens: u64,
528528
}
529529

530+
#[derive(Debug, Clone, Deserialize, Serialize)]
531+
pub struct TokenUsageInfo {
532+
pub total_token_usage: TokenUsage,
533+
pub last_token_usage: TokenUsage,
534+
pub model_context_window: Option<u64>,
535+
}
536+
537+
impl TokenUsageInfo {
538+
pub fn new_or_append(
539+
info: &Option<TokenUsageInfo>,
540+
last: &Option<TokenUsage>,
541+
model_context_window: Option<u64>,
542+
) -> Option<Self> {
543+
if info.is_none() && last.is_none() {
544+
return None;
545+
}
546+
547+
let mut info = match info {
548+
Some(info) => info.clone(),
549+
None => Self {
550+
total_token_usage: TokenUsage::default(),
551+
last_token_usage: TokenUsage::default(),
552+
model_context_window,
553+
},
554+
};
555+
if let Some(last) = last {
556+
info.append_last_usage(last);
557+
}
558+
Some(info)
559+
}
560+
561+
pub fn append_last_usage(&mut self, last: &TokenUsage) {
562+
self.total_token_usage.add_assign(last);
563+
self.last_token_usage = last.clone();
564+
}
565+
}
566+
567+
#[derive(Debug, Clone, Deserialize, Serialize)]
568+
pub struct TokenCountEvent {
569+
pub info: Option<TokenUsageInfo>,
570+
}
571+
530572
// Includes prompts, tools and space to call compact.
531573
const BASELINE_TOKENS: u64 = 12000;
532574

@@ -536,7 +578,7 @@ impl TokenUsage {
536578
}
537579

538580
pub fn cached_input(&self) -> u64 {
539-
self.cached_input_tokens.unwrap_or(0)
581+
self.cached_input_tokens
540582
}
541583

542584
pub fn non_cached_input(&self) -> u64 {
@@ -554,7 +596,7 @@ impl TokenUsage {
554596
/// This will be off for the current turn and pending function calls.
555597
pub fn tokens_in_context_window(&self) -> u64 {
556598
self.total_tokens
557-
.saturating_sub(self.reasoning_output_tokens.unwrap_or(0))
599+
.saturating_sub(self.reasoning_output_tokens)
558600
}
559601

560602
/// Estimate the remaining user-controllable percentage of the model's context window.
@@ -579,6 +621,15 @@ impl TokenUsage {
579621
let remaining = effective_window.saturating_sub(used);
580622
((remaining as f32 / effective_window as f32) * 100.0).clamp(0.0, 100.0) as u8
581623
}
624+
625+
/// In-place element-wise sum of token counts.
626+
pub fn add_assign(&mut self, other: &TokenUsage) {
627+
self.input_tokens += other.input_tokens;
628+
self.cached_input_tokens += other.cached_input_tokens;
629+
self.output_tokens += other.output_tokens;
630+
self.reasoning_output_tokens += other.reasoning_output_tokens;
631+
self.total_tokens += other.total_tokens;
632+
}
582633
}
583634

584635
#[derive(Debug, Clone, Deserialize, Serialize)]
@@ -606,10 +657,11 @@ impl fmt::Display for FinalOutput {
606657
String::new()
607658
},
608659
token_usage.output_tokens,
609-
token_usage
610-
.reasoning_output_tokens
611-
.map(|r| format!(" (reasoning {r})"))
612-
.unwrap_or_default()
660+
if token_usage.reasoning_output_tokens > 0 {
661+
format!(" (reasoning {})", token_usage.reasoning_output_tokens)
662+
} else {
663+
String::new()
664+
}
613665
)
614666
}
615667
}

‎codex-rs/tui/src/bottom_pane/chat_composer.rs‎

Lines changed: 12 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use codex_core::protocol::TokenUsage;
1+
use codex_core::protocol::TokenUsageInfo;
22
use crossterm::event::KeyCode;
33
use crossterm::event::KeyEvent;
44
use crossterm::event::KeyEventKind;
@@ -63,12 +63,6 @@ struct AttachedImage {
6363
path: PathBuf,
6464
}
6565

66-
struct TokenUsageInfo {
67-
total_token_usage: TokenUsage,
68-
last_token_usage: TokenUsage,
69-
model_context_window: Option<u64>,
70-
}
71-
7266
pub(crate) struct ChatComposer {
7367
textarea: TextArea,
7468
textarea_state: RefCell<TextAreaState>,
@@ -166,17 +160,8 @@ impl ChatComposer {
166160
/// Update the cached *context-left* percentage and refresh the placeholder
167161
/// text. The UI relies on the placeholder to convey the remaining
168162
/// context when the composer is empty.
169-
pub(crate) fn set_token_usage(
170-
&mut self,
171-
total_token_usage: TokenUsage,
172-
last_token_usage: TokenUsage,
173-
model_context_window: Option<u64>,
174-
) {
175-
self.token_usage_info = Some(TokenUsageInfo {
176-
total_token_usage,
177-
last_token_usage,
178-
model_context_window,
179-
});
163+
pub(crate) fn set_token_usage(&mut self, token_info: Option<TokenUsageInfo>) {
164+
self.token_usage_info = token_info;
180165
}
181166

182167
/// Record the history metadata advertised by `SessionConfiguredEvent` so
@@ -1290,11 +1275,16 @@ impl WidgetRef for ChatComposer {
12901275
} else {
12911276
100
12921277
};
1278+
let context_style = if percent_remaining < 20 {
1279+
Style::default().fg(Color::Yellow)
1280+
} else {
1281+
Style::default().add_modifier(Modifier::DIM)
1282+
};
12931283
hint.push(" ".into());
1294-
hint.push(
1295-
Span::from(format!("{percent_remaining}% context left"))
1296-
.style(Style::default().add_modifier(Modifier::DIM)),
1297-
);
1284+
hint.push(Span::styled(
1285+
format!("{percent_remaining}% context left"),
1286+
context_style,
1287+
));
12981288
}
12991289
}
13001290

‎codex-rs/tui/src/bottom_pane/mod.rs‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use crate::app_event_sender::AppEventSender;
55
use crate::tui::FrameRequester;
66
use crate::user_approval_widget::ApprovalRequest;
77
use bottom_pane_view::BottomPaneView;
8-
use codex_core::protocol::TokenUsage;
8+
use codex_core::protocol::TokenUsageInfo;
99
use codex_file_search::FileMatch;
1010
use crossterm::event::KeyEvent;
1111
use ratatui::buffer::Buffer;
@@ -358,14 +358,8 @@ impl BottomPane {
358358

359359
/// Update the *context-window remaining* indicator in the composer. This
360360
/// is forwarded directly to the underlying `ChatComposer`.
361-
pub(crate) fn set_token_usage(
362-
&mut self,
363-
total_token_usage: TokenUsage,
364-
last_token_usage: TokenUsage,
365-
model_context_window: Option<u64>,
366-
) {
367-
self.composer
368-
.set_token_usage(total_token_usage, last_token_usage, model_context_window);
361+
pub(crate) fn set_token_usage(&mut self, token_info: Option<TokenUsageInfo>) {
362+
self.composer.set_token_usage(token_info);
369363
self.request_redraw();
370364
}
371365

0 commit comments

Comments
 (0)