codex/codex-rs/core/src/state/session.rs

//! Session-wide mutable state.

use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;

use crate::conversation_history::ConversationHistory;
use crate::protocol::RateLimitSnapshot;
use crate::protocol::TokenUsage;
use crate::protocol::TokenUsageInfo;
use crate::state::TaskKind;

/// Persistent, session-scoped state previously stored directly on `Session`.
#[derive(Default)]
pub(crate) struct SessionState {
    pub(crate) history: ConversationHistory,
    pub(crate) token_info: Option<TokenUsageInfo>,
    pub(crate) latest_rate_limits: Option<RateLimitSnapshot>,
}

impl SessionState {
    /// Create a new session state mirroring previous `State::default()` semantics.
    pub(crate) fn new() -> Self {
        Self {
            history: ConversationHistory::new(),
            ..Default::default()
        }
    }

    // History helpers
    pub(crate) fn record_items<I>(&mut self, items: I, task_kind: TaskKind)
    where
        I: IntoIterator<Item = ResponseItem>,
    {
        self.history.record_items(items, task_kind)
    }

    pub(crate) fn history_snapshot(&self) -> Vec<ResponseItem> {
        self.history.contents()
    }

    pub(crate) fn replace_history(&mut self, items: Vec<ResponseItem>) {
        self.history.replace(items);
    }

    pub(crate) fn clear_review_thread(&mut self) {
        self.history.clear_review_thread();
    }

    pub(crate) fn initialize_review_history(
        &mut self,
        response_input: &ResponseInputItem,
        initial_context: Vec<ResponseItem>,
    ) {
        self.history
            .initialize_review_history(response_input, initial_context);
    }

    pub(crate) fn prepare_prompt_input(
        &mut self,
        task_kind: TaskKind,
        pending_input: Vec<ResponseItem>,
    ) -> Vec<ResponseItem> {
        if !pending_input.is_empty() {
            self.history.add_pending_input(pending_input, task_kind);
        }
        self.history.handle_missing_tool_call_output(task_kind);
        self.history.prompt(task_kind)
    }

    // Token/rate limit helpers
    pub(crate) fn update_token_info_from_usage(
        &mut self,
        usage: &TokenUsage,
        model_context_window: Option<u64>,
    ) {
        self.token_info = TokenUsageInfo::new_or_append(
            &self.token_info,
            &Some(usage.clone()),
            model_context_window,
        );
    }

    pub(crate) fn set_rate_limits(&mut self, snapshot: RateLimitSnapshot) {
        self.latest_rate_limits = Some(snapshot);
    }

    pub(crate) fn token_info_and_rate_limits(
        &self,
    ) -> (Option<TokenUsageInfo>, Option<RateLimitSnapshot>) {
        (self.token_info.clone(), self.latest_rate_limits.clone())
    }

    pub(crate) fn set_token_usage_full(&mut self, context_window: u64) {
        match &mut self.token_info {
            Some(info) => info.fill_to_context_window(context_window),
            None => {
                self.token_info = Some(TokenUsageInfo::full_context_window(context_window));
            }
        }
    }

    // Pending input/approval moved to TurnState.
}