From 25ec9ab031f5c8bf861ccb9f0803f985d36c2d95 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Mon, 6 Jul 2026 16:19:08 -0700 Subject: [PATCH 1/2] Extract remote compaction request attempts --- codex-rs/core/src/compact_remote.rs | 109 +++------------ codex-rs/core/src/compact_remote_request.rs | 119 ++++++++++++++++ codex-rs/core/src/compact_remote_v2.rs | 110 ++++----------- .../core/src/compact_remote_v2_attempt.rs | 130 ++++++++++++++++++ 4 files changed, 294 insertions(+), 174 deletions(-) create mode 100644 codex-rs/core/src/compact_remote_request.rs create mode 100644 codex-rs/core/src/compact_remote_v2_attempt.rs diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 6b251562f798..d9bf8673fb3d 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"; @@ -182,89 +180,26 @@ async fn run_remote_compact_task_inner_impl( ) -> CodexResult<()> { let turn_context = &step_context.turn; 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_id = context_compaction_item.id.clone(); 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 active_context_tokens_before = analytics_details.active_context_tokens_before; + let RemoteCompactAttempt { + turn_context, + new_history, + trace_input_history, + compaction_trace, + } = run_remote_compact_attempt( + sess, + step_context, + turn_state, + compaction_item_id.as_str(), + compaction_metadata, + active_context_tokens_before, + 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(), @@ -303,9 +238,9 @@ async fn run_remote_compact_task_inner_impl( compacted_item, ) .await; - sess.recompute_token_usage(turn_context).await; + sess.recompute_token_usage(&turn_context).await; - sess.emit_turn_item_completed(turn_context, compaction_item) + sess.emit_turn_item_completed(&turn_context, compaction_item) .await; Ok(()) } 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..95318551b318 --- /dev/null +++ b/codex-rs/core/src/compact_remote_request.rs @@ -0,0 +1,119 @@ +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 crate::session::turn_context::TurnContext; +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) turn_context: Arc, + pub(super) new_history: Vec, + pub(super) trace_input_history: Vec, + pub(super) compaction_trace: CompactionTraceContext, +} + +pub(super) async fn run_remote_compact_attempt( + sess: &Arc, + step_context: &Arc, + turn_state: Option>>, + compaction_item_id: &str, + compaction_metadata: CompactionTurnMetadata, + active_context_tokens_before: Option, + analytics_details: &mut CompactionAnalyticsDetails, +) -> CodexResult { + let turn_context = Arc::clone(&step_context.turn); + let compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context( + turn_context.sub_id.as_str(), + compaction_item_id, + turn_context.model_info.slug.as_str(), + turn_context.provider.info().name.as_str(), + ); + 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" + ); + } + analytics_details.active_context_tokens_before = active_context_tokens_before; + 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 = + 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 { + turn_context, + new_history, + trace_input_history, + compaction_trace, + }) +} diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 6e110045bbe4..ca13fb8c4e87 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. @@ -193,75 +193,11 @@ async fn run_remote_compact_task_inner_impl( ) -> CodexResult<()> { let turn_context = &step_context.turn; 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_id = context_compaction_item.id.clone(); 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(), - 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, @@ -270,24 +206,24 @@ async fn run_remote_compact_task_inner_impl( &mut owned_client_session } }; - let compaction_output_result = run_remote_compaction_request_v2( - sess, + let active_context_tokens_before = analytics_details.active_context_tokens_before; + let RemoteCompactV2Attempt { 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 { + trace_input_history, + prompt_input, compaction_output, token_usage, - } = compaction_output_result?; + compaction_trace, + } = run_remote_compact_v2_attempt( + sess, + step_context, + client_session, + compaction_item_id.as_str(), + compaction_metadata, + active_context_tokens_before, + analytics_details, + ) + .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); @@ -332,9 +268,9 @@ async fn run_remote_compact_task_inner_impl( compacted_item, ) .await; - sess.recompute_token_usage(turn_context).await; + sess.recompute_token_usage(&turn_context).await; - sess.emit_turn_item_completed(turn_context, compaction_item) + sess.emit_turn_item_completed(&turn_context, compaction_item) .await; Ok(()) } 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..46ec5675c28b --- /dev/null +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -0,0 +1,130 @@ +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 crate::session::turn_context::TurnContext; +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) turn_context: Arc, + pub(super) trace_input_history: Vec, + pub(super) prompt_input: Vec, + pub(super) compaction_output: ResponseItem, + pub(super) token_usage: Option, + pub(super) compaction_trace: CompactionTraceContext, +} + +pub(super) async fn run_remote_compact_v2_attempt( + sess: &Arc, + step_context: &Arc, + client_session: &mut ModelClientSession, + compaction_item_id: &str, + compaction_metadata: CompactionTurnMetadata, + active_context_tokens_before: Option, + analytics_details: &mut CompactionAnalyticsDetails, +) -> CodexResult { + let turn_context = Arc::clone(&step_context.turn); + let compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context( + turn_context.sub_id.as_str(), + compaction_item_id, + turn_context.model_info.slug.as_str(), + turn_context.provider.info().name.as_str(), + ); + 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" + ); + } + analytics_details.active_context_tokens_before = active_context_tokens_before; + 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 = + 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 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 { + turn_context, + trace_input_history, + prompt_input, + compaction_output, + token_usage, + compaction_trace, + }) +} From 1f3200acfaf2062f353ee11186eb77d54a66376b Mon Sep 17 00:00:00 2001 From: celia-oai Date: Mon, 6 Jul 2026 17:39:17 -0700 Subject: [PATCH 2/2] Preserve remote compaction behavior --- codex-rs/core/src/compact_remote.rs | 19 ++++++----- codex-rs/core/src/compact_remote_request.rs | 24 ++++---------- codex-rs/core/src/compact_remote_v2.rs | 26 ++++++--------- .../core/src/compact_remote_v2_attempt.rs | 32 ++++++++----------- 4 files changed, 41 insertions(+), 60 deletions(-) diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index d9bf8673fb3d..8964b9aa425d 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -180,23 +180,26 @@ async fn run_remote_compact_task_inner_impl( ) -> CodexResult<()> { let turn_context = &step_context.turn; let context_compaction_item = ContextCompactionItem::new(); - let compaction_item_id = context_compaction_item.id.clone(); + // 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 active_context_tokens_before = analytics_details.active_context_tokens_before; let RemoteCompactAttempt { - turn_context, new_history, trace_input_history, - compaction_trace, } = run_remote_compact_attempt( sess, step_context, turn_state, - compaction_item_id.as_str(), + &compaction_trace, compaction_metadata, - active_context_tokens_before, analytics_details, ) .await?; @@ -238,9 +241,9 @@ async fn run_remote_compact_task_inner_impl( compacted_item, ) .await; - sess.recompute_token_usage(&turn_context).await; + sess.recompute_token_usage(turn_context).await; - sess.emit_turn_item_completed(&turn_context, compaction_item) + sess.emit_turn_item_completed(turn_context, compaction_item) .await; Ok(()) } diff --git a/codex-rs/core/src/compact_remote_request.rs b/codex-rs/core/src/compact_remote_request.rs index 95318551b318..bbf2f044466d 100644 --- a/codex-rs/core/src/compact_remote_request.rs +++ b/codex-rs/core/src/compact_remote_request.rs @@ -10,7 +10,6 @@ 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_protocol::auth::AuthMode; use codex_protocol::error::Result as CodexResult; use codex_protocol::models::ResponseItem; @@ -19,28 +18,19 @@ use tokio_util::sync::CancellationToken; use tracing::info; pub(super) struct RemoteCompactAttempt { - pub(super) turn_context: Arc, pub(super) new_history: Vec, pub(super) trace_input_history: Vec, - pub(super) compaction_trace: CompactionTraceContext, } pub(super) async fn run_remote_compact_attempt( sess: &Arc, step_context: &Arc, turn_state: Option>>, - compaction_item_id: &str, + compaction_trace: &CompactionTraceContext, compaction_metadata: CompactionTurnMetadata, - active_context_tokens_before: Option, analytics_details: &mut CompactionAnalyticsDetails, ) -> CodexResult { - let turn_context = Arc::clone(&step_context.turn); - let compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context( - turn_context.sub_id.as_str(), - compaction_item_id, - turn_context.model_info.slug.as_str(), - turn_context.provider.info().name.as_str(), - ); + 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) = @@ -56,13 +46,13 @@ pub(super) async fn run_remote_compact_attempt( "rewrote history outputs before remote compaction" ); } - analytics_details.active_context_tokens_before = active_context_tokens_before; 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 = - active_context_tokens_before.map(|active_context_tokens_before| { + 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)) }); @@ -106,14 +96,12 @@ pub(super) async fn run_remote_compact_attempt( }, }, &turn_context.session_telemetry, - &compaction_trace, + compaction_trace, &responses_metadata, ) .await?; Ok(RemoteCompactAttempt { - turn_context, new_history, trace_input_history, - compaction_trace, }) } diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index ca13fb8c4e87..8105f63b7afa 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -193,34 +193,28 @@ async fn run_remote_compact_task_inner_impl( ) -> CodexResult<()> { let turn_context = &step_context.turn; let context_compaction_item = ContextCompactionItem::new(); - let compaction_item_id = context_compaction_item.id.clone(); + 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 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 active_context_tokens_before = analytics_details.active_context_tokens_before; let RemoteCompactV2Attempt { - turn_context, trace_input_history, prompt_input, compaction_output, token_usage, - compaction_trace, + owned_client_session: _owned_client_session, } = run_remote_compact_v2_attempt( sess, step_context, client_session, - compaction_item_id.as_str(), + &compaction_trace, compaction_metadata, - active_context_tokens_before, analytics_details, ) .await?; @@ -268,9 +262,9 @@ async fn run_remote_compact_task_inner_impl( compacted_item, ) .await; - sess.recompute_token_usage(&turn_context).await; + sess.recompute_token_usage(turn_context).await; - sess.emit_turn_item_completed(&turn_context, compaction_item) + sess.emit_turn_item_completed(turn_context, compaction_item) .await; Ok(()) } diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index 46ec5675c28b..a145d00ae668 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -11,7 +11,6 @@ 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_protocol::error::Result as CodexResult; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::TokenUsage; @@ -20,30 +19,23 @@ use tokio_util::sync::CancellationToken; use tracing::info; pub(super) struct RemoteCompactV2Attempt { - pub(super) turn_context: Arc, pub(super) trace_input_history: Vec, pub(super) prompt_input: Vec, pub(super) compaction_output: ResponseItem, pub(super) token_usage: Option, - pub(super) compaction_trace: CompactionTraceContext, + /// 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: &mut ModelClientSession, - compaction_item_id: &str, + client_session: Option<&mut ModelClientSession>, + compaction_trace: &CompactionTraceContext, compaction_metadata: CompactionTurnMetadata, - active_context_tokens_before: Option, analytics_details: &mut CompactionAnalyticsDetails, ) -> CodexResult { - let turn_context = Arc::clone(&step_context.turn); - let compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context( - turn_context.sub_id.as_str(), - compaction_item_id, - turn_context.model_info.slug.as_str(), - turn_context.provider.info().name.as_str(), - ); + 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) = @@ -59,13 +51,13 @@ pub(super) async fn run_remote_compact_v2_attempt( "rewrote history outputs before remote compaction v2" ); } - analytics_details.active_context_tokens_before = active_context_tokens_before; 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 = - active_context_tokens_before.map(|active_context_tokens_before| { + 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)) }); @@ -102,6 +94,11 @@ pub(super) async fn run_remote_compact_v2_attempt( "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(), @@ -120,11 +117,10 @@ pub(super) async fn run_remote_compact_v2_attempt( token_usage, } = compaction_output_result?; Ok(RemoteCompactV2Attempt { - turn_context, trace_input_history, prompt_input, compaction_output, token_usage, - compaction_trace, + owned_client_session, }) }