Skip to content

chore: extract remote compaction request attempts #31316

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
celia-oai merged 2 commits into from
Jul 7, 2026
+265 −164
Merged

chore: extract remote compaction request attempts #31316

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
25ec9ab
Extract remote compaction request attempts
celia-oai Jul 6, 2026
1f3200a
Preserve remote compaction behavior
celia-oai Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 15 additions & 77 deletions codex-rs/core/src/compact_remote.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand All @@ -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";
Expand Down Expand Up @@ -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(),
Expand Down
107 changes: 107 additions & 0 deletions codex-rs/core/src/compact_remote_request.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<ResponseItem>,
pub(super) trace_input_history: Vec<ResponseItem>,
}

pub(super) async fn run_remote_compact_attempt(
sess: &Arc<Session>,
step_context: &Arc<StepContext>,
turn_state: Option<Arc<OnceLock<String>>>,
compaction_trace: &CompactionTraceContext,
compaction_metadata: CompactionTurnMetadata,
analytics_details: &mut CompactionAnalyticsDetails,
) -> CodexResult<RemoteCompactAttempt> {
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,
})
}
104 changes: 17 additions & 87 deletions codex-rs/core/src/compact_remote_v2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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.
Expand Down Expand Up @@ -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);
Expand Down
Loading
Loading