diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index e7a533d76b..fab38a9a7e 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -5115,6 +5115,13 @@ }, "turnId": { "type": "string" + }, + "windowNumber": { + "default": 0, + "description": "Current context-window generation. This starts at zero and advances after compaction.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 53bd7cef4e..c0a4915e5c 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -19821,6 +19821,13 @@ }, "turnId": { "type": "string" + }, + "windowNumber": { + "default": 0, + "description": "Current context-window generation. This starts at zero and advances after compaction.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 11540d40c7..7e89da3ac2 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -17600,6 +17600,13 @@ }, "turnId": { "type": "string" + }, + "windowNumber": { + "default": 0, + "description": "Current context-window generation. This starts at zero and advances after compaction.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadTokenUsageUpdatedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadTokenUsageUpdatedNotification.json index 111de85c62..1c213f326b 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadTokenUsageUpdatedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadTokenUsageUpdatedNotification.json @@ -65,6 +65,13 @@ }, "turnId": { "type": "string" + }, + "windowNumber": { + "default": 0, + "description": "Current context-window generation. This starts at zero and advances after compaction.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadTokenUsageUpdatedNotification.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadTokenUsageUpdatedNotification.ts index 1be282500c..6e077945a2 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadTokenUsageUpdatedNotification.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadTokenUsageUpdatedNotification.ts @@ -3,4 +3,8 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { ThreadTokenUsage } from "./ThreadTokenUsage"; -export type ThreadTokenUsageUpdatedNotification = { threadId: string, turnId: string, tokenUsage: ThreadTokenUsage, }; +export type ThreadTokenUsageUpdatedNotification = { threadId: string, turnId: string, tokenUsage: ThreadTokenUsage, +/** + * Current context-window generation. This starts at zero and advances after compaction. + */ +windowNumber: number, }; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index 367a4a64c8..bea41256f5 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -1373,6 +1373,10 @@ pub struct ThreadTokenUsageUpdatedNotification { pub thread_id: String, pub turn_id: String, pub token_usage: ThreadTokenUsage, + /// Current context-window generation. This starts at zero and advances after compaction. + #[serde(default)] + #[ts(type = "number")] + pub window_number: u64, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index d52f1bd168..3283859050 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -313,7 +313,7 @@ Start a fresh thread when you need a new Codex conversation. Valid `personality` values are `"friendly"`, `"pragmatic"`, and `"none"`. When `"none"` is selected, the personality placeholder is replaced with an empty string. -To continue a stored session, call `thread/resume` with the `thread.id` you previously recorded. The response shape matches `thread/start`. When the stored session includes persisted token usage, the server emits `thread/tokenUsage/updated` immediately after the response so clients can render restored usage before the next turn starts. You can also pass the same configuration overrides supported by `thread/start`, including `approvalsReviewer`. +To continue a stored session, call `thread/resume` with the `thread.id` you previously recorded. The response shape matches `thread/start`. When the stored session includes persisted token usage, the server emits `thread/tokenUsage/updated` immediately after the response so clients can render restored usage before the next turn starts. The notification's `windowNumber` is the current zero-based context-window generation, so it also gives the number of compactions in the surviving thread history. You can also pass the same configuration overrides supported by `thread/start`, including `approvalsReviewer`. By default, `thread/resume` includes the reconstructed turn history in `thread.turns`. Experimental clients can pass `excludeTurns: true` to return only thread metadata and live resume state, then call `thread/turns/list` separately if they want to page the turn history over the network. In that mode the server also skips replaying restored `thread/tokenUsage/updated`, which avoids rebuilding turns just to attribute historical usage. diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 863aeac79b..aac1f6dc84 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -901,8 +901,15 @@ pub(crate) async fn apply_bespoke_event_handling( .await; } EventMsg::TokenCount(token_count_event) => { - handle_token_count_event(conversation_id, event_turn_id, token_count_event, &outgoing) - .await; + let window_number = conversation.context_window_number().await; + handle_token_count_event( + conversation_id, + event_turn_id, + token_count_event, + window_number, + &outgoing, + ) + .await; } EventMsg::Error(ev) => { thread_watch_manager @@ -1610,6 +1617,7 @@ async fn handle_token_count_event( conversation_id: ThreadId, turn_id: String, token_count_event: TokenCountEvent, + window_number: u64, outgoing: &ThreadScopedOutgoingMessageSender, ) { let TokenCountEvent { info, rate_limits } = token_count_event; @@ -1618,6 +1626,7 @@ async fn handle_token_count_event( thread_id: conversation_id.to_string(), turn_id, token_usage, + window_number, }; outgoing .send_server_notification(ServerNotification::ThreadTokenUsageUpdated(notification)) @@ -3829,6 +3838,7 @@ mod tests { info: Some(info), rate_limits: Some(rate_limits), }, + /*window_number*/ 7, &outgoing, ) .await; @@ -3840,6 +3850,7 @@ mod tests { ) => { assert_eq!(payload.thread_id, conversation_id.to_string()); assert_eq!(payload.turn_id, turn_id); + assert_eq!(payload.window_number, 7); let usage = payload.token_usage; assert_eq!(usage.total.total_tokens, 200); assert_eq!(usage.total.cached_input_tokens, 25); @@ -3886,6 +3897,7 @@ mod tests { info: None, rate_limits: None, }, + /*window_number*/ 0, &outgoing, ) .await; diff --git a/codex-rs/app-server/src/request_processors/token_usage_replay.rs b/codex-rs/app-server/src/request_processors/token_usage_replay.rs index d4d2228d4b..b12089a239 100644 --- a/codex-rs/app-server/src/request_processors/token_usage_replay.rs +++ b/codex-rs/app-server/src/request_processors/token_usage_replay.rs @@ -44,10 +44,12 @@ pub(super) async fn send_thread_token_usage_update_to_connection( let Some(info) = conversation.token_usage_info().await else { return; }; + let window_number = conversation.context_window_number().await; let notification = ThreadTokenUsageUpdatedNotification { thread_id: thread_id.to_string(), turn_id: token_usage_turn_id.unwrap_or_else(|| latest_token_usage_turn_id(thread)), token_usage: ThreadTokenUsage::from(info), + window_number, }; outgoing .send_server_notification_to_connections( diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index c80268e2f9..f04c52b3e9 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -450,6 +450,11 @@ impl CodexThread { self.codex.session.token_usage_info().await } + /// Returns the current context-window generation for this thread. + pub async fn context_window_number(&self) -> u64 { + self.codex.session.context_window_number().await + } + /// Records a user-role session-prefix message without creating a new user turn boundary. pub(crate) async fn inject_user_message_without_turn(&self, message: String) { let item = ResponseItem::Message { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index af2412e5c4..f203b53945 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1273,6 +1273,12 @@ impl Session { state.token_info() } + /// Returns the current context-window generation for this session. + pub(crate) async fn context_window_number(&self) -> u64 { + let state = self.state.lock().await; + state.auto_compact_window_number() + } + pub(crate) async fn get_estimated_token_count( &self, turn_context: &TurnContext, @@ -3487,9 +3493,8 @@ impl Session { } pub(crate) async fn current_window_id(&self) -> String { - let state = self.state.lock().await; let thread_id = self.thread_id; - let window_number = state.auto_compact_window_number(); + let window_number = self.context_window_number().await; format!("{thread_id}:{window_number}") } diff --git a/codex-rs/exec/tests/event_processor_with_json_output.rs b/codex-rs/exec/tests/event_processor_with_json_output.rs index 1f44ca31a8..3c26f2272c 100644 --- a/codex-rs/exec/tests/event_processor_with_json_output.rs +++ b/codex-rs/exec/tests/event_processor_with_json_output.rs @@ -1239,6 +1239,7 @@ fn token_usage_update_is_emitted_on_turn_completion() { }, model_context_window: Some(128_000), }, + window_number: 0, }, )); assert_eq!( diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index d8d63575bb..ce25a9fb5b 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -1135,6 +1135,7 @@ async fn token_usage_update_refreshes_status_line_with_runtime_context_window() ThreadId::new(), "turn-1", Some(950_000), + /*window_number*/ 0, ))); assert_eq!( @@ -4651,10 +4652,12 @@ fn token_usage_notification( thread_id: ThreadId, turn_id: &str, model_context_window: Option, + window_number: u64, ) -> ServerNotification { ServerNotification::ThreadTokenUsageUpdated(ThreadTokenUsageUpdatedNotification { thread_id: thread_id.to_string(), turn_id: turn_id.to_string(), + window_number, token_usage: ThreadTokenUsage { total: TokenUsageBreakdown { total_tokens: 10, diff --git a/codex-rs/tui/src/app/thread_events.rs b/codex-rs/tui/src/app/thread_events.rs index 8d41f1c052..4c090d7eba 100644 --- a/codex-rs/tui/src/app/thread_events.rs +++ b/codex-rs/tui/src/app/thread_events.rs @@ -45,6 +45,7 @@ pub(super) struct ThreadEventStore { pub(super) pending_interactive_replay: PendingInteractiveReplayState, pub(super) active_turn_id: Option, pub(super) input_state: Option, + latest_token_usage: Option, pub(super) capacity: usize, pub(super) active: bool, } @@ -69,6 +70,7 @@ impl ThreadEventStore { pending_interactive_replay: PendingInteractiveReplayState::default(), active_turn_id: None, input_state: None, + latest_token_usage: None, capacity, active: false, } @@ -107,6 +109,9 @@ impl ThreadEventStore { pub(super) fn push_notification(&mut self, notification: ServerNotification) { self.pending_interactive_replay .note_server_notification(¬ification); + if let ServerNotification::ThreadTokenUsageUpdated(notification) = ¬ification { + self.latest_token_usage = Some(notification.clone()); + } match ¬ification { ServerNotification::TurnStarted(turn) => { self.active_turn_id = Some(turn.turn.id.clone()); @@ -206,6 +211,7 @@ impl ThreadEventStore { } pub(super) fn snapshot(&self) -> ThreadEventSnapshot { + let latest_token_usage = self.latest_token_usage.as_ref(); ThreadEventSnapshot { session: self.session.clone(), turns: self.turns.clone(), @@ -218,11 +224,19 @@ impl ThreadEventStore { ThreadBufferedEvent::Request(request) => self .pending_interactive_replay .should_replay_snapshot_request(request), + ThreadBufferedEvent::Notification( + ServerNotification::ThreadTokenUsageUpdated(_), + ) => latest_token_usage.is_none(), ThreadBufferedEvent::Notification(_) | ThreadBufferedEvent::HistoryEntryResponse(_) | ThreadBufferedEvent::FeedbackSubmission(_) => true, }) .cloned() + .chain(latest_token_usage.cloned().map(|notification| { + ThreadBufferedEvent::Notification(ServerNotification::ThreadTokenUsageUpdated( + notification, + )) + })) .collect(), input_state: self.input_state.clone(), } @@ -349,6 +363,9 @@ mod tests { use codex_app_server_protocol::HookScope as AppServerHookScope; use codex_app_server_protocol::HookStartedNotification; use codex_app_server_protocol::RequestId as AppServerRequestId; + use codex_app_server_protocol::ThreadTokenUsage; + use codex_app_server_protocol::ThreadTokenUsageUpdatedNotification; + use codex_app_server_protocol::TokenUsageBreakdown; use codex_app_server_protocol::TurnCompletedNotification; use codex_app_server_protocol::TurnStartedNotification; use codex_config::types::ApprovalsReviewer; @@ -419,6 +436,31 @@ mod tests { }) } + fn token_usage_notification(thread_id: ThreadId, window_number: u64) -> ServerNotification { + ServerNotification::ThreadTokenUsageUpdated(ThreadTokenUsageUpdatedNotification { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + window_number, + token_usage: ThreadTokenUsage { + total: TokenUsageBreakdown { + total_tokens: 10, + input_tokens: 4, + cached_input_tokens: 1, + output_tokens: 5, + reasoning_output_tokens: 0, + }, + last: TokenUsageBreakdown { + total_tokens: 10, + input_tokens: 4, + cached_input_tokens: 1, + output_tokens: 5, + reasoning_output_tokens: 0, + }, + model_context_window: Some(100), + }, + }) + } + fn hook_started_notification(thread_id: ThreadId, turn_id: &str) -> ServerNotification { ServerNotification::HookStarted(HookStartedNotification { thread_id: thread_id.to_string(), @@ -555,6 +597,35 @@ mod tests { assert_eq!(store.active_turn_id(), None); } + #[test] + fn thread_event_store_snapshot_preserves_latest_token_usage() { + let thread_id = ThreadId::new(); + let mut store = ThreadEventStore::new(/*capacity*/ 1); + store.push_notification(token_usage_notification( + thread_id, /*window_number*/ 3, + )); + store.push_notification(ServerNotification::ConfigWarning( + codex_app_server_protocol::ConfigWarningNotification { + summary: "unrelated warning".to_string(), + details: None, + path: None, + range: None, + }, + )); + store.rebase_buffer_after_session_refresh(); + + let snapshot = store.snapshot(); + let [ + ThreadBufferedEvent::Notification(ServerNotification::ThreadTokenUsageUpdated( + notification, + )), + ] = snapshot.events.as_slice() + else { + panic!("expected the latest token usage notification"); + }; + assert_eq!(notification.window_number, 3); + } + #[test] fn thread_event_store_rebase_preserves_resolved_request_state() { let thread_id = ThreadId::new(); diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index fa745c652e..c1f758f33c 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -551,6 +551,8 @@ pub(crate) struct ChatWidget { runtime_model_provider_base_url: Option, pub(crate) remote_connection: Option, token_info: Option, + #[cfg_attr(not(test), allow(dead_code))] + window_number: u64, rate_limit_snapshots_by_limit_id: BTreeMap, refreshing_status_outputs: Vec<(u64, StatusHistoryHandle)>, next_status_refresh_request_id: u64, diff --git a/codex-rs/tui/src/chatwidget/constructor.rs b/codex-rs/tui/src/chatwidget/constructor.rs index 7fb6fe54e9..a34a118aba 100644 --- a/codex-rs/tui/src/chatwidget/constructor.rs +++ b/codex-rs/tui/src/chatwidget/constructor.rs @@ -123,6 +123,7 @@ impl ChatWidget { runtime_model_provider_base_url, remote_connection: None, token_info: None, + window_number: 0, rate_limit_snapshots_by_limit_id: BTreeMap::new(), refreshing_status_outputs: Vec::new(), next_status_refresh_request_id: 0, diff --git a/codex-rs/tui/src/chatwidget/protocol.rs b/codex-rs/tui/src/chatwidget/protocol.rs index 0c5b0ed3ee..47a71e6407 100644 --- a/codex-rs/tui/src/chatwidget/protocol.rs +++ b/codex-rs/tui/src/chatwidget/protocol.rs @@ -30,6 +30,7 @@ impl ChatWidget { } match notification { ServerNotification::ThreadTokenUsageUpdated(notification) => { + self.window_number = notification.window_number; self.set_token_info(Some(token_usage_info_from_app_server( notification.token_usage, ))); diff --git a/codex-rs/tui/src/chatwidget/session_flow.rs b/codex-rs/tui/src/chatwidget/session_flow.rs index 177700d4f4..99fdf79228 100644 --- a/codex-rs/tui/src/chatwidget/session_flow.rs +++ b/codex-rs/tui/src/chatwidget/session_flow.rs @@ -27,6 +27,7 @@ impl ChatWidget { } self.refresh_plan_mode_nudge(); self.turn_lifecycle.reset_thread(); + self.window_number = 0; self.clear_safety_buffering(); self.thread_name = session.thread_name.clone(); self.current_goal_status_indicator = None; diff --git a/codex-rs/tui/src/chatwidget/tests/helpers.rs b/codex-rs/tui/src/chatwidget/tests/helpers.rs index 5f890ab000..d9d64a003c 100644 --- a/codex-rs/tui/src/chatwidget/tests/helpers.rs +++ b/codex-rs/tui/src/chatwidget/tests/helpers.rs @@ -382,6 +382,14 @@ fn token_usage_breakdown(usage: TokenUsage) -> codex_app_server_protocol::TokenU } pub(super) fn handle_token_count(chat: &mut ChatWidget, info: Option) { + handle_token_count_with_window_number(chat, info, /*window_number*/ 0); +} + +pub(super) fn handle_token_count_with_window_number( + chat: &mut ChatWidget, + info: Option, + window_number: u64, +) { match info { Some(info) => { chat.handle_server_notification( @@ -393,6 +401,7 @@ pub(super) fn handle_token_count(chat: &mut ChatWidget, info: Option