Compare commits

...

1 Commits

Author SHA1 Message Date
Michael Chen
0a9a2ab1c7 [Core] Add model to MCP turn metadata 2026-05-05 09:41:44 -07:00
10 changed files with 138 additions and 32 deletions

View File

@@ -184,7 +184,9 @@ async fn run_compact_task_inner_impl(
personality: turn_context.personality,
..Default::default()
};
let turn_metadata_header = turn_context.turn_metadata_state.current_header_value();
let turn_metadata_header = turn_context
.turn_metadata_state
.current_header_value_for_model(turn_context.model_info.slug.as_str());
let attempt_result = drain_to_completed(
&sess,
turn_context.as_ref(),

View File

@@ -177,7 +177,9 @@ async fn run_remote_compact_task_inner_impl(
output_schema_strict: true,
};
let turn_metadata_header = turn_context.turn_metadata_state.current_header_value();
let turn_metadata_header = turn_context
.turn_metadata_state
.current_header_value_for_model(turn_context.model_info.slug.as_str());
let trace_attempt = compaction_trace.start_attempt(&serde_json::json!({
"model": turn_context.model_info.slug.as_str(),
"instructions": prompt.base_instructions.text.as_str(),

View File

@@ -895,7 +895,11 @@ fn build_mcp_tool_call_request_meta(
) -> Option<serde_json::Value> {
let mut request_meta = serde_json::Map::new();
if let Some(turn_metadata) = turn_context.turn_metadata_state.current_meta_value() {
if let Some(turn_metadata) = turn_context
.turn_metadata_state
.current_header_value_for_model(turn_context.model_info.slug.as_str())
.and_then(|header| serde_json::from_str(&header).ok())
{
request_meta.insert(
crate::X_CODEX_TURN_METADATA_HEADER.to_string(),
turn_metadata,

View File

@@ -911,7 +911,7 @@ async fn mcp_tool_call_request_meta_includes_turn_metadata_for_custom_server() {
let expected_turn_metadata = serde_json::from_str::<serde_json::Value>(
&turn_context
.turn_metadata_state
.current_header_value()
.current_header_value_for_model(turn_context.model_info.slug.as_str())
.expect("turn metadata header"),
)
.expect("turn metadata json");
@@ -950,6 +950,12 @@ async fn mcp_tool_call_request_meta_includes_turn_started_at_unix_ms() {
.get(crate::X_CODEX_TURN_METADATA_HEADER)
.expect("turn metadata should be present");
assert_eq!(
turn_metadata
.get("model")
.and_then(serde_json::Value::as_str),
Some(turn_context.model_info.slug.as_str())
);
assert_eq!(
turn_metadata
.get("turn_started_at_unix_ms")
@@ -964,7 +970,7 @@ async fn codex_apps_tool_call_request_meta_includes_turn_metadata_and_codex_apps
let expected_turn_metadata = serde_json::from_str::<serde_json::Value>(
&turn_context
.turn_metadata_state
.current_header_value()
.current_header_value_for_model(turn_context.model_info.slug.as_str())
.expect("turn metadata header"),
)
.expect("turn metadata json");
@@ -1014,7 +1020,7 @@ async fn codex_apps_tool_call_request_meta_includes_call_id_without_existing_cod
let expected_turn_metadata = serde_json::from_str::<serde_json::Value>(
&turn_context
.turn_metadata_state
.current_header_value()
.current_header_value_for_model(turn_context.model_info.slug.as_str())
.expect("turn metadata header"),
)
.expect("turn metadata json");

View File

@@ -443,7 +443,9 @@ pub(crate) async fn run_turn(
})
.map(|user_message| user_message.message())
.collect::<Vec<String>>();
let turn_metadata_header = turn_context.turn_metadata_state.current_header_value();
let turn_metadata_header = turn_context
.turn_metadata_state
.current_header_value_for_model(turn_context.model_info.slug.as_str());
match run_sampling_request(
Arc::clone(&sess),
Arc::clone(&turn_context),

View File

@@ -223,7 +223,7 @@ async fn schedule_startup_prewarm_inner(
);
let startup_turn_metadata_header = startup_turn_context
.turn_metadata_state
.current_header_value();
.current_header_value_for_model(startup_turn_context.model_info.slug.as_str());
let mut client_session = session.services.model_client.new_session();
client_session
.prewarm_websocket(

View File

@@ -20,6 +20,14 @@ use codex_protocol::protocol::SessionSource;
use codex_utils_absolute_path::AbsolutePathBuf;
const TURN_STARTED_AT_UNIX_MS_KEY: &str = "turn_started_at_unix_ms";
const MODEL_KEY: &str = "model";
#[derive(Clone, Copy, Debug, Default)]
struct TurnMetadataOverlay<'a> {
model: Option<&'a str>,
turn_started_at_unix_ms: Option<i64>,
responsesapi_client_metadata: Option<&'a HashMap<String, String>>,
}
#[derive(Clone, Debug, Default)]
struct WorkspaceGitMetadata {
@@ -76,25 +84,27 @@ impl TurnMetadataBag {
}
}
fn merge_turn_metadata(
header: &str,
turn_started_at_unix_ms: Option<i64>,
responsesapi_client_metadata: Option<&HashMap<String, String>>,
) -> Option<String> {
if turn_started_at_unix_ms.is_none() && responsesapi_client_metadata.is_none() {
fn merge_turn_metadata(header: &str, overlay: TurnMetadataOverlay<'_>) -> Option<String> {
if overlay.model.is_none()
&& overlay.turn_started_at_unix_ms.is_none()
&& overlay.responsesapi_client_metadata.is_none()
{
return None;
}
let mut metadata = serde_json::from_str::<serde_json::Map<String, Value>>(header).ok()?;
if let Some(turn_started_at_unix_ms) = turn_started_at_unix_ms {
if let Some(model) = overlay.model {
metadata.insert(MODEL_KEY.to_string(), Value::String(model.to_string()));
}
if let Some(turn_started_at_unix_ms) = overlay.turn_started_at_unix_ms {
metadata.insert(
TURN_STARTED_AT_UNIX_MS_KEY.to_string(),
Value::Number(turn_started_at_unix_ms.into()),
);
}
if let Some(responsesapi_client_metadata) = responsesapi_client_metadata {
if let Some(responsesapi_client_metadata) = overlay.responsesapi_client_metadata {
for (key, value) in responsesapi_client_metadata {
if key == TURN_STARTED_AT_UNIX_MS_KEY {
if matches!(key.as_str(), MODEL_KEY | TURN_STARTED_AT_UNIX_MS_KEY) {
continue;
}
metadata
@@ -219,7 +229,11 @@ impl TurnMetadataState {
}
}
pub(crate) fn current_header_value(&self) -> Option<String> {
pub(crate) fn current_header_value_for_model(&self, model: &str) -> Option<String> {
self.build_current_header_value(Some(model))
}
fn build_current_header_value(&self, model: Option<&str>) -> Option<String> {
let header = if let Some(header) = self
.enriched_header
.read()
@@ -242,17 +256,15 @@ impl TurnMetadataState {
.clone();
merge_turn_metadata(
&header,
turn_started_at_unix_ms,
responsesapi_client_metadata.as_ref(),
TurnMetadataOverlay {
model,
turn_started_at_unix_ms,
responsesapi_client_metadata: responsesapi_client_metadata.as_ref(),
},
)
.or(Some(header))
}
pub(crate) fn current_meta_value(&self) -> Option<serde_json::Value> {
self.current_header_value()
.and_then(|header| serde_json::from_str(&header).ok())
}
pub(crate) fn set_responsesapi_client_metadata(
&self,
responsesapi_client_metadata: HashMap<String, String>,

View File

@@ -95,7 +95,9 @@ fn turn_metadata_state_uses_platform_sandbox_tag() {
/*enforce_managed_network*/ false,
);
let header = state.current_header_value().expect("header");
let header = state
.current_header_value_for_model("gpt-5.4")
.expect("header");
let json: Value = serde_json::from_str(&header).expect("json");
let sandbox_name = json.get("sandbox").and_then(Value::as_str);
let session_id = json.get("session_id").and_then(Value::as_str);
@@ -125,7 +127,9 @@ fn turn_metadata_state_classifies_subagent_thread_source() {
/*enforce_managed_network*/ false,
);
let header = state.current_header_value().expect("header");
let header = state
.current_header_value_for_model("gpt-5.4")
.expect("header");
let json: Value = serde_json::from_str(&header).expect("json");
assert_eq!(json["thread_source"].as_str(), Some("subagent"));
@@ -149,7 +153,9 @@ fn turn_metadata_state_includes_turn_started_at_unix_ms_after_start() {
);
state.set_turn_started_at_unix_ms(/*turn_started_at_unix_ms*/ 1_700_000_000_123);
let header = state.current_header_value().expect("header");
let header = state
.current_header_value_for_model("gpt-5.4")
.expect("header");
let json: Value = serde_json::from_str(&header).expect("json");
assert_eq!(
@@ -158,6 +164,30 @@ fn turn_metadata_state_includes_turn_started_at_unix_ms_after_start() {
);
}
#[test]
fn turn_metadata_state_includes_model_when_requested() {
let temp_dir = TempDir::new().expect("temp dir");
let cwd = temp_dir.path().abs();
let permission_profile = PermissionProfile::read_only();
let state = TurnMetadataState::new(
"session-a".to_string(),
&SessionSource::Exec,
"turn-a".to_string(),
cwd,
&permission_profile,
WindowsSandboxLevel::Disabled,
/*enforce_managed_network*/ false,
);
let header = state
.current_header_value_for_model("gpt-5.4")
.expect("header");
let json: Value = serde_json::from_str(&header).expect("json");
assert_eq!(json["model"].as_str(), Some("gpt-5.4"));
}
#[test]
fn turn_metadata_state_ignores_client_turn_started_at_unix_ms_before_start() {
let temp_dir = TempDir::new().expect("temp dir");
@@ -178,7 +208,9 @@ fn turn_metadata_state_ignores_client_turn_started_at_unix_ms_before_start() {
"client-supplied".to_string(),
)]));
let header = state.current_header_value().expect("header");
let header = state
.current_header_value_for_model("gpt-5.4")
.expect("header");
let json: Value = serde_json::from_str(&header).expect("json");
assert!(json.get("turn_started_at_unix_ms").is_none());
@@ -202,6 +234,7 @@ fn turn_metadata_state_merges_client_metadata_without_replacing_reserved_fields(
state.set_responsesapi_client_metadata(HashMap::from([
("fiber_run_id".to_string(), "fiber-123".to_string()),
("origin".to_string(), "東京".to_string()),
("model".to_string(), "client-supplied".to_string()),
("session_id".to_string(), "client-supplied".to_string()),
("thread_source".to_string(), "client-supplied".to_string()),
(
@@ -211,13 +244,16 @@ fn turn_metadata_state_merges_client_metadata_without_replacing_reserved_fields(
]));
state.set_turn_started_at_unix_ms(/*turn_started_at_unix_ms*/ 1_700_000_000_123);
let header = state.current_header_value().expect("header");
let header = state
.current_header_value_for_model("gpt-5.4")
.expect("header");
assert!(header.is_ascii());
assert!(!header.contains("東京"));
let json: Value = serde_json::from_str(&header).expect("json");
assert_eq!(json["fiber_run_id"].as_str(), Some("fiber-123"));
assert_eq!(json["origin"].as_str(), Some("東京"));
assert_eq!(json["model"].as_str(), Some("gpt-5.4"));
assert_eq!(json["session_id"].as_str(), Some("session-a"));
assert_eq!(json["thread_source"].as_str(), Some("user"));
assert_eq!(json["turn_id"].as_str(), Some("turn-a"));

View File

@@ -419,8 +419,12 @@ async fn responses_stream_includes_turn_metadata_header_for_git_workspace_e2e()
test.submit_turn("hello")
.await
.expect("submit first turn prompt");
let initial_header = first_request
.single_request()
let initial_request = first_request.single_request();
let initial_request_body = initial_request.body_json();
let initial_model = initial_request_body["model"]
.as_str()
.expect("initial request should include model");
let initial_header = initial_request
.header("x-codex-turn-metadata")
.expect("x-codex-turn-metadata header should be present");
let initial_parsed: serde_json::Value =
@@ -442,6 +446,12 @@ async fn responses_stream_includes_turn_metadata_header_for_git_workspace_e2e()
initial_turn_started_at_unix_ms > 0,
"turn_started_at_unix_ms should be positive"
);
assert_eq!(
initial_parsed
.get("model")
.and_then(serde_json::Value::as_str),
Some(initial_model)
);
assert_eq!(
initial_parsed
.get("sandbox")
@@ -561,6 +571,26 @@ async fn responses_stream_includes_turn_metadata_header_for_git_workspace_e2e()
first_turn_started_at_unix_ms, second_turn_started_at_unix_ms,
"requests in the same turn should share turn_started_at_unix_ms"
);
let first_request_model = requests[0].body_json()["model"]
.as_str()
.expect("first request should include model")
.to_string();
let second_request_model = requests[1].body_json()["model"]
.as_str()
.expect("second request should include model")
.to_string();
assert_eq!(
first_parsed
.get("model")
.and_then(serde_json::Value::as_str),
Some(first_request_model.as_str())
);
assert_eq!(
second_parsed
.get("model")
.and_then(serde_json::Value::as_str),
Some(second_request_model.as_str())
);
assert_eq!(
first_parsed
.get("thread_source")

View File

@@ -597,6 +597,10 @@ async fn tool_search_returns_deferred_tools_without_follow_up_tool_injection() -
apps_tool_call.pointer("/params/_meta/x-codex-turn-metadata/session_id"),
Some(&json!(test.session_configured.session_id.to_string()))
);
assert_eq!(
apps_tool_call.pointer("/params/_meta/x-codex-turn-metadata/model"),
Some(&json!("gpt-5.4"))
);
assert!(
apps_tool_call
.pointer("/params/_meta/x-codex-turn-metadata/turn_id")
@@ -619,6 +623,14 @@ async fn tool_search_returns_deferred_tools_without_follow_up_tool_injection() -
.expect("first response request should include turn metadata"),
)
.expect("first response request turn metadata should be valid JSON");
assert_eq!(
first_request_turn_metadata
.get("model")
.and_then(Value::as_str),
apps_tool_call
.pointer("/params/_meta/x-codex-turn-metadata/model")
.and_then(Value::as_str)
);
assert_eq!(
first_request_turn_metadata
.get("turn_started_at_unix_ms")