Codex Handbook
core/src/compact_remote_v2.rs 843 lines
use std::sync::Arc;use crate::Prompt;use crate::ResponseStream;use crate::client::ModelClientSession;use crate::client_common::ResponseEvent;use crate::compact::CompactionAnalyticsAttempt;use crate::compact::CompactionAnalyticsDetails;use crate::compact::InitialContextInjection;use crate::compact::compaction_status_from_result;use crate::compact_remote::process_compacted_history;use crate::compact_remote::should_keep_compacted_history_item;use crate::compact_remote::trim_function_call_history_to_fit_context_window;use crate::hook_runtime::PostCompactHookOutcome;use crate::hook_runtime::PreCompactHookOutcome;use crate::hook_runtime::run_post_compact_hooks;use crate::hook_runtime::run_pre_compact_hooks;use crate::responses_metadata::CodexResponsesMetadata;use crate::responses_metadata::CodexResponsesRequestKind;use crate::responses_metadata::CompactionTurnMetadata;use crate::responses_retry::ResponsesStreamRequest;use crate::responses_retry::handle_retryable_response_stream_error;use crate::session::session::Session;use crate::session::turn::built_tools;use crate::session::turn_context::TurnContext;use codex_analytics::CompactionImplementation;use codex_analytics::CompactionPhase;use codex_analytics::CompactionReason;use codex_analytics::CompactionTrigger;use codex_protocol::error::CodexErr;use codex_protocol::error::Result as CodexResult;use codex_protocol::items::ContextCompactionItem;use codex_protocol::items::TurnItem;use codex_protocol::models::ContentItem;use codex_protocol::models::ResponseItem;use codex_protocol::protocol::CompactedItem;use codex_protocol::protocol::EventMsg;use codex_protocol::protocol::TokenUsage;use codex_protocol::protocol::TruncationPolicy;use codex_protocol::protocol::TurnStartedEvent;use codex_rollout_trace::CompactionCheckpointTracePayload;use codex_rollout_trace::InferenceTraceContext;use codex_utils_output_truncation::approx_token_count;use codex_utils_output_truncation::truncate_text;use futures::StreamExt;use tokio_util::sync::CancellationToken;use tracing::info;// Mirror the current /responses/compact retained-message default while the// server-side path remains the reference implementation.const RETAINED_MESSAGE_TOKEN_BUDGET: usize = 64_000;// Compact attempts can run much longer than normal turns, so keep the per-transport// retry budget smaller than the general Responses stream retry budget.const MAX_REMOTE_COMPACTION_V2_STREAM_RETRIES: u64 = 2;pub(crate) async fn run_inline_remote_auto_compact_task(    sess: Arc<Session>,    turn_context: Arc<TurnContext>,    client_session: &mut ModelClientSession,    initial_context_injection: InitialContextInjection,    reason: CompactionReason,    phase: CompactionPhase,) -> CodexResult<()> {    run_remote_compact_task_inner(        &sess,        &turn_context,        Some(client_session),        initial_context_injection,        CompactionTrigger::Auto,        reason,        phase,    )    .await}pub(crate) async fn run_remote_compact_task(    sess: Arc<Session>,    turn_context: Arc<TurnContext>,) -> CodexResult<()> {    let start_event = EventMsg::TurnStarted(TurnStartedEvent {        turn_id: turn_context.sub_id.clone(),        trace_id: turn_context.trace_id.clone(),        started_at: turn_context.turn_timing_state.started_at_unix_secs().await,        model_context_window: turn_context.model_context_window(),        collaboration_mode_kind: turn_context.collaboration_mode.mode,    });    sess.send_event(&turn_context, start_event).await;    run_remote_compact_task_inner(        &sess,        &turn_context,        /*client_session*/ None,        InitialContextInjection::DoNotInject,        CompactionTrigger::Manual,        CompactionReason::UserRequested,        CompactionPhase::StandaloneTurn,    )    .await}async fn run_remote_compact_task_inner(    sess: &Arc<Session>,    turn_context: &Arc<TurnContext>,    client_session: Option<&mut ModelClientSession>,    initial_context_injection: InitialContextInjection,    trigger: CompactionTrigger,    reason: CompactionReason,    phase: CompactionPhase,) -> CodexResult<()> {    let compaction_metadata = CompactionTurnMetadata::new(        trigger,        reason,        CompactionImplementation::ResponsesCompactionV2,        phase,    );    let mut analytics_details = CompactionAnalyticsDetails {        active_context_tokens_before: Some(sess.get_total_token_usage().await),        ..Default::default()    };    let attempt = CompactionAnalyticsAttempt::begin(        sess.as_ref(),        turn_context.as_ref(),        trigger,        reason,        CompactionImplementation::ResponsesCompactionV2,        phase,    )    .await;    let pre_compact_outcome = run_pre_compact_hooks(sess, turn_context, trigger).await;    match pre_compact_outcome {        PreCompactHookOutcome::Continue => {}        PreCompactHookOutcome::Stopped => {            let error = CodexErr::TurnAborted;            attempt                .track(                    sess.as_ref(),                    codex_analytics::CompactionStatus::Interrupted,                    Some(&error),                    analytics_details,                )                .await;            return Err(error);        }    }    let result = run_remote_compact_task_inner_impl(        sess,        turn_context,        client_session,        initial_context_injection,        compaction_metadata,        &mut analytics_details,    )    .await;    let status = compaction_status_from_result(&result);    let codex_error = result.as_ref().err();    if result.is_ok() {        let post_compact_outcome = run_post_compact_hooks(sess, turn_context, trigger).await;        if let PostCompactHookOutcome::Stopped = post_compact_outcome {            attempt                .track(sess.as_ref(), status, codex_error, analytics_details)                .await;            return Err(CodexErr::TurnAborted);        }    }    attempt        .track(sess.as_ref(), status, codex_error, analytics_details)        .await;    if let Err(err) = result {        sess.track_turn_codex_error(turn_context, &err);        let event = EventMsg::Error(            err.to_error_event(Some("Error running remote compact task".to_string())),        );        sess.send_event(turn_context, event).await;        return Err(err);    }    Ok(())}async fn run_remote_compact_task_inner_impl(    sess: &Arc<Session>,    turn_context: &Arc<TurnContext>,    client_session: Option<&mut ModelClientSession>,    initial_context_injection: InitialContextInjection,    compaction_metadata: CompactionTurnMetadata,    analytics_details: &mut CompactionAnalyticsDetails,) -> CodexResult<()> {    let context_compaction_item = ContextCompactionItem::new();    let compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context(        turn_context.sub_id.as_str(),        context_compaction_item.id.as_str(),        turn_context.model_info.slug.as_str(),        turn_context.provider.info().name.as_str(),    );    let compaction_item = TurnItem::ContextCompaction(context_compaction_item);    sess.emit_turn_item_started(turn_context, &compaction_item)        .await;    let mut history = sess.clone_history().await;    let base_instructions = sess.get_base_instructions().await;    let (rewritten_outputs, estimated_deleted_tokens) =        trim_function_call_history_to_fit_context_window(            &mut history,            turn_context.as_ref(),            &base_instructions,        );    if rewritten_outputs > 0 {        info!(            turn_id = %turn_context.sub_id,            rewritten_outputs,            "rewrote history outputs before remote compaction v2"        );    }    if estimated_deleted_tokens > 0 {        let max_local_deleted_tokens = sess            .estimated_tokens_after_last_model_generated_item()            .await;        analytics_details.active_context_tokens_before = analytics_details            .active_context_tokens_before            .map(|active_context_tokens_before| {                active_context_tokens_before                    .saturating_sub(estimated_deleted_tokens.min(max_local_deleted_tokens))            });    }    let trace_input_history = history.raw_items().to_vec();    let prompt_input = history.for_prompt(&turn_context.model_info.input_modalities);    let tool_router = built_tools(        sess.as_ref(),        turn_context.as_ref(),        &CancellationToken::new(),    )    .await?;    let mut input = prompt_input.clone();    input.push(ResponseItem::CompactionTrigger { metadata: None });    let prompt = Prompt {        input,        tools: tool_router.model_visible_specs(),        parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls,        base_instructions,        personality: turn_context.personality,        output_schema: None,        output_schema_strict: true,    };    let window_id = sess.current_window_id().await;    let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata(        sess.installation_id.clone(),        window_id,        CodexResponsesRequestKind::Compaction(compaction_metadata),    );    let trace_attempt = compaction_trace.start_attempt(&serde_json::json!({        "model": turn_context.model_info.slug.as_str(),        "instructions": prompt.base_instructions.text.as_str(),        "input": &prompt.input,        "parallel_tool_calls": prompt.parallel_tool_calls,    }));    let mut owned_client_session;    let client_session = match client_session {        Some(client_session) => client_session,        None => {            owned_client_session = sess.services.model_client.new_session();            &mut owned_client_session        }    };    let compaction_output_result = run_remote_compaction_request_v2(        sess,        turn_context,        client_session,        &prompt,        &responses_metadata,    )    .await;    trace_attempt.record_result(        compaction_output_result            .as_ref()            .map(|output| std::slice::from_ref(&output.compaction_output)),    );    let RemoteCompactionV2Output {        compaction_output,        token_usage,    } = compaction_output_result?;    if let Some(token_usage) = token_usage {        analytics_details.active_context_tokens_before = Some(token_usage.input_tokens);        analytics_details.compaction_summary_tokens = Some(token_usage.output_tokens);        analytics_details.cached_input_tokens = Some(token_usage.cached_input_tokens);    }    let (compacted_history, retained_images) =        build_v2_compacted_history(&prompt_input, compaction_output);    analytics_details.retained_image_count = Some(retained_images);    let new_window_id = sess.advance_auto_compact_window_id().await;    let new_history = process_compacted_history(        sess.as_ref(),        turn_context.as_ref(),        compacted_history,        initial_context_injection,    )    .await;    let reference_context_item = match initial_context_injection {        InitialContextInjection::DoNotInject => None,        InitialContextInjection::BeforeLastUserMessage => Some(turn_context.to_turn_context_item()),    };    let compacted_item = CompactedItem {        message: String::new(),        replacement_history: Some(new_history.clone()),        window_id: Some(new_window_id),    };    compaction_trace.record_installed(&CompactionCheckpointTracePayload {        input_history: &trace_input_history,        replacement_history: &new_history,    });    sess.replace_compacted_history(new_history, reference_context_item, compacted_item)        .await;    sess.recompute_token_usage(turn_context).await;    sess.emit_turn_item_completed(turn_context, compaction_item)        .await;    Ok(())}struct RemoteCompactionV2Output {    compaction_output: ResponseItem,    token_usage: Option<TokenUsage>,}async fn run_remote_compaction_request_v2(    sess: &Session,    turn_context: &TurnContext,    client_session: &mut ModelClientSession,    prompt: &Prompt,    responses_metadata: &CodexResponsesMetadata,) -> CodexResult<RemoteCompactionV2Output> {    let max_retries = turn_context        .provider        .info()        .stream_max_retries()        .min(MAX_REMOTE_COMPACTION_V2_STREAM_RETRIES);    let mut retries = 0;    loop {        let result = match client_session            .stream(                prompt,                &turn_context.model_info,                &turn_context.session_telemetry,                turn_context.reasoning_effort.clone(),                turn_context.reasoning_summary,                turn_context.config.service_tier.clone(),                responses_metadata,                &InferenceTraceContext::disabled(),            )            .await        {            Ok(stream) => collect_compaction_output(stream).await,            Err(err) => Err(err),        };        match result {            Ok(compaction_output) => return Ok(compaction_output),            Err(err) if !err.is_retryable() => return Err(err),            Err(err) => {                handle_retryable_response_stream_error(                    &mut retries,                    max_retries,                    err,                    client_session,                    sess,                    turn_context,                    ResponsesStreamRequest::RemoteCompactionV2,                )                .await?;            }        }    }}async fn collect_compaction_output(    mut stream: ResponseStream,) -> CodexResult<RemoteCompactionV2Output> {    let mut output_item_count = 0usize;    let mut compaction_count = 0usize;    let mut compaction_output = None;    let mut saw_completed = false;    let mut completed_token_usage = None;    while let Some(event) = stream.next().await {        match event? {            ResponseEvent::OutputItemDone(item) => {                output_item_count += 1;                if let ResponseItem::Compaction { .. } = item {                    compaction_count += 1;                    if compaction_output.is_none() {                        compaction_output = Some(item);                    }                }            }            ResponseEvent::Completed { token_usage, .. } => {                saw_completed = true;                completed_token_usage = token_usage;                break;            }            _ => {}        }    }    if !saw_completed {        return Err(CodexErr::Stream(            "remote compaction v2 stream closed before response.completed".to_string(),            None,        ));    }    if compaction_count != 1 {        return Err(CodexErr::Fatal(format!(            "remote compaction v2 expected exactly one compaction output item, got {compaction_count} from {output_item_count} output items"        )));    }    let Some(compaction_output) = compaction_output else {        unreachable!("compaction output must exist when count is exactly one");    };    Ok(RemoteCompactionV2Output {        compaction_output,        token_usage: completed_token_usage,    })}fn build_v2_compacted_history(    prompt_input: &[ResponseItem],    compaction_output: ResponseItem,) -> (Vec<ResponseItem>, usize) {    let retained = prompt_input        .iter()        .filter(|item| is_retained_for_remote_compaction_v2(item))        .filter(|item| should_keep_compacted_history_item(item))        .cloned()        .collect::<Vec<_>>();    let mut retained =        truncate_retained_messages_for_remote_compaction(retained, RETAINED_MESSAGE_TOKEN_BUDGET);    let retained_image_count = retained        .iter()        .map(retained_input_image_count)        .sum::<usize>();    retained.push(compaction_output);    (retained, retained_image_count)}fn is_retained_for_remote_compaction_v2(item: &ResponseItem) -> bool {    let ResponseItem::Message { role, .. } = item else {        return false;    };    matches!(role.as_str(), "user" | "developer" | "system")}fn retained_input_image_count(item: &ResponseItem) -> usize {    let ResponseItem::Message { content, .. } = item else {        return 0;    };    content        .iter()        .filter(|item| matches!(item, ContentItem::InputImage { .. }))        .count()}fn truncate_retained_messages_for_remote_compaction(    items: Vec<ResponseItem>,    max_tokens: usize,) -> Vec<ResponseItem> {    let mut remaining = max_tokens;    let mut truncated_reversed = Vec::with_capacity(items.len());    for item in items.into_iter().rev() {        if remaining == 0 {            continue;        }        let token_count = message_text_token_count(&item).max(1);        if token_count <= remaining {            truncated_reversed.push(item);            remaining = remaining.saturating_sub(token_count);        } else if let Some(truncated_item) =            truncate_message_text_to_token_budget(item, /*max_tokens*/ remaining)        {            truncated_reversed.push(truncated_item);            remaining = 0;        }    }    truncated_reversed.reverse();    truncated_reversed}fn message_text_token_count(item: &ResponseItem) -> usize {    let ResponseItem::Message { content, .. } = item else {        return 0;    };    content        .iter()        .map(|item| match item {            ContentItem::InputText { text } | ContentItem::OutputText { text } => {                approx_token_count(text)            }            ContentItem::InputImage { .. } => 0,        })        .sum()}fn truncate_message_text_to_token_budget(    item: ResponseItem,    max_tokens: usize,) -> Option<ResponseItem> {    let ResponseItem::Message {        id,        role,        content,        phase,        metadata,    } = item    else {        return Some(item);    };    let mut remaining = max_tokens;    let mut truncated_content = Vec::with_capacity(content.len());    for mut content_item in content {        match &mut content_item {            ContentItem::InputText { text } | ContentItem::OutputText { text } => {                if remaining == 0 {                    continue;                }                let token_count = approx_token_count(text);                if token_count <= remaining {                    remaining = remaining.saturating_sub(token_count);                } else {                    *text = truncate_text(text, TruncationPolicy::Tokens(remaining));                    remaining = 0;                }                if !text.is_empty() {                    truncated_content.push(content_item);                }            }            ContentItem::InputImage { .. } => truncated_content.push(content_item),        }    }    if truncated_content.is_empty() {        return None;    }    Some(ResponseItem::Message {        id,        role,        content: truncated_content,        phase,        metadata,    })}#[cfg(test)]mod tests {    use super::*;    use codex_protocol::models::ContentItem;    use codex_protocol::models::MessagePhase;    use pretty_assertions::assert_eq;    use tokio::sync::mpsc;    use tokio_util::sync::CancellationToken;    fn message(role: &str, text: &str, phase: Option<MessagePhase>) -> ResponseItem {        ResponseItem::Message {            id: None,            role: role.to_string(),            content: vec![ContentItem::InputText {                text: text.to_string(),            }],            phase,            metadata: None,        }    }    fn response_stream(events: Vec<CodexResult<ResponseEvent>>) -> ResponseStream {        let (tx_event, rx_event) = mpsc::channel(events.len().max(1));        for event in events {            tx_event                .try_send(event)                .expect("response stream test channel should have capacity");        }        drop(tx_event);        ResponseStream {            rx_event,            consumer_dropped: CancellationToken::new(),        }    }    #[test]    fn build_v2_compacted_history_filters_to_installed_retention_shape() {        let input = vec![            message("developer", "dev", /*phase*/ None),            message("system", "sys", /*phase*/ None),            message("user", "user", /*phase*/ None),            message("assistant", "commentary", Some(MessagePhase::Commentary)),            message("assistant", "final", Some(MessagePhase::FinalAnswer)),            ResponseItem::FunctionCall {                id: None,                name: "shell_command".to_string(),                namespace: None,                arguments: "{}".to_string(),                call_id: "call_1".to_string(),                metadata: None,            },            ResponseItem::Compaction {                encrypted_content: "old".to_string(),                metadata: None,            },        ];        let output = ResponseItem::Compaction {            encrypted_content: "new".to_string(),            metadata: None,        };        let (history, _) = build_v2_compacted_history(&input, output.clone());        assert_eq!(            history,            vec![message("user", "user", /*phase*/ None), output]        );    }    #[test]    fn build_v2_compacted_history_discards_messages_before_truncating() {        let old = message("user", "old", /*phase*/ None);        let new = message("user", "new", /*phase*/ None);        let huge_developer_message = "d".repeat((RETAINED_MESSAGE_TOKEN_BUDGET + 1) * 4);        let huge_contextual_message = format!(            "<environment_context>\n{}\n</environment_context>",            "c".repeat((RETAINED_MESSAGE_TOKEN_BUDGET + 1) * 4)        );        let input = vec![            old.clone(),            message("developer", &huge_developer_message, /*phase*/ None),            message("user", &huge_contextual_message, /*phase*/ None),            new.clone(),        ];        let output = ResponseItem::Compaction {            encrypted_content: "new".to_string(),            metadata: None,        };        let (history, _) = build_v2_compacted_history(&input, output.clone());        assert_eq!(history, vec![old, new, output]);    }    #[test]    fn build_v2_compacted_history_counts_retained_input_images() {        let input = vec![ResponseItem::Message {            id: None,            role: "user".to_string(),            content: vec![                ContentItem::InputText {                    text: "user".to_string(),                },                ContentItem::InputImage {                    image_url: "data:image/png;base64,abc".to_string(),                    detail: None,                },                ContentItem::InputImage {                    image_url: "data:image/png;base64,def".to_string(),                    detail: None,                },            ],            phase: None,            metadata: None,        }];        let output = ResponseItem::Compaction {            encrypted_content: "new".to_string(),            metadata: None,        };        let (_, retained_image_count) = build_v2_compacted_history(&input, output);        assert_eq!(retained_image_count, 2);    }    #[test]    fn retained_history_truncation_keeps_newest_messages_first() {        let middle = message("user", "middle1234", /*phase*/ None);        let new = message("user", "new", /*phase*/ None);        let retained = vec![            message("user", "old-old", /*phase*/ None),            middle,            new.clone(),        ];        let truncated =            truncate_retained_messages_for_remote_compaction(retained, /*max_tokens*/ 3);        assert_eq!(            truncated,            vec![                message("user", "midd…1 tokens truncated…1234", /*phase*/ None),                new,            ]        );    }    #[test]    fn retained_history_truncation_preserves_images_and_truncates_later_text_parts() {        let item = ResponseItem::Message {            id: None,            role: "user".to_string(),            content: vec![                ContentItem::InputText {                    text: "abcdef".to_string(),                },                ContentItem::InputImage {                    image_url: "data:image/png;base64,abc".to_string(),                    detail: None,                },                ContentItem::OutputText {                    text: "uvwxyz".to_string(),                },            ],            phase: None,            metadata: None,        };        let truncated =            truncate_retained_messages_for_remote_compaction(vec![item], /*max_tokens*/ 3);        assert_eq!(            truncated,            vec![ResponseItem::Message {                id: None,                role: "user".to_string(),                content: vec![                    ContentItem::InputText {                        text: "abcdef".to_string(),                    },                    ContentItem::InputImage {                        image_url: "data:image/png;base64,abc".to_string(),                        detail: None,                    },                    ContentItem::OutputText {                        text: "uv…1 tokens truncated…yz".to_string(),                    },                ],                phase: None,                metadata: None,            }]        );    }    #[test]    fn retained_history_truncation_charges_image_only_messages() {        let image_only_message = ResponseItem::Message {            id: None,            role: "user".to_string(),            content: vec![ContentItem::InputImage {                image_url: "data:image/png;base64,abc".to_string(),                detail: None,            }],            phase: None,            metadata: None,        };        let newest = message("user", "new", /*phase*/ None);        let retained = vec![            message("user", "old", /*phase*/ None),            image_only_message.clone(),            newest.clone(),        ];        let truncated =            truncate_retained_messages_for_remote_compaction(retained, /*max_tokens*/ 2);        assert_eq!(truncated, vec![image_only_message, newest]);    }    #[test]    fn retained_history_truncation_drops_image_only_messages_after_budget_is_spent() {        let image_only_message = ResponseItem::Message {            id: None,            role: "user".to_string(),            content: vec![ContentItem::InputImage {                image_url: "data:image/png;base64,abc".to_string(),                detail: None,            }],            phase: None,            metadata: None,        };        let newest = message("user", "new", /*phase*/ None);        let retained = vec![image_only_message, newest.clone()];        let truncated =            truncate_retained_messages_for_remote_compaction(retained, /*max_tokens*/ 1);        assert_eq!(truncated, vec![newest]);    }    #[tokio::test]    async fn collect_compaction_output_accepts_additional_output_items() {        let compaction = ResponseItem::Compaction {            encrypted_content: "encrypted".to_string(),            metadata: None,        };        let stream = response_stream(vec![            Ok(ResponseEvent::OutputItemDone(message(                "assistant",                "IGNORED_COMPACT_REPLY",                Some(MessagePhase::FinalAnswer),            ))),            Ok(ResponseEvent::OutputItemDone(compaction.clone())),            Ok(ResponseEvent::Completed {                response_id: "resp-compact".to_string(),                token_usage: Some(TokenUsage {                    input_tokens: 123_456,                    cached_input_tokens: 7_890,                    output_tokens: 42,                    reasoning_output_tokens: 5,                    total_tokens: 123_498,                }),                end_turn: Some(true),            }),        ]);        let output = collect_compaction_output(stream)            .await            .expect("compaction should be collected");        assert_eq!(output.compaction_output, compaction);        assert_eq!(            output.token_usage,            Some(TokenUsage {                input_tokens: 123_456,                cached_input_tokens: 7_890,                output_tokens: 42,                reasoning_output_tokens: 5,                total_tokens: 123_498,            })        );    }}