[codex] add context-window lineage headers (#16758)

This change adds client-owned context-window and parent thread id
headers to all requests to responses api.
This commit is contained in:
Thibault Sottiaux
2026-04-03 19:54:31 -10:00
committed by GitHub
Unverified
parent 39097ab65d
commit 9e19004bc2
7 changed files with 338 additions and 37 deletions
+111 -27
View File
@@ -28,6 +28,7 @@ use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use codex_api::CompactClient as ApiCompactClient;
@@ -119,6 +120,9 @@ use codex_response_debug_context::telemetry_transport_error_message;
pub const OPENAI_BETA_HEADER: &str = "OpenAI-Beta";
pub const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state";
pub const X_CODEX_TURN_METADATA_HEADER: &str = "x-codex-turn-metadata";
pub const X_CODEX_PARENT_THREAD_ID_HEADER: &str = "x-codex-parent-thread-id";
pub const X_CODEX_WINDOW_ID_HEADER: &str = "x-codex-window-id";
pub const X_OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent";
pub const X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER: &str =
"x-responsesapi-include-timing-metrics";
const RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06";
@@ -137,6 +141,7 @@ pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration =
struct ModelClientState {
auth_manager: Option<Arc<AuthManager>>,
conversation_id: ThreadId,
window_generation: AtomicU64,
provider: ModelProviderInfo,
auth_env_telemetry: AuthEnvTelemetry,
session_source: SessionSource,
@@ -274,6 +279,7 @@ impl ModelClient {
state: Arc::new(ModelClientState {
auth_manager,
conversation_id,
window_generation: AtomicU64::new(0),
provider,
auth_env_telemetry,
session_source,
@@ -303,6 +309,24 @@ impl ModelClient {
self.state.auth_manager.clone()
}
pub(crate) fn set_window_generation(&self, window_generation: u64) {
self.state
.window_generation
.store(window_generation, Ordering::Relaxed);
self.store_cached_websocket_session(WebsocketSession::default());
}
pub(crate) fn advance_window_generation(&self) {
self.state.window_generation.fetch_add(1, Ordering::Relaxed);
self.store_cached_websocket_session(WebsocketSession::default());
}
fn current_window_id(&self) -> String {
let conversation_id = self.state.conversation_id;
let window_generation = self.state.window_generation.load(Ordering::Relaxed);
format!("{conversation_id}:{window_generation}")
}
fn take_cached_websocket_session(&self) -> WebsocketSession {
let mut cached_websocket_session = self
.state
@@ -401,7 +425,7 @@ impl ModelClient {
text,
};
let mut extra_headers = self.build_subagent_headers();
let mut extra_headers = self.build_responses_identity_headers();
extra_headers.extend(build_conversation_headers(Some(
self.state.conversation_id.to_string(),
)));
@@ -461,21 +485,56 @@ impl ModelClient {
fn build_subagent_headers(&self) -> ApiHeaderMap {
let mut extra_headers = ApiHeaderMap::new();
if let SessionSource::SubAgent(sub) = &self.state.session_source {
let subagent = match sub {
SubAgentSource::Review => "review".to_string(),
SubAgentSource::Compact => "compact".to_string(),
SubAgentSource::MemoryConsolidation => "memory_consolidation".to_string(),
SubAgentSource::ThreadSpawn { .. } => "collab_spawn".to_string(),
SubAgentSource::Other(label) => label.clone(),
};
if let Ok(val) = HeaderValue::from_str(&subagent) {
extra_headers.insert("x-openai-subagent", val);
}
if let Some(subagent) = subagent_header_value(&self.state.session_source)
&& let Ok(val) = HeaderValue::from_str(&subagent)
{
extra_headers.insert(X_OPENAI_SUBAGENT_HEADER, val);
}
extra_headers
}
fn build_responses_identity_headers(&self) -> ApiHeaderMap {
let mut extra_headers = self.build_subagent_headers();
if let Some(parent_thread_id) = parent_thread_id_header_value(&self.state.session_source)
&& let Ok(val) = HeaderValue::from_str(&parent_thread_id)
{
extra_headers.insert(X_CODEX_PARENT_THREAD_ID_HEADER, val);
}
if let Ok(val) = HeaderValue::from_str(&self.current_window_id()) {
extra_headers.insert(X_CODEX_WINDOW_ID_HEADER, val);
}
extra_headers
}
fn build_ws_client_metadata(
&self,
turn_metadata_header: Option<&str>,
) -> HashMap<String, String> {
let mut client_metadata = HashMap::new();
client_metadata.insert(
X_CODEX_WINDOW_ID_HEADER.to_string(),
self.current_window_id(),
);
if let Some(subagent) = subagent_header_value(&self.state.session_source) {
client_metadata.insert(X_OPENAI_SUBAGENT_HEADER.to_string(), subagent);
}
if let Some(parent_thread_id) = parent_thread_id_header_value(&self.state.session_source) {
client_metadata.insert(
X_CODEX_PARENT_THREAD_ID_HEADER.to_string(),
parent_thread_id,
);
}
if let Some(turn_metadata_header) = parse_turn_metadata_header(turn_metadata_header)
&& let Ok(turn_metadata) = turn_metadata_header.to_str()
{
client_metadata.insert(
X_CODEX_TURN_METADATA_HEADER.to_string(),
turn_metadata.to_string(),
);
}
client_metadata
}
/// Builds request telemetry for unary API calls (e.g., Compact endpoint).
fn build_request_telemetry(
session_telemetry: &SessionTelemetry,
@@ -655,6 +714,7 @@ impl ModelClient {
headers.insert("x-client-request-id", header_value);
}
headers.extend(build_conversation_headers(Some(conversation_id)));
headers.extend(self.build_responses_identity_headers());
headers.insert(
OPENAI_BETA_HEADER,
HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE),
@@ -678,7 +738,7 @@ impl Drop for ModelClientSession {
}
impl ModelClientSession {
fn reset_websocket_session(&mut self) {
pub(crate) fn reset_websocket_session(&mut self) {
self.websocket_session.connection = None;
self.websocket_session.last_request = None;
self.websocket_session.last_response_rx = None;
@@ -769,11 +829,15 @@ impl ModelClientSession {
ApiResponsesOptions {
conversation_id: Some(conversation_id),
session_source: Some(self.client.state.session_source.clone()),
extra_headers: build_responses_headers(
self.client.state.beta_features_header.as_deref(),
Some(&self.turn_state),
turn_metadata_header.as_ref(),
),
extra_headers: {
let mut headers = build_responses_headers(
self.client.state.beta_features_header.as_deref(),
Some(&self.turn_state),
turn_metadata_header.as_ref(),
);
headers.extend(self.client.build_responses_identity_headers());
headers
},
compression,
turn_state: Some(Arc::clone(&self.turn_state)),
}
@@ -1137,7 +1201,7 @@ impl ModelClientSession {
)?;
let mut ws_payload = ResponseCreateWsRequest {
client_metadata: response_create_client_metadata(
build_ws_client_metadata(turn_metadata_header),
Some(self.client.build_ws_client_metadata(turn_metadata_header)),
request_trace.as_ref(),
),
..ResponseCreateWsRequest::from(&request)
@@ -1370,14 +1434,6 @@ fn parse_turn_metadata_header(turn_metadata_header: Option<&str>) -> Option<Head
turn_metadata_header.and_then(|value| HeaderValue::from_str(value).ok())
}
fn build_ws_client_metadata(turn_metadata_header: Option<&str>) -> Option<HashMap<String, String>> {
let turn_metadata_header = parse_turn_metadata_header(turn_metadata_header)?;
let turn_metadata = turn_metadata_header.to_str().ok()?.to_string();
let mut client_metadata = HashMap::new();
client_metadata.insert(X_CODEX_TURN_METADATA_HEADER.to_string(), turn_metadata);
Some(client_metadata)
}
/// Builds the extra headers attached to Responses API requests.
///
/// These headers implement Codex-specific conventions:
@@ -1409,6 +1465,34 @@ fn build_responses_headers(
headers
}
fn subagent_header_value(session_source: &SessionSource) -> Option<String> {
let SessionSource::SubAgent(subagent_source) = session_source else {
return None;
};
match subagent_source {
SubAgentSource::Review => Some("review".to_string()),
SubAgentSource::Compact => Some("compact".to_string()),
SubAgentSource::MemoryConsolidation => Some("memory_consolidation".to_string()),
SubAgentSource::ThreadSpawn { .. } => Some("collab_spawn".to_string()),
SubAgentSource::Other(label) => Some(label.clone()),
}
}
fn parent_thread_id_header_value(session_source: &SessionSource) -> Option<String> {
match session_source {
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id, ..
}) => Some(parent_thread_id.to_string()),
SessionSource::Cli
| SessionSource::VSCode
| SessionSource::Exec
| SessionSource::Mcp
| SessionSource::Custom(_)
| SessionSource::SubAgent(_)
| SessionSource::Unknown => None,
}
}
fn map_response_stream<S>(
api_stream: S,
session_telemetry: SessionTelemetry,
+43 -1
View File
@@ -2,6 +2,10 @@ use super::AuthRequestTelemetryContext;
use super::ModelClient;
use super::PendingUnauthorizedRetry;
use super::UnauthorizedRecoveryExecution;
use super::X_CODEX_PARENT_THREAD_ID_HEADER;
use super::X_CODEX_TURN_METADATA_HEADER;
use super::X_CODEX_WINDOW_ID_HEADER;
use super::X_OPENAI_SUBAGENT_HEADER;
use codex_api::api_bridge::CoreAuthProvider;
use codex_app_server_protocol::AuthMode;
use codex_model_provider_info::WireApi;
@@ -80,11 +84,49 @@ fn build_subagent_headers_sets_other_subagent_label() {
)));
let headers = client.build_subagent_headers();
let value = headers
.get("x-openai-subagent")
.get(X_OPENAI_SUBAGENT_HEADER)
.and_then(|value| value.to_str().ok());
assert_eq!(value, Some("memory_consolidation"));
}
#[test]
fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() {
let parent_thread_id = ThreadId::new();
let client = test_model_client(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: 2,
agent_path: None,
agent_nickname: None,
agent_role: None,
}));
client.advance_window_generation();
let client_metadata = client.build_ws_client_metadata(Some(r#"{"turn_id":"turn-123"}"#));
let conversation_id = client.state.conversation_id;
assert_eq!(
client_metadata,
std::collections::HashMap::from([
(
X_CODEX_WINDOW_ID_HEADER.to_string(),
format!("{conversation_id}:1"),
),
(
X_OPENAI_SUBAGENT_HEADER.to_string(),
"collab_spawn".to_string(),
),
(
X_CODEX_PARENT_THREAD_ID_HEADER.to_string(),
parent_thread_id.to_string(),
),
(
X_CODEX_TURN_METADATA_HEADER.to_string(),
r#"{"turn_id":"turn-123"}"#.to_string(),
),
])
);
}
#[tokio::test]
async fn summarize_memories_returns_empty_for_empty_input() {
let client = test_model_client(SessionSource::Cli);
+30 -9
View File
@@ -1531,6 +1531,17 @@ impl Session {
),
),
};
let window_generation = match &initial_history {
InitialHistory::Resumed(resumed_history) => u64::try_from(
resumed_history
.history
.iter()
.filter(|item| matches!(item, RolloutItem::Compacted(_)))
.count(),
)
.unwrap_or(u64::MAX),
InitialHistory::New | InitialHistory::Forked(_) => 0,
};
let state_builder = match &initial_history {
InitialHistory::Resumed(resumed) => metadata::builder_from_items(
resumed.history.as_slice(),
@@ -1920,6 +1931,9 @@ impl Session {
),
environment: environment_manager.current().await?,
};
services
.model_client
.set_window_generation(window_generation);
let js_repl = Arc::new(JsReplHandle::with_node_path(
config.js_repl_node_path.clone(),
config.js_repl_node_module_dirs.clone(),
@@ -3514,6 +3528,7 @@ impl Session {
self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item)])
.await;
}
self.services.model_client.advance_window_generation();
}
async fn persist_rollout_response_items(&self, items: &[ResponseItem]) {
@@ -5733,16 +5748,20 @@ pub(crate) async fn run_turn(
let model_info = turn_context.model_info.clone();
let auto_compact_limit = model_info.auto_compact_token_limit().unwrap_or(i64::MAX);
let mut prewarmed_client_session = prewarmed_client_session;
// TODO(ccunningham): Pre-turn compaction runs before context updates and the
// new user message are recorded. Estimate pending incoming items (context
// diffs/full reinjection + user input) and trigger compaction preemptively
// when they would push the thread over the compaction threshold.
if run_pre_sampling_compact(&sess, &turn_context)
.await
.is_err()
{
error!("Failed to run pre-sampling compact");
return None;
let pre_sampling_compacted = match run_pre_sampling_compact(&sess, &turn_context).await {
Ok(pre_sampling_compacted) => pre_sampling_compacted,
Err(_) => {
error!("Failed to run pre-sampling compact");
return None;
}
};
if pre_sampling_compacted && let Some(mut client_session) = prewarmed_client_session.take() {
client_session.reset_websocket_session();
}
let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref());
@@ -6055,6 +6074,7 @@ pub(crate) async fn run_turn(
{
return None;
}
client_session.reset_websocket_session();
continue;
}
@@ -6215,9 +6235,9 @@ pub(crate) async fn run_turn(
async fn run_pre_sampling_compact(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
) -> CodexResult<()> {
) -> CodexResult<bool> {
let total_usage_tokens_before_compaction = sess.get_total_token_usage().await;
maybe_run_previous_model_inline_compact(
let mut pre_sampling_compacted = maybe_run_previous_model_inline_compact(
sess,
turn_context,
total_usage_tokens_before_compaction,
@@ -6231,8 +6251,9 @@ async fn run_pre_sampling_compact(
// Compact if the total usage tokens are greater than the auto compact limit
if total_usage_tokens >= auto_compact_limit {
run_auto_compact(sess, turn_context, InitialContextInjection::DoNotInject).await?;
pre_sampling_compacted = true;
}
Ok(())
Ok(pre_sampling_compacted)
}
/// Runs pre-sampling compaction against the previous model when switching to a smaller
+1
View File
@@ -220,6 +220,7 @@ async fn run_compact_task_inner(
};
sess.replace_compacted_history(new_history, reference_context_item, compacted_item)
.await;
client_session.reset_websocket_session();
sess.recompute_token_usage(&turn_context).await;
sess.emit_turn_item_completed(&turn_context, compaction_item)
+6
View File
@@ -129,10 +129,16 @@ async fn responses_stream_includes_subagent_header_on_review() {
}
let request = request_recorder.single_request();
let expected_window_id = format!("{conversation_id}:0");
assert_eq!(
request.header("x-openai-subagent").as_deref(),
Some("review")
);
assert_eq!(
request.header("x-codex-window-id").as_deref(),
Some(expected_window_id.as_str())
);
assert_eq!(request.header("x-codex-parent-thread-id"), None);
assert_eq!(request.header("x-codex-sandbox"), None);
}
+1
View File
@@ -154,3 +154,4 @@ mod user_shell_cmd;
mod view_image;
mod web_search;
mod websocket_fallback;
mod window_headers;
+146
View File
@@ -0,0 +1,146 @@
#![allow(clippy::expect_used)]
use super::compact::COMPACT_WARNING_MESSAGE;
use anyhow::Result;
use codex_core::CodexThread;
use codex_core::compact::SUMMARIZATION_PROMPT;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::WarningEvent;
use codex_protocol::user_input::UserInput;
use core_test_support::responses::ResponsesRequest;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::mount_sse_sequence;
use core_test_support::responses::sse;
use core_test_support::responses::start_mock_server;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use std::sync::Arc;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_fork() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let request_log = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_assistant_message("msg-1", "first reply"),
ev_completed("resp-1"),
]),
sse(vec![
ev_assistant_message("msg-2", "summary"),
ev_completed("resp-2"),
]),
sse(vec![ev_completed("resp-3")]),
sse(vec![ev_completed("resp-4")]),
sse(vec![ev_completed("resp-5")]),
],
)
.await;
let mut builder = test_codex().with_config(|config| {
config.model_provider.name = "Non-OpenAI Model provider".to_string();
config.compact_prompt = Some(SUMMARIZATION_PROMPT.to_string());
});
let initial = builder.build(&server).await?;
let initial_thread = Arc::clone(&initial.codex);
let rollout_path = initial
.session_configured
.rollout_path
.clone()
.expect("rollout path");
submit_user_turn(&initial_thread, "before compact").await?;
submit_compact_turn(&initial_thread).await?;
submit_user_turn(&initial_thread, "after compact").await?;
shutdown_thread(&initial_thread).await?;
let resumed = builder
.resume(&server, initial.home.clone(), rollout_path.clone())
.await?;
submit_user_turn(&resumed.codex, "after resume").await?;
shutdown_thread(&resumed.codex).await?;
let forked = resumed
.thread_manager
.fork_thread(
/*snapshot*/ 0usize,
resumed.config.clone(),
rollout_path,
/*persist_extended_history*/ false,
/*parent_trace*/ None,
)
.await?;
submit_user_turn(&forked.thread, "after fork").await?;
shutdown_thread(&forked.thread).await?;
let requests = request_log.requests();
assert_eq!(requests.len(), 5, "expected five model requests");
let (initial_thread_id, first_generation) = window_id_parts(&requests[0]);
let (compact_thread_id, compact_generation) = window_id_parts(&requests[1]);
let (after_compact_thread_id, after_compact_generation) = window_id_parts(&requests[2]);
let (after_resume_thread_id, after_resume_generation) = window_id_parts(&requests[3]);
let (after_fork_thread_id, after_fork_generation) = window_id_parts(&requests[4]);
assert_eq!(first_generation, 0);
assert_eq!(compact_thread_id, initial_thread_id);
assert_eq!(compact_generation, 0);
assert_eq!(after_compact_thread_id, initial_thread_id);
assert_eq!(after_compact_generation, 1);
assert_eq!(after_resume_thread_id, initial_thread_id);
assert_eq!(after_resume_generation, 1);
assert_ne!(after_fork_thread_id, initial_thread_id);
assert_eq!(after_fork_generation, 0);
Ok(())
}
async fn submit_user_turn(codex: &Arc<CodexThread>, text: &str) -> Result<()> {
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: text.to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
})
.await?;
wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await;
Ok(())
}
async fn submit_compact_turn(codex: &Arc<CodexThread>) -> Result<()> {
codex.submit(Op::Compact).await?;
let warning_event = wait_for_event(codex, |event| matches!(event, EventMsg::Warning(_))).await;
let EventMsg::Warning(WarningEvent { message }) = warning_event else {
panic!("expected warning event after compact");
};
assert_eq!(message, COMPACT_WARNING_MESSAGE);
wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await;
Ok(())
}
async fn shutdown_thread(codex: &Arc<CodexThread>) -> Result<()> {
codex.submit(Op::Shutdown).await?;
wait_for_event(codex, |event| matches!(event, EventMsg::ShutdownComplete)).await;
Ok(())
}
fn window_id_parts(request: &ResponsesRequest) -> (String, u64) {
let window_id = request
.header("x-codex-window-id")
.expect("missing x-codex-window-id header");
let (thread_id, generation) = window_id
.rsplit_once(':')
.unwrap_or_else(|| panic!("invalid window id header: {window_id}"));
let generation = generation
.parse::<u64>()
.unwrap_or_else(|err| panic!("invalid window generation in {window_id}: {err}"));
(thread_id.to_string(), generation)
}