Codex Handbook
tools/src/response_history.rs 150 lines
use codex_protocol::models::ContentItem;use codex_protocol::models::ResponseItem;use codex_utils_output_truncation::TruncationPolicy;use codex_utils_output_truncation::approx_token_count;use codex_utils_output_truncation::truncate_text;/// Retains items from the earliest of the last `user_message_count` user/// messages through the latest user message.pub fn retain_tail_from_last_n_user_messages(    items: &mut Vec<ResponseItem>,    user_message_count: usize,) {    if user_message_count == 0 {        items.clear();        return;    }    let Some(latest_user_idx) = items.iter().rposition(ResponseItem::is_user_message) else {        items.clear();        return;    };    items.truncate(latest_user_idx + 1);    let earliest_retained_user_idx = items        .iter()        .enumerate()        .rev()        .filter(|(_, item)| item.is_user_message())        .take(user_message_count)        .last()        .map(|(idx, _)| idx)        .unwrap_or(latest_user_idx);    items.drain(..earliest_retained_user_idx);}/// Truncates assistant output text to a shared token budget across items.pub fn truncate_assistant_output_text_to_token_budget(    items: &mut Vec<ResponseItem>,    max_tokens: usize,) {    let mut remaining_budget = max_tokens;    items.retain_mut(|item| {        let ResponseItem::Message { role, content, .. } = item else {            return true;        };        if role != "assistant" {            return true;        }        content.retain_mut(|content_item| {            let ContentItem::OutputText { text } = content_item else {                return true;            };            if remaining_budget == 0 {                return false;            }            let token_count = approx_token_count(text);            if token_count <= remaining_budget {                remaining_budget = remaining_budget.saturating_sub(token_count);                return true;            }            *text = truncate_text(text, TruncationPolicy::Tokens(remaining_budget));            remaining_budget = 0;            true        });        !content.is_empty()    });}#[cfg(test)]mod tests {    use codex_protocol::models::ContentItem;    use codex_protocol::models::ResponseItem;    use codex_utils_output_truncation::TruncationPolicy;    use codex_utils_output_truncation::truncate_text;    use pretty_assertions::assert_eq;    use super::retain_tail_from_last_n_user_messages;    use super::truncate_assistant_output_text_to_token_budget;    fn message(role: &str, text: &str) -> ResponseItem {        ResponseItem::Message {            id: None,            role: role.to_string(),            content: vec![if role == "assistant" {                ContentItem::OutputText {                    text: text.to_string(),                }            } else {                ContentItem::InputText {                    text: text.to_string(),                }            }],            phase: None,            metadata: None,        }    }    #[test]    fn retains_tail_through_latest_user_message() {        let mut items = vec![            message("system", "system"),            message("user", "old user"),            message("assistant", "old assistant"),            message("user", "previous user"),            message("assistant", "previous assistant"),            message("user", "current user"),            message("assistant", "later assistant"),        ];        retain_tail_from_last_n_user_messages(&mut items, /*user_message_count*/ 2);        assert_eq!(            items,            vec![                message("user", "previous user"),                message("assistant", "previous assistant"),                message("user", "current user"),            ]        );    }    #[test]    fn truncates_assistant_output_text_across_items() {        let long_assistant = "a".repeat(16);        let mut items = vec![            message("user", "previous user"),            message("assistant", &long_assistant),            message("assistant", "after budget"),            message("user", "current user"),        ];        truncate_assistant_output_text_to_token_budget(&mut items, /*max_tokens*/ 2);        assert_eq!(            items,            vec![                message("user", "previous user"),                message(                    "assistant",                    &truncate_text(&long_assistant, TruncationPolicy::Tokens(2)),                ),                message("user", "current user"),            ]        );    }}