From 4626154ba596e8016311732bd59411e2eedcabd2 Mon Sep 17 00:00:00 2001 From: miaolin Date: Wed, 10 Jun 2026 17:09:37 -0700 Subject: [PATCH] Pass agent path metadata to MCP tools --- codex-rs/core/src/mcp_tool_call.rs | 23 ++++++- codex-rs/core/src/mcp_tool_call_tests.rs | 37 +++++++++- codex-rs/core/src/turn_metadata.rs | 27 ++++++++ codex-rs/core/src/turn_metadata_tests.rs | 88 ++++++++++++++++++++++++ 4 files changed, 172 insertions(+), 3 deletions(-) diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index e5273a6669..b646bec795 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -339,8 +339,14 @@ async fn handle_approved_mcp_tool_call( }; let result = async { let rewritten_arguments = rewrite?; - let request_meta = - build_mcp_tool_call_request_meta(turn_context, &server, call_id, metadata); + let has_spawned_subagent = has_spawned_subagent(sess).await; + let request_meta = build_mcp_tool_call_request_meta( + turn_context, + &server, + call_id, + metadata, + has_spawned_subagent, + ); let result = execute_mcp_tool_call( sess, turn_context, @@ -1026,6 +1032,7 @@ fn build_mcp_tool_call_request_meta( server: &str, call_id: &str, metadata: Option<&McpToolApprovalMetadata>, + has_spawned_subagent: bool, ) -> Option { let mut request_meta = serde_json::Map::new(); @@ -1034,6 +1041,7 @@ fn build_mcp_tool_call_request_meta( .current_meta_value_for_mcp_request(McpTurnMetadataContext { model: turn_context.model_info.slug.as_str(), reasoning_effort: turn_context.effective_reasoning_effort(), + has_spawned_subagent, }) { request_meta.insert( @@ -1065,6 +1073,17 @@ fn build_mcp_tool_call_request_meta( (!request_meta.is_empty()).then_some(serde_json::Value::Object(request_meta)) } +async fn has_spawned_subagent(sess: &Session) -> bool { + let Some(state_db) = sess.services.state_db.as_ref() else { + return false; + }; + + state_db + .list_thread_spawn_children(sess.thread_id) + .await + .is_ok_and(|children| !children.is_empty()) +} + fn with_mcp_tool_call_thread_id_meta( meta: Option, thread_id: &str, diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 723948b7be..145c2df783 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -91,6 +91,7 @@ fn mcp_turn_metadata_context(turn_context: &TurnContext) -> McpTurnMetadataConte McpTurnMetadataContext { model: turn_context.model_info.slug.as_str(), reasoning_effort: turn_context.effective_reasoning_effort(), + has_spawned_subagent: false, } } @@ -1035,6 +1036,7 @@ async fn mcp_tool_call_request_meta_includes_turn_metadata_for_custom_server() { "custom_server", "call-custom", /*metadata*/ None, + /*has_spawned_subagent*/ false, ) .expect("custom servers should receive turn metadata"); let turn_metadata = meta @@ -1077,6 +1079,7 @@ async fn mcp_tool_call_request_meta_includes_turn_started_at_unix_ms() { "custom_server", "call-custom", /*metadata*/ None, + /*has_spawned_subagent*/ false, ) .expect("custom servers should receive turn metadata"); let turn_metadata = meta @@ -1091,6 +1094,30 @@ async fn mcp_tool_call_request_meta_includes_turn_started_at_unix_ms() { ); } +#[tokio::test] +async fn mcp_tool_call_request_meta_includes_has_spawned_subagent() { + let (_, turn_context) = make_session_and_context().await; + + let meta = build_mcp_tool_call_request_meta( + &turn_context, + "custom_server", + "call-custom", + /*metadata*/ None, + /*has_spawned_subagent*/ true, + ) + .expect("custom servers should receive turn metadata"); + let turn_metadata = meta + .get(crate::X_CODEX_TURN_METADATA_HEADER) + .expect("turn metadata should be present"); + + assert_eq!( + turn_metadata + .get("has_spawned_subagent") + .and_then(serde_json::Value::as_bool), + Some(true) + ); +} + #[tokio::test] async fn plugin_mcp_tool_call_request_meta_includes_plugin_id() { let (_, turn_context) = make_session_and_context().await; @@ -1106,7 +1133,13 @@ async fn plugin_mcp_tool_call_request_meta_includes_plugin_id() { metadata.plugin_id = Some("sample@test".to_string()); assert_eq!( - build_mcp_tool_call_request_meta(&turn_context, "sample", "call-plugin", Some(&metadata),), + build_mcp_tool_call_request_meta( + &turn_context, + "sample", + "call-plugin", + Some(&metadata), + /*has_spawned_subagent*/ false, + ), Some(serde_json::json!({ crate::X_CODEX_TURN_METADATA_HEADER: expected_turn_metadata, MCP_TOOL_PLUGIN_ID_META_KEY: "sample@test", @@ -1183,6 +1216,7 @@ async fn codex_apps_tool_call_request_meta_includes_turn_metadata_and_codex_apps CODEX_APPS_MCP_SERVER_NAME, "call_abc123xyz789", Some(&metadata), + /*has_spawned_subagent*/ false, ), Some(serde_json::json!({ crate::X_CODEX_TURN_METADATA_HEADER: expected_turn_metadata, @@ -1210,6 +1244,7 @@ async fn codex_apps_tool_call_request_meta_includes_call_id_without_existing_cod CODEX_APPS_MCP_SERVER_NAME, "call_abc123xyz789", /*metadata*/ None, + /*has_spawned_subagent*/ false, ), Some(serde_json::json!({ crate::X_CODEX_TURN_METADATA_HEADER: expected_turn_metadata, diff --git a/codex-rs/core/src/turn_metadata.rs b/codex-rs/core/src/turn_metadata.rs index e66a1fbb50..baa456b4be 100644 --- a/codex-rs/core/src/turn_metadata.rs +++ b/codex-rs/core/src/turn_metadata.rs @@ -21,6 +21,7 @@ use codex_git_utils::get_git_remote_urls_assume_git_repo; use codex_git_utils::get_git_repo_root; use codex_git_utils::get_has_changes; use codex_git_utils::get_head_commit_hash; +use codex_protocol::AgentPath; use codex_protocol::ThreadId; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::models::PermissionProfile; @@ -30,6 +31,8 @@ use codex_protocol::protocol::ThreadSource; use codex_utils_absolute_path::AbsolutePathBuf; const MODEL_KEY: &str = "model"; +const AGENT_PATH_KEY: &str = "agent_path"; +const HAS_SPAWNED_SUBAGENT_KEY: &str = "has_spawned_subagent"; const REASONING_EFFORT_KEY: &str = "reasoning_effort"; const TURN_STARTED_AT_UNIX_MS_KEY: &str = "turn_started_at_unix_ms"; const USER_INPUT_REQUESTED_DURING_TURN_KEY: &str = "user_input_requested_during_turn"; @@ -41,6 +44,7 @@ const WINDOW_ID_KEY: &str = "window_id"; pub(crate) struct McpTurnMetadataContext<'a> { pub(crate) model: &'a str, pub(crate) reasoning_effort: Option, + pub(crate) has_spawned_subagent: bool, } /// Metadata present only on outbound model requests that perform compaction. @@ -194,6 +198,8 @@ fn merge_turn_metadata( | "forked_from_thread_id" | "parent_thread_id" | "subagent_kind" + | AGENT_PATH_KEY + | HAS_SPAWNED_SUBAGENT_KEY | REQUEST_KIND_KEY | COMPACTION_KEY | WINDOW_ID_KEY @@ -249,6 +255,7 @@ pub(crate) struct TurnMetadataState { repo_root: Option, base_metadata: TurnMetadataBag, base_header: Option, + agent_path: Option, enriched_header: Arc>>, turn_started_at_unix_ms: Arc>>, responsesapi_client_metadata: Arc>>>, @@ -290,6 +297,13 @@ impl TurnMetadataState { | SessionSource::Internal(_) | SessionSource::Unknown => None, }; + let agent_path = session_source + .get_agent_path() + .map(String::from) + .or_else(|| { + (!matches!(session_source, SessionSource::SubAgent(_))) + .then(|| AgentPath::ROOT.to_string()) + }); let base_metadata = TurnMetadataBag { request_kind: None, session_id: Some(session_id), @@ -309,6 +323,7 @@ impl TurnMetadataState { repo_root, base_metadata, base_header, + agent_path, enriched_header: Arc::new(RwLock::new(None)), turn_started_at_unix_ms: Arc::new(RwLock::new(None)), responsesapi_client_metadata: Arc::new(RwLock::new(None)), @@ -353,10 +368,22 @@ impl TurnMetadataState { let header = self.current_header_value()?; let mut metadata = serde_json::from_str::>(&header).ok()?; metadata.remove(REQUEST_KIND_KEY); + if let Some(agent_path) = self.agent_path.as_ref() { + metadata.insert( + AGENT_PATH_KEY.to_string(), + Value::String(agent_path.clone()), + ); + } else { + metadata.remove(AGENT_PATH_KEY); + } metadata.insert( MODEL_KEY.to_string(), Value::String(context.model.to_string()), ); + metadata.insert( + HAS_SPAWNED_SUBAGENT_KEY.to_string(), + Value::Bool(context.has_spawned_subagent), + ); match context.reasoning_effort { Some(reasoning_effort) => { metadata.insert( diff --git a/codex-rs/core/src/turn_metadata_tests.rs b/codex-rs/core/src/turn_metadata_tests.rs index 32a559848a..8a4c411fe6 100644 --- a/codex-rs/core/src/turn_metadata_tests.rs +++ b/codex-rs/core/src/turn_metadata_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::sandbox_tags::permission_profile_sandbox_tag; +use codex_protocol::AgentPath; use codex_protocol::models::PermissionProfile; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::SessionSource; @@ -20,6 +21,7 @@ fn test_mcp_turn_metadata_context() -> McpTurnMetadataContext<'static> { McpTurnMetadataContext { model: "gpt-5.4", reasoning_effort: Some(ReasoningEffortConfig::High), + has_spawned_subagent: false, } } @@ -404,6 +406,7 @@ fn turn_metadata_state_includes_model_and_reasoning_effort_only_in_request_meta( .current_meta_value_for_mcp_request(McpTurnMetadataContext { model: "gpt-5.4", reasoning_effort: None, + has_spawned_subagent: false, }) .expect("turn metadata should be present"); assert_eq!( @@ -417,6 +420,84 @@ fn turn_metadata_state_includes_model_and_reasoning_effort_only_in_request_meta( ); } +#[test] +fn turn_metadata_state_includes_agent_path_only_in_request_meta() { + let temp_dir = TempDir::new().expect("temp dir"); + let cwd = temp_dir.path().abs(); + let permission_profile = PermissionProfile::read_only(); + let parent_thread_id = ThreadId::new(); + let agent_path = AgentPath::try_from("/root/worker").expect("agent path"); + + let root_state = TurnMetadataState::new( + "session-root".to_string(), + "thread-root".to_string(), + /*forked_from_thread_id*/ None, + /*parent_thread_id*/ None, + &SessionSource::Exec, + /*thread_source*/ None, + "turn-root".to_string(), + cwd.clone(), + &permission_profile, + WindowsSandboxLevel::Disabled, + /*enforce_managed_network*/ false, + ); + let subagent_state = TurnMetadataState::new( + "session-child".to_string(), + "thread-child".to_string(), + /*forked_from_thread_id*/ None, + Some(parent_thread_id), + &SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: Some(agent_path), + agent_nickname: None, + agent_role: None, + }), + /*thread_source*/ None, + "turn-child".to_string(), + cwd, + &permission_profile, + WindowsSandboxLevel::Disabled, + /*enforce_managed_network*/ false, + ); + + let root_header: Value = + serde_json::from_str(&root_state.current_header_value().expect("root header")) + .expect("valid json"); + assert!(root_header.get(AGENT_PATH_KEY).is_none()); + + let root_meta = root_state + .current_meta_value_for_mcp_request(test_mcp_turn_metadata_context()) + .expect("root mcp metadata"); + assert_eq!( + serde_json::json!({ + "agent_path": "/root", + "has_spawned_subagent": false, + }), + serde_json::json!({ + "agent_path": root_meta.get(AGENT_PATH_KEY), + "has_spawned_subagent": root_meta.get(HAS_SPAWNED_SUBAGENT_KEY), + }) + ); + + let subagent_meta = subagent_state + .current_meta_value_for_mcp_request(McpTurnMetadataContext { + has_spawned_subagent: true, + ..test_mcp_turn_metadata_context() + }) + .expect("subagent mcp metadata"); + assert_eq!( + serde_json::json!({ + "agent_path": "/root/worker", + "has_spawned_subagent": true, + }), + serde_json::json!({ + "agent_path": subagent_meta.get(AGENT_PATH_KEY), + "has_spawned_subagent": subagent_meta.get(HAS_SPAWNED_SUBAGENT_KEY), + }) + ); +} + #[test] fn turn_metadata_state_marks_user_input_requested_during_turn_only_for_mcp_request_meta() { let temp_dir = TempDir::new().expect("temp dir"); @@ -503,6 +584,11 @@ fn turn_metadata_state_ignores_client_reserved_metadata_before_start() { "client-supplied".to_string(), ), ("subagent_kind".to_string(), "client-supplied".to_string()), + (AGENT_PATH_KEY.to_string(), "client-supplied".to_string()), + ( + HAS_SPAWNED_SUBAGENT_KEY.to_string(), + "client-supplied".to_string(), + ), ])); let header = state.current_header_value().expect("header"); @@ -512,6 +598,8 @@ fn turn_metadata_state_ignores_client_reserved_metadata_before_start() { assert!(json.get("forked_from_thread_id").is_none()); assert!(json.get("parent_thread_id").is_none()); assert!(json.get("subagent_kind").is_none()); + assert!(json.get(AGENT_PATH_KEY).is_none()); + assert!(json.get(HAS_SPAWNED_SUBAGENT_KEY).is_none()); } #[test]