diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 642c2080c6..ada43aa6f9 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -830,7 +830,7 @@ impl Session { async fn get_total_token_usage(&self) -> i64 { let state = self.state.lock().await; - state.get_total_token_usage(state.server_reasoning_included()) + state.get_total_token_usage() } pub(crate) async fn get_base_instructions(&self) -> BaseInstructions { diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 0502d538a0..7225721c93 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -25,6 +25,7 @@ pub(crate) struct ContextManager { /// The oldest items are at the beginning of the vector. items: Vec, token_info: Option, + server_reasoning_included: bool, codex_home: PathBuf, } @@ -33,6 +34,7 @@ impl ContextManager { Self { items: Vec::new(), token_info: TokenUsageInfo::new_or_append(&None, &None, None), + server_reasoning_included: false, codex_home: codex_home.to_path_buf(), } } @@ -54,6 +56,14 @@ impl ContextManager { } } + pub(crate) fn set_server_reasoning_included(&mut self, included: bool) { + self.server_reasoning_included = included; + } + + pub(crate) fn total_token_usage(&self) -> i64 { + self.get_total_token_usage(self.server_reasoning_included) + } + /// `items` is ordered from oldest to newest. pub(crate) fn record_items(&mut self, items: I, policy: TruncationPolicy) where @@ -321,7 +331,8 @@ impl ContextManager { return item.clone(); }; - let mut remaining = context_window.saturating_sub(self.get_total_token_usage()); + let mut remaining = context_window + .saturating_sub(self.get_total_token_usage(self.server_reasoning_included)); let offloader = ContextOffloader::new(&self.codex_home); let mut new_content = Vec::with_capacity(content.len().saturating_add(1)); for item in content { diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index cc02ee365e..59d53f84c4 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -14,7 +14,6 @@ pub(crate) struct SessionState { pub(crate) session_configuration: SessionConfiguration, pub(crate) history: ContextManager, pub(crate) latest_rate_limits: Option, - pub(crate) server_reasoning_included: bool, } impl SessionState { @@ -25,7 +24,6 @@ impl SessionState { session_configuration, history, latest_rate_limits: None, - server_reasoning_included: false, } } @@ -80,17 +78,12 @@ impl SessionState { self.history.set_token_usage_full(context_window); } - pub(crate) fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { - self.history - .get_total_token_usage(server_reasoning_included) + pub(crate) fn get_total_token_usage(&self) -> i64 { + self.history.total_token_usage() } pub(crate) fn set_server_reasoning_included(&mut self, included: bool) { - self.server_reasoning_included = included; - } - - pub(crate) fn server_reasoning_included(&self) -> bool { - self.server_reasoning_included + self.history.set_server_reasoning_included(included); } }