From 15ea598c6e7e0914a7ae8c881ac05dacea2f7902 Mon Sep 17 00:00:00 2001 From: Rasmus Rygaard Date: Wed, 5 Aug 2026 21:34:47 +0000 Subject: [PATCH] Expose session sources to MCP contributors (#37167) ## What changed - Add `session_source()` to `McpServerContributionContext` for thread-scoped MCP resolution. - Propagate each thread's `SessionSource` through initial setup, per-step resolution, and runtime refreshes. ## Testing - Verify global contexts omit the source while thread-scoped contexts retain it. - Verify MCP contributors observe distinct sources for a root thread and its spawned subagent. GitOrigin-RevId: 23da1509858e023fbd9e2d179034ddc83553f17b --- codex-rs/core/src/mcp.rs | 15 ++- codex-rs/core/src/session/mcp.rs | 33 +++-- codex-rs/core/src/session/mcp_runtime.rs | 3 + codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/session.rs | 11 +- codex-rs/core/src/thread_manager_tests.rs | 32 +++-- .../core/tests/suite/mcp_tool_exposure.rs | 124 ++++++++++++++++++ .../ext/extension-api/src/contributors/mcp.rs | 16 +++ codex-rs/ext/extension-api/tests/registry.rs | 28 ++++ 9 files changed, 239 insertions(+), 24 deletions(-) diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 5ae3315e70..84b3859084 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -26,6 +26,7 @@ use codex_mcp::configured_mcp_servers; use codex_mcp::effective_mcp_servers; use codex_plugin::AppConnectorId; use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::protocol::SessionSource; const LEGACY_CODEX_APPS_REGISTRATION_ID: &str = "legacy_codex_apps"; @@ -36,6 +37,11 @@ pub(crate) struct McpRuntimeProjection { pub(crate) plugins_available: bool, } +pub(crate) struct McpThreadIdentity<'a> { + pub(crate) session_source: &'a SessionSource, + pub(crate) originator: &'a str, +} + enum OrderedMcpOverlay { Set { contributor_id: &'static str, @@ -109,7 +115,7 @@ impl McpManager { config: &Config, thread_init: &ExtensionDataInit, thread_store: &ExtensionData, - originator: &str, + identity: McpThreadIdentity<'_>, ready_selected_capability_roots: &[SelectedCapabilityRoot], executor_capability_discovery: Option<&ExecutorCapabilityDiscoverySnapshot>, ) -> McpRuntimeProjection { @@ -118,11 +124,12 @@ impl McpManager { config, thread_init, thread_store, - originator, + identity.originator, ready_selected_capability_roots, executor_capability_discovery, - ), - Some(originator), + ) + .with_session_source(identity.session_source), + Some(identity.originator), ) .await } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 81edaa4f43..eb70a4cd06 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -95,12 +95,13 @@ impl Session { config: &Config, ) -> (McpConfig, McpRuntimeContext) { let originator = self.originator().await; - let windows_sandbox_level = self - .state - .lock() - .await - .session_configuration - .windows_sandbox_level; + let (windows_sandbox_level, session_source) = { + let state = self.state.lock().await; + ( + state.session_configuration.windows_sandbox_level, + state.session_configuration.session_source.clone(), + ) + }; let environments = self.services.turn_environments.snapshot().await; let selected_capability_roots = self .resolve_selected_capability_roots_for_step(&environments) @@ -122,7 +123,10 @@ impl Session { config, &self.services.mcp_thread_init, &self.services.thread_extension_data, - &originator, + McpThreadIdentity { + session_source: &session_source, + originator: &originator, + }, &ready_selected_capability_roots, executor_capability_discovery.as_deref(), ) @@ -199,7 +203,10 @@ impl Session { &desired.config, &self.services.mcp_thread_init, &self.services.thread_extension_data, - &desired.originator, + McpThreadIdentity { + session_source: &desired.session_source, + originator: &desired.originator, + }, &ready_selected_capability_roots, executor_capability_discovery.as_deref(), ) @@ -253,7 +260,10 @@ impl Session { &desired.config, &self.services.mcp_thread_init, &self.services.thread_extension_data, - &desired.originator, + McpThreadIdentity { + session_source: &desired.session_source, + originator: &desired.originator, + }, &ready_selected_capability_roots, executor_capability_discovery.as_deref(), ) @@ -616,7 +626,10 @@ impl Session { refresh_config, &self.services.mcp_thread_init, &self.services.thread_extension_data, - &turn_context.originator, + McpThreadIdentity { + session_source: &turn_context.session_source, + originator: &turn_context.originator, + }, &ready_selected_capability_roots, executor_capability_discovery.as_deref(), ) diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index b2dd2dc61b..77efaae831 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -16,6 +16,7 @@ pub(super) struct McpDesiredState { pub(super) auth: Option, pub(super) submit_id: String, pub(super) originator: String, + pub(super) session_source: SessionSource, pub(super) environments: TurnEnvironmentSnapshot, pub(super) windows_sandbox_level: WindowsSandboxLevel, } @@ -74,6 +75,7 @@ impl Session { auth, submit_id: self.next_internal_sub_id(), originator: session_configuration.originator.clone(), + session_source: session_configuration.session_source.clone(), environments, windows_sandbox_level: session_configuration.windows_sandbox_level, } @@ -95,6 +97,7 @@ impl Session { auth, submit_id: INITIAL_SUBMIT_ID.to_owned(), originator: session_configuration.originator.clone(), + session_source: session_configuration.session_source.clone(), environments: resolved_environments.clone(), windows_sandbox_level: session_configuration.windows_sandbox_level, }; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 7b02bd51c9..a7e511f28f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -304,6 +304,7 @@ use crate::exec_policy::ExecPolicyUpdateError; use crate::guardian::GuardianReviewOptions; use crate::guardian::GuardianReviewSessionManager; use crate::mcp::McpManager; +use crate::mcp::McpThreadIdentity; use crate::network_policy_decision::execpolicy_network_rule_amendment; use crate::rollout::map_session_init_error; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8c7f36af44..e5de254b07 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -738,6 +738,7 @@ impl Session { let mcp_thread_init_for_startup = &mcp_thread_init; let thread_extension_data_for_mcp = &thread_extension_data; let mcp_originator = session_configuration.originator.clone(); + let mcp_session_source = session_configuration.session_source.clone(); let mcp_runtime_cwd = session_configuration .environment_selections() .first() @@ -751,7 +752,10 @@ impl Session { &config_for_mcp, mcp_thread_init_for_startup, thread_extension_data_for_mcp, - &mcp_originator, + McpThreadIdentity { + session_source: &mcp_session_source, + originator: &mcp_originator, + }, /*ready_selected_capability_roots*/ &[], /*executor_capability_discovery*/ None, ) @@ -1271,7 +1275,10 @@ impl Session { config.as_ref(), &sess.services.mcp_thread_init, &sess.services.thread_extension_data, - &session_configuration.originator, + McpThreadIdentity { + session_source: &session_configuration.session_source, + originator: &session_configuration.originator, + }, /*ready_selected_capability_roots*/ &[], /*executor_capability_discovery*/ None, ) diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index c1e4b49dec..3e1c0422aa 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::config::test_config; use crate::init_state_db; use crate::installation_id::INSTALLATION_ID_FILENAME; +use crate::mcp::McpThreadIdentity; use crate::rollout::RolloutRecorder; use crate::session::session::SessionSettingsUpdate; use crate::session::tests::build_world_state_from_turn_context; @@ -27,6 +28,7 @@ use codex_protocol::protocol::ResumedHistory; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::UserMessageEvent; @@ -666,7 +668,7 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() { struct InitialDataRecorder { lifecycle_observed: Arc>>, - mcp_observed: Arc>>, + mcp_observed: Arc>>, } impl codex_extension_api::ThreadLifecycleContributor for InitialDataRecorder { @@ -712,7 +714,13 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() self.mcp_observed .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(selected_root.id.clone()); + .push(( + selected_root.id.clone(), + context + .session_source() + .expect("thread-scoped MCP resolution should identify its source") + .clone(), + )); let mut server = codex_mcp::codex_apps_mcp_server_config( "https://selected.invalid", /*apps_mcp_product_sku*/ None, @@ -792,9 +800,11 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() }) .await .expect("start first thread"); + let second_session_source = SessionSource::SubAgent(SubAgentSource::Review); let second_thread = manager .start_thread(StartThreadOptions { environments: Some(Vec::new()), + session_source: Some(second_session_source.clone()), thread_extension_init: selected_root_init("selected-b", "env-b"), ..StartThreadOptions::new(config.clone()) }) @@ -809,7 +819,10 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() &config, &first_session.services.mcp_thread_init, &first_session.services.thread_extension_data, - &first_originator, + McpThreadIdentity { + session_source: &SessionSource::Exec, + originator: &first_originator, + }, /*ready_selected_capability_roots*/ &[], /*executor_capability_discovery*/ None, ) @@ -823,7 +836,10 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() &config, &second_session.services.mcp_thread_init, &second_session.services.thread_extension_data, - &second_originator, + McpThreadIdentity { + session_source: &second_session_source, + originator: &second_originator, + }, /*ready_selected_capability_roots*/ &[], /*executor_capability_discovery*/ None, ) @@ -846,10 +862,10 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() .lock() .unwrap_or_else(std::sync::PoisonError::into_inner), vec![ - "selected-a".to_string(), - "selected-b".to_string(), - "selected-a".to_string(), - "selected-b".to_string(), + ("selected-a".to_string(), SessionSource::Exec), + ("selected-b".to_string(), second_session_source.clone()), + ("selected-a".to_string(), SessionSource::Exec), + ("selected-b".to_string(), second_session_source), ] ); let selected_servers = |config: &codex_mcp::McpConfig| { diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 38fa96e775..4235a39aa4 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -17,6 +17,8 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::request_user_input::RequestUserInputAnswer; use codex_protocol::request_user_input::RequestUserInputResponse; use codex_protocol::user_input::UserInput; @@ -32,7 +34,9 @@ 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::ev_function_call; +use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::namespace_child_tool; use core_test_support::responses::sse; @@ -74,6 +78,10 @@ struct AppsMcpServerContributor { url: String, } +struct SessionSourceMcpContributor { + observed_sources: Arc>>, +} + impl CoalescingMcpContributor { fn new() -> Self { Self { @@ -142,6 +150,30 @@ impl McpServerContributor for AppsMcpServerContributor { } } +impl McpServerContributor for SessionSourceMcpContributor { + fn id(&self) -> &'static str { + "session_source_mcp_test" + } + + fn contribute<'a>( + &'a self, + context: McpServerContributionContext<'a, Config>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + self.observed_sources + .lock() + .expect("observed sources lock should not be poisoned") + .push( + context + .session_source() + .expect("thread-scoped MCP resolution should identify its session source") + .clone(), + ); + Vec::new() + }) + } +} + fn format_labeled_requests_snapshot( scenario: &str, sections: &[(&str, &ResponsesRequest)], @@ -215,6 +247,98 @@ fn config_with_mcp_marker(base: &Config, marker: &str) -> Config { config } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn root_and_spawned_subagent_receive_distinct_mcp_session_sources() -> Result<()> { + skip_if_no_network!(Ok(())); + + const PARENT_PROMPT: &str = "spawn an agent to verify its MCP session source"; + const CHILD_PROMPT: &str = "child: report that MCP configuration completed"; + const SPAWN_CALL_ID: &str = "mcp-session-source-spawn"; + + let server = responses::start_mock_server().await; + let spawn_args = serde_json::to_string(&json!({ "message": CHILD_PROMPT }))?; + mount_sse_once_match( + &server, + |request: &Request| { + std::str::from_utf8(&request.body).is_ok_and(|body| body.contains(PARENT_PROMPT)) + && !request.headers.contains_key("x-openai-subagent") + }, + sse(vec![ + ev_response_created("resp-parent-spawn"), + ev_function_call_with_namespace( + SPAWN_CALL_ID, + "multi_agent_v1", + "spawn_agent", + &spawn_args, + ), + ev_completed("resp-parent-spawn"), + ]), + ) + .await; + let child_response = mount_sse_once_match( + &server, + |request: &Request| { + std::str::from_utf8(&request.body).is_ok_and(|body| body.contains(CHILD_PROMPT)) + && request + .headers + .get("x-openai-subagent") + .and_then(|value| value.to_str().ok()) + == Some("collab_spawn") + }, + sse(vec![ + ev_response_created("resp-child"), + ev_assistant_message("msg-child", "child done"), + ev_completed("resp-child"), + ]), + ) + .await; + mount_sse_once_match( + &server, + |request: &Request| { + std::str::from_utf8(&request.body).is_ok_and(|body| body.contains(SPAWN_CALL_ID)) + && !request.headers.contains_key("x-openai-subagent") + }, + sse(vec![ + ev_response_created("resp-parent-complete"), + ev_assistant_message("msg-parent", "parent done"), + ev_completed("resp-parent-complete"), + ]), + ) + .await; + + let observed_sources = Arc::new(Mutex::new(Vec::new())); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.mcp_server_contributor(Arc::new(SessionSourceMcpContributor { + observed_sources: observed_sources.clone(), + })); + let test = core_test_support::test_codex::test_codex() + .with_extensions(Arc::new(extensions.build())) + .build(&server) + .await?; + + test.submit_turn(PARENT_PROMPT).await?; + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async { + while child_response.requests().is_empty() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await?; + + let observed_sources = observed_sources + .lock() + .expect("observed sources lock should not be poisoned"); + assert!(observed_sources.contains(&SessionSource::Exec)); + assert!(observed_sources.iter().any(|source| matches!( + source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + .. + }) if *parent_thread_id == test.session_configured.thread_id + ))); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn rapid_mcp_refreshes_coalesce_to_the_latest_config() -> Result<()> { let server = responses::start_mock_server().await; diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index 66a716d2c2..d32715a403 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -1,6 +1,7 @@ use codex_config::McpServerConfig; use codex_exec_server_protocol::ExecutorCapabilityDiscoverySnapshot; use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::protocol::SessionSource; use crate::ExtensionData; use crate::ExtensionDataInit; @@ -17,6 +18,8 @@ pub struct McpServerContributionContext<'a, C> { thread_store: Option<&'a ExtensionData>, /// Stable host inputs for the active thread, when resolution is thread-scoped. thread_init: Option<&'a ExtensionDataInit>, + /// Source of the active thread, when supplied by the host runtime. + session_source: Option<&'a SessionSource>, /// Effective request originator for the active thread, when resolution is thread-scoped. originator: Option<&'a str>, /// Selected roots resolved against ready environments for this exact step. @@ -40,6 +43,7 @@ impl<'a, C> McpServerContributionContext<'a, C> { config, thread_store: None, thread_init: None, + session_source: None, originator: None, ready_selected_capability_roots: None, executor_capability_discovery: None, @@ -59,12 +63,19 @@ impl<'a, C> McpServerContributionContext<'a, C> { config, thread_store: Some(thread_store), thread_init: Some(thread_init), + session_source: None, originator: Some(originator), ready_selected_capability_roots: Some(ready_selected_capability_roots), executor_capability_discovery, } } + /// Attaches the stable source of the active thread to this contribution. + pub fn with_session_source(mut self, session_source: &'a SessionSource) -> Self { + self.session_source = Some(session_source); + self + } + /// Returns the host configuration visible during resolution. pub fn config(&self) -> &'a C { self.config @@ -80,6 +91,11 @@ impl<'a, C> McpServerContributionContext<'a, C> { self.thread_init } + /// Returns the active thread's source when supplied by the host runtime. + pub fn session_source(&self) -> Option<&'a SessionSource> { + self.session_source + } + /// Returns the effective request originator when resolving for a running thread. pub fn originator(&self) -> Option<&'a str> { self.originator diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index 6761a239f5..e5579e5453 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/codex-rs/ext/extension-api/tests/registry.rs @@ -8,11 +8,13 @@ use codex_extension_api::ConfigContributor; use codex_extension_api::ContextContributor; use codex_extension_api::ContextualUserFragment; use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionMetrics; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionWarning; +use codex_extension_api::McpServerContributionContext; use codex_extension_api::PromptFragment; use codex_extension_api::PromptSlot; use codex_extension_api::SkillInvocationContributor; @@ -33,11 +35,37 @@ use codex_protocol::items::TurnItem; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ReviewDecision; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::WarningEvent; use pretty_assertions::assert_eq; struct AllContributors; +#[test] +fn mcp_contribution_context_identifies_the_running_thread() { + let config = (); + let thread_init = ExtensionDataInit::new(); + let thread_store = ExtensionData::new("child-thread"); + let session_source = SessionSource::SubAgent(SubAgentSource::Review); + + let thread_context = McpServerContributionContext::for_step( + &config, + &thread_init, + &thread_store, + "codex_work_cca", + &[], + /*executor_capability_discovery*/ None, + ) + .with_session_source(&session_source); + + assert_eq!(thread_context.session_source(), Some(&session_source)); + assert_eq!( + McpServerContributionContext::global(&config).session_source(), + None + ); +} + impl ContextContributor for AllContributors { fn contribute_thread_context<'a>( &'a self,