From be5296d4672f59205c1b8a76f46e74eae4f10d51 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Tue, 13 Oct 2026 08:04:25 +0000 Subject: [PATCH] Prepare compaction helpers for background requests (#53606) ## What changed - Extract an HTTP Responses executor that can run without owning a turn's WebSocket state, with optional authentication-owner checks. - Add cancellation-aware compaction output collection that records background token usage against the rollout budget. - Support installing a compacted history prefix while preserving later history items and placing initial context around real user or agent messages. - Allow callers to supply a context window for history trimming and a captured window for compaction request metadata. GitOrigin-RevId: 29f3ca0ae8841f295058feb207e3a225af9a1487 --- codex-rs/core/src/client.rs | 77 ++++++++++--- codex-rs/core/src/compact.rs | 40 ++++--- codex-rs/core/src/compact_remote_history.rs | 3 +- codex-rs/core/src/compact_remote_v2.rs | 65 +++++++++-- .../core/src/compact_remote_v2_attempt.rs | 3 +- codex-rs/core/src/session/mod.rs | 107 ++++++++++++++++-- codex-rs/core/src/session/session.rs | 14 ++- 7 files changed, 248 insertions(+), 61 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index ab83a91a61..89eb0e1edc 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -309,6 +309,11 @@ pub struct ModelClientSession { turn_state: Arc>, } +/// Executes HTTP Responses requests without owning a turn's WebSocket state. +pub(crate) struct ResponsesHttpExecutor<'a> { + client: &'a ModelClient, +} + #[derive(Debug, Clone)] struct LastResponse { response_id: String, @@ -685,7 +690,7 @@ impl ModelClient { self.state.provider.auth_manager() } - fn auth_owner_generation(&self) -> Option { + pub(crate) fn auth_owner_generation(&self) -> Option { self.auth_manager().map(|manager| { manager .auth_change_state_receiver() @@ -1428,6 +1433,12 @@ impl ModelClientSession { .replace(Arc::clone(tools)); previous.is_some_and(|previous| previous != *tools) } +} + +impl<'a> ResponsesHttpExecutor<'a> { + pub(crate) fn new(client: &'a ModelClient) -> Self { + Self { client } + } #[allow(clippy::too_many_arguments)] /// Builds shared Responses API transport options and request-body options. @@ -1436,6 +1447,7 @@ impl ModelClientSession { /// regardless of transport choice. async fn build_responses_options( &self, + turn_state: &Arc>, responses_metadata: &CodexResponsesMetadata, compression: Compression, use_responses_lite: bool, @@ -1447,7 +1459,7 @@ impl ModelClientSession { extra_headers: { let mut headers = build_responses_headers( self.client.state.beta_features_header.as_deref(), - Some(&self.turn_state), + Some(turn_state), ); add_originator_header(&mut headers, self.client.state.originator.as_str()); headers.extend( @@ -1461,9 +1473,20 @@ impl ModelClientSession { headers }, compression, - turn_state: Some(Arc::clone(&self.turn_state)), + turn_state: Some(Arc::clone(turn_state)), } } +} + +impl ModelClientSession { + #[allow(dead_code)] + pub(crate) fn turn_state_with_auth_owner(&self) -> (Option, Option) { + let owner = self.client.auth_owner_generation(); + let turn_state = (self.websocket_session.auth_owner_generation == owner) + .then(|| self.turn_state.get().cloned()) + .flatten(); + (turn_state, owner) + } /// Checks whether the current request is an incremental extension of the previous request. /// We only reuse an incremental input delta when non-input request fields are unchanged and @@ -1700,7 +1723,9 @@ impl ModelClientSession { "websocket connection is unavailable".to_string(), )) } +} +impl ResponsesHttpExecutor<'_> { fn responses_request_compression(&self, auth: Option<&CodexAuth>) -> Compression { if self.client.state.enable_request_compression && auth.is_some_and(CodexAuth::uses_codex_backend) @@ -1712,7 +1737,7 @@ impl ModelClientSession { } } - /// Streams a turn via the OpenAI Responses API. + /// Streams a turn via the OpenAI Responses API over HTTP. /// /// Handles reasoning summaries, verbosity, and the `text` controls used for output schemas. #[allow(clippy::too_many_arguments)] @@ -1729,8 +1754,11 @@ impl ModelClientSession { turn.has_metadata_header = responses_metadata.has_turn_metadata() ) )] - async fn stream_responses_api( + pub(crate) async fn stream_responses_http( &self, + turn_state: &Arc>, + expected_auth_owner: Option, + event_sender: Option<&Sender>, prompt: &Prompt, model_info: &ModelInfo, session_telemetry: &SessionTelemetry, @@ -1751,6 +1779,12 @@ impl ModelClientSession { .client .current_client_setup(ClientRouting::Workspace) .await?; + if expected_auth_owner.is_some_and(|owner| { + Some(owner) != client_setup.auth_owner_generation + || Some(owner) != self.client.auth_owner_generation() + }) { + return Err(CodexErr::TurnAborted); + } let include_internal = self .client .state @@ -1771,7 +1805,7 @@ impl ModelClientSession { client_setup.agent_identity_telemetry.clone(), pending_retry, ); - let (request_telemetry, sse_telemetry) = Self::build_streaming_telemetry( + let (request_telemetry, sse_telemetry) = ModelClientSession::build_streaming_telemetry( session_telemetry, request_auth_context, RequestRouteTelemetry::for_endpoint("/responses"), @@ -1780,6 +1814,7 @@ impl ModelClientSession { let compression = self.responses_request_compression(client_setup.auth.as_ref()); let mut options = self .build_responses_options( + turn_state, responses_metadata, compression, model_info.use_responses_lite, @@ -1883,7 +1918,7 @@ impl ModelClientSession { &mut provider_auth_recovery_attempted, session_telemetry, &self.client.state.provider, - self.client.event_sender.as_ref(), + event_sender, responses_metadata.turn_id.as_deref(), ) .await?, @@ -1904,7 +1939,9 @@ impl ModelClientSession { } } } +} +impl ModelClientSession { /// Streams a turn via the Responses API over WebSocket transport. #[allow(clippy::too_many_arguments)] #[instrument( @@ -2356,17 +2393,21 @@ impl ModelClientSession { } } - self.stream_responses_api( - prompt, - model_info, - session_telemetry, - effort, - summary, - service_tier, - responses_metadata, - inference_trace, - ) - .await + ResponsesHttpExecutor::new(&self.client) + .stream_responses_http( + &self.turn_state, + /*expected_auth_owner*/ None, + self.client.event_sender.as_ref(), + prompt, + model_info, + session_telemetry, + effort, + summary, + service_tier, + responses_metadata, + inference_trace, + ) + .await } } } diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index a3c3165fdb..9379204c5c 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -608,6 +608,21 @@ pub(crate) fn is_summary_message(message: &str) -> bool { message.starts_with(format!("{SUMMARY_PREFIX}\n").as_str()) } +pub(crate) fn is_real_user_or_agent_message(item: &ResponseItem) -> bool { + if let ResponseItem::AgentMessage { content, .. } = item { + return !matches!( + content.first(), + Some(AgentMessageInputContent::InputText { text }) + if text.starts_with("Message Type: FINAL_ANSWER\n") + ); + } + matches!(item, ResponseItem::Message { role, .. } if role == "user") + && matches!( + crate::event_mapping::parse_turn_item(item), + Some(TurnItem::UserMessage(user)) if !is_summary_message(&user.message()) + ) +} + /// Inserts canonical initial context into compacted replacement history at the /// model-expected boundary. /// @@ -625,31 +640,20 @@ pub(crate) fn insert_initial_context_before_last_real_user_or_summary( let mut last_user_or_summary_index = None; let mut last_real_user_index = None; for (i, item) in compacted_history.iter().enumerate().rev() { - if crate::compact_retained_tool_outputs::is_retained_tool_output(item) { - last_real_user_index = Some(i); - break; - } - if let ResponseItem::AgentMessage { content, .. } = &item.item - && !matches!( - content.first(), - Some(AgentMessageInputContent::InputText { text }) - if text.starts_with("Message Type: FINAL_ANSWER\n") - ) + if crate::compact_retained_tool_outputs::is_retained_tool_output(item) + || is_real_user_or_agent_message(&item.item) { last_real_user_index = Some(i); break; } - let Some(TurnItem::UserMessage(user)) = crate::event_mapping::parse_turn_item(&item.item) - else { - continue; - }; // Compaction summaries are encoded as user messages, so track both: // the last real user message (preferred insertion point) and the last // user-message-like item (fallback summary insertion point). - last_user_or_summary_index.get_or_insert(i); - if !is_summary_message(&user.message()) { - last_real_user_index = Some(i); - break; + if matches!( + crate::event_mapping::parse_turn_item(&item.item), + Some(TurnItem::UserMessage(_)) + ) { + last_user_or_summary_index.get_or_insert(i); } } let last_compaction_index = compacted_history diff --git a/codex-rs/core/src/compact_remote_history.rs b/codex-rs/core/src/compact_remote_history.rs index b411dab2f1..82fc5e3f0e 100644 --- a/codex-rs/core/src/compact_remote_history.rs +++ b/codex-rs/core/src/compact_remote_history.rs @@ -70,9 +70,10 @@ fn is_attached_notice(notice: &ResponseItem) -> bool { pub(crate) fn trim_function_call_history_to_fit_context_window( history: &mut ContextManager, turn_context: &TurnContext, + context_window: Option, base_instructions: &BaseInstructions, ) -> (usize, i64) { - let Some(context_window) = turn_context.model_context_window() else { + let Some(context_window) = context_window else { return (0, 0); }; // Keep the unclamped total so replacing an item cannot lose an overflow hidden by i64 diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index de9da20f8e..316751f827 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -60,6 +60,7 @@ 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; #[path = "compact_remote_v2_attempt.rs"] mod attempt; @@ -81,6 +82,13 @@ const MAX_RETAINED_AGENT_MESSAGE_TOKENS: i64 = 10_000; // retry budget smaller than the general Responses stream retry budget. const MAX_REMOTE_COMPACTION_V2_STREAM_RETRIES: u64 = 2; +#[derive(Clone, Copy)] +enum CompactionRequestKind<'a> { + Foreground, + #[allow(dead_code)] + Background(&'a CancellationToken), +} + pub(crate) async fn run_inline_remote_auto_compact_task( sess: Arc, step_context: Arc, @@ -410,7 +418,15 @@ async fn run_remote_compaction_request_v2( ) .await { - Ok(stream) => collect_compaction_output(sess, turn_context, stream).await, + Ok(stream) => { + collect_compaction_output( + sess, + turn_context, + stream, + CompactionRequestKind::Foreground, + ) + .await + } Err(err) => Err(err), }; @@ -436,13 +452,24 @@ async fn collect_compaction_output( sess: &Session, turn_context: &TurnContext, mut stream: ResponseStream, + request_kind: CompactionRequestKind<'_>, ) -> CodexResult { let mut output_item_count = 0usize; let mut compaction_count = 0usize; let mut compaction_output = None; let mut completed_response_id = None; let mut completed_token_usage = None; - while let Some(event) = stream.next().await { + loop { + let event = match request_kind { + CompactionRequestKind::Foreground => stream.next().await, + CompactionRequestKind::Background(cancellation) => tokio::select! { + biased; + // Charge completed usage even when cancellation arrives at the same time. + event = stream.next() => event, + _ = cancellation.cancelled() => return Err(CodexErr::TurnAborted), + }, + }; + let Some(event) = event else { break }; match event? { ResponseEvent::OutputItemDone(item) => { output_item_count += 1; @@ -459,13 +486,22 @@ async fn collect_compaction_output( usage_metadata, .. } => { - sess.record_observed_response_completed( - turn_context, - &response_id, - token_usage.as_ref(), - usage_metadata.as_ref(), - ) - .await; + match request_kind { + CompactionRequestKind::Background(_) => { + if let Some(usage) = token_usage.as_ref() { + sess.record_rollout_budget_usage(usage).await?; + } + } + CompactionRequestKind::Foreground => { + sess.record_observed_response_completed( + turn_context, + &response_id, + token_usage.as_ref(), + usage_metadata.as_ref(), + ) + .await; + } + } completed_response_id = Some(response_id); completed_token_usage = token_usage; break; @@ -1294,9 +1330,14 @@ mod tests { let (sess, turn_context, rx) = crate::session::tests::make_session_and_context_with_rx().await; - let output = collect_compaction_output(&sess, &turn_context, stream) - .await - .expect("compaction should be collected"); + let output = collect_compaction_output( + &sess, + &turn_context, + stream, + CompactionRequestKind::Foreground, + ) + .await + .expect("compaction should be collected"); assert_eq!(output.compaction_output, compaction); assert_eq!(output.response_id, "resp-compact"); diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index 4df467edf4..97a70670bc 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -58,7 +58,8 @@ pub(super) async fn run_remote_compact_v2_attempt( let (rewritten_outputs, estimated_deleted_tokens) = trim_function_call_history_to_fit_context_window( &mut history, - turn_context.as_ref(), + turn_context, + turn_context.model_context_window(), &base_instructions, ); if rewritten_outputs > 0 { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index e844a5cdb1..1a55f1cfc9 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -441,6 +441,32 @@ pub(crate) enum ForkPersistence { }, } +/// Normal appends keep this generation; compaction and rollback start a new one. +#[derive(Clone, Copy)] +pub(crate) struct CompactionPrefix { + pub(crate) len: usize, +} + +pub(crate) enum CompactedHistoryInstallation<'a> { + Complete(CompactedHistoryMetadata), + #[allow(dead_code)] + Prefix { + prefix: CompactionPrefix, + prefix_context: Vec, + initial_context: Vec, + model_info: &'a ModelInfo, + compaction_response_id: String, + reviewer_compaction_hash: Option, + started_item: (&'a TurnContext, &'a TurnItem), + }, +} + +impl From for CompactedHistoryInstallation<'_> { + fn from(metadata: CompactedHistoryMetadata) -> Self { + Self::Complete(metadata) + } +} + pub(crate) struct SessionSpawnArgs { pub(crate) startup: Option>, pub(crate) config: Config, @@ -4122,16 +4148,18 @@ impl Session { state.replace_history(items, reference_context_item); } - pub(crate) async fn replace_compacted_history( + pub(crate) async fn replace_compacted_history<'a>( &self, mut items: Vec, reference_context_item: TurnContextItem, world_state_baseline: WorldStateSnapshot, - metadata: CompactedHistoryMetadata, + installation: impl Into>, ) { - for envelope in &mut items { - Self::assign_missing_response_item_id(&mut envelope.item); - } + let installation = installation.into(); + let compaction_model_hash = match &installation { + CompactedHistoryInstallation::Complete(metadata) => &metadata.compaction_model_hash, + CompactedHistoryInstallation::Prefix { model_info, .. } => &model_info.comp_hash, + }; let mcp_revision = self .services .executed_tool_calls @@ -4151,14 +4179,63 @@ impl Session { checkpoint .metadata .get_or_insert_default() - .compaction_model_hash = metadata.compaction_model_hash; + .compaction_model_hash = compaction_model_hash.clone(); } // Wait for accepted updates to finish persisting, then keep later updates from // overtaking the current settings snapshot while its checkpoint is written. let _settings_guard = thread_settings::acquire_persistence_lock(self).await; let world_state_item = WorldStateItem::full(world_state_baseline.clone().into_object()); - let compacted_item = { + let (compacted_item, started_item) = { let mut state = self.state.lock().await; + let (metadata, started_item) = match installation { + CompactedHistoryInstallation::Complete(metadata) => (metadata, None), + CompactedHistoryInstallation::Prefix { + prefix, + prefix_context, + initial_context, + model_info, + compaction_response_id, + reviewer_compaction_hash, + started_item: (turn_context, started_item), + } => { + let current = state.history.annotated_items(); + let suffix = ¤t[prefix.len..]; + if suffix + .iter() + .any(|envelope| compact::is_real_user_or_agent_message(&envelope.item)) + { + items.extend(initial_context); + } else { + items = compact::insert_initial_context_before_last_real_user_or_summary( + items, + initial_context, + ); + } + items.splice( + 0..0, + prefix_context.into_iter().map(ResponseItemEnvelope::new), + ); + items.extend_from_slice(suffix); + let input_goal_ids = + crate::context::UserGoalUpdate::message_ids(state.history.raw_items()); + let (window_number, window_ids) = state.advance_auto_compact_window(); + ( + CompactedHistoryMetadata { + input_goal_ids, + message: String::new(), + window_number, + window_ids, + compaction_response_id: Some(compaction_response_id), + compaction_model_hash: model_info.comp_hash.clone(), + reviewer_compaction_hash, + }, + Some((turn_context, started_item)), + ) + } + }; + for envelope in &mut items { + Self::assign_missing_response_item_id(&mut envelope.item); + } // Goal edits are published outside the running task. Keep edits accepted after // the compaction input snapshot, in their original order, after its replacement. let replacement_goal_ids = crate::context::UserGoalUpdate::message_ids( @@ -4188,7 +4265,7 @@ impl Session { ); state.reasoning_effort_pin = ReasoningEffortPin::Compacted; state.history.set_world_state_baseline(world_state_baseline); - CompactedItem { + let compacted_item = CompactedItem { message: metadata.message, replacement_history: Some(replacement_history), guardian_history: state.history.guardian_history_checkpoint(), @@ -4209,9 +4286,14 @@ impl Session { turn_attribution: state.turn_attribution.clone(), previous_turn_settings: state.previous_turn_settings(), }), - } + }; + (compacted_item, started_item) }; + if let Some((turn_context, item)) = started_item { + // The prefix is now known to be valid; announce the install before checkpoint I/O. + self.emit_turn_item_started(turn_context, item).await; + } let rollout_items = [ RolloutItem::Compacted(compacted_item), RolloutItem::WorldState(world_state_item), @@ -4314,7 +4396,8 @@ impl Session { .collect() } - /// `step_context` and `world_state` must come from the same captured step. + /// Normally both inputs come from the same step. Async compaction uses the prefix's + /// world state so later state updates can stay in history unchanged. /// Returns rendered items with placement preserved for the caller's history layout. pub(crate) async fn build_initial_context_with_world_state( &self, @@ -4630,6 +4713,10 @@ impl Session { pub(crate) async fn current_window(&self) -> (String, u64, Uuid) { let state = self.state.lock().await; + self.current_window_from_state(&state) + } + + pub(crate) fn current_window_from_state(&self, state: &SessionState) -> (String, u64, Uuid) { let thread_id = self.thread_id; let window_number = state.auto_compact_window_number(); let context_window_id = state.auto_compact_window_ids().window_id; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 5cd89a8672..a617d9367b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -719,7 +719,19 @@ impl Session { turn_context: &TurnContext, compaction_metadata: CompactionTurnMetadata, ) -> CodexResponsesMetadata { - let (window_id, window_number, context_window_id) = self.current_window().await; + self.compaction_responses_metadata_for_window( + turn_context, + compaction_metadata, + self.current_window().await, + ) + } + + pub(crate) fn compaction_responses_metadata_for_window( + &self, + turn_context: &TurnContext, + compaction_metadata: CompactionTurnMetadata, + (window_id, window_number, context_window_id): (String, u64, uuid::Uuid), + ) -> CodexResponsesMetadata { let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( self.installation_id.clone(), window_id,