openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Expose session sources to MCP contributors (#37167)

Browse Source 
## 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
This commit is contained in:
Rasmus Rygaard
2026-08-05 21:45:33 +00:00
committed by copyberry
parent ad6e48ddd3
commit 15ea598c6e
9 changed files with 239 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/mcp.rs
+11 -4
View File
@@ -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
}
codex-rs/core/src/session/mcp.rs
+23 -10
View File
@@ -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(),
)
codex-rs/core/src/session/mcp_runtime.rs
+3
View File
@@ -16,6 +16,7 @@ pub(super) struct McpDesiredState {
pub(super) auth: Option<CodexAuth>,
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,
};
codex-rs/core/src/session/mod.rs
+1
View File
@@ -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;
codex-rs/core/src/session/session.rs
+9 -2
View File
@@ -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,
)
codex-rs/core/src/thread_manager_tests.rs
+24 -8
View File
@@ -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<std::sync::Mutex<Vec<(String, String)>>>,
mcp_observed: Arc<std::sync::Mutex<Vec<String>>>,
mcp_observed: Arc<std::sync::Mutex<Vec<(String, SessionSource)>>>,
}
impl codex_extension_api::ThreadLifecycleContributor<Config> 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| {
codex-rs/core/tests/suite/mcp_tool_exposure.rs
+124
View File
@@ -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<Mutex<Vec<SessionSource>>>,
}
impl CoalescingMcpContributor {
fn new() -> Self {
Self {
@@ -142,6 +150,30 @@ impl McpServerContributor<Config> for AppsMcpServerContributor {
}
}
impl McpServerContributor<Config> for SessionSourceMcpContributor {
fn id(&self) -> &'static str {
"session_source_mcp_test"
}
fn contribute<'a>(
&'a self,
context: McpServerContributionContext<'a, Config>,
) -> ExtensionFuture<'a, Vec<McpServerContribution>> {
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::<Config>::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;
codex-rs/ext/extension-api/src/contributors/mcp.rs
+16
View File
@@ -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
codex-rs/ext/extension-api/tests/registry.rs
+28
View File
@@ -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,