diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 6b251562f798..8964b9aa425d 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -1,8 +1,6 @@ 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; @@ -15,17 +13,14 @@ 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::step_context::StepContext; -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::auth::AuthMode; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_protocol::items::ContextCompactionItem; @@ -38,8 +33,11 @@ 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; + +#[path = "compact_remote_request.rs"] +mod request; +use request::RemoteCompactAttempt; +use request::run_remote_compact_attempt; const CONTEXT_WINDOW_TRUNCATED_OUTPUT_MESSAGE: &str = "Output exceeded the available model context and was truncated"; @@ -193,78 +191,18 @@ async fn run_remote_compact_task_inner_impl( 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(), - step_context.as_ref(), - &CancellationToken::new(), + let RemoteCompactAttempt { + new_history, + trace_input_history, + } = run_remote_compact_attempt( + sess, + step_context, + turn_state, + &compaction_trace, + compaction_metadata, + analytics_details, ) .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, - 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 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_number, new_window_ids) = sess.advance_auto_compact_window().await; let (new_history, world_state_baseline) = process_compacted_history( sess.as_ref(), diff --git a/codex-rs/core/src/compact_remote_request.rs b/codex-rs/core/src/compact_remote_request.rs new file mode 100644 index 000000000000..bbf2f044466d --- /dev/null +++ b/codex-rs/core/src/compact_remote_request.rs @@ -0,0 +1,107 @@ +use std::sync::Arc; +use std::sync::OnceLock; + +use super::trim_function_call_history_to_fit_context_window; +use crate::Prompt; +use crate::client::CompactConversationRequestSettings; +use crate::compact::CompactionAnalyticsDetails; +use crate::responses_metadata::CodexResponsesRequestKind; +use crate::responses_metadata::CompactionTurnMetadata; +use crate::session::session::Session; +use crate::session::step_context::StepContext; +use crate::session::turn::built_tools; +use codex_protocol::auth::AuthMode; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::models::ResponseItem; +use codex_rollout_trace::CompactionTraceContext; +use tokio_util::sync::CancellationToken; +use tracing::info; + +pub(super) struct RemoteCompactAttempt { + pub(super) new_history: Vec, + pub(super) trace_input_history: Vec, +} + +pub(super) async fn run_remote_compact_attempt( + sess: &Arc, + step_context: &Arc, + turn_state: Option>>, + compaction_trace: &CompactionTraceContext, + compaction_metadata: CompactionTurnMetadata, + analytics_details: &mut CompactionAnalyticsDetails, +) -> CodexResult { + let turn_context = &step_context.turn; + 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)) + }); + } + 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(), + step_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, + 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 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?; + Ok(RemoteCompactAttempt { + new_history, + trace_input_history, + }) +} diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 6e110045bbe4..8105f63b7afa 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -10,19 +10,16 @@ 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::step_context::StepContext; -use crate::session::turn::built_tools; use crate::session::turn_context::TurnContext; use codex_analytics::CompactionImplementation; use codex_analytics::CompactionPhase; @@ -44,8 +41,11 @@ 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; + +#[path = "compact_remote_v2_attempt.rs"] +mod attempt; +use attempt::RemoteCompactV2Attempt; +use attempt::run_remote_compact_v2_attempt; // Mirror the current /responses/compact retained-message default while the // server-side path remains the reference implementation. @@ -203,91 +203,21 @@ async fn run_remote_compact_task_inner_impl( 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(), - step_context.as_ref(), - &CancellationToken::new(), - ) - .await?; - let mut input = prompt_input.clone(); - input.push(ResponseItem::CompactionTrigger {}); - let prompt = Prompt { - input, - tools: tool_router.model_visible_specs(), - parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, - base_instructions, - 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( + let RemoteCompactV2Attempt { + trace_input_history, + prompt_input, + compaction_output, + token_usage, + owned_client_session: _owned_client_session, + } = run_remote_compact_v2_attempt( sess, - turn_context, + step_context, client_session, - &prompt, - &responses_metadata, + &compaction_trace, + compaction_metadata, + analytics_details, ) - .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?; + .await?; if let Some(token_usage) = token_usage { sess.record_rollout_budget_usage(&token_usage)?; analytics_details.active_context_tokens_before = Some(token_usage.input_tokens); diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs new file mode 100644 index 000000000000..a145d00ae668 --- /dev/null +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -0,0 +1,126 @@ +use std::sync::Arc; + +use super::RemoteCompactionV2Output; +use super::run_remote_compaction_request_v2; +use crate::Prompt; +use crate::client::ModelClientSession; +use crate::compact::CompactionAnalyticsDetails; +use crate::compact_remote::trim_function_call_history_to_fit_context_window; +use crate::responses_metadata::CodexResponsesRequestKind; +use crate::responses_metadata::CompactionTurnMetadata; +use crate::session::session::Session; +use crate::session::step_context::StepContext; +use crate::session::turn::built_tools; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::TokenUsage; +use codex_rollout_trace::CompactionTraceContext; +use tokio_util::sync::CancellationToken; +use tracing::info; + +pub(super) struct RemoteCompactV2Attempt { + pub(super) trace_input_history: Vec, + pub(super) prompt_input: Vec, + pub(super) compaction_output: ResponseItem, + pub(super) token_usage: Option, + /// Keeps a session created for standalone compaction alive through lifecycle completion. + pub(super) owned_client_session: Option, +} + +pub(super) async fn run_remote_compact_v2_attempt( + sess: &Arc, + step_context: &Arc, + client_session: Option<&mut ModelClientSession>, + compaction_trace: &CompactionTraceContext, + compaction_metadata: CompactionTurnMetadata, + analytics_details: &mut CompactionAnalyticsDetails, +) -> CodexResult { + let turn_context = &step_context.turn; + 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(), + step_context.as_ref(), + &CancellationToken::new(), + ) + .await?; + let mut input = prompt_input.clone(); + input.push(ResponseItem::CompactionTrigger {}); + let prompt = Prompt { + input, + tools: tool_router.model_visible_specs(), + parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, + base_instructions, + 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 = None; + let client_session = match client_session { + Some(client_session) => client_session, + None => owned_client_session.insert(sess.services.model_client.new_session()), + }; + let compaction_output_result = run_remote_compaction_request_v2( + sess, + turn_context.as_ref(), + 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?; + Ok(RemoteCompactV2Attempt { + trace_input_history, + prompt_input, + compaction_output, + token_usage, + owned_client_session, + }) +}