Codex Handbook
core/src/state/auto_compact_window.rs 162 lines
use codex_protocol::protocol::TokenUsage;#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub(crate) struct AutoCompactWindowSnapshot {    pub(crate) prefill_input_tokens: Option<i64>,}#[derive(Debug, Clone, Copy, PartialEq, Eq)]enum AutoCompactWindowPrefill {    ServerObserved(i64),    Estimated(i64),}#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub(super) struct AutoCompactWindow {    window_id: u64,    new_context_window_requested: bool,    /// Absolute input-token baseline for the current compaction window.    ///    /// `body_after_prefix` subtracts this from later active-context usage. It is    /// not the growth itself; server-observed usage replaces estimated    /// resume/recompute baselines when available.    prefill_input_tokens: Option<AutoCompactWindowPrefill>,}impl AutoCompactWindow {    pub(super) fn new() -> Self {        Self {            window_id: 0,            new_context_window_requested: false,            prefill_input_tokens: None,        }    }    pub(super) fn clear_prefill(&mut self) {        self.prefill_input_tokens = None;    }    pub(super) fn window_id(&self) -> u64 {        self.window_id    }    pub(super) fn set_window_id(&mut self, window_id: u64) {        self.window_id = window_id;    }    pub(super) fn advance_window_id(&mut self) -> u64 {        self.window_id = self.window_id.saturating_add(1);        self.new_context_window_requested = false;        self.window_id    }    pub(super) fn request_new_context_window(&mut self) {        self.new_context_window_requested = true;    }    pub(super) fn take_new_context_window_request(&mut self) -> bool {        let requested = self.new_context_window_requested;        self.new_context_window_requested = false;        requested    }    /// Records the request-input side of the first server usage sample. The    /// sampled output from that response is body growth and should remain    /// counted against the scoped auto-compact budget.    pub(super) fn ensure_server_observed_prefill_from_usage(&mut self, usage: &TokenUsage) {        if matches!(            self.prefill_input_tokens,            Some(AutoCompactWindowPrefill::ServerObserved(_))        ) {            return;        }        self.prefill_input_tokens = Some(AutoCompactWindowPrefill::ServerObserved(            usage.input_tokens.max(0),        ));    }    pub(super) fn set_estimated_prefill(&mut self, tokens: i64) {        if matches!(            self.prefill_input_tokens,            Some(AutoCompactWindowPrefill::ServerObserved(_))        ) {            return;        }        self.prefill_input_tokens = Some(AutoCompactWindowPrefill::Estimated(tokens.max(0)));    }    pub(super) fn snapshot(&self) -> AutoCompactWindowSnapshot {        let prefill_input_tokens = match self.prefill_input_tokens {            Some(AutoCompactWindowPrefill::ServerObserved(tokens))            | Some(AutoCompactWindowPrefill::Estimated(tokens)) => Some(tokens),            None => None,        };        AutoCompactWindowSnapshot {            prefill_input_tokens,        }    }}#[cfg(test)]mod tests {    use super::*;    use pretty_assertions::assert_eq;    #[test]    fn tracks_prefill_and_window_boundaries() {        let mut window = AutoCompactWindow::new();        assert_eq!(window.window_id(), 0);        window.set_window_id(/*window_id*/ 3);        assert_eq!(window.window_id(), 3);        window.request_new_context_window();        assert!(window.take_new_context_window_request());        assert!(!window.take_new_context_window_request());        window.request_new_context_window();        assert_eq!(window.advance_window_id(), 4);        assert_eq!(window.window_id(), 4);        assert!(!window.take_new_context_window_request());        assert_eq!(            window.snapshot(),            AutoCompactWindowSnapshot {                prefill_input_tokens: None,            }        );        window.set_estimated_prefill(/*tokens*/ 150);        assert_eq!(            window.snapshot(),            AutoCompactWindowSnapshot {                prefill_input_tokens: Some(150),            }        );        window.ensure_server_observed_prefill_from_usage(&TokenUsage {            input_tokens: 120,            total_tokens: 170,            ..Default::default()        });        assert_eq!(            window.snapshot(),            AutoCompactWindowSnapshot {                prefill_input_tokens: Some(120),            }        );        window.ensure_server_observed_prefill_from_usage(&TokenUsage {            input_tokens: 130,            total_tokens: 180,            ..Default::default()        });        window.set_estimated_prefill(/*tokens*/ 90);        assert_eq!(            window.snapshot(),            AutoCompactWindowSnapshot {                prefill_input_tokens: Some(120),            }        );    }}