mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Make MCP server contributions thread-scoped (#27670)
## Why `selectedCapabilityRoots` belongs to one thread, but MCP contributors previously received only the global Codex config. That left no clean way for a selected executor capability to contribute MCP servers to its own thread. ## What this PR does - Gives MCP contributors a small context containing the config and, for a running thread, its frozen host-seeded inputs. - Uses the same thread inputs during startup, status queries, refreshes, and skill dependency checks. - Keeps threadless MCP operations and the existing hosted Apps behavior unchanged. - Adds coverage showing that two threads resolve independent registrations and that later lifecycle mutations do not change the frozen MCP inputs. This PR does not discover plugin manifests, add MCP servers, or launch anything new. It only establishes the thread-scoped registration boundary. ## Follow-ups - Resolve selected executor plugin roots through their owning environment filesystem. - Convert their stdio MCP declarations into environment-bound registrations and add an executor MCP end-to-end test. ## Verification - `just fmt` - `cargo check --tests -p codex-protocol -p codex-extension-api -p codex-mcp-extension -p codex-core -p codex-app-server` Tests and Clippy were not run.
This commit is contained in:
@@ -1,7 +1,6 @@
|
||||
use crate::config_manager::ConfigManager;
|
||||
use codex_core::CodexThread;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::McpServerRefreshConfig;
|
||||
use codex_protocol::protocol::Op;
|
||||
@@ -22,8 +21,7 @@ pub(crate) async fn queue_strict_refresh(
|
||||
.get_thread(thread_id)
|
||||
.await
|
||||
.map_err(|err| io::Error::other(format!("failed to load thread {thread_id}: {err}")))?;
|
||||
let config =
|
||||
build_refresh_config(thread_manager, config_manager, thread.config().await).await?;
|
||||
let config = build_refresh_config(thread.as_ref(), config_manager).await?;
|
||||
refreshes.push((thread_id, thread, config));
|
||||
}
|
||||
for (thread_id, thread, config) in refreshes {
|
||||
@@ -44,15 +42,13 @@ pub(crate) async fn queue_best_effort_refresh(
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let config =
|
||||
match build_refresh_config(thread_manager, config_manager, thread.config().await).await
|
||||
{
|
||||
Ok(config) => config,
|
||||
Err(err) => {
|
||||
warn!("failed to build MCP refresh config for thread {thread_id}: {err}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let config = match build_refresh_config(thread.as_ref(), config_manager).await {
|
||||
Ok(config) => config,
|
||||
Err(err) => {
|
||||
warn!("failed to build MCP refresh config for thread {thread_id}: {err}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if let Err(err) = queue_refresh(thread_id, thread, config).await {
|
||||
warn!("{err}");
|
||||
}
|
||||
@@ -60,14 +56,15 @@ pub(crate) async fn queue_best_effort_refresh(
|
||||
}
|
||||
|
||||
async fn build_refresh_config(
|
||||
thread_manager: &ThreadManager,
|
||||
thread: &CodexThread,
|
||||
config_manager: &ConfigManager,
|
||||
thread_config: Arc<Config>,
|
||||
) -> io::Result<McpServerRefreshConfig> {
|
||||
let thread_config = thread.config().await;
|
||||
let config = config_manager
|
||||
.load_latest_config_for_thread(thread_config.as_ref())
|
||||
.await?;
|
||||
let mcp_servers = thread_manager.mcp_manager().runtime_servers(&config).await;
|
||||
let mcp_config = thread.runtime_mcp_config(&config).await;
|
||||
let mcp_servers = codex_mcp::configured_mcp_servers(&mcp_config);
|
||||
Ok(McpServerRefreshConfig {
|
||||
mcp_servers: serde_json::to_value(mcp_servers).map_err(io::Error::other)?,
|
||||
mcp_oauth_credentials_store_mode: serde_json::to_value(
|
||||
|
||||
@@ -203,22 +203,28 @@ impl McpRequestProcessor {
|
||||
let request = request_id.clone();
|
||||
|
||||
let outgoing = Arc::clone(&self.outgoing);
|
||||
let config = match params.thread_id.as_deref() {
|
||||
let (config, thread) = match params.thread_id.as_deref() {
|
||||
Some(thread_id) => {
|
||||
let (_, thread) = self.load_thread(thread_id).await?;
|
||||
let thread_config = thread.config().await;
|
||||
self.config_manager
|
||||
let config = self
|
||||
.config_manager
|
||||
.load_latest_config_for_thread(thread_config.as_ref())
|
||||
.await
|
||||
.map_err(|err| internal_error(format!("failed to reload config: {err}")))?
|
||||
.map_err(|err| internal_error(format!("failed to reload config: {err}")))?;
|
||||
(config, Some(thread))
|
||||
}
|
||||
None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None),
|
||||
};
|
||||
let mcp_config = match thread {
|
||||
Some(thread) => thread.runtime_mcp_config(&config).await,
|
||||
None => {
|
||||
self.thread_manager
|
||||
.mcp_manager()
|
||||
.runtime_config(&config)
|
||||
.await
|
||||
}
|
||||
None => self.load_latest_config(/*fallback_cwd*/ None).await?,
|
||||
};
|
||||
let mcp_config = self
|
||||
.thread_manager
|
||||
.mcp_manager()
|
||||
.runtime_config(&config)
|
||||
.await;
|
||||
let auth = self.auth_manager.auth().await;
|
||||
let environment_manager = self.thread_manager.environment_manager();
|
||||
// This status path has no turn-selected environment. Use config cwd
|
||||
|
||||
@@ -550,6 +550,11 @@ impl CodexThread {
|
||||
self.codex.session.get_config().await
|
||||
}
|
||||
|
||||
/// Resolves the MCP runtime configuration using this thread's extension data.
|
||||
pub async fn runtime_mcp_config(&self, config: &crate::config::Config) -> codex_mcp::McpConfig {
|
||||
self.codex.session.runtime_mcp_config(config).await
|
||||
}
|
||||
|
||||
pub fn multi_agent_version(&self) -> Option<MultiAgentVersion> {
|
||||
self.codex.session.multi_agent_version()
|
||||
}
|
||||
|
||||
@@ -4,8 +4,10 @@ use std::sync::Arc;
|
||||
use crate::config::Config;
|
||||
use codex_config::McpServerConfig;
|
||||
use codex_core_plugins::PluginsManager;
|
||||
use codex_extension_api::ExtensionDataInit;
|
||||
use codex_extension_api::ExtensionRegistry;
|
||||
use codex_extension_api::McpServerContribution;
|
||||
use codex_extension_api::McpServerContributionContext;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
|
||||
use codex_mcp::EffectiveMcpServer;
|
||||
@@ -45,6 +47,24 @@ impl McpManager {
|
||||
/// Returns the MCP config after applying compatibility built-ins and
|
||||
/// runtime-only extension overlays.
|
||||
pub async fn runtime_config(&self, config: &Config) -> McpConfig {
|
||||
self.runtime_config_with_context(config, /*thread_init*/ None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn runtime_config_for_thread(
|
||||
&self,
|
||||
config: &Config,
|
||||
thread_init: &ExtensionDataInit,
|
||||
) -> McpConfig {
|
||||
self.runtime_config_with_context(config, Some(thread_init))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn runtime_config_with_context(
|
||||
&self,
|
||||
config: &Config,
|
||||
thread_init: Option<&ExtensionDataInit>,
|
||||
) -> McpConfig {
|
||||
let mut mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await;
|
||||
let mut catalog = mcp_config.mcp_server_catalog.to_builder();
|
||||
if mcp_config.apps_enabled {
|
||||
@@ -63,9 +83,13 @@ impl McpManager {
|
||||
);
|
||||
}
|
||||
|
||||
let context = match thread_init {
|
||||
Some(thread_init) => McpServerContributionContext::for_thread(config, thread_init),
|
||||
None => McpServerContributionContext::global(config),
|
||||
};
|
||||
let mut contribution_order = 0;
|
||||
for contributor in self.extensions.mcp_server_contributors() {
|
||||
for contribution in contributor.contribute(config).await {
|
||||
for contribution in contributor.contribute(context).await {
|
||||
match contribution {
|
||||
McpServerContribution::Set {
|
||||
name,
|
||||
|
||||
@@ -53,11 +53,7 @@ pub(crate) async fn maybe_prompt_and_install_mcp_dependencies(
|
||||
return;
|
||||
}
|
||||
|
||||
let installed = sess
|
||||
.services
|
||||
.mcp_manager
|
||||
.runtime_servers(config.as_ref())
|
||||
.await;
|
||||
let installed = sess.runtime_mcp_servers(config.as_ref()).await;
|
||||
let missing = collect_missing_mcp_dependencies(mentioned_skills, &installed);
|
||||
if missing.is_empty() {
|
||||
return;
|
||||
@@ -98,7 +94,7 @@ pub(crate) async fn maybe_install_mcp_dependencies(
|
||||
}
|
||||
|
||||
let codex_home = config.codex_home.clone();
|
||||
let installed = sess.services.mcp_manager.runtime_servers(config).await;
|
||||
let installed = sess.runtime_mcp_servers(config).await;
|
||||
let missing = collect_missing_mcp_dependencies(mentioned_skills, &installed);
|
||||
if missing.is_empty() {
|
||||
return;
|
||||
@@ -201,11 +197,7 @@ pub(crate) async fn maybe_install_mcp_dependencies(
|
||||
warn!("failed to refresh MCP dependencies for mentioned skills: {err}");
|
||||
return;
|
||||
}
|
||||
let refresh_servers = sess
|
||||
.services
|
||||
.mcp_manager
|
||||
.runtime_servers(&refresh_config)
|
||||
.await;
|
||||
let refresh_servers = sess.runtime_mcp_servers(&refresh_config).await;
|
||||
sess.refresh_mcp_servers_now(
|
||||
turn_context,
|
||||
refresh_servers,
|
||||
|
||||
@@ -74,6 +74,20 @@ impl ElicitationReviewer for GuardianMcpElicitationReviewer {
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub(crate) async fn runtime_mcp_config(&self, config: &Config) -> McpConfig {
|
||||
self.services
|
||||
.mcp_manager
|
||||
.runtime_config_for_thread(config, &self.services.mcp_thread_init)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn runtime_mcp_servers(
|
||||
&self,
|
||||
config: &Config,
|
||||
) -> HashMap<String, McpServerConfig> {
|
||||
codex_mcp::configured_mcp_servers(&self.runtime_mcp_config(config).await)
|
||||
}
|
||||
|
||||
pub(crate) fn mcp_elicitation_reviewer(self: &Arc<Self>) -> ElicitationReviewerHandle {
|
||||
Arc::new(GuardianMcpElicitationReviewer::new(self))
|
||||
}
|
||||
@@ -289,11 +303,7 @@ impl Session {
|
||||
) {
|
||||
let auth = self.services.auth_manager.auth().await;
|
||||
let config = self.get_config().await;
|
||||
let mcp_config = self
|
||||
.services
|
||||
.mcp_manager
|
||||
.runtime_config(config.as_ref())
|
||||
.await;
|
||||
let mcp_config = self.runtime_mcp_config(config.as_ref()).await;
|
||||
let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config);
|
||||
let mcp_servers =
|
||||
effective_mcp_servers_from_configured(mcp_servers, &mcp_config, auth.as_ref());
|
||||
|
||||
@@ -323,6 +323,7 @@ use crate::unified_exec::UnifiedExecProcessManager;
|
||||
use crate::windows_sandbox::WindowsSandboxLevelExt;
|
||||
use codex_core_plugins::PluginsManager;
|
||||
use codex_git_utils::get_git_repo_root;
|
||||
use codex_mcp::McpConfig;
|
||||
use codex_mcp::compute_auth_statuses;
|
||||
use codex_mcp::effective_mcp_servers_from_configured;
|
||||
use codex_mcp::host_owned_codex_apps_enabled;
|
||||
|
||||
@@ -516,6 +516,11 @@ impl Session {
|
||||
}
|
||||
InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id,
|
||||
};
|
||||
let mcp_thread_init = thread_extension_init.clone();
|
||||
let thread_extension_data = codex_extension_api::ExtensionData::new_with_init(
|
||||
thread_id.to_string(),
|
||||
thread_extension_init,
|
||||
);
|
||||
// Kick off independent async setup tasks in parallel to reduce startup latency.
|
||||
//
|
||||
// - initialize thread persistence with new or resumed session info
|
||||
@@ -598,9 +603,12 @@ impl Session {
|
||||
let auth_manager_clone = Arc::clone(&auth_manager);
|
||||
let config_for_mcp = Arc::clone(&config);
|
||||
let mcp_manager_for_mcp = Arc::clone(&mcp_manager);
|
||||
let mcp_thread_init_for_startup = &mcp_thread_init;
|
||||
let auth_and_mcp_fut = async move {
|
||||
let auth = auth_manager_clone.auth().await;
|
||||
let mcp_config = mcp_manager_for_mcp.runtime_config(&config_for_mcp).await;
|
||||
let mcp_config = mcp_manager_for_mcp
|
||||
.runtime_config_for_thread(&config_for_mcp, mcp_thread_init_for_startup)
|
||||
.await;
|
||||
let mcp_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref());
|
||||
let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config);
|
||||
let auth_statuses = compute_auth_statuses(
|
||||
@@ -955,10 +963,6 @@ impl Session {
|
||||
session_extension_data.insert(McpResourceClient::new(Arc::clone(
|
||||
&mcp_connection_manager,
|
||||
)));
|
||||
let thread_extension_data = codex_extension_api::ExtensionData::new_with_init(
|
||||
thread_id.to_string(),
|
||||
thread_extension_init,
|
||||
);
|
||||
for contributor in extensions.thread_lifecycle_contributors() {
|
||||
contributor.on_thread_start(codex_extension_api::ThreadStartInput {
|
||||
config: config.as_ref(),
|
||||
@@ -1005,6 +1009,7 @@ impl Session {
|
||||
// TODO(jif): extract session to share between sub-agents
|
||||
session_extension_data,
|
||||
thread_extension_data,
|
||||
mcp_thread_init,
|
||||
agent_control,
|
||||
network_proxy: arc_swap::ArcSwapOption::from(network_proxy.map(Arc::new)),
|
||||
network_proxy_audit_metadata,
|
||||
|
||||
@@ -4978,6 +4978,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
|
||||
agent_control.session_id().to_string(),
|
||||
),
|
||||
thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()),
|
||||
mcp_thread_init: codex_extension_api::ExtensionDataInit::default(),
|
||||
agent_control,
|
||||
network_proxy: arc_swap::ArcSwapOption::from(None),
|
||||
network_proxy_audit_metadata: crate::config::NetworkProxyAuditMetadata::default(),
|
||||
@@ -6977,6 +6978,7 @@ where
|
||||
agent_control.session_id().to_string(),
|
||||
),
|
||||
thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()),
|
||||
mcp_thread_init: codex_extension_api::ExtensionDataInit::default(),
|
||||
agent_control,
|
||||
network_proxy: arc_swap::ArcSwapOption::from(None),
|
||||
network_proxy_audit_metadata: crate::config::NetworkProxyAuditMetadata::default(),
|
||||
|
||||
@@ -22,6 +22,7 @@ use codex_analytics::AnalyticsEventsClient;
|
||||
use codex_core_plugins::PluginsManager;
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use codex_extension_api::ExtensionData;
|
||||
use codex_extension_api::ExtensionDataInit;
|
||||
use codex_extension_api::ExtensionRegistry;
|
||||
use codex_hooks::Hooks;
|
||||
use codex_login::AuthManager;
|
||||
@@ -67,6 +68,7 @@ pub(crate) struct SessionServices {
|
||||
pub(crate) extensions: Arc<ExtensionRegistry<crate::config::Config>>,
|
||||
pub(crate) session_extension_data: ExtensionData,
|
||||
pub(crate) thread_extension_data: ExtensionData,
|
||||
pub(crate) mcp_thread_init: ExtensionDataInit,
|
||||
pub(crate) agent_control: AgentControl,
|
||||
pub(crate) network_proxy: ArcSwapOption<StartedNetworkProxy>,
|
||||
pub(crate) network_proxy_audit_metadata: NetworkProxyAuditMetadata,
|
||||
|
||||
@@ -9,6 +9,8 @@ use crate::tasks::InterruptedTurnHistoryMarker;
|
||||
use crate::tasks::interrupted_turn_history_marker;
|
||||
use codex_extension_api::empty_extension_registry;
|
||||
use codex_models_manager::manager::RefreshStrategy;
|
||||
use codex_protocol::capabilities::CapabilityRootLocation;
|
||||
use codex_protocol::capabilities::SelectedCapabilityRoot;
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::ReasoningItemReasoningSummary;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
@@ -387,11 +389,10 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() {
|
||||
struct InitialMarker(&'static str);
|
||||
|
||||
async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() {
|
||||
struct InitialDataRecorder {
|
||||
observed: Arc<std::sync::Mutex<Option<(String, String)>>>,
|
||||
lifecycle_observed: Arc<std::sync::Mutex<Vec<(String, String)>>>,
|
||||
mcp_observed: Arc<std::sync::Mutex<Vec<String>>>,
|
||||
}
|
||||
|
||||
impl codex_extension_api::ThreadLifecycleContributor<Config> for InitialDataRecorder {
|
||||
@@ -400,17 +401,56 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() {
|
||||
input: codex_extension_api::ThreadStartInput<'a, Config>,
|
||||
) -> codex_extension_api::ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let marker = input
|
||||
let selected_root = input
|
||||
.thread_store
|
||||
.get::<InitialMarker>()
|
||||
.expect("initial extension data should be available");
|
||||
*self
|
||||
.observed
|
||||
.get::<Vec<SelectedCapabilityRoot>>()
|
||||
.and_then(|roots| roots.first().cloned())
|
||||
.expect("selected root should be available");
|
||||
self.lifecycle_observed
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some((
|
||||
input.thread_store.level_id().to_string(),
|
||||
marker.0.to_string(),
|
||||
));
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.push((input.thread_store.level_id().to_string(), selected_root.id));
|
||||
input
|
||||
.thread_store
|
||||
.insert(Vec::<SelectedCapabilityRoot>::new());
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl codex_extension_api::McpServerContributor<Config> for InitialDataRecorder {
|
||||
fn id(&self) -> &'static str {
|
||||
"selected_root_test"
|
||||
}
|
||||
|
||||
fn contribute<'a>(
|
||||
&'a self,
|
||||
context: codex_extension_api::McpServerContributionContext<'a, Config>,
|
||||
) -> codex_extension_api::ExtensionFuture<'a, Vec<codex_extension_api::McpServerContribution>>
|
||||
{
|
||||
Box::pin(async move {
|
||||
let thread_init = context
|
||||
.thread_init()
|
||||
.expect("initial MCP resolution should be thread-scoped");
|
||||
let selected_root = thread_init
|
||||
.get::<Vec<SelectedCapabilityRoot>>()
|
||||
.and_then(|roots| roots.first().cloned())
|
||||
.expect("selected root should be available");
|
||||
self.mcp_observed
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.push(selected_root.id.clone());
|
||||
let mut server = codex_mcp::codex_apps_mcp_server_config(
|
||||
"https://selected.invalid",
|
||||
/*apps_mcp_product_sku*/ None,
|
||||
);
|
||||
let CapabilityRootLocation::Environment { environment_id, .. } =
|
||||
&selected_root.location;
|
||||
server.environment_id = environment_id.clone();
|
||||
server.enabled = false;
|
||||
vec![codex_extension_api::McpServerContribution::Set {
|
||||
name: selected_root.id,
|
||||
config: Box::new(server),
|
||||
}]
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -421,11 +461,15 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() {
|
||||
config.cwd = config.codex_home.abs();
|
||||
std::fs::create_dir_all(&config.codex_home).expect("create codex home");
|
||||
|
||||
let observed = Arc::new(std::sync::Mutex::new(None));
|
||||
let lifecycle_observed = Arc::new(std::sync::Mutex::new(Vec::new()));
|
||||
let mcp_observed = Arc::new(std::sync::Mutex::new(Vec::new()));
|
||||
let recorder = Arc::new(InitialDataRecorder {
|
||||
lifecycle_observed: Arc::clone(&lifecycle_observed),
|
||||
mcp_observed: Arc::clone(&mcp_observed),
|
||||
});
|
||||
let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new();
|
||||
extensions.thread_lifecycle_contributor(Arc::new(InitialDataRecorder {
|
||||
observed: Arc::clone(&observed),
|
||||
}));
|
||||
extensions.thread_lifecycle_contributor(recorder.clone());
|
||||
extensions.mcp_server_contributor(recorder);
|
||||
let manager = ThreadManager::new(
|
||||
&config,
|
||||
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()),
|
||||
@@ -439,12 +483,21 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() {
|
||||
TEST_INSTALLATION_ID.to_string(),
|
||||
/*attestation_provider*/ None,
|
||||
);
|
||||
let mut thread_extension_init = codex_extension_api::ExtensionDataInit::new();
|
||||
thread_extension_init.insert(InitialMarker("seeded"));
|
||||
let selected_root_init = |id: &str, environment_id: &str| {
|
||||
let mut init = codex_extension_api::ExtensionDataInit::new();
|
||||
init.insert(vec![SelectedCapabilityRoot {
|
||||
id: id.to_string(),
|
||||
location: CapabilityRootLocation::Environment {
|
||||
environment_id: environment_id.to_string(),
|
||||
path: format!("/plugins/{id}"),
|
||||
},
|
||||
}]);
|
||||
init
|
||||
};
|
||||
|
||||
let thread = manager
|
||||
let first_thread = manager
|
||||
.start_thread_with_options(StartThreadOptions {
|
||||
config,
|
||||
config: config.clone(),
|
||||
initial_history: InitialHistory::New,
|
||||
session_source: None,
|
||||
thread_source: None,
|
||||
@@ -452,17 +505,64 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() {
|
||||
metrics_service_name: None,
|
||||
parent_trace: None,
|
||||
environments: Vec::new(),
|
||||
thread_extension_init,
|
||||
thread_extension_init: selected_root_init("selected-a", "env-a"),
|
||||
})
|
||||
.await
|
||||
.expect("start thread");
|
||||
.expect("start first thread");
|
||||
let second_thread = manager
|
||||
.start_thread_with_options(StartThreadOptions {
|
||||
config: config.clone(),
|
||||
initial_history: InitialHistory::New,
|
||||
session_source: None,
|
||||
thread_source: None,
|
||||
dynamic_tools: Vec::new(),
|
||||
metrics_service_name: None,
|
||||
parent_trace: None,
|
||||
environments: Vec::new(),
|
||||
thread_extension_init: selected_root_init("selected-b", "env-b"),
|
||||
})
|
||||
.await
|
||||
.expect("start second thread");
|
||||
let first_resolved = first_thread.thread.runtime_mcp_config(&config).await;
|
||||
let second_resolved = second_thread.thread.runtime_mcp_config(&config).await;
|
||||
|
||||
assert_eq!(
|
||||
observed
|
||||
*lifecycle_observed
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone(),
|
||||
Some((thread.thread_id.to_string(), "seeded".to_string()))
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner),
|
||||
vec![
|
||||
(first_thread.thread_id.to_string(), "selected-a".to_string()),
|
||||
(
|
||||
second_thread.thread_id.to_string(),
|
||||
"selected-b".to_string()
|
||||
),
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
*mcp_observed
|
||||
.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(),
|
||||
]
|
||||
);
|
||||
let selected_servers = |config: &codex_mcp::McpConfig| {
|
||||
codex_mcp::configured_mcp_servers(config)
|
||||
.into_iter()
|
||||
.filter(|(name, _)| name.starts_with("selected-"))
|
||||
.map(|(name, server)| (name, server.environment_id))
|
||||
.collect::<std::collections::BTreeMap<_, _>>()
|
||||
};
|
||||
assert_eq!(
|
||||
selected_servers(&first_resolved),
|
||||
std::collections::BTreeMap::from([("selected-a".to_string(), "env-a".to_string())])
|
||||
);
|
||||
assert_eq!(
|
||||
selected_servers(&second_resolved),
|
||||
std::collections::BTreeMap::from([("selected-b".to_string(), "env-b".to_string())])
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ mod turn_input;
|
||||
mod turn_lifecycle;
|
||||
|
||||
pub use mcp::McpServerContribution;
|
||||
pub use mcp::McpServerContributionContext;
|
||||
pub use prompt::PromptFragment;
|
||||
pub use prompt::PromptSlot;
|
||||
pub use thread_lifecycle::ThreadIdleInput;
|
||||
@@ -45,13 +46,19 @@ pub type ExtensionFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
|
||||
/// Contributors run in registration order. Later contributions for the same
|
||||
/// name replace earlier ones. Implementations must contribute only names they
|
||||
/// own and must apply any source-specific policy before returning a server.
|
||||
/// Thread-scoped resolution exposes the host-seeded thread inputs; global
|
||||
/// resolution exposes none and must not imply a local fallback. Thread inputs
|
||||
/// are frozen for the runtime and do not include lifecycle-contributor state.
|
||||
/// Plugin-owned servers and their provenance continue to be resolved by the
|
||||
/// plugin manager until that ownership moves into an extension explicitly.
|
||||
pub trait McpServerContributor<C: Sync>: Send + Sync {
|
||||
/// Stable identity used for registration provenance and conflict diagnostics.
|
||||
fn id(&self) -> &'static str;
|
||||
|
||||
fn contribute<'a>(&'a self, config: &'a C) -> ExtensionFuture<'a, Vec<McpServerContribution>>;
|
||||
fn contribute<'a>(
|
||||
&'a self,
|
||||
context: McpServerContributionContext<'a, C>,
|
||||
) -> ExtensionFuture<'a, Vec<McpServerContribution>>;
|
||||
}
|
||||
|
||||
/// Extension contribution that adds prompt fragments during prompt assembly.
|
||||
@@ -69,8 +76,7 @@ pub trait ContextContributor: Send + Sync {
|
||||
/// extension-private thread state. Heavy dependencies belong on the extension
|
||||
/// value created by the host, not in these inputs.
|
||||
pub trait ThreadLifecycleContributor<C: Sync>: Send + Sync {
|
||||
/// Called after thread-scoped extension stores are created, before later
|
||||
/// contributors can read from them.
|
||||
/// Called after host startup has initialized the thread-scoped store.
|
||||
fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let _self = self;
|
||||
|
||||
@@ -1,5 +1,55 @@
|
||||
use codex_config::McpServerConfig;
|
||||
|
||||
use crate::ExtensionDataInit;
|
||||
|
||||
/// Input supplied while resolving MCP server contributions.
|
||||
///
|
||||
/// Thread-scoped implementations can read the immutable host-seeded inputs
|
||||
/// through [`Self::thread_init`]. Implementations should not retain borrowed
|
||||
/// context after contribution completes.
|
||||
pub struct McpServerContributionContext<'a, C> {
|
||||
/// Host configuration visible during MCP resolution.
|
||||
config: &'a C,
|
||||
/// Initial inputs for the active thread, when resolution is thread-scoped.
|
||||
thread_init: Option<&'a ExtensionDataInit>,
|
||||
}
|
||||
|
||||
impl<C> Clone for McpServerContributionContext<'_, C> {
|
||||
fn clone(&self) -> Self {
|
||||
*self
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> Copy for McpServerContributionContext<'_, C> {}
|
||||
|
||||
impl<'a, C> McpServerContributionContext<'a, C> {
|
||||
/// Creates context for resolution that is not associated with a running thread.
|
||||
pub fn global(config: &'a C) -> Self {
|
||||
Self {
|
||||
config,
|
||||
thread_init: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates context for one active thread runtime.
|
||||
pub fn for_thread(config: &'a C, thread_init: &'a ExtensionDataInit) -> Self {
|
||||
Self {
|
||||
config,
|
||||
thread_init: Some(thread_init),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the host configuration visible during resolution.
|
||||
pub fn config(&self) -> &'a C {
|
||||
self.config
|
||||
}
|
||||
|
||||
/// Returns the frozen initial inputs when resolving for a running thread.
|
||||
pub fn thread_init(&self) -> Option<&'a ExtensionDataInit> {
|
||||
self.thread_init
|
||||
}
|
||||
}
|
||||
|
||||
/// One extension-owned overlay for the runtime MCP server configuration.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum McpServerContribution {
|
||||
|
||||
@@ -36,6 +36,7 @@ pub use contributors::ConfigContributor;
|
||||
pub use contributors::ContextContributor;
|
||||
pub use contributors::ExtensionFuture;
|
||||
pub use contributors::McpServerContribution;
|
||||
pub use contributors::McpServerContributionContext;
|
||||
pub use contributors::McpServerContributor;
|
||||
pub use contributors::PromptFragment;
|
||||
pub use contributors::PromptSlot;
|
||||
|
||||
@@ -9,9 +9,11 @@ type ErasedData = Arc<dyn Any + Send + Sync>;
|
||||
|
||||
/// Typed values supplied before an [`ExtensionData`] scope is created.
|
||||
///
|
||||
/// Hosts consume this value once to seed a scope before lifecycle contributors
|
||||
/// run. It does not install extensions or provide persistence.
|
||||
#[derive(Debug, Default)]
|
||||
/// Hosts may retain a clone when later operations must use the same initial
|
||||
/// inputs. Cloning freezes the attachment map and shares each value by `Arc`;
|
||||
/// values with interior mutability remain shared. This type does not install
|
||||
/// extensions or provide persistence.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct ExtensionDataInit {
|
||||
entries: HashMap<TypeId, ErasedData>,
|
||||
}
|
||||
@@ -31,6 +33,15 @@ impl ExtensionDataInit {
|
||||
.insert(TypeId::of::<T>(), Arc::new(value))
|
||||
.map(downcast_data)
|
||||
}
|
||||
|
||||
/// Returns a host-supplied initial attachment without creating a mutable scope.
|
||||
pub fn get<T>(&self) -> Option<Arc<T>>
|
||||
where
|
||||
T: Any + Send + Sync,
|
||||
{
|
||||
let value = self.entries.get(&TypeId::of::<T>())?.clone();
|
||||
Some(downcast_data(value))
|
||||
}
|
||||
}
|
||||
|
||||
/// Typed extension-owned data attached to one host object.
|
||||
|
||||
@@ -2,6 +2,7 @@ use codex_core::config::Config;
|
||||
use codex_extension_api::ExtensionFuture;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_extension_api::McpServerContribution;
|
||||
use codex_extension_api::McpServerContributionContext;
|
||||
use codex_extension_api::McpServerContributor;
|
||||
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
|
||||
use codex_mcp::hosted_plugin_runtime_mcp_server_config;
|
||||
@@ -15,9 +16,10 @@ impl McpServerContributor<Config> for HostedPluginRuntimeExtension {
|
||||
|
||||
fn contribute<'a>(
|
||||
&'a self,
|
||||
config: &'a Config,
|
||||
context: McpServerContributionContext<'a, Config>,
|
||||
) -> ExtensionFuture<'a, Vec<McpServerContribution>> {
|
||||
Box::pin(async move {
|
||||
let config = context.config();
|
||||
let name = CODEX_APPS_MCP_SERVER_NAME.to_string();
|
||||
if !config.features.enabled(codex_features::Feature::Apps) {
|
||||
return vec![McpServerContribution::Remove { name }];
|
||||
|
||||
@@ -7,6 +7,7 @@ use codex_core::config::ConfigBuilder;
|
||||
use codex_core_plugins::PluginsManager;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_extension_api::McpServerContribution;
|
||||
use codex_extension_api::McpServerContributionContext;
|
||||
use codex_extension_api::McpServerContributor;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
|
||||
@@ -191,7 +192,7 @@ impl McpServerContributor<Config> for RemoveCodexApps {
|
||||
|
||||
fn contribute<'a>(
|
||||
&'a self,
|
||||
_config: &'a Config,
|
||||
_context: McpServerContributionContext<'a, Config>,
|
||||
) -> codex_extension_api::ExtensionFuture<'a, Vec<McpServerContribution>> {
|
||||
Box::pin(async move {
|
||||
vec![McpServerContribution::Remove {
|
||||
|
||||
@@ -177,7 +177,9 @@ pub struct W3cTraceContext {
|
||||
/// Config payload for refreshing MCP servers.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct McpServerRefreshConfig {
|
||||
/// Complete runtime server map after source and thread-scoped resolution.
|
||||
pub mcp_servers: Value,
|
||||
/// OAuth credential store mode to use with this server snapshot.
|
||||
pub mcp_oauth_credentials_store_mode: Value,
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user