diff --git a/src/providers/codex/compaction.rs b/src/providers/codex/compaction.rs index 8597c85..2db3f36 100644 --- a/src/providers/codex/compaction.rs +++ b/src/providers/codex/compaction.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; use crate::anthropic::sse::parse_sse_events; @@ -33,17 +34,27 @@ impl std::fmt::Display for CompactionError { } enum CompactionPhase { + Preparing, Unconfirmed, Anchored { portable_summary: String }, } struct CompactionState { + attempt: CompactionAttempt, model: String, native_history: Vec, phase: CompactionPhase, updated_at: u64, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CompactionAttempt(u64); + +pub struct CompactionReplay { + pub request: ResponsesRequest, + pub attempt: CompactionAttempt, +} + #[derive(Default)] struct CompactionRegistry { states: HashMap, @@ -51,6 +62,7 @@ struct CompactionRegistry { } static REGISTRY: Mutex> = Mutex::new(None); +static NEXT_ATTEMPT_ID: AtomicU64 = AtomicU64::new(1); pub async fn request_compaction( client: &CodexHttpClient, @@ -78,41 +90,67 @@ pub async fn request_compaction( Ok(build_compacted_history(&conversation, compaction)) } -pub fn store_compaction( - session_id: &str, - model: &str, - native_history: Vec, -) -> bool { +pub fn begin_compaction(session_id: &str, model: &str) -> CompactionAttempt { + let attempt = CompactionAttempt(NEXT_ATTEMPT_ID.fetch_add(1, Ordering::Relaxed)); let state = CompactionState { + attempt, model: model.to_string(), - native_history, - phase: CompactionPhase::Unconfirmed, + native_history: Vec::new(), + phase: CompactionPhase::Preparing, updated_at: now_ms(), }; - if state_size(session_id, &state) > MAX_STATE_BYTES { - clear_compaction(session_id); - return false; - } - let now = state.updated_at; let mut guard = REGISTRY.lock().unwrap(); let registry = guard.get_or_insert_with(CompactionRegistry::default); evict_states(registry, now); registry.states.insert(session_id.to_string(), state); evict_states(registry, now); - registry.states.contains_key(session_id) + attempt +} + +pub fn store_compaction( + session_id: &str, + attempt: CompactionAttempt, + native_history: Vec, +) -> bool { + let now = now_ms(); + let mut guard = REGISTRY.lock().unwrap(); + let Some(registry) = guard.as_mut() else { + return false; + }; + evict_states(registry, now); + let Some(state) = registry.states.get_mut(session_id) else { + return false; + }; + if state.attempt != attempt || !matches!(state.phase, CompactionPhase::Preparing) { + return false; + } + state.native_history = native_history; + state.phase = CompactionPhase::Unconfirmed; + state.updated_at = now; + if state_size(session_id, state) > MAX_STATE_BYTES { + registry.states.remove(session_id); + update_total_bytes(registry); + return false; + } + evict_states(registry, now); + registry + .states + .get(session_id) + .is_some_and(|state| state.attempt == attempt) } pub fn activate_compaction( session_id: Option<&str>, + attempt: Option, model: &str, output: &[ResponsesInputItem], ) -> bool { - let Some(session_id) = session_id else { + let (Some(session_id), Some(attempt)) = (session_id, attempt) else { return false; }; let Some(portable_summary) = portable_summary_text(output) else { - clear_compaction(session_id); + abort_compaction_attempt(Some(session_id), Some(attempt)); return false; }; @@ -125,6 +163,9 @@ pub fn activate_compaction( let Some(state) = registry.states.get_mut(session_id) else { return false; }; + if state.attempt != attempt { + return false; + } if state.model != model || !matches!(state.phase, CompactionPhase::Unconfirmed) { registry.states.remove(session_id); update_total_bytes(registry); @@ -144,13 +185,16 @@ pub fn activate_compaction( pub fn apply_compaction_replay( session_id: Option<&str>, request: &ResponsesRequest, -) -> Option { +) -> Option { let session_id = session_id?; let now = now_ms(); let mut guard = REGISTRY.lock().unwrap(); let registry = guard.as_mut()?; evict_states(registry, now); let state = registry.states.get_mut(session_id)?; + if !matches!(state.phase, CompactionPhase::Anchored { .. }) { + return None; + } if state.model != request.model { registry.states.remove(session_id); update_total_bytes(registry); @@ -189,28 +233,30 @@ pub fn apply_compaction_replay( return None; } state.updated_at = now; - Some(replay) + Some(CompactionReplay { + request: replay, + attempt: state.attempt, + }) } -pub fn abort_compaction_attempt( - session_id: Option<&str>, - compact_boundary: bool, - request: &ResponsesRequest, -) { - if (compact_boundary || request_contains_compaction(request)) - && let Some(session_id) = session_id +pub fn abort_compaction_attempt(session_id: Option<&str>, attempt: Option) { + let (Some(session_id), Some(attempt)) = (session_id, attempt) else { + return; + }; + let mut guard = REGISTRY.lock().unwrap(); + let Some(registry) = guard.as_mut() else { + return; + }; + if registry + .states + .get(session_id) + .is_some_and(|state| state.attempt == attempt) { - clear_compaction(session_id); + registry.states.remove(session_id); + update_total_bytes(registry); } } -pub fn request_contains_compaction(request: &ResponsesRequest) -> bool { - request - .input - .iter() - .any(|item| matches!(item, ResponsesInputItem::Compaction { .. })) -} - pub fn clear_compaction(session_id: &str) { let mut guard = REGISTRY.lock().unwrap(); if let Some(registry) = guard.as_mut() { @@ -461,7 +507,7 @@ fn serialized_size(items: &[ResponsesInputItem]) -> usize { fn state_size(session_id: &str, state: &CompactionState) -> usize { let summary_len = match &state.phase { - CompactionPhase::Unconfirmed => 0, + CompactionPhase::Preparing | CompactionPhase::Unconfirmed => 0, CompactionPhase::Anchored { portable_summary } => portable_summary.len(), }; session_id.len() + state.model.len() + summary_len + serialized_size(&state.native_history) @@ -512,6 +558,8 @@ mod tests { const SUMMARY: &str = "portable summary with enough detail to identify this compacted conversation"; + const STALE_SUMMARY: &str = + "stale portable summary from an older overlapping compaction attempt"; static TEST_REGISTRY_LOCK: Mutex<()> = Mutex::new(()); fn request(input: serde_json::Value) -> ResponsesRequest { @@ -536,6 +584,15 @@ mod tests { .unwrap() } + fn stored_compaction( + session_id: &str, + native_history: Vec, + ) -> CompactionAttempt { + let attempt = begin_compaction(session_id, "gpt-5.6-sol"); + assert!(store_compaction(session_id, attempt, native_history)); + attempt + } + #[test] fn parses_exactly_one_completed_compaction_item() { let body = b"data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"compaction\",\"encrypted_content\":\"opaque\"}}\n\ndata: {\"type\":\"response.completed\",\"response\":{}}\n\n"; @@ -567,9 +624,8 @@ mod tests { fn replay_requires_activation_and_wrapped_summary_anchor() { let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); clear_all_compactions_for_tests(); - store_compaction( + let attempt = stored_compaction( "session", - "gpt-5.6-sol", vec![ResponsesInputItem::Compaction { encrypted_content: "opaque".to_string(), }], @@ -583,13 +639,16 @@ mod tests { assert!(apply_compaction_replay(Some("session"), &next).is_none()); assert!(activate_compaction( Some("session"), + Some(attempt), "gpt-5.6-sol", &output(&format!( "summary preparation\n\n{SUMMARY}\n" )) )); - let replay = apply_compaction_replay(Some("session"), &next).unwrap(); + let replay = apply_compaction_replay(Some("session"), &next) + .unwrap() + .request; assert!(matches!( replay.input[0], ResponsesInputItem::AdditionalTools { .. } @@ -604,6 +663,195 @@ mod tests { assert_eq!(replay.client_metadata, next.client_metadata); } + #[test] + fn stale_activation_cannot_anchor_newer_native_history() { + let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); + clear_all_compactions_for_tests(); + let older = stored_compaction( + "session", + vec![ResponsesInputItem::Compaction { + encrypted_content: "older-native-history".to_string(), + }], + ); + let newer = stored_compaction( + "session", + vec![ResponsesInputItem::Compaction { + encrypted_content: "newer-native-history".to_string(), + }], + ); + + assert!(!activate_compaction( + Some("session"), + Some(older), + "gpt-5.6-sol", + &output(STALE_SUMMARY), + )); + assert!(activate_compaction( + Some("session"), + Some(newer), + "gpt-5.6-sol", + &output(SUMMARY), + )); + } + + #[test] + fn stale_store_cannot_replace_newer_compaction() { + let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); + clear_all_compactions_for_tests(); + let older = begin_compaction("session", "gpt-5.6-sol"); + let newer = stored_compaction( + "session", + vec![ResponsesInputItem::Compaction { + encrypted_content: "newer-native-history".to_string(), + }], + ); + + assert!(!store_compaction( + "session", + older, + vec![ResponsesInputItem::Compaction { + encrypted_content: "late-older-native-history".to_string(), + }], + )); + assert!(activate_compaction( + Some("session"), + Some(newer), + "gpt-5.6-sol", + &output(SUMMARY), + )); + let next = request(json!([ + {"type":"message","role":"user","content":[{"type":"input_text","text":SUMMARY}]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]} + ])); + let replay = apply_compaction_replay(Some("session"), &next) + .unwrap() + .request; + assert!(replay.input.iter().any(|item| matches!( + item, + ResponsesInputItem::Compaction { encrypted_content } + if encrypted_content == "newer-native-history" + ))); + } + + #[test] + fn preparing_compaction_survives_model_mismatched_replay_check() { + let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); + clear_all_compactions_for_tests(); + let attempt = begin_compaction("session", "gpt-5.6-sol"); + let mut mismatched = request(json!([ + {"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]} + ])); + mismatched.model = "gpt-5.6-terra".to_string(); + + assert!(apply_compaction_replay(Some("session"), &mismatched).is_none()); + assert!(store_compaction( + "session", + attempt, + vec![ResponsesInputItem::Compaction { + encrypted_content: "native-history".to_string(), + }], + )); + } + + #[test] + fn stale_abort_cannot_clear_newer_compaction() { + let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); + clear_all_compactions_for_tests(); + let older = begin_compaction("session", "gpt-5.6-sol"); + let newer = stored_compaction( + "session", + vec![ResponsesInputItem::Compaction { + encrypted_content: "newer-native-history".to_string(), + }], + ); + + abort_compaction_attempt(Some("session"), Some(older)); + + assert!(activate_compaction( + Some("session"), + Some(newer), + "gpt-5.6-sol", + &output(SUMMARY), + )); + } + + #[test] + fn stale_replay_abort_cannot_clear_newer_compaction() { + let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); + clear_all_compactions_for_tests(); + let older = stored_compaction( + "session", + vec![ResponsesInputItem::Compaction { + encrypted_content: "older-native-history".to_string(), + }], + ); + assert!(activate_compaction( + Some("session"), + Some(older), + "gpt-5.6-sol", + &output(STALE_SUMMARY), + )); + let older_next = request(json!([ + {"type":"message","role":"user","content":[{"type":"input_text","text":STALE_SUMMARY}]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"continue old"}]} + ])); + let older_replay = apply_compaction_replay(Some("session"), &older_next).unwrap(); + + let newer = stored_compaction( + "session", + vec![ResponsesInputItem::Compaction { + encrypted_content: "newer-native-history".to_string(), + }], + ); + assert!(activate_compaction( + Some("session"), + Some(newer), + "gpt-5.6-sol", + &output(SUMMARY), + )); + + abort_compaction_attempt(Some("session"), Some(older_replay.attempt)); + + let newer_next = request(json!([ + {"type":"message","role":"user","content":[{"type":"input_text","text":SUMMARY}]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"continue new"}]} + ])); + let replay = apply_compaction_replay(Some("session"), &newer_next) + .unwrap() + .request; + assert!(replay.input.iter().any(|item| matches!( + item, + ResponsesInputItem::Compaction { encrypted_content } + if encrypted_content == "newer-native-history" + ))); + } + + #[test] + fn invalid_stale_summary_cannot_clear_newer_compaction() { + let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); + clear_all_compactions_for_tests(); + let older = begin_compaction("session", "gpt-5.6-sol"); + let newer = stored_compaction( + "session", + vec![ResponsesInputItem::Compaction { + encrypted_content: "newer-native-history".to_string(), + }], + ); + + assert!(!activate_compaction( + Some("session"), + Some(older), + "gpt-5.6-sol", + &output("too short"), + )); + assert!(activate_compaction( + Some("session"), + Some(newer), + "gpt-5.6-sol", + &output(SUMMARY), + )); + } + #[test] fn replay_clears_on_missing_or_duplicate_anchor() { let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); @@ -612,14 +860,18 @@ mod tests { format!("{SUMMARY} and {SUMMARY}"), ] { clear_all_compactions_for_tests(); - store_compaction( + let attempt = stored_compaction( "session", - "gpt-5.6-sol", vec![ResponsesInputItem::Compaction { encrypted_content: "opaque".to_string(), }], ); - activate_compaction(Some("session"), "gpt-5.6-sol", &output(SUMMARY)); + activate_compaction( + Some("session"), + Some(attempt), + "gpt-5.6-sol", + &output(SUMMARY), + ); let changed = request(json!([ {"type":"message","role":"user","content":[{"type":"input_text","text":text}]}, {"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]} @@ -653,21 +905,25 @@ mod tests { fn failed_replay_clears_anchored_state() { let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); clear_all_compactions_for_tests(); - store_compaction( + let attempt = stored_compaction( "failed-replay", - "gpt-5.6-sol", vec![ResponsesInputItem::Compaction { encrypted_content: "opaque".to_string(), }], ); - activate_compaction(Some("failed-replay"), "gpt-5.6-sol", &output(SUMMARY)); + activate_compaction( + Some("failed-replay"), + Some(attempt), + "gpt-5.6-sol", + &output(SUMMARY), + ); let next = request(json!([ {"type":"message","role":"user","content":[{"type":"input_text","text":SUMMARY}]}, {"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]} ])); let replay = apply_compaction_replay(Some("failed-replay"), &next).unwrap(); - abort_compaction_attempt(Some("failed-replay"), false, &replay); + abort_compaction_attempt(Some("failed-replay"), Some(replay.attempt)); assert!(apply_compaction_replay(Some("failed-replay"), &next).is_none()); } @@ -676,14 +932,18 @@ mod tests { fn replay_clears_on_model_change() { let _guard = TEST_REGISTRY_LOCK.lock().unwrap(); clear_all_compactions_for_tests(); - store_compaction( + let attempt = stored_compaction( "session", - "gpt-5.6-sol", vec![ResponsesInputItem::Compaction { encrypted_content: "opaque".to_string(), }], ); - activate_compaction(Some("session"), "gpt-5.6-sol", &output(SUMMARY)); + activate_compaction( + Some("session"), + Some(attempt), + "gpt-5.6-sol", + &output(SUMMARY), + ); let mut changed = request(json!([ {"type":"message","role":"user","content":[{"type":"input_text","text":SUMMARY}]}, {"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]} diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index c3b1aba..b98447d 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -37,8 +37,8 @@ use self::auth::manager::CodexAuthManager; use self::auth::token_store::file_store; use self::client::CodexHttpClient; use self::compaction::{ - abort_compaction_attempt, activate_compaction, apply_compaction_replay, request_compaction, - store_compaction, + CompactionAttempt, abort_compaction_attempt, activate_compaction, apply_compaction_replay, + begin_compaction, request_compaction, store_compaction, }; use self::continuation::{ ContinuationCandidate, abort_continuation, continuation_candidate, record_continuation, @@ -231,6 +231,7 @@ impl Provider for CodexProvider { let compact_boundary = is_compact_messages_request(&body); let server_compaction_enabled = config::codex_server_compaction(); + let mut compaction_attempt = None; if !server_compaction_enabled && let Some(session_id) = ctx.session_id.as_deref() { compaction::clear_compaction(session_id); } @@ -238,7 +239,8 @@ impl Provider for CodexProvider { && compact_boundary && let Some(session_id) = ctx.session_id.as_deref() { - compaction::clear_compaction(session_id); + let attempt = begin_compaction(session_id, &translated.model); + compaction_attempt = Some(attempt); log_compaction_event( "server_compaction_triggered", &ctx, @@ -252,7 +254,7 @@ impl Provider for CodexProvider { compaction_ctx.monitor = None; match request_compaction(self.client.as_ref(), &translated, &compaction_ctx).await { Ok(native_history) => { - if store_compaction(session_id, &translated.model, native_history) { + if store_compaction(session_id, attempt, native_history) { log_compaction_event( "server_compaction_completed", &ctx, @@ -264,12 +266,12 @@ impl Provider for CodexProvider { "server_compaction_failed", &ctx, translated.input.len(), - Some("compaction state exceeded the in-memory limit"), + Some("compaction state was superseded or exceeded the in-memory limit"), ); } } Err(error) => { - compaction::clear_compaction(session_id); + abort_compaction_attempt(Some(session_id), Some(attempt)); log_compaction_event( "server_compaction_failed", &ctx, @@ -282,7 +284,8 @@ impl Provider for CodexProvider { && !compact_boundary && let Some(replay) = apply_compaction_replay(ctx.session_id.as_deref(), &translated) { - translated = replay; + translated = replay.request; + compaction_attempt = Some(replay.attempt); } // Check continuation @@ -308,7 +311,10 @@ impl Provider for CodexProvider { ctx, stream_request, continuation, - compact_boundary, + LiveStreamCompaction { + compact_boundary, + attempt: compaction_attempt, + }, ) .await; } @@ -322,11 +328,7 @@ impl Provider for CodexProvider { { Ok(r) => r, Err(e) => { - abort_compaction_attempt( - ctx.session_id.as_deref(), - compact_boundary, - &translated, - ); + abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); abort_continuation(ctx.session_id.as_deref(), turn_id); return map_codex_error_to_response(&e); } @@ -339,13 +341,13 @@ impl Provider for CodexProvider { let error = empty_buffered_completion_error(); drop_live_continuation_for_retry(&mut continuation); if attempt >= MAX_EMPTY_COMPLETION_RETRIES { - abort_compaction_attempt(ctx.session_id.as_deref(), compact_boundary, &translated); + abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); abort_continuation(ctx.session_id.as_deref(), turn_id); return map_codex_error_to_response(&error); } let delay = compute_backoff_delay(attempt, None); if delay.exceeds_budget { - abort_compaction_attempt(ctx.session_id.as_deref(), compact_boundary, &translated); + abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); abort_continuation(ctx.session_id.as_deref(), turn_id); return map_codex_error_to_response(&error); } @@ -364,11 +366,7 @@ impl Provider for CodexProvider { ) { Ok(b) => b, Err(e) => { - abort_compaction_attempt( - ctx.session_id.as_deref(), - compact_boundary, - &translated, - ); + abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); abort_continuation(ctx.session_id.as_deref(), turn_id); return map_codex_failure_to_response(&format!( "Stream translation error: {e}" @@ -388,6 +386,7 @@ impl Provider for CodexProvider { update_continuation_from_upstream( ctx.session_id.as_deref(), turn_id, + compaction_attempt, &translated, &upstream.body, compact_boundary, @@ -418,6 +417,7 @@ impl Provider for CodexProvider { update_continuation_from_upstream( ctx.session_id.as_deref(), turn_id, + compaction_attempt, &translated, &upstream.body, compact_boundary, @@ -425,11 +425,7 @@ impl Provider for CodexProvider { (StatusCode::OK, Json(json)).into_response() } Err(e) => { - abort_compaction_attempt( - ctx.session_id.as_deref(), - compact_boundary, - &translated, - ); + abort_compaction_attempt(ctx.session_id.as_deref(), compaction_attempt); abort_continuation(ctx.session_id.as_deref(), turn_id); map_codex_failure_to_response(&format!("Accumulation error: {e}")) } @@ -525,10 +521,9 @@ fn log_compaction_event( fn abort_request_state( session_id: Option<&str>, turn_id: Option, - compact_boundary: bool, - request: &translate::request::ResponsesRequest, + compaction_attempt: Option, ) { - abort_compaction_attempt(session_id, compact_boundary, request); + abort_compaction_attempt(session_id, compaction_attempt); abort_continuation(session_id, turn_id); } @@ -537,6 +532,12 @@ enum LiveStreamStart { Retry { error: client::CodexError }, } +#[derive(Clone, Copy)] +struct LiveStreamCompaction { + compact_boundary: bool, + attempt: Option, +} + async fn live_stream_response( client: Arc, message_id: String, @@ -544,7 +545,7 @@ async fn live_stream_response( ctx: RequestContext, request_body: translate::request::ResponsesRequest, continuation: ContinuationCandidate, - compact_boundary: bool, + compaction: LiveStreamCompaction, ) -> Response { let model = model.to_string(); let turn_id = continuation.turn_id; @@ -564,22 +565,12 @@ async fn live_stream_response( continue; } if attempt >= MAX_RETRYABLE_LIVE_STREAM_RETRIES { - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); return map_codex_error_to_response(&err); } let delay = compute_backoff_delay(attempt, err.retry_after.as_deref()); if delay.exceeds_budget { - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); return map_codex_error_to_response(&err); } attempt += 1; @@ -587,12 +578,7 @@ async fn live_stream_response( continue; } Err(err) => { - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); return map_codex_error_to_response(&err); } }; @@ -604,7 +590,7 @@ async fn live_stream_response( ctx.clone(), turn_id, request_body.clone(), - compact_boundary, + compaction, ) .await { @@ -616,22 +602,12 @@ async fn live_stream_response( continue; } if attempt >= MAX_RETRYABLE_LIVE_STREAM_RETRIES { - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); return map_codex_error_to_response(&error); } let delay = compute_backoff_delay(attempt, error.retry_after.as_deref()); if delay.exceeds_budget { - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); return map_codex_error_to_response(&error); } attempt += 1; @@ -648,7 +624,7 @@ async fn live_stream_response_once( ctx: RequestContext, turn_id: Option, request_body: translate::request::ResponsesRequest, - compact_boundary: bool, + compaction: LiveStreamCompaction, ) -> LiveStreamStart { let estimated_input_tokens = count_translated_tokens(&request_body); let mut translator = LiveStreamTranslator::with_estimated_input_tokens( @@ -669,12 +645,7 @@ async fn live_stream_response_once( if retryable_live_start_codex_error(&err) { return LiveStreamStart::Retry { error: err }; } - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); return LiveStreamStart::Response(map_codex_error_to_response(&err)); } }; @@ -722,12 +693,7 @@ async fn live_stream_response_once( }, }; } - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); return LiveStreamStart::Response(map_codex_failure_to_response(&message)); } }; @@ -747,9 +713,10 @@ async fn live_stream_response_once( update_continuation_from_upstream( ctx.session_id.as_deref(), turn_id, + compaction.attempt, &request_body, &upstream_sse_body, - compact_boundary, + compaction.compact_boundary, ); return LiveStreamStart::Response(single_live_stream_response(pending_chunk)); } @@ -761,16 +728,17 @@ async fn live_stream_response_once( turn_id, request_body, upstream_sse_body, - compact_boundary, + compaction, )); } if terminal { update_continuation_from_upstream( ctx.session_id.as_deref(), turn_id, + compaction.attempt, &request_body, &upstream_sse_body, - compact_boundary, + compaction.compact_boundary, ); if pending_chunk.is_empty() { return LiveStreamStart::Response(empty_live_stream_response()); @@ -869,17 +837,12 @@ fn remaining_live_stream_response( turn_id: Option, request_body: translate::request::ResponsesRequest, mut upstream_sse_body: Vec, - compact_boundary: bool, + compaction: LiveStreamCompaction, ) -> Response { let (tx, rx) = tokio::sync::mpsc::channel::>(64); tokio::spawn(async move { if tx.send(Ok(Bytes::from(first_chunk))).await.is_err() { - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); return; } while let Some(item) = upstream_events.recv().await { @@ -896,8 +859,7 @@ fn remaining_live_stream_response( abort_request_state( ctx.session_id.as_deref(), turn_id, - compact_boundary, - &request_body, + compaction.attempt, ); let chunk = translator.error_chunk( &message, @@ -917,8 +879,7 @@ fn remaining_live_stream_response( abort_request_state( ctx.session_id.as_deref(), turn_id, - compact_boundary, - &request_body, + compaction.attempt, ); return; } @@ -927,20 +888,16 @@ fn remaining_live_stream_response( update_continuation_from_upstream( ctx.session_id.as_deref(), turn_id, + compaction.attempt, &request_body, &upstream_sse_body, - compact_boundary, + compaction.compact_boundary, ); return; } } Err(err) => { - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); let chunk = translator.finish_after_closed_completed_tool_call(ctx.traffic.as_deref()); if !chunk.is_empty() { @@ -963,12 +920,7 @@ fn remaining_live_stream_response( } } - abort_request_state( - ctx.session_id.as_deref(), - turn_id, - compact_boundary, - &request_body, - ); + abort_request_state(ctx.session_id.as_deref(), turn_id, compaction.attempt); let chunk = translator.finish_after_closed_completed_tool_call(ctx.traffic.as_deref()); if !chunk.is_empty() { record_live_stream_progress(&ctx, &chunk); @@ -1144,6 +1096,7 @@ fn codex_stream_error_type(err: &client::CodexError) -> &'static str { fn update_continuation_from_upstream( session_id: Option<&str>, turn_id: Option, + compaction_attempt: Option, request_body: &translate::request::ResponsesRequest, upstream_body: &[u8], compact_boundary: bool, @@ -1151,7 +1104,12 @@ fn update_continuation_from_upstream( match finish_metadata_from_upstream(upstream_body) { Ok(Some(finish)) if finish.continuation_eligible => { if compact_boundary { - activate_compaction(session_id, &request_body.model, &finish.output_items); + activate_compaction( + session_id, + compaction_attempt, + &request_body.model, + &finish.output_items, + ); } record_continuation( session_id, @@ -1162,7 +1120,7 @@ fn update_continuation_from_upstream( ); } _ => { - abort_compaction_attempt(session_id, compact_boundary, request_body); + abort_compaction_attempt(session_id, compaction_attempt); abort_continuation(session_id, turn_id); } }