diff --git a/codex-rs/codex-api/src/endpoint/compact.rs b/codex-rs/codex-api/src/endpoint/compact.rs index 939bc7d85..c8730235f 100644 --- a/codex-rs/codex-api/src/endpoint/compact.rs +++ b/codex-rs/codex-api/src/endpoint/compact.rs @@ -11,8 +11,11 @@ use http::Method; use serde::Deserialize; use serde_json::to_value; use std::sync::Arc; +use std::sync::OnceLock; use std::time::Duration; +const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state"; + pub struct CompactClient { session: EndpointSession, } @@ -39,6 +42,7 @@ impl CompactClient { body: serde_json::Value, extra_headers: HeaderMap, request_timeout: Duration, + turn_state: Option<&OnceLock>, ) -> Result, ApiError> { let resp = self .session @@ -52,6 +56,14 @@ impl CompactClient { }, ) .await?; + if let Some(turn_state) = turn_state + && let Some(header_value) = resp + .headers + .get(X_CODEX_TURN_STATE_HEADER) + .and_then(|value| value.to_str().ok()) + { + let _ = turn_state.set(header_value.to_string()); + } let parsed: CompactHistoryResponse = serde_json::from_slice(&resp.body).map_err(|e| ApiError::Stream(e.to_string()))?; Ok(parsed.output) @@ -62,10 +74,12 @@ impl CompactClient { input: &CompactionInput<'_>, extra_headers: HeaderMap, request_timeout: Duration, + turn_state: Option<&OnceLock>, ) -> Result, ApiError> { let body = to_value(input) .map_err(|e| ApiError::Stream(format!("failed to encode compaction input: {e}")))?; - self.compact(body, extra_headers, request_timeout).await + self.compact(body, extra_headers, request_timeout, turn_state) + .await } } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 94ddc036a..f98b21299 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -433,6 +433,7 @@ impl ModelClient { &self, prompt: &Prompt, model_info: &ModelInfo, + turn_state: Option>>, settings: CompactConversationRequestSettings, session_telemetry: &SessionTelemetry, compaction_trace: &CompactionTraceContext, @@ -492,7 +493,7 @@ impl ModelClient { } extra_headers.extend(build_responses_headers( self.state.beta_features_header.as_deref(), - /*turn_state*/ None, + turn_state.as_ref(), )); extra_headers.extend(self.build_responses_compatibility_headers(responses_metadata)); extra_headers.extend(build_session_headers( @@ -512,7 +513,12 @@ impl ModelClient { .with_telemetry(Some(request_telemetry)); let trace_attempt = compaction_trace.start_attempt(&payload); let result = client - .compact_input(&payload, extra_headers, compact_request_timeout) + .compact_input( + &payload, + extra_headers, + compact_request_timeout, + turn_state.as_deref(), + ) .await .map_err(map_api_error); trace_attempt.record_result(result.as_deref()); @@ -925,6 +931,10 @@ impl Drop for ModelClientSession { } impl ModelClientSession { + pub(crate) fn turn_state(&self) -> Arc> { + Arc::clone(&self.turn_state) + } + fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; self.websocket_session.last_request = None; diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 0822d0f0a..a3da98cc8 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::OnceLock; use crate::Prompt; use crate::client::CompactConversationRequestSettings; @@ -43,6 +44,7 @@ const CONTEXT_WINDOW_TRUNCATED_OUTPUT_MESSAGE: &str = pub(crate) async fn run_inline_remote_auto_compact_task( sess: Arc, turn_context: Arc, + turn_state: Arc>, initial_context_injection: InitialContextInjection, reason: CompactionReason, phase: CompactionPhase, @@ -50,6 +52,7 @@ pub(crate) async fn run_inline_remote_auto_compact_task( run_remote_compact_task_inner( &sess, &turn_context, + Some(turn_state), initial_context_injection, CompactionTrigger::Auto, reason, @@ -75,6 +78,7 @@ pub(crate) async fn run_remote_compact_task( run_remote_compact_task_inner( &sess, &turn_context, + /*turn_state*/ None, InitialContextInjection::DoNotInject, CompactionTrigger::Manual, CompactionReason::UserRequested, @@ -87,6 +91,7 @@ pub(crate) async fn run_remote_compact_task( async fn run_remote_compact_task_inner( sess: &Arc, turn_context: &Arc, + turn_state: Option>>, initial_context_injection: InitialContextInjection, trigger: CompactionTrigger, reason: CompactionReason, @@ -130,6 +135,7 @@ async fn run_remote_compact_task_inner( let result = run_remote_compact_task_inner_impl( sess, turn_context, + turn_state, initial_context_injection, compaction_metadata, &mut analytics_details, @@ -163,6 +169,7 @@ async fn run_remote_compact_task_inner( async fn run_remote_compact_task_inner_impl( sess: &Arc, turn_context: &Arc, + turn_state: Option>>, initial_context_injection: InitialContextInjection, compaction_metadata: CompactionTurnMetadata, analytics_details: &mut CompactionAnalyticsDetails, @@ -237,6 +244,7 @@ async fn run_remote_compact_task_inner_impl( .compact_conversation_history( &prompt, &turn_context.model_info, + turn_state, CompactConversationRequestSettings { effort: turn_context.reasoning_effort.clone(), summary: turn_context.reasoning_summary, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 6244274a4..5857d8f32 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -922,6 +922,7 @@ async fn run_auto_compact( run_inline_remote_auto_compact_task( Arc::clone(sess), Arc::clone(turn_context), + client_session.turn_state(), initial_context_injection, reason, phase, diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 3345cd72f..781592875 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -131,6 +131,7 @@ fn canonical_json(value: &Value) -> Value { const PRETURN_CONTEXT_DIFF_CWD: &str = "/tmp/PRETURN_CONTEXT_DIFF_CWD"; const DUMMY_FUNCTION_NAME: &str = "test_tool"; +const TURN_STATE_HEADER: &str = "x-codex-turn-state"; const REMOTE_COMPACT_TURN_COMPLETE_TIMEOUT: Duration = Duration::from_secs(30); fn summary_with_prefix(summary: &str) -> String { @@ -3584,6 +3585,355 @@ async fn snapshot_request_shape_remote_pre_turn_compaction_context_window_exceed Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_pre_turn_compact_response_seeds_turn_state() -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = TestCodexHarness::with_builder( + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + config.model_auto_compact_token_limit = Some(200); + }), + ) + .await?; + let codex = harness.test().codex.clone(); + + let responses_mock = responses::mount_response_sequence( + harness.server(), + vec![ + responses::sse_response(responses::sse(vec![ + responses::ev_assistant_message("m1", "BEFORE_COMPACT_REPLY"), + responses::ev_completed_with_tokens("r1", /*total_tokens*/ 500), + ])), + responses::sse_response(responses::sse(vec![ + responses::ev_assistant_message("m2", "AFTER_COMPACT_REPLY"), + responses::ev_completed_with_tokens("r2", /*total_tokens*/ 80), + ])), + ], + ) + .await; + let compact_mock = responses::mount_compact_response_once( + harness.server(), + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .insert_header(TURN_STATE_HEADER, "compact-state") + .set_body_json(json!({ + "output": compacted_summary_only_output("PRE_TURN_COMPACT_SUMMARY"), + })), + ) + .await; + + // Phase 1: the first turn raises usage above the pre-turn compact threshold. + // Phase 2: the next turn compacts before sampling and establishes turn state. + for text in ["BEFORE_COMPACT_USER", "AFTER_COMPACT_USER"] { + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_turn_complete(&codex).await; + } + + // Phase 3: compact starts empty, and its returned state is sent to the first sample. + assert_eq!( + compact_mock.single_request().header(TURN_STATE_HEADER), + None + ); + let requests = responses_mock.requests(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].header(TURN_STATE_HEADER), None); + assert_eq!( + requests[1].header(TURN_STATE_HEADER).as_deref(), + Some("compact-state") + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_mid_turn_compact_v1_sends_turn_state_over_http() -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = TestCodexHarness::with_builder( + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + config.model_auto_compact_token_limit = Some(200); + }), + ) + .await?; + let codex = harness.test().codex.clone(); + let responses_mock = responses::mount_response_sequence( + harness.server(), + vec![ + responses::sse_response(responses::sse(vec![ + responses::ev_function_call("call-before-compact", DUMMY_FUNCTION_NAME, "{}"), + responses::ev_completed_with_tokens("r1", /*total_tokens*/ 500), + ])) + .insert_header(TURN_STATE_HEADER, "sampling-state"), + responses::sse_response(responses::sse(vec![ + responses::ev_function_call("call-after-compact", DUMMY_FUNCTION_NAME, "{}"), + responses::ev_completed_with_tokens("r2", /*total_tokens*/ 80), + ])) + .insert_header(TURN_STATE_HEADER, "continuation-state"), + responses::sse_response(responses::sse(vec![ + responses::ev_assistant_message("m1", "FINAL_REPLY"), + responses::ev_completed_with_tokens("r3", /*total_tokens*/ 80), + ])), + ], + ) + .await; + let compact_mock = responses::mount_compact_response_once( + harness.server(), + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .insert_header(TURN_STATE_HEADER, "compact-state") + .set_body_json(json!({ + "output": compacted_summary_only_output("MID_TURN_COMPACT_SUMMARY"), + })), + ) + .await; + + // Phase 1: sampling mints state and crosses the token limit with a pending tool follow-up. + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "RUN_WITH_MID_TURN_COMPACT".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_turn_complete(&codex).await; + + // Phase 2: v1 compact receives the state established by sampling. + let compact_request = compact_mock.single_request(); + assert_eq!(compact_request.path(), "/v1/responses/compact"); + assert_eq!( + compact_request.header(TURN_STATE_HEADER).as_deref(), + Some("sampling-state") + ); + + // Phase 3: every remaining request keeps replaying that first value. + let requests = responses_mock.requests(); + assert_eq!(requests.len(), 3); + assert_eq!(requests[0].header(TURN_STATE_HEADER), None); + assert_eq!( + requests[1].header(TURN_STATE_HEADER).as_deref(), + Some("sampling-state") + ); + assert_eq!( + requests[2].header(TURN_STATE_HEADER).as_deref(), + Some("sampling-state") + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_mid_turn_compact_v2_sends_turn_state_over_http() -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = TestCodexHarness::with_builder( + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + let _ = config.features.enable(Feature::RemoteCompactionV2); + config.model_auto_compact_token_limit = Some(200); + }), + ) + .await?; + let codex = harness.test().codex.clone(); + let responses_mock = responses::mount_response_sequence( + harness.server(), + vec![ + responses::sse_response(responses::sse(vec![ + responses::ev_function_call("call-before-compact", DUMMY_FUNCTION_NAME, "{}"), + responses::ev_completed_with_tokens("r1", /*total_tokens*/ 500), + ])) + .insert_header(TURN_STATE_HEADER, "sampling-state"), + responses::sse_response(responses::sse(vec![ + json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "V2_COMPACT_SUMMARY", + } + }), + responses::ev_completed("r-compact"), + ])) + .insert_header(TURN_STATE_HEADER, "compact-state"), + responses::sse_response(responses::sse(vec![ + responses::ev_function_call("call-after-compact", DUMMY_FUNCTION_NAME, "{}"), + responses::ev_completed_with_tokens("r2", /*total_tokens*/ 80), + ])) + .insert_header(TURN_STATE_HEADER, "continuation-state"), + responses::sse_response(responses::sse(vec![ + responses::ev_assistant_message("m1", "FINAL_REPLY"), + responses::ev_completed_with_tokens("r3", /*total_tokens*/ 80), + ])), + ], + ) + .await; + + // Phase 1: sampling mints state and schedules inline v2 compaction. + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "RUN_WITH_MID_TURN_COMPACT_V2".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_turn_complete(&codex).await; + + let requests = responses_mock.requests(); + assert_eq!(requests.len(), 4); + assert!( + requests + .iter() + .all(|request| request.path() == "/v1/responses") + ); + assert_eq!(requests[0].header(TURN_STATE_HEADER), None); + + // Phase 2: the v2 compaction request replays the state already established by sampling. + assert!( + requests[1] + .body_json() + .to_string() + .contains("\"type\":\"compaction_trigger\"") + ); + assert_eq!( + requests[1].header(TURN_STATE_HEADER).as_deref(), + Some("sampling-state") + ); + + // Phase 3: later response headers do not replace the first value in the OnceLock. + assert_eq!( + requests[2].header(TURN_STATE_HEADER).as_deref(), + Some("sampling-state") + ); + assert_eq!( + requests[3].header(TURN_STATE_HEADER).as_deref(), + Some("sampling-state") + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_mid_turn_compact_v2_sends_turn_state_over_websocket() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_websocket_server(vec![vec![ + vec![ + responses::ev_response_created("warm-1"), + responses::ev_completed("warm-1"), + ], + vec![ + json!({ + "type": "response.metadata", + "headers": {(TURN_STATE_HEADER): "sampling-state"}, + }), + responses::ev_function_call("call-before-compact", DUMMY_FUNCTION_NAME, "{}"), + responses::ev_completed_with_tokens("r1", /*total_tokens*/ 500), + ], + vec![ + json!({ + "type": "response.metadata", + "headers": {(TURN_STATE_HEADER): "compact-state"}, + }), + json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "V2_WS_COMPACT_SUMMARY", + } + }), + responses::ev_completed("r-compact"), + ], + vec![ + json!({ + "type": "response.metadata", + "headers": {(TURN_STATE_HEADER): "continuation-state"}, + }), + responses::ev_function_call("call-after-compact", DUMMY_FUNCTION_NAME, "{}"), + responses::ev_completed_with_tokens("r2", /*total_tokens*/ 80), + ], + vec![ + responses::ev_assistant_message("m1", "FINAL_REPLY"), + responses::ev_completed_with_tokens("r3", /*total_tokens*/ 80), + ], + ]]) + .await; + let mut builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + let _ = config.features.enable(Feature::RemoteCompactionV2); + config.model_auto_compact_token_limit = Some(200); + }); + let test = builder.build_with_websocket_server(&server).await?; + + // Phase 1: startup prewarm stays empty, then WebSocket sampling mints state and schedules + // inline v2 compaction. + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "RUN_WITH_WS_MID_TURN_COMPACT_V2".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_turn_complete(&test.codex).await; + + let requests = server.single_connection(); + assert_eq!(requests.len(), 5); + assert_eq!(requests[0].body_json()["generate"].as_bool(), Some(false)); + // Phase 2: the v2 compact request replays the state already established by sampling. + assert!( + requests[2] + .body_json() + .to_string() + .contains("\"type\":\"compaction_trigger\"") + ); + // Phase 3: both post-compact requests keep replaying that first value. + assert_eq!( + requests + .iter() + .map(|request| request.body_json()["client_metadata"][TURN_STATE_HEADER].clone()) + .collect::>(), + vec![ + json!(null), + json!(null), + json!("sampling-state"), + json!("sampling-state"), + json!("sampling-state"), + ] + ); + + server.shutdown().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn snapshot_request_shape_remote_mid_turn_continuation_compaction() -> Result<()> { skip_if_no_network!(Ok(()));