agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[codex] Send turn state through compact requests (#28002)

Browse Source 
## Context

Inline compaction is part of the active logical turn. Compact requests
and the sampling requests around them should use the same turn state,
including when compaction is the first request to establish it.

## Change

Pass the turn-scoped `OnceLock` directly to inline v1 compaction so
`/responses/compact` includes an established value in the existing HTTP
header. Capture `x-codex-turn-state` from the compact response into that
same lock, allowing pre-turn compact to establish the value that
subsequent sampling reuses.

V2 compact already uses the normal Responses HTTP/WebSocket path and
continues to share the same `OnceLock` without separate plumbing. The
first returned value wins for the logical turn.

## Test plan

Integration coverage verifies that:

- pre-turn v1 compact can establish state for the first sampling request
- inline v1 compact receives established state over HTTP
- inline v2 compact reuses established state over HTTP
- inline v2 compact reuses established state over WebSocket

CI validates the full change.
This commit is contained in:
Ahmed Ibrahim
2026-06-13 01:27:58 -07:00
committed by GitHub
Unverified
parent 640d61b121
commit f297b9f07d
5 changed files with 386 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-api/src/endpoint/compact.rs
+15 -1
View File
@@ -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<T: HttpTransport> {
session: EndpointSession<T>,
}
@@ -39,6 +42,7 @@ impl<T: HttpTransport> CompactClient<T> {
body: serde_json::Value,
extra_headers: HeaderMap,
request_timeout: Duration,
turn_state: Option<&OnceLock<String>>,
) -> Result<Vec<ResponseItem>, ApiError> {
let resp = self
.session
@@ -52,6 +56,14 @@ impl<T: HttpTransport> CompactClient<T> {
},
)
.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<T: HttpTransport> CompactClient<T> {
input: &CompactionInput<'_>,
extra_headers: HeaderMap,
request_timeout: Duration,
turn_state: Option<&OnceLock<String>>,
) -> Result<Vec<ResponseItem>, 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
}
}
codex-rs/core/src/client.rs
+12 -2
View File
@@ -433,6 +433,7 @@ impl ModelClient {
&self,
prompt: &Prompt,
model_info: &ModelInfo,
turn_state: Option<Arc<OnceLock<String>>>,
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<OnceLock<String>> {
Arc::clone(&self.turn_state)
}
fn reset_websocket_session(&mut self) {
self.websocket_session.connection = None;
self.websocket_session.last_request = None;
codex-rs/core/src/compact_remote.rs
+8
View File
@@ -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<Session>,
turn_context: Arc<TurnContext>,
turn_state: Arc<OnceLock<String>>,
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<Session>,
turn_context: &Arc<TurnContext>,
turn_state: Option<Arc<OnceLock<String>>>,
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<Session>,
turn_context: &Arc<TurnContext>,
turn_state: Option<Arc<OnceLock<String>>>,
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,
codex-rs/core/src/session/turn.rs
+1
View File
@@ -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,
codex-rs/core/tests/suite/compact_remote.rs
+350
View File
@@ -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<_>>(),
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(()));