From 715aac0c0022f2f1119d65ac311b0eb39cab6941 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Wed, 25 Jun 2025 16:33:04 -0700 Subject: [PATCH] feat: show number of tokens remaining in UI --- codex-rs/core/src/chat_completions.rs | 18 +++++++-- codex-rs/core/src/client.rs | 39 ++++++++++++++++--- codex-rs/core/src/client_common.rs | 5 ++- codex-rs/core/src/codex.rs | 16 +++++++- codex-rs/core/src/protocol.rs | 10 +++++ codex-rs/exec/src/event_processor.rs | 4 ++ codex-rs/mcp-server/src/codex_tool_runner.rs | 1 + codex-rs/tui/src/bottom_pane/chat_composer.rs | 31 +++++++++++++++ codex-rs/tui/src/bottom_pane/mod.rs | 7 ++++ codex-rs/tui/src/chatwidget.rs | 11 ++++++ 10 files changed, 132 insertions(+), 10 deletions(-) diff --git a/codex-rs/core/src/chat_completions.rs b/codex-rs/core/src/chat_completions.rs index f381c72e51..053129da0b 100644 --- a/codex-rs/core/src/chat_completions.rs +++ b/codex-rs/core/src/chat_completions.rs @@ -215,6 +215,7 @@ where let _ = tx_event .send(Ok(ResponseEvent::Completed { response_id: String::new(), + total_tokens: None, })) .await; return; @@ -232,6 +233,7 @@ where let _ = tx_event .send(Ok(ResponseEvent::Completed { response_id: String::new(), + total_tokens: None, })) .await; return; @@ -317,6 +319,7 @@ where let _ = tx_event .send(Ok(ResponseEvent::Completed { response_id: String::new(), + total_tokens: None, })) .await; @@ -394,7 +397,10 @@ where // Not an assistant message – forward immediately. return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(item)))); } - Poll::Ready(Some(Ok(ResponseEvent::Completed { response_id }))) => { + Poll::Ready(Some(Ok(ResponseEvent::Completed { + response_id, + total_tokens, + }))) => { if !this.cumulative.is_empty() { let aggregated_item = crate::models::ResponseItem::Message { role: "assistant".to_string(), @@ -404,7 +410,10 @@ where }; // Buffer Completed so it is returned *after* the aggregated message. - this.pending_completed = Some(ResponseEvent::Completed { response_id }); + this.pending_completed = Some(ResponseEvent::Completed { + response_id, + total_tokens, + }); return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone( aggregated_item, @@ -412,7 +421,10 @@ where } // Nothing aggregated – forward Completed directly. - return Poll::Ready(Some(Ok(ResponseEvent::Completed { response_id }))); + return Poll::Ready(Some(Ok(ResponseEvent::Completed { + response_id, + total_tokens, + }))); } // No other `Ok` variants exist at the moment, continue polling. } } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index aff838887a..3311e94185 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -210,6 +210,29 @@ struct SseEvent { #[derive(Debug, Deserialize)] struct ResponseCompleted { id: String, + usage: Option, +} + +#[derive(Debug, Deserialize)] +#[allow(dead_code)] // Fields should print in debug output. +struct ResponseCompletedUsage { + input_tokens: u64, + input_tokens_details: Option, + output_tokens: u64, + output_tokens_details: Option, + total_tokens: u64, +} + +#[derive(Debug, Deserialize)] +struct ResponseCompletedInputTokensDetails { + #[allow(dead_code)] // Fields should print in debug output. + cached_tokens: u64, +} + +#[derive(Debug, Deserialize)] +struct ResponseCompletedOutputTokensDetails { + #[allow(dead_code)] // Fields should print in debug output. + reasoning_tokens: u64, } async fn process_sse(stream: S, tx_event: mpsc::Sender>) @@ -221,7 +244,7 @@ where // If the stream stays completely silent for an extended period treat it as disconnected. let idle_timeout = *OPENAI_STREAM_IDLE_TIMEOUT_MS; // The response id returned from the "complete" message. - let mut response_id = None; + let mut response_completed: Option = None; loop { let sse = match timeout(idle_timeout, stream.next()).await { @@ -233,9 +256,15 @@ where return; } Ok(None) => { - match response_id { - Some(response_id) => { - let event = ResponseEvent::Completed { response_id }; + match response_completed { + Some(ResponseCompleted { + id: response_id, + usage, + }) => { + let event = ResponseEvent::Completed { + response_id, + total_tokens: usage.map(|u| u.total_tokens), + }; let _ = tx_event.send(Ok(event)).await; } None => { @@ -301,7 +330,7 @@ where if let Some(resp_val) = event.response { match serde_json::from_value::(resp_val) { Ok(r) => { - response_id = Some(r.id); + response_completed = Some(r); } Err(e) => { debug!("failed to parse ResponseCompleted: {e}"); diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index a2633475df..5aeee12f57 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -51,7 +51,10 @@ impl Prompt { #[derive(Debug)] pub enum ResponseEvent { OutputItemDone(ResponseItem), - Completed { response_id: String }, + Completed { + response_id: String, + total_tokens: Option, + }, } #[derive(Debug, Serialize)] diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index e12a3a600b..eb79a0fdc9 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -81,6 +81,7 @@ use crate::protocol::SandboxPolicy; use crate::protocol::SessionConfiguredEvent; use crate::protocol::Submission; use crate::protocol::TaskCompleteEvent; +use crate::protocol::TokenCountEvent; use crate::rollout::RolloutRecorder; use crate::safety::SafetyCheck; use crate::safety::assess_command_safety; @@ -1078,7 +1079,20 @@ async fn try_run_turn( let response = handle_response_item(sess, sub_id, item.clone()).await?; output.push(ProcessedResponseItem { item, response }); } - ResponseEvent::Completed { response_id } => { + ResponseEvent::Completed { + response_id, + total_tokens, + } => { + if let Some(total_tokens) = total_tokens { + sess.tx_event + .send(Event { + id: sub_id.to_string(), + msg: EventMsg::TokenCount(TokenCountEvent { total_tokens }), + }) + .await + .ok(); + } + let mut state = sess.state.lock().unwrap(); state.previous_response_id = Some(response_id); break; diff --git a/codex-rs/core/src/protocol.rs b/codex-rs/core/src/protocol.rs index d4aa769852..c161f57210 100644 --- a/codex-rs/core/src/protocol.rs +++ b/codex-rs/core/src/protocol.rs @@ -275,6 +275,10 @@ pub enum EventMsg { /// Agent has completed all actions TaskComplete(TaskCompleteEvent), + /// Token count event, sent periodically to report the number of tokens + /// used in the current session. + TokenCount(TokenCountEvent), + /// Agent text output message AgentMessage(AgentMessageEvent), @@ -322,6 +326,12 @@ pub struct TaskCompleteEvent { pub last_agent_message: Option, } +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct TokenCountEvent { + /// Total number of tokens used in the current session. + pub total_tokens: u64, +} + #[derive(Debug, Clone, Deserialize, Serialize)] pub struct AgentMessageEvent { pub message: String, diff --git a/codex-rs/exec/src/event_processor.rs b/codex-rs/exec/src/event_processor.rs index e2a8bbb20a..7b641cfeb4 100644 --- a/codex-rs/exec/src/event_processor.rs +++ b/codex-rs/exec/src/event_processor.rs @@ -16,6 +16,7 @@ use codex_core::protocol::McpToolCallEndEvent; use codex_core::protocol::PatchApplyBeginEvent; use codex_core::protocol::PatchApplyEndEvent; use codex_core::protocol::SessionConfiguredEvent; +use codex_core::protocol::TokenCountEvent; use owo_colors::OwoColorize; use owo_colors::Style; use shlex::try_join; @@ -180,6 +181,9 @@ impl EventProcessor { EventMsg::TaskStarted | EventMsg::TaskComplete(_) => { // Ignore. } + EventMsg::TokenCount(TokenCountEvent { total_tokens }) => { + ts_println!(self, "tokens used: {total_tokens}"); + } EventMsg::AgentMessage(AgentMessageEvent { message }) => { ts_println!( self, diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index 67c990b00c..796a119e5c 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -162,6 +162,7 @@ pub async fn run_codex_tool_session( } EventMsg::Error(_) | EventMsg::TaskStarted + | EventMsg::TokenCount(_) | EventMsg::AgentReasoning(_) | EventMsg::McpToolCallBegin(_) | EventMsg::McpToolCallEnd(_) diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index 1218f76ec7..8342922df5 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -35,6 +35,11 @@ pub(crate) struct ChatComposer<'a> { command_popup: Option, app_event_tx: AppEventSender, history: ChatComposerHistory, + + /// Percentage of context window remaining for the currently selected + /// model. Stored as an integer 0-100 so we can easily embed it in the + /// placeholder text without additional formatting each render. + context_left_percent: Option, } impl ChatComposer<'_> { @@ -48,11 +53,37 @@ impl ChatComposer<'_> { command_popup: None, app_event_tx, history: ChatComposerHistory::new(), + context_left_percent: None, }; this.update_border(has_input_focus); this } + /// Update the cached *context-left* percentage and refresh the placeholder + /// text. The UI relies on the placeholder to convey the remaining + /// context when the composer is empty (mirroring the behaviour of the + /// TypeScript CLI). + pub(crate) fn set_context_left_percent(&mut self, percent: u8) { + // Update only when the value actually changed to avoid unnecessary + // redraws. + if self.context_left_percent == Some(percent) { + return; + } + + self.context_left_percent = Some(percent); + + // Build placeholder string similar to the JS CLI. We include the + // context indicator only when there is *enough* space so the hint + // remains concise. + let placeholder = if percent > 25 { + format!("send a message — {percent}% context left") + } else { + format!("send a message — {percent}% context left (consider /compact)") + }; + + self.textarea.set_placeholder_text(placeholder); + } + /// Record the history metadata advertised by `SessionConfiguredEvent` so /// that the composer can navigate cross-session history. pub(crate) fn set_history_metadata(&mut self, log_id: u64, entry_count: usize) { diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index c654581ccd..d5bd64f026 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -129,6 +129,13 @@ impl BottomPane<'_> { } } + /// Update the *context-window remaining* indicator in the composer. This + /// is forwarded directly to the underlying `ChatComposer`. + pub(crate) fn set_context_left_percent(&mut self, percent: u8) { + self.composer.set_context_left_percent(percent); + self.request_redraw(); + } + /// Called when the agent requests user approval. pub fn push_approval_request(&mut self, request: ApprovalRequest) { let request = if let Some(view) = self.active_view.as_mut() { diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index bd5197c73b..32cd375354 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -18,6 +18,7 @@ use codex_core::protocol::McpToolCallEndEvent; use codex_core::protocol::Op; use codex_core::protocol::PatchApplyBeginEvent; use codex_core::protocol::TaskCompleteEvent; +use codex_core::protocol::TokenCountEvent; use crossterm::event::KeyEvent; use ratatui::buffer::Buffer; use ratatui::layout::Constraint; @@ -231,6 +232,7 @@ impl ChatWidget<'_> { EventMsg::AgentMessage(AgentMessageEvent { message }) => { self.conversation_history .add_agent_message(&self.config, message); + self.request_redraw(); } EventMsg::AgentReasoning(AgentReasoningEvent { text }) => { @@ -250,6 +252,15 @@ impl ChatWidget<'_> { self.bottom_pane.set_task_running(false); self.request_redraw(); } + EventMsg::TokenCount(TokenCountEvent { total_tokens }) => { + let max_tokens = 128_000; + let percent: u8 = if total_tokens > 0 { + ((1.0 - total_tokens as f32 / max_tokens as f32) * 100.0) as u8 + } else { + 100 + }; + self.bottom_pane.set_context_left_percent(percent); + } EventMsg::Error(ErrorEvent { message }) => { self.conversation_history.add_error(message); self.bottom_pane.set_task_running(false);