Codex Handbook
core/src/compact_remote.rs 448 lines
use std::sync::Arc;use std::sync::OnceLock;use crate::Prompt;use crate::client::CompactConversationRequestSettings;use crate::compact::CompactionAnalyticsAttempt;use crate::compact::CompactionAnalyticsDetails;use crate::compact::InitialContextInjection;use crate::compact::compaction_status_from_result;use crate::compact::insert_initial_context_before_last_real_user_or_summary;use crate::context_manager::ContextManager;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::CodexResponsesRequestKind;use crate::responses_metadata::CompactionTurnMetadata;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_app_server_protocol::AuthMode;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::BaseInstructions;use codex_protocol::models::FunctionCallOutputBody;use codex_protocol::models::FunctionCallOutputPayload;use codex_protocol::models::ResponseItem;use codex_protocol::protocol::CompactedItem;use codex_protocol::protocol::EventMsg;use codex_protocol::protocol::TurnStartedEvent;use codex_rollout_trace::CompactionCheckpointTracePayload;use tokio_util::sync::CancellationToken;use tracing::info;const CONTEXT_WINDOW_TRUNCATED_OUTPUT_MESSAGE: &str =    "Output exceeded the available model context and was truncated";pub(crate) async fn run_inline_remote_auto_compact_task(    sess: Arc<Session>,    turn_context: Arc<TurnContext>,    turn_state: Arc<OnceLock<String>>,    initial_context_injection: InitialContextInjection,    reason: CompactionReason,    phase: CompactionPhase,) -> CodexResult<()> {    run_remote_compact_task_inner(        &sess,        &turn_context,        Some(turn_state),        initial_context_injection,        CompactionTrigger::Auto,        reason,        phase,    )    .await?;    Ok(())}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,        /*turn_state*/ None,        InitialContextInjection::DoNotInject,        CompactionTrigger::Manual,        CompactionReason::UserRequested,        CompactionPhase::StandaloneTurn,    )    .await?;    Ok(())}async fn run_remote_compact_task_inner(    sess: &Arc<Session>,    turn_context: &Arc<TurnContext>,    turn_state: Option<Arc<OnceLock<String>>>,    initial_context_injection: InitialContextInjection,    trigger: CompactionTrigger,    reason: CompactionReason,    phase: CompactionPhase,) -> CodexResult<()> {    let compaction_metadata = CompactionTurnMetadata::new(        trigger,        reason,        CompactionImplementation::ResponsesCompact,        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::ResponsesCompact,        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,        turn_state,        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>,    turn_state: Option<Arc<OnceLock<String>>>,    initial_context_injection: InitialContextInjection,    compaction_metadata: CompactionTurnMetadata,    analytics_details: &mut CompactionAnalyticsDetails,) -> CodexResult<()> {    let context_compaction_item = ContextCompactionItem::new();    // Use the UI compaction item ID as the trace compaction ID so protocol lifecycle events,    // endpoint attempts, and the installed history checkpoint all have one join key.    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"        );    }    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))            });    }    // This is the history selected for remote compaction, after any output rewriting required to    // fit the compact endpoint. The checkpoint below records it separately from the next sampling    // request, whose prompt will repeat current developer/context prefix items.    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 prompt = Prompt {        input: 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 mut new_history = sess        .services        .model_client        .compact_conversation_history(            &prompt,            &turn_context.model_info,            turn_state,            CompactConversationRequestSettings {                effort: turn_context.reasoning_effort.clone(),                summary: turn_context.reasoning_summary,                service_tier: if sess.services.auth_manager.auth_mode() == Some(AuthMode::ApiKey) {                    None                } else {                    turn_context.config.service_tier.clone()                },            },            &turn_context.session_telemetry,            &compaction_trace,            &responses_metadata,        )        .await?;    let new_window_id = sess.advance_auto_compact_window_id().await;    new_history = process_compacted_history(        sess.as_ref(),        turn_context.as_ref(),        new_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),    };    // Install is the semantic boundary where the compact endpoint's output becomes live    // thread history. Keep it distinct from the later inference request so the reducer can    // still represent repeated developer/context prefix items exactly as the model saw them.    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(())}pub(crate) async fn process_compacted_history(    sess: &Session,    turn_context: &TurnContext,    mut compacted_history: Vec<ResponseItem>,    initial_context_injection: InitialContextInjection,) -> Vec<ResponseItem> {    // Mid-turn compaction is the only path that must inject initial context above the last user    // message in the replacement history. Pre-turn compaction instead injects context after the    // compaction item, but mid-turn compaction keeps the compaction item last for model training.    let initial_context = if matches!(        initial_context_injection,        InitialContextInjection::BeforeLastUserMessage    ) {        sess.build_initial_context(turn_context).await    } else {        Vec::new()    };    compacted_history.retain(should_keep_compacted_history_item);    insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context)}/// Returns whether an item from remote compaction output should be preserved.////// Called while processing the model-provided compacted transcript, before we/// append fresh canonical context from the current session.////// We drop:/// - `developer` messages because remote output can include stale/duplicated///   instruction content./// - non-user-content `user` messages (session prefix/instruction wrappers),///   while preserving real user messages and persisted hook prompts.////// This intentionally keeps:/// - `assistant` messages (future remote compaction models may emit them)/// - `user`-role warnings that parse as `TurnItem::UserMessage` and compaction-generated summary///   messages. Legacy warning fragments are filtered by `parse_turn_item` before they reach this///   check.pub(crate) fn should_keep_compacted_history_item(item: &ResponseItem) -> bool {    match item {        ResponseItem::Message { role, .. } if role == "developer" => false,        ResponseItem::Message { role, .. } if role == "user" => {            matches!(                crate::event_mapping::parse_turn_item(item),                Some(TurnItem::UserMessage(_) | TurnItem::HookPrompt(_))            )        }        ResponseItem::Message { role, .. } if role == "assistant" => true,        ResponseItem::Message { .. } => false,        ResponseItem::AgentMessage { .. } => true,        ResponseItem::Compaction { .. } | ResponseItem::ContextCompaction { .. } => true,        ResponseItem::CompactionTrigger { .. } => false,        ResponseItem::Reasoning { .. }        | ResponseItem::LocalShellCall { .. }        | ResponseItem::FunctionCall { .. }        | ResponseItem::ToolSearchCall { .. }        | ResponseItem::FunctionCallOutput { .. }        | ResponseItem::ToolSearchOutput { .. }        | ResponseItem::CustomToolCall { .. }        | ResponseItem::CustomToolCallOutput { .. }        | ResponseItem::WebSearchCall { .. }        | ResponseItem::ImageGenerationCall { .. }        | ResponseItem::Other => false,    }}pub(crate) fn trim_function_call_history_to_fit_context_window(    history: &mut ContextManager,    turn_context: &TurnContext,    base_instructions: &BaseInstructions,) -> (usize, i64) {    let Some(context_window) = turn_context.model_context_window() else {        return (0, 0);    };    let mut rewritten_outputs = 0usize;    let mut estimated_deleted_tokens = 0i64;    let item_count = history.raw_items().len();    for index in (0..item_count).rev() {        let Some(estimated_tokens_before) =            history.estimate_token_count_with_base_instructions(base_instructions)        else {            break;        };        if estimated_tokens_before <= context_window {            break;        }        let Some(rewritten_item) = history            .raw_items()            .get(index)            .and_then(rewritten_output_for_context_window)        else {            break;        };        let mut items = history.raw_items().to_vec();        items[index] = rewritten_item;        history.replace(items);        let estimated_tokens_after = history            .estimate_token_count_with_base_instructions(base_instructions)            .unwrap_or_default();        rewritten_outputs += 1;        estimated_deleted_tokens = estimated_deleted_tokens            .saturating_add(estimated_tokens_before.saturating_sub(estimated_tokens_after));    }    (rewritten_outputs, estimated_deleted_tokens)}fn rewritten_output_for_context_window(item: &ResponseItem) -> Option<ResponseItem> {    Some(match item {        ResponseItem::FunctionCallOutput {            call_id,            output,            metadata,        } => ResponseItem::FunctionCallOutput {            call_id: call_id.clone(),            output: truncated_output_payload(output),            metadata: metadata.clone(),        },        ResponseItem::CustomToolCallOutput {            call_id,            name,            output,            metadata,        } => ResponseItem::CustomToolCallOutput {            call_id: call_id.clone(),            name: name.clone(),            output: truncated_output_payload(output),            metadata: metadata.clone(),        },        ResponseItem::ToolSearchOutput {            call_id,            status,            execution,            metadata,            ..        } => ResponseItem::ToolSearchOutput {            call_id: call_id.clone(),            status: status.clone(),            execution: execution.clone(),            tools: Vec::new(),            metadata: metadata.clone(),        },        _ => return None,    })}fn truncated_output_payload(output: &FunctionCallOutputPayload) -> FunctionCallOutputPayload {    FunctionCallOutputPayload {        body: FunctionCallOutputBody::Text(CONTEXT_WINDOW_TRUNCATED_OUTPUT_MESSAGE.to_string()),        success: output.success,    }}