mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Add turn-start cloud plugin discovery to the MCP extension (#46572)
## What changed - Support an optional cloud provider alongside executor plugins, refreshing its thread-scoped catalog at each regular turn when plugins are enabled. - Use cloud connector metadata for app attribution without starting cloud-declared MCP servers. - Scope cached catalogs to authentication and Apps availability. Retain the previous catalog on discovery failure within the same scope, clear stale catalogs on auth changes, and reject results completed under outdated credentials. - Add a cancellable regular-task startup phase that reconciles MCP before discovery and reapplies contributor state afterward without invalidating resource caches. ## Testing Add regression coverage for per-turn refresh, empty catalogs, discovery failures, plugin disabling, auth changes during discovery, MCP reconciliation and cancellation, and cloud app attribution alongside executor MCP tools. GitOrigin-RevId: 8469eb109ea316e21cbc4122c52ad8000caeef03
This commit is contained in:
committed by
copyberry
parent
d84ebff592
commit
394dadb93d
Generated
+2
@@ -3935,6 +3935,7 @@ dependencies = [
|
||||
"codex-protocol",
|
||||
"codex-utils-absolute-path",
|
||||
"codex-utils-path-uri",
|
||||
"core_test_support",
|
||||
"futures",
|
||||
"pretty_assertions",
|
||||
"serde_json",
|
||||
@@ -3943,6 +3944,7 @@ dependencies = [
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"tracing",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -16,6 +16,7 @@ pub use resource_client::McpEventDefinition;
|
||||
pub use resource_client::McpEventNotification;
|
||||
pub use resource_client::McpEventStream;
|
||||
pub use resource_client::McpResourceClient;
|
||||
pub use resource_client::McpResourceClientAuthKey;
|
||||
pub use resource_client::McpResourceClientCacheKey;
|
||||
pub use resource_client::McpResourcePage;
|
||||
pub use resource_client::McpResourceReadResult;
|
||||
|
||||
@@ -219,6 +219,25 @@ impl PartialEq for McpResourceServerCacheKey {
|
||||
|
||||
impl Eq for McpResourceServerCacheKey {}
|
||||
|
||||
/// Opaque auth scope and server availability, independent of connection-set rebuilds.
|
||||
/// This is for fixed, host-owned sources; it does not fingerprint arbitrary server config.
|
||||
#[derive(Clone)]
|
||||
pub struct McpResourceClientAuthKey {
|
||||
pub(crate) generation: Arc<()>,
|
||||
pub(crate) server: String,
|
||||
pub(crate) available: bool,
|
||||
}
|
||||
|
||||
impl PartialEq for McpResourceClientAuthKey {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
Arc::ptr_eq(&self.generation, &other.generation)
|
||||
&& self.server == other.server
|
||||
&& self.available == other.available
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for McpResourceClientAuthKey {}
|
||||
|
||||
impl std::fmt::Debug for McpResourceClient {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
@@ -243,6 +262,12 @@ impl McpResourceClient {
|
||||
self.runtime.resource_cache_key(server)
|
||||
}
|
||||
|
||||
/// Tracks published auth and source availability without invalidating on unrelated
|
||||
/// environment changes. This key contains no credentials.
|
||||
pub fn auth_cache_key_for_server(&self, server: &str) -> McpResourceClientAuthKey {
|
||||
self.runtime.auth_cache_key_for_server(server)
|
||||
}
|
||||
|
||||
/// Returns whether this client can address the named server.
|
||||
///
|
||||
/// This does not wait for server startup.
|
||||
|
||||
@@ -119,6 +119,7 @@ struct PublishedMcpRuntime {
|
||||
config: Option<Arc<McpConfig>>,
|
||||
auth: Option<CodexAuth>,
|
||||
auth_token: Option<String>,
|
||||
auth_generation: Arc<()>,
|
||||
plugins_available: bool,
|
||||
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
|
||||
environment_selections: Arc<[TurnEnvironmentSelection]>,
|
||||
@@ -145,6 +146,23 @@ fn ensure_host_owned_apps_registration(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl PublishedMcpRuntime {
|
||||
// Shared by dirty detection and cache-key publication so their auth rules agree.
|
||||
fn auth_matches(&self, auth: Option<&CodexAuth>) -> bool {
|
||||
match (self.auth.as_ref(), auth) {
|
||||
(Some(previous), Some(latest)) => {
|
||||
previous == latest
|
||||
&& previous.get_account_id() == latest.get_account_id()
|
||||
&& previous.get_chatgpt_user_id() == latest.get_chatgpt_user_id()
|
||||
&& previous.is_fedramp_account() == latest.is_fedramp_account()
|
||||
&& self.auth_token == latest.get_token().ok()
|
||||
}
|
||||
(None, None) => true,
|
||||
(Some(_), None) | (None, Some(_)) => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct CachedMcpBinding {
|
||||
catalog_revisions: HashMap<String, BindingCatalogRevision>,
|
||||
// Reuse a frozen binding while a model step or caller still needs it.
|
||||
@@ -211,6 +229,7 @@ impl McpRuntime {
|
||||
config: None,
|
||||
auth: None,
|
||||
auth_token: None,
|
||||
auth_generation: Arc::new(()),
|
||||
plugins_available: false,
|
||||
ready_selected_capability_roots: Vec::new(),
|
||||
environment_selections: Arc::default(),
|
||||
@@ -314,6 +333,12 @@ impl McpRuntime {
|
||||
let config = Arc::clone(&input.config);
|
||||
let auth = input.auth.clone();
|
||||
let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok());
|
||||
let current = self.current.load_full();
|
||||
let auth_generation = if current.auth_matches(auth.as_ref()) {
|
||||
Arc::clone(¤t.auth_generation)
|
||||
} else {
|
||||
Arc::new(())
|
||||
};
|
||||
let plugins_available = input.plugins_available;
|
||||
let ready_selected_capability_roots = input.ready_selected_capability_roots.clone();
|
||||
let environment_selections = Arc::clone(&input.runtime_context.environment_selections);
|
||||
@@ -345,6 +370,7 @@ impl McpRuntime {
|
||||
config: Some(config),
|
||||
auth,
|
||||
auth_token,
|
||||
auth_generation,
|
||||
plugins_available,
|
||||
ready_selected_capability_roots,
|
||||
environment_selections,
|
||||
@@ -463,17 +489,18 @@ impl McpRuntime {
|
||||
|
||||
/// Returns whether the published snapshot still belongs to the current credentials.
|
||||
pub fn current_auth_matches(&self, auth: Option<&CodexAuth>) -> bool {
|
||||
self.current.load().auth_matches(auth)
|
||||
}
|
||||
|
||||
pub(crate) fn auth_cache_key_for_server(
|
||||
&self,
|
||||
server: &str,
|
||||
) -> crate::McpResourceClientAuthKey {
|
||||
let current = self.current.load();
|
||||
match (current.auth.as_ref(), auth) {
|
||||
(Some(previous), Some(latest)) => {
|
||||
previous == latest
|
||||
&& previous.get_account_id() == latest.get_account_id()
|
||||
&& previous.get_chatgpt_user_id() == latest.get_chatgpt_user_id()
|
||||
&& previous.is_fedramp_account() == latest.is_fedramp_account()
|
||||
&& current.auth_token == latest.get_token().ok()
|
||||
}
|
||||
(None, None) => true,
|
||||
(Some(_), None) | (None, Some(_)) => false,
|
||||
crate::McpResourceClientAuthKey {
|
||||
generation: Arc::clone(¤t.auth_generation),
|
||||
server: server.to_string(),
|
||||
available: current.connections.contains_server(server),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -982,6 +1009,7 @@ mod tests {
|
||||
config: Some(Arc::new(config)),
|
||||
auth: None,
|
||||
auth_token: None,
|
||||
auth_generation: Arc::new(()),
|
||||
plugins_available: false,
|
||||
ready_selected_capability_roots: Vec::new(),
|
||||
environment_selections: Arc::default(),
|
||||
@@ -1055,6 +1083,7 @@ mod tests {
|
||||
))),
|
||||
auth: None,
|
||||
auth_token: None,
|
||||
auth_generation: Arc::new(()),
|
||||
plugins_available: false,
|
||||
ready_selected_capability_roots: Vec::new(),
|
||||
environment_selections: Arc::default(),
|
||||
|
||||
@@ -82,6 +82,7 @@ impl<'a> McpEnvironmentScope<'a> {
|
||||
}
|
||||
|
||||
pub(crate) struct McpThreadIdentity<'a> {
|
||||
pub(crate) auth_changed: bool,
|
||||
pub(crate) session_source: &'a SessionSource,
|
||||
pub(crate) originator: &'a str,
|
||||
pub(crate) environments: McpEnvironmentScope<'a>,
|
||||
@@ -171,7 +172,8 @@ impl McpManager {
|
||||
ready_selected_capability_roots,
|
||||
executor_capability_discovery,
|
||||
)
|
||||
.with_session_source(identity.session_source),
|
||||
.with_session_source(identity.session_source)
|
||||
.with_auth_changed(identity.auth_changed),
|
||||
Some(identity.originator),
|
||||
identity.environments,
|
||||
identity.disabled_plugin_ids,
|
||||
|
||||
@@ -121,23 +121,26 @@ impl Session {
|
||||
&environments,
|
||||
)
|
||||
.await;
|
||||
let mcp_projection = self
|
||||
.services
|
||||
.mcp_manager
|
||||
.runtime_config_for_step(
|
||||
config,
|
||||
&self.services.mcp_thread_init,
|
||||
&self.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
session_source: &session_source,
|
||||
originator: &originator,
|
||||
disabled_plugin_ids: &disabled_plugin_ids,
|
||||
environments: McpEnvironmentScope::Selected(&environment_selections),
|
||||
},
|
||||
&ready_selected_capability_roots,
|
||||
executor_capability_discovery.as_deref(),
|
||||
)
|
||||
.await;
|
||||
let mcp_projection =
|
||||
self.services
|
||||
.mcp_manager
|
||||
.runtime_config_for_step(
|
||||
config,
|
||||
&self.services.mcp_thread_init,
|
||||
&self.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
auth_changed: !self.services.mcp_runtime.current_auth_matches(
|
||||
self.services.auth_manager.auth_cached().as_ref(),
|
||||
),
|
||||
session_source: &session_source,
|
||||
originator: &originator,
|
||||
disabled_plugin_ids: &disabled_plugin_ids,
|
||||
environments: McpEnvironmentScope::Selected(&environment_selections),
|
||||
},
|
||||
&ready_selected_capability_roots,
|
||||
executor_capability_discovery.as_deref(),
|
||||
)
|
||||
.await;
|
||||
let mcp_config = self
|
||||
.project_selected_environment_mcp_servers(config, &environments, mcp_projection)
|
||||
.await
|
||||
@@ -219,6 +222,10 @@ impl Session {
|
||||
&self.services.mcp_thread_init,
|
||||
&self.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
auth_changed: !self
|
||||
.services
|
||||
.mcp_runtime
|
||||
.current_auth_matches(desired.auth.as_ref()),
|
||||
session_source: &desired.session_source,
|
||||
originator: &desired.originator,
|
||||
disabled_plugin_ids: &desired.disabled_plugin_ids,
|
||||
@@ -288,6 +295,10 @@ impl Session {
|
||||
&self.services.mcp_thread_init,
|
||||
&self.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
auth_changed: !self
|
||||
.services
|
||||
.mcp_runtime
|
||||
.current_auth_matches(desired.auth.as_ref()),
|
||||
session_source: &desired.session_source,
|
||||
originator: &desired.originator,
|
||||
disabled_plugin_ids: &desired.disabled_plugin_ids,
|
||||
@@ -688,6 +699,10 @@ impl Session {
|
||||
&self.services.mcp_thread_init,
|
||||
&self.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
auth_changed: !self
|
||||
.services
|
||||
.mcp_runtime
|
||||
.current_auth_matches(desired.auth.as_ref()),
|
||||
session_source: &turn_context.session_source,
|
||||
originator: &turn_context.originator,
|
||||
disabled_plugin_ids: &disabled_plugin_ids,
|
||||
|
||||
@@ -9,6 +9,11 @@ impl Session {
|
||||
pub(crate) fn request_mcp_runtime_refresh(&self) {
|
||||
// Plugin changes can reuse connections but still change their skill resources.
|
||||
self.services.mcp_runtime.invalidate_resource_caches();
|
||||
self.request_mcp_runtime_reprojection();
|
||||
}
|
||||
|
||||
/// Reproject contributor state without invalidating cached MCP resources.
|
||||
pub(crate) fn request_mcp_runtime_reprojection(&self) {
|
||||
self.mark_mcp_runtime_dirty();
|
||||
self.schedule_mcp_prewarm();
|
||||
}
|
||||
|
||||
@@ -1125,6 +1125,7 @@ impl Session {
|
||||
mcp_thread_init_for_startup,
|
||||
thread_extension_data_for_mcp,
|
||||
McpThreadIdentity {
|
||||
auth_changed: false,
|
||||
session_source: &mcp_session_source,
|
||||
originator: &mcp_originator,
|
||||
disabled_plugin_ids: &mcp_disabled_plugin_ids,
|
||||
@@ -1818,6 +1819,10 @@ impl Session {
|
||||
&sess.services.mcp_thread_init,
|
||||
&sess.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
auth_changed: !sess
|
||||
.services
|
||||
.mcp_runtime
|
||||
.current_auth_matches(latest_auth.as_ref()),
|
||||
session_source: &session_configuration.session_source,
|
||||
originator: &session_configuration.originator,
|
||||
disabled_plugin_ids: &session_configuration.disabled_plugin_ids,
|
||||
|
||||
@@ -3,6 +3,9 @@ mod notification_tests;
|
||||
|
||||
use super::mcp_refresh::McpRefresh;
|
||||
use super::step_context::StepInputs;
|
||||
|
||||
#[path = "turn_start_mcp_tests.rs"]
|
||||
mod turn_start_mcp_tests;
|
||||
use super::step_settings::ResolvedStepSettings;
|
||||
use super::step_settings::StepSettings;
|
||||
use super::step_settings::StepSettingsUpdate;
|
||||
@@ -3285,7 +3288,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() {
|
||||
turn_level_id: String,
|
||||
turn_id: String,
|
||||
collaboration_mode: CollaborationMode,
|
||||
token_usage_at_turn_start: TokenUsage,
|
||||
token_usage_at_turn_start: Option<TokenUsage>,
|
||||
saw_session_store: bool,
|
||||
saw_thread_store: bool,
|
||||
}
|
||||
@@ -3309,7 +3312,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() {
|
||||
turn_level_id: input.turn_store.level_id().to_string(),
|
||||
turn_id: input.turn_id.to_string(),
|
||||
collaboration_mode: input.collaboration_mode.clone(),
|
||||
token_usage_at_turn_start: input.token_usage_at_turn_start.clone(),
|
||||
token_usage_at_turn_start: input.token_usage_at_turn_start.cloned(),
|
||||
saw_session_store: input
|
||||
.session_store
|
||||
.get::<SessionTurnStartMarker>()
|
||||
@@ -3356,7 +3359,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() {
|
||||
turn_level_id: turn_context.sub_id.clone(),
|
||||
turn_id: turn_context.sub_id.clone(),
|
||||
collaboration_mode: turn_context.collaboration_mode(),
|
||||
token_usage_at_turn_start,
|
||||
token_usage_at_turn_start: Some(token_usage_at_turn_start),
|
||||
saw_session_store: true,
|
||||
saw_thread_store: true,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
//! Turn-start MCP reconciliation and cooperative discovery cancellation regressions.
|
||||
|
||||
use super::*;
|
||||
use codex_mcp::McpResourceClient;
|
||||
use codex_mcp::McpResourceClientCacheKey;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
struct Recorder {
|
||||
client: McpResourceClient,
|
||||
requires_mcp: std::sync::atomic::AtomicBool,
|
||||
keys: std::sync::Mutex<Vec<McpResourceClientCacheKey>>,
|
||||
discovered: tokio::sync::Notify,
|
||||
}
|
||||
|
||||
impl codex_extension_api::TurnLifecycleContributor for Recorder {
|
||||
fn turn_start_phase(
|
||||
&self,
|
||||
_thread_store: &codex_extension_api::ExtensionData,
|
||||
) -> codex_extension_api::TurnStartPhase {
|
||||
codex_extension_api::TurnStartPhase::RegularTaskStart
|
||||
}
|
||||
|
||||
fn requires_mcp_runtime(&self, _thread_store: &codex_extension_api::ExtensionData) -> bool {
|
||||
self.requires_mcp.load(std::sync::atomic::Ordering::Relaxed)
|
||||
}
|
||||
|
||||
fn on_turn_start<'a>(
|
||||
&'a self,
|
||||
_input: codex_extension_api::TurnStartInput<'a>,
|
||||
) -> codex_extension_api::ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
self.keys
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.push(self.client.cache_key());
|
||||
self.discovered.notify_one();
|
||||
std::future::pending::<()>().await;
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn turn_start_refreshes_dirty_mcp_reuses_clean_runtime_and_cancels_discovery() {
|
||||
let (mut session, turn_context) = make_session_and_context().await;
|
||||
let recorder = Arc::new(Recorder {
|
||||
client: McpResourceClient::new(Arc::clone(&session.services.mcp_runtime)),
|
||||
requires_mcp: Default::default(),
|
||||
keys: Default::default(),
|
||||
discovered: Default::default(),
|
||||
});
|
||||
let initial_key = recorder.client.cache_key();
|
||||
let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
|
||||
builder.turn_lifecycle_contributor(recorder.clone());
|
||||
session.services.extensions = Arc::new(builder.build());
|
||||
let session = Arc::new(session);
|
||||
let turn_context = Arc::new(turn_context);
|
||||
|
||||
for turn in 0..4 {
|
||||
recorder
|
||||
.requires_mcp
|
||||
.store(turn > 0, std::sync::atomic::Ordering::Relaxed);
|
||||
if turn == 0 {
|
||||
session.mark_mcp_runtime_dirty();
|
||||
} else if turn == 3 {
|
||||
session.services.mcp_runtime.reconnect_on_next_refresh();
|
||||
session.request_mcp_runtime_refresh();
|
||||
}
|
||||
session
|
||||
.spawn_task(
|
||||
Arc::clone(&turn_context),
|
||||
Vec::new(),
|
||||
crate::tasks::RegularTask::new(),
|
||||
)
|
||||
.await;
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
recorder.discovered.notified(),
|
||||
)
|
||||
.await
|
||||
.expect("regular-task contributor should run");
|
||||
assert_eq!(session.mcp_refresh.is_pending(), turn == 0);
|
||||
assert_eq!(
|
||||
recorder
|
||||
.keys
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.len(),
|
||||
turn + 1
|
||||
);
|
||||
if turn > 0 {
|
||||
session.refresh_mcp_if_dirty().await;
|
||||
assert!(
|
||||
recorder.client.cache_key()
|
||||
== recorder
|
||||
.keys
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)[turn],
|
||||
"a clean refresh must reuse the published MCP runtime"
|
||||
);
|
||||
}
|
||||
let (cancellation_token, done) = {
|
||||
let active_turn = session.active_turn.lock().await;
|
||||
let task = active_turn
|
||||
.as_ref()
|
||||
.and_then(|turn| turn.task.as_ref())
|
||||
.expect("discovery should run inside the registered task");
|
||||
(task.cancellation_token.clone(), Arc::clone(&task.done))
|
||||
};
|
||||
let completed = done.notified();
|
||||
tokio::pin!(completed);
|
||||
completed.as_mut().enable();
|
||||
cancellation_token.cancel();
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), completed)
|
||||
.await
|
||||
.expect("discovery cancellation must finish the task without forced abort");
|
||||
assert!(
|
||||
session.mcp_refresh.is_pending(),
|
||||
"cancelled discovery must request an MCP refresh before the task completes"
|
||||
);
|
||||
session.abort_all_tasks(TurnAbortReason::Interrupted).await;
|
||||
}
|
||||
let keys = recorder
|
||||
.keys
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
assert!(
|
||||
keys[0] == initial_key,
|
||||
"unrelated contributors must not reconcile MCP"
|
||||
);
|
||||
assert!(
|
||||
keys[1] != initial_key,
|
||||
"discovery must observe the reconciled generation"
|
||||
);
|
||||
assert!(
|
||||
keys[2] != keys[1],
|
||||
"cancelled cloud preparation must request reprojection on the next turn"
|
||||
);
|
||||
assert!(
|
||||
keys[3] != keys[2],
|
||||
"the next turn must consume a pending reconnect before discovery"
|
||||
);
|
||||
}
|
||||
@@ -1,5 +1,8 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_extension_api::ExtensionData;
|
||||
use codex_extension_api::ThreadIdleCause;
|
||||
use codex_extension_api::TurnStartPhase;
|
||||
use codex_protocol::protocol::CodexErrorInfo;
|
||||
use codex_protocol::protocol::TokenUsage;
|
||||
use codex_protocol::protocol::TurnAbortReason;
|
||||
@@ -9,12 +12,21 @@ use crate::session::turn_context::TurnContext;
|
||||
|
||||
impl Session {
|
||||
pub(super) async fn emit_turn_start_lifecycle(
|
||||
&self,
|
||||
self: &Arc<Self>,
|
||||
turn_context: &TurnContext,
|
||||
token_usage_at_turn_start: &TokenUsage,
|
||||
token_usage_at_turn_start: Option<&TokenUsage>,
|
||||
phase: TurnStartPhase,
|
||||
) {
|
||||
let collaboration_mode = turn_context.collaboration_mode();
|
||||
for contributor in self.services.extensions.turn_lifecycle_contributors() {
|
||||
if contributor.turn_start_phase(&self.services.thread_extension_data) != phase {
|
||||
continue;
|
||||
}
|
||||
if phase == TurnStartPhase::RegularTaskStart
|
||||
&& contributor.requires_mcp_runtime(&self.services.thread_extension_data)
|
||||
{
|
||||
self.refresh_mcp_if_dirty().await;
|
||||
}
|
||||
contributor
|
||||
.on_turn_start(codex_extension_api::TurnStartInput {
|
||||
turn_id: turn_context.sub_id.as_str(),
|
||||
|
||||
@@ -322,8 +322,12 @@ impl Session {
|
||||
self.input_queue
|
||||
.extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items)
|
||||
.await;
|
||||
self.emit_turn_start_lifecycle(turn_context.as_ref(), &token_usage_at_turn_start)
|
||||
.await;
|
||||
self.emit_turn_start_lifecycle(
|
||||
turn_context.as_ref(),
|
||||
Some(&token_usage_at_turn_start),
|
||||
codex_extension_api::TurnStartPhase::BeforeTaskRegistration,
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut active = self.active_turn.lock().await;
|
||||
let turn = active.get_or_insert_with(ActiveTurn::default);
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_async_utils::OrCancelExt;
|
||||
use codex_extension_api::TurnStartPhase;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::session::TurnInput;
|
||||
@@ -47,6 +49,32 @@ impl SessionTask for RegularTask {
|
||||
// not wait on startup prewarm resolution.
|
||||
let prewarmed_client_session = async {
|
||||
sess.emit_turn_started(&ctx).await;
|
||||
// Regular-start contributors run once, after the task is visible and interruptible.
|
||||
let prepares_mcp = sess
|
||||
.services
|
||||
.extensions
|
||||
.turn_lifecycle_contributors()
|
||||
.iter()
|
||||
.any(|contributor| {
|
||||
contributor.turn_start_phase(&sess.services.thread_extension_data)
|
||||
== TurnStartPhase::RegularTaskStart
|
||||
&& contributor.requires_mcp_runtime(&sess.services.thread_extension_data)
|
||||
});
|
||||
let preparation = sess
|
||||
.emit_turn_start_lifecycle(
|
||||
&ctx,
|
||||
/*token_usage_at_turn_start*/ None,
|
||||
TurnStartPhase::RegularTaskStart,
|
||||
)
|
||||
.or_cancel(&cancellation_token)
|
||||
.await;
|
||||
// Even cancelled discovery may have cleared the previous account's catalog.
|
||||
if prepares_mcp {
|
||||
sess.request_mcp_runtime_reprojection();
|
||||
}
|
||||
if preparation.is_err() {
|
||||
return SessionStartupPrewarmResolution::Cancelled;
|
||||
}
|
||||
sess.set_server_reasoning_included(/*included*/ false).await;
|
||||
sess.consume_startup_prewarm_for_regular_turn(&cancellation_token)
|
||||
.await
|
||||
|
||||
@@ -1704,6 +1704,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
|
||||
&first_session.services.mcp_thread_init,
|
||||
&first_session.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
auth_changed: false,
|
||||
session_source: &SessionSource::Exec,
|
||||
originator: &first_originator,
|
||||
disabled_plugin_ids: &[],
|
||||
@@ -1725,6 +1726,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
|
||||
&second_session.services.mcp_thread_init,
|
||||
&second_session.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
auth_changed: false,
|
||||
session_source: &second_session_source,
|
||||
originator: &second_originator,
|
||||
disabled_plugin_ids: &[],
|
||||
@@ -1799,6 +1801,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
|
||||
&first_session.services.mcp_thread_init,
|
||||
&first_session.services.thread_extension_data,
|
||||
McpThreadIdentity {
|
||||
auth_changed: false,
|
||||
session_source: &SessionSource::Exec,
|
||||
originator: &first_originator,
|
||||
disabled_plugin_ids: &disabled_plugin_ids,
|
||||
|
||||
@@ -102,6 +102,7 @@ enum AppsTestToolsListBehavior {
|
||||
AlwaysAvailable,
|
||||
AvailableWhen(Arc<AtomicBool>),
|
||||
AlwaysUnavailable,
|
||||
Custom(Arc<Mutex<Vec<Value>>>),
|
||||
}
|
||||
|
||||
impl AppsTestServer {
|
||||
@@ -226,6 +227,13 @@ impl AppsTestServer {
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn mount_with_tools(
|
||||
server: &MockServer,
|
||||
tools: Arc<Mutex<Vec<Value>>>,
|
||||
) -> Result<Self> {
|
||||
Self::mount_with_tools_list_behavior(server, AppsTestToolsListBehavior::Custom(tools)).await
|
||||
}
|
||||
|
||||
pub async fn mount_without_tools(server: &MockServer) -> Result<Self> {
|
||||
Self::mount_with_tools_list_behavior(server, AppsTestToolsListBehavior::AlwaysUnavailable)
|
||||
.await
|
||||
@@ -551,7 +559,8 @@ impl Respond for CodexAppsJsonRpcResponder {
|
||||
"notifications/initialized" => ResponseTemplate::new(202),
|
||||
"tools/list" => {
|
||||
let tools_available = match &self.tools_list_behavior {
|
||||
AppsTestToolsListBehavior::AlwaysAvailable => true,
|
||||
AppsTestToolsListBehavior::AlwaysAvailable
|
||||
| AppsTestToolsListBehavior::Custom(_) => true,
|
||||
AppsTestToolsListBehavior::AvailableWhen(tools_available) => {
|
||||
tools_available.load(Ordering::SeqCst)
|
||||
}
|
||||
@@ -660,6 +669,14 @@ impl Respond for CodexAppsJsonRpcResponder {
|
||||
"nextCursor": null
|
||||
}
|
||||
});
|
||||
if let AppsTestToolsListBehavior::Custom(tools) = &self.tools_list_behavior {
|
||||
response["result"]["tools"] = json!(
|
||||
tools
|
||||
.lock()
|
||||
.expect("Apps test tools lock should not be poisoned")
|
||||
.clone()
|
||||
);
|
||||
}
|
||||
if !tools_available
|
||||
&& let Some(tools) = response
|
||||
.pointer_mut("/result/tools")
|
||||
|
||||
@@ -1426,6 +1426,31 @@ pub fn test_codex() -> TestCodexBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn run_test_with_large_stack<F, Fut>(name: &str, test: F) -> Result<()>
|
||||
where
|
||||
F: FnOnce() -> Fut + Send + 'static,
|
||||
Fut: Future<Output = Result<()>> + Send + 'static,
|
||||
{
|
||||
const WORKER_THREADS: usize = 2;
|
||||
const TEST_STACK_SIZE_BYTES: usize = 8 * 1024 * 1024;
|
||||
|
||||
let handle = std::thread::Builder::new()
|
||||
.name(name.to_string())
|
||||
.stack_size(TEST_STACK_SIZE_BYTES)
|
||||
.spawn(move || -> Result<()> {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.worker_threads(WORKER_THREADS)
|
||||
.thread_stack_size(TEST_STACK_SIZE_BYTES)
|
||||
.enable_all()
|
||||
.build()?;
|
||||
runtime.block_on(Box::pin(test()))
|
||||
})?;
|
||||
|
||||
handle
|
||||
.join()
|
||||
.map_err(|_| anyhow!("{name} thread panicked"))?
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@@ -55,6 +55,7 @@ pub use turn_input::TurnInputEnvironment;
|
||||
pub use turn_lifecycle::TurnAbortInput;
|
||||
pub use turn_lifecycle::TurnErrorInput;
|
||||
pub use turn_lifecycle::TurnStartInput;
|
||||
pub use turn_lifecycle::TurnStartPhase;
|
||||
pub use turn_lifecycle::TurnStopInput;
|
||||
pub use world_state::PreviousWorldStateSection;
|
||||
pub use world_state::RenderedWorldStateFragment;
|
||||
@@ -185,8 +186,21 @@ pub trait ThreadLifecycleContributor<C: Sync>: Send + Sync {
|
||||
/// extension-private turn state. The host exposes stable identifiers and
|
||||
/// extension stores instead of core runtime objects.
|
||||
pub trait TurnLifecycleContributor: Send + Sync {
|
||||
/// Called after turn-scoped extension stores are created, before the task
|
||||
/// for the turn starts running.
|
||||
/// Selects the start phase using the current thread policy. Disabled callbacks
|
||||
/// may retain the default phase and return without doing any work.
|
||||
fn turn_start_phase(&self, _thread_store: &ExtensionData) -> TurnStartPhase {
|
||||
TurnStartPhase::BeforeTaskRegistration
|
||||
}
|
||||
|
||||
/// Whether regular-task startup must reconcile MCP before this callback runs.
|
||||
/// Ignored before task registration; contributors unrelated to MCP keep the default.
|
||||
fn requires_mcp_runtime(&self, _thread_store: &ExtensionData) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Called once in the selected phase, before task-specific work begins.
|
||||
/// Regular-task preparation may be cancelled before this callback completes;
|
||||
/// stop and abort callbacks must tolerate missing or partial initialization.
|
||||
fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let _self = self;
|
||||
|
||||
@@ -14,6 +14,8 @@ use crate::ExtensionDataInit;
|
||||
pub struct McpServerContributionContext<'a, C> {
|
||||
/// Host configuration visible during MCP resolution.
|
||||
config: &'a C,
|
||||
/// Whether pending auth differs from published auth; invalidate catalogs before projection.
|
||||
auth_changed: bool,
|
||||
/// Extension-owned data for the active thread, when resolution is thread-scoped.
|
||||
thread_store: Option<&'a ExtensionData>,
|
||||
/// Stable host inputs for the active thread, when resolution is thread-scoped.
|
||||
@@ -41,6 +43,7 @@ impl<'a, C> McpServerContributionContext<'a, C> {
|
||||
pub fn global(config: &'a C) -> Self {
|
||||
Self {
|
||||
config,
|
||||
auth_changed: false,
|
||||
thread_store: None,
|
||||
thread_init: None,
|
||||
session_source: None,
|
||||
@@ -61,6 +64,7 @@ impl<'a, C> McpServerContributionContext<'a, C> {
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
auth_changed: false,
|
||||
thread_store: Some(thread_store),
|
||||
thread_init: Some(thread_init),
|
||||
session_source: None,
|
||||
@@ -70,6 +74,17 @@ impl<'a, C> McpServerContributionContext<'a, C> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Marks whether this projection replaces the published MCP authentication.
|
||||
pub fn with_auth_changed(mut self, auth_changed: bool) -> Self {
|
||||
self.auth_changed = auth_changed;
|
||||
self
|
||||
}
|
||||
|
||||
/// Returns whether auth-scoped catalogs must be invalidated before projection.
|
||||
pub fn auth_changed(&self) -> bool {
|
||||
self.auth_changed
|
||||
}
|
||||
|
||||
/// 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);
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
//! Turn lifecycle inputs and scheduling phases for host-owned contributors.
|
||||
|
||||
use codex_protocol::config_types::CollaborationMode;
|
||||
use codex_protocol::protocol::CodexErrorInfo;
|
||||
use codex_protocol::protocol::TokenUsage;
|
||||
@@ -5,6 +7,13 @@ use codex_protocol::protocol::TurnAbortReason;
|
||||
|
||||
use crate::ExtensionData;
|
||||
|
||||
/// Runs before task registration or during cancellable regular-task startup.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum TurnStartPhase {
|
||||
BeforeTaskRegistration,
|
||||
RegularTaskStart,
|
||||
}
|
||||
|
||||
/// Input supplied when the host starts a turn.
|
||||
pub struct TurnStartInput<'a> {
|
||||
/// Stable host-owned turn identifier.
|
||||
@@ -12,7 +21,9 @@ pub struct TurnStartInput<'a> {
|
||||
/// Effective collaboration mode for this turn.
|
||||
pub collaboration_mode: &'a CollaborationMode,
|
||||
/// Total token usage snapshot captured when the turn started.
|
||||
pub token_usage_at_turn_start: &'a TokenUsage,
|
||||
/// Present before task registration; absent during regular-task preparation.
|
||||
/// Token-accounting contributors must remain in `BeforeTaskRegistration`.
|
||||
pub token_usage_at_turn_start: Option<&'a TokenUsage>,
|
||||
/// Store scoped to the host session runtime.
|
||||
pub session_store: &'a ExtensionData,
|
||||
/// Store scoped to this thread runtime.
|
||||
|
||||
@@ -91,6 +91,7 @@ pub use contributors::TurnInputEnvironment;
|
||||
pub use contributors::TurnItemContributor;
|
||||
pub use contributors::TurnLifecycleContributor;
|
||||
pub use contributors::TurnStartInput;
|
||||
pub use contributors::TurnStartPhase;
|
||||
pub use contributors::TurnStopInput;
|
||||
pub use contributors::WorldStateContributionInput;
|
||||
pub use contributors::WorldStateSectionContribution;
|
||||
|
||||
@@ -234,6 +234,11 @@ where
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(token_usage_at_turn_start) = input.token_usage_at_turn_start else {
|
||||
tracing::warn!("skipping goal turn accounting: token baseline unavailable");
|
||||
return;
|
||||
};
|
||||
|
||||
if let Err(err) = self
|
||||
.state_dbs
|
||||
.thread_goals()
|
||||
@@ -247,7 +252,7 @@ where
|
||||
accounting.start_turn(
|
||||
input.turn_id,
|
||||
input.collaboration_mode.mode,
|
||||
input.token_usage_at_turn_start,
|
||||
token_usage_at_turn_start,
|
||||
);
|
||||
if matches!(
|
||||
input.collaboration_mode.mode,
|
||||
|
||||
@@ -263,6 +263,19 @@ async fn create_goal_resets_baseline_before_turn_stop_accounting() -> anyhow::Re
|
||||
let thread_id = test_thread_id()?;
|
||||
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
||||
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
||||
// A missing baseline must not panic or seed accounting with invented token usage.
|
||||
for contributor in harness.registry.turn_lifecycle_contributors() {
|
||||
contributor
|
||||
.on_turn_start(TurnStartInput {
|
||||
turn_id: "missing-baseline",
|
||||
collaboration_mode: &default_collaboration_mode(),
|
||||
token_usage_at_turn_start: None,
|
||||
session_store: &harness.session_store,
|
||||
thread_store: &harness.thread_store,
|
||||
turn_store: &ExtensionData::new("missing-baseline"),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
harness
|
||||
.start_turn(
|
||||
"turn-1",
|
||||
@@ -1766,7 +1779,7 @@ impl GoalExtensionHarness {
|
||||
.on_turn_start(TurnStartInput {
|
||||
turn_id,
|
||||
collaboration_mode: &collaboration_mode,
|
||||
token_usage_at_turn_start: usage,
|
||||
token_usage_at_turn_start: Some(usage),
|
||||
session_store: &self.session_store,
|
||||
thread_store: &self.thread_store,
|
||||
turn_store: &turn_store,
|
||||
|
||||
@@ -33,6 +33,7 @@ tracing = { workspace = true }
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
|
||||
[dev-dependencies]
|
||||
core_test_support = { workspace = true }
|
||||
axum = { workspace = true }
|
||||
codex-exec-server-test-support = { workspace = true }
|
||||
codex-login = { workspace = true }
|
||||
@@ -41,3 +42,4 @@ pretty_assertions = { workspace = true }
|
||||
tempfile = { workspace = true }
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
|
||||
tokio-util = { workspace = true, features = ["rt"] }
|
||||
wiremock = { workspace = true }
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
//! Refreshes the cloud catalog each regular turn and projects Apps-only plugin packages.
|
||||
|
||||
use crate::PluginListQuery;
|
||||
use crate::PluginsThreadState;
|
||||
use crate::plugin_contributor::PluginContributor;
|
||||
use crate::plugin_contributor_state::SelectedPluginMetadata;
|
||||
use codex_core::config::Config;
|
||||
use codex_core_plugins::PluginCatalog;
|
||||
use codex_core_plugins::PluginIdentity;
|
||||
use codex_extension_api::ConfigContributor;
|
||||
use codex_extension_api::ExtensionData;
|
||||
use codex_extension_api::ExtensionFuture;
|
||||
use codex_extension_api::ThreadLifecycleContributor;
|
||||
use codex_extension_api::ThreadStartInput;
|
||||
use codex_extension_api::TurnLifecycleContributor;
|
||||
use codex_extension_api::TurnStartInput;
|
||||
use codex_extension_api::TurnStartPhase;
|
||||
use codex_features::Feature;
|
||||
use codex_mcp::McpResourceClient;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
const MAX_PROJECTED_CLOUD_PLUGINS: usize = 100;
|
||||
const MAX_PLUGIN_DISPLAY_NAME_CHARS: usize = 64;
|
||||
const MAX_CONNECTOR_IDS_PER_PLUGIN: usize = 64;
|
||||
|
||||
impl ThreadLifecycleContributor<Config> for PluginContributor {
|
||||
fn on_thread_start<'a>(
|
||||
&'a self,
|
||||
input: ThreadStartInput<'a, Config>,
|
||||
) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let state = input.thread_store.get_or_init(PluginsThreadState::default);
|
||||
state.contributor_state().plugins_enabled =
|
||||
input.config.features.enabled(Feature::Plugins);
|
||||
if let Some(resources) = &input.mcp_resource_client {
|
||||
input.session_store.insert(resources.as_ref().clone());
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ConfigContributor<Config> for PluginContributor {
|
||||
fn on_config_changed(
|
||||
&self,
|
||||
_session_store: &ExtensionData,
|
||||
thread_store: &ExtensionData,
|
||||
_previous_config: &Config,
|
||||
new_config: &Config,
|
||||
) {
|
||||
let plugins_enabled = new_config.features.enabled(Feature::Plugins);
|
||||
let state = thread_store.get_or_init(PluginsThreadState::default);
|
||||
let mut state = state.contributor_state();
|
||||
state.plugins_enabled = plugins_enabled;
|
||||
if self.providers.cloud.is_none() || !plugins_enabled {
|
||||
state.cloud_generation = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TurnLifecycleContributor for PluginContributor {
|
||||
fn turn_start_phase(&self, thread_store: &ExtensionData) -> TurnStartPhase {
|
||||
if self.providers.cloud.is_some()
|
||||
&& thread_store
|
||||
.get::<PluginsThreadState>()
|
||||
.is_some_and(|state| state.contributor_state().plugins_enabled)
|
||||
{
|
||||
TurnStartPhase::RegularTaskStart
|
||||
} else {
|
||||
TurnStartPhase::BeforeTaskRegistration
|
||||
}
|
||||
}
|
||||
|
||||
fn requires_mcp_runtime(&self, thread_store: &ExtensionData) -> bool {
|
||||
self.turn_start_phase(thread_store) == TurnStartPhase::RegularTaskStart
|
||||
}
|
||||
|
||||
fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
let Some(provider) = &self.providers.cloud else {
|
||||
return;
|
||||
};
|
||||
if !input
|
||||
.thread_store
|
||||
.get::<PluginsThreadState>()
|
||||
.is_some_and(|state| state.contributor_state().plugins_enabled)
|
||||
{
|
||||
return;
|
||||
}
|
||||
let resources = input
|
||||
.session_store
|
||||
.get::<McpResourceClient>()
|
||||
.map(|resources| resources.as_ref().clone());
|
||||
let state = input.thread_store.get_or_init(PluginsThreadState::default);
|
||||
// Change generations before the fallible lookup so only failures within the
|
||||
// same auth scope may retain a previously published catalog.
|
||||
let auth_cache_key = state
|
||||
.contributor_state()
|
||||
.prepare_cloud_generation(resources.clone());
|
||||
match provider
|
||||
.list(PluginListQuery {
|
||||
thread_id: input.thread_store.level_id().to_string(),
|
||||
turn_id: input.turn_id.to_string(),
|
||||
mcp_resources: resources.map(Arc::new),
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(catalog) => {
|
||||
let mut state = state.contributor_state();
|
||||
if !state.publish_cloud_catalog(auth_cache_key.as_ref(), catalog) {
|
||||
tracing::warn!("cloud plugin discovery finished after auth changed");
|
||||
}
|
||||
}
|
||||
Err(_) => tracing::warn!("cloud plugin discovery failed"),
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Cloud packages contribute Apps attribution, never executor MCP declarations.
|
||||
pub(crate) struct CloudPluginMetadata {
|
||||
pub(crate) selection_order: usize,
|
||||
pub(crate) selected_root_id: String,
|
||||
pub(crate) metadata: SelectedPluginMetadata,
|
||||
}
|
||||
|
||||
pub(crate) fn catalog_to_metadata(catalog: &PluginCatalog) -> Vec<CloudPluginMetadata> {
|
||||
catalog
|
||||
.entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(selection_order, entry)| {
|
||||
let PluginIdentity::Remote { remote_plugin_id } = &entry.id else {
|
||||
return None;
|
||||
};
|
||||
let mut seen_connector_ids = HashSet::new();
|
||||
Some(CloudPluginMetadata {
|
||||
selection_order,
|
||||
selected_root_id: format!("cloud:{remote_plugin_id}"),
|
||||
metadata: SelectedPluginMetadata {
|
||||
plugin_id: remote_plugin_id.clone(),
|
||||
plugin_display_name: entry
|
||||
.display_name
|
||||
.chars()
|
||||
.take(MAX_PLUGIN_DISPLAY_NAME_CHARS)
|
||||
.collect(),
|
||||
servers: Vec::new(),
|
||||
connector_ids: entry
|
||||
.connector_ids
|
||||
.iter()
|
||||
.filter(|connector_id| seen_connector_ids.insert((*connector_id).clone()))
|
||||
.take(MAX_CONNECTOR_IDS_PER_PLUGIN)
|
||||
.cloned()
|
||||
.collect(),
|
||||
},
|
||||
})
|
||||
})
|
||||
.take(MAX_PROJECTED_CLOUD_PLUGINS)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "cloud_plugin_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,583 @@
|
||||
//! Tests cloud plugin lifecycle, tool calls, scheduling, and plugin policy.
|
||||
|
||||
use super::*;
|
||||
use crate::PluginProvider;
|
||||
use crate::PluginProviderError;
|
||||
use crate::PluginProviderFuture;
|
||||
use crate::PluginProviders;
|
||||
use crate::PluginsThreadState;
|
||||
use crate::install_plugin_providers;
|
||||
use anyhow::Context;
|
||||
use codex_core_plugins::ExecutorPluginProvider;
|
||||
use codex_core_plugins::PluginCatalog;
|
||||
use codex_core_plugins::PluginCatalogEntry;
|
||||
use codex_core_plugins::PluginIdentity;
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use codex_extension_api::ContextContributor;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_extension_api::SelectedPluginSnapshot;
|
||||
use codex_extension_api::WorldStateContributionInput;
|
||||
use codex_extension_api::WorldStateSectionContribution;
|
||||
use codex_protocol::config_types::CollaborationMode;
|
||||
use codex_protocol::config_types::ModeKind;
|
||||
use codex_protocol::config_types::Settings;
|
||||
use core_test_support::apps_test_server::AppsTestServer;
|
||||
use core_test_support::apps_test_server::recorded_apps_tool_calls;
|
||||
use core_test_support::responses;
|
||||
use core_test_support::responses::ev_completed;
|
||||
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_sequence;
|
||||
use core_test_support::responses::sse;
|
||||
use core_test_support::test_codex::run_test_with_large_stack;
|
||||
use core_test_support::test_codex::test_codex;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Notify;
|
||||
|
||||
#[tokio::test]
|
||||
async fn cloud_discovery_selects_regular_task_phase() {
|
||||
struct ExistingContributor;
|
||||
impl TurnLifecycleContributor for ExistingContributor {}
|
||||
|
||||
let mut builder = ExtensionRegistryBuilder::new();
|
||||
builder.turn_lifecycle_contributor(Arc::new(ExistingContributor));
|
||||
let executor = Arc::new(ExecutorPluginProvider::new(Arc::new(
|
||||
EnvironmentManager::default_for_tests(),
|
||||
)));
|
||||
install_plugin_providers(&mut builder, PluginProviders::new(executor.clone()));
|
||||
install_plugin_providers(
|
||||
&mut builder,
|
||||
PluginProviders::new(executor).with_cloud_provider(Arc::new(CloudProvider::default())),
|
||||
);
|
||||
let registry = builder.build();
|
||||
let thread = ExtensionData::new("thread");
|
||||
for plugins_enabled in [false, true] {
|
||||
thread
|
||||
.get_or_init(PluginsThreadState::default)
|
||||
.contributor_state()
|
||||
.plugins_enabled = plugins_enabled;
|
||||
let dependencies = registry
|
||||
.turn_lifecycle_contributors()
|
||||
.iter()
|
||||
.map(|contributor| contributor.requires_mcp_runtime(&thread))
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(dependencies, vec![false, false, plugins_enabled]);
|
||||
let phases = registry
|
||||
.turn_lifecycle_contributors()
|
||||
.iter()
|
||||
.map(|contributor| contributor.turn_start_phase(&thread))
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
phases,
|
||||
vec![
|
||||
TurnStartPhase::BeforeTaskRegistration,
|
||||
TurnStartPhase::BeforeTaskRegistration,
|
||||
if plugins_enabled {
|
||||
TurnStartPhase::RegularTaskStart
|
||||
} else {
|
||||
TurnStartPhase::BeforeTaskRegistration
|
||||
},
|
||||
],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabling_plugins_clears_cloud_catalog_and_skips_discovery() -> anyhow::Result<()> {
|
||||
let config_home = tempfile::tempdir()?;
|
||||
let mut config = codex_core::config::ConfigBuilder::default()
|
||||
.codex_home(config_home.path().to_path_buf())
|
||||
.fallback_cwd(Some(config_home.path().to_path_buf()))
|
||||
.build()
|
||||
.await?;
|
||||
let mode = CollaborationMode {
|
||||
mode: ModeKind::Default,
|
||||
settings: Settings {
|
||||
model: "test".to_string(),
|
||||
reasoning_effort: None,
|
||||
developer_instructions: None,
|
||||
},
|
||||
};
|
||||
let provider = Arc::new(CloudProvider::default());
|
||||
let mut builder = ExtensionRegistryBuilder::new();
|
||||
let providers = PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new(
|
||||
EnvironmentManager::default_for_tests(),
|
||||
))))
|
||||
.with_cloud_provider(provider.clone());
|
||||
install_plugin_providers(&mut builder, providers);
|
||||
let registry = builder.build();
|
||||
let session = ExtensionData::new("session");
|
||||
let thread = ExtensionData::new("thread");
|
||||
let thread_init = codex_extension_api::ExtensionDataInit::default();
|
||||
for (turn_id, plugins_enabled, reply) in [
|
||||
("off", false, Ok(cloud_catalog())),
|
||||
("on", true, Ok(cloud_catalog())),
|
||||
("plugins-off", false, Ok(cloud_catalog())),
|
||||
(
|
||||
"re-enabled-failed",
|
||||
true,
|
||||
Err("discovery unavailable".to_string()),
|
||||
),
|
||||
("re-enabled", true, Ok(cloud_catalog())),
|
||||
] {
|
||||
let expected_catalog = if plugins_enabled {
|
||||
reply.clone().ok()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let previous_discoveries = provider.calls.load(Ordering::SeqCst);
|
||||
provider.set_reply(reply);
|
||||
let previous = config.clone();
|
||||
config
|
||||
.features
|
||||
.set_enabled(Feature::Plugins, plugins_enabled)?;
|
||||
for contributor in registry.config_contributors() {
|
||||
contributor.on_config_changed(&session, &thread, &previous, &config);
|
||||
}
|
||||
if !plugins_enabled && let Some(state) = thread.get::<PluginsThreadState>() {
|
||||
assert_eq!(state.cloud_catalog(), None);
|
||||
}
|
||||
for contributor in registry.turn_lifecycle_contributors() {
|
||||
contributor
|
||||
.on_turn_start(TurnStartInput {
|
||||
turn_id,
|
||||
collaboration_mode: &mode,
|
||||
token_usage_at_turn_start: None,
|
||||
session_store: &session,
|
||||
thread_store: &thread,
|
||||
turn_store: &ExtensionData::new(turn_id),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
registry.mcp_server_contributors()[0]
|
||||
.contribute(codex_extension_api::McpServerContributionContext::for_step(
|
||||
&config,
|
||||
&thread_init,
|
||||
&thread,
|
||||
"test",
|
||||
&[],
|
||||
/*executor_capability_discovery*/ None,
|
||||
))
|
||||
.await;
|
||||
if let Some(state) = thread.get::<PluginsThreadState>() {
|
||||
assert_eq!(state.cloud_catalog(), expected_catalog, "turn {turn_id}");
|
||||
}
|
||||
assert_eq!(
|
||||
provider.calls.load(Ordering::SeqCst),
|
||||
previous_discoveries + usize::from(plugins_enabled),
|
||||
"turn {turn_id}",
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct CloudProvider {
|
||||
reply: Mutex<Result<PluginCatalog, String>>,
|
||||
calls: AtomicUsize,
|
||||
resources: Mutex<Option<Arc<McpResourceClient>>>,
|
||||
pause_discovery: AtomicBool,
|
||||
discovery_started: Notify,
|
||||
resume_discovery: Notify,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct StepPackages(Mutex<Option<Vec<String>>>);
|
||||
|
||||
impl ContextContributor for StepPackages {
|
||||
fn contribute_world_state<'a>(
|
||||
&'a self,
|
||||
input: WorldStateContributionInput<'a>,
|
||||
) -> ExtensionFuture<'a, Vec<WorldStateSectionContribution>> {
|
||||
Box::pin(async move {
|
||||
*self
|
||||
.0
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = input
|
||||
.thread_store
|
||||
.get::<SelectedPluginSnapshot>()
|
||||
.map(|snapshot| {
|
||||
snapshot
|
||||
.plugins
|
||||
.iter()
|
||||
.map(|plugin| plugin.selected_root_id.clone())
|
||||
.collect()
|
||||
});
|
||||
Vec::new()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for CloudProvider {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
reply: Mutex::new(Ok(PluginCatalog::default())),
|
||||
calls: AtomicUsize::new(0),
|
||||
resources: Mutex::new(None),
|
||||
pause_discovery: AtomicBool::new(false),
|
||||
discovery_started: Notify::new(),
|
||||
resume_discovery: Notify::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CloudProvider {
|
||||
fn set_reply(&self, reply: Result<PluginCatalog, String>) {
|
||||
*self
|
||||
.reply
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = reply;
|
||||
}
|
||||
}
|
||||
|
||||
impl PluginProvider for CloudProvider {
|
||||
fn list(&self, query: PluginListQuery) -> PluginProviderFuture<'_, PluginCatalog> {
|
||||
*self
|
||||
.resources
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = query.mcp_resources;
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
let reply = self
|
||||
.reply
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone();
|
||||
let pause = self.pause_discovery.swap(false, Ordering::SeqCst);
|
||||
Box::pin(async move {
|
||||
if pause {
|
||||
self.discovery_started.notify_one();
|
||||
self.resume_discovery.notified().await;
|
||||
}
|
||||
reply.map_err(PluginProviderError::Message)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn cloud_catalog() -> PluginCatalog {
|
||||
PluginCatalog {
|
||||
entries: vec![PluginCatalogEntry {
|
||||
id: PluginIdentity::Remote {
|
||||
remote_plugin_id: "plugin_remote".to_string(),
|
||||
},
|
||||
display_name: "Cloud Demo".to_string(),
|
||||
version: Some("1.0.0".to_string()),
|
||||
mcp_servers: Default::default(),
|
||||
connector_ids: vec!["calendar".to_string()],
|
||||
locations: vec![codex_core_plugins::PluginSourceLocation::Cloud {
|
||||
resource_uri: "plugin://plugin_remote".to_string(),
|
||||
bundle_uri: None,
|
||||
}],
|
||||
}],
|
||||
warnings: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn app_tool(connector: &str, action: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"name": format!("{connector}_{action}"),
|
||||
"description": format!("Run {action} in {connector}."),
|
||||
"annotations": {"readOnlyHint": true},
|
||||
"inputSchema": {"type": "object", "properties": {}, "additionalProperties": false},
|
||||
"_meta": {
|
||||
"connector_id": connector,
|
||||
"connector_name": connector,
|
||||
"_codex_apps": {
|
||||
"resource_uri": format!("connector://{connector}/tools/{connector}_{action}"),
|
||||
"contains_mcp_source": true,
|
||||
"connector_id": connector,
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cloud_catalog_refreshes_per_turn_and_invalidates_on_auth_change() -> anyhow::Result<()> {
|
||||
run_test_with_large_stack(
|
||||
"cloud-plugin-lifecycle",
|
||||
run_cloud_catalog_refresh_lifecycle,
|
||||
)
|
||||
}
|
||||
|
||||
async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> {
|
||||
let server = responses::start_mock_server().await;
|
||||
// Each account has a fixed inventory; only the auth switch changes the server's tools.
|
||||
let apps_tools = Arc::new(Mutex::new(vec![app_tool("calendar", "list_events")]));
|
||||
let apps_server = AppsTestServer::mount_with_tools(&server, apps_tools.clone()).await?;
|
||||
let mut gmail_catalog = cloud_catalog();
|
||||
gmail_catalog.entries[0].id = PluginIdentity::Remote {
|
||||
remote_plugin_id: "plugin_gmail".to_string(),
|
||||
};
|
||||
gmail_catalog.entries[0].display_name = "Gmail Demo".to_string();
|
||||
gmail_catalog.entries[0].connector_ids = vec!["gmail".to_string()];
|
||||
let provider = Arc::new(CloudProvider::default());
|
||||
let step_packages = Arc::new(StepPackages::default());
|
||||
let mut extensions = ExtensionRegistryBuilder::new();
|
||||
extensions.prompt_contributor(step_packages.clone());
|
||||
install_plugin_providers(
|
||||
&mut extensions,
|
||||
PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new(
|
||||
EnvironmentManager::default_for_tests(),
|
||||
))))
|
||||
.with_cloud_provider(provider.clone()),
|
||||
);
|
||||
let mut builder = test_codex()
|
||||
.with_exec_server_url("none")
|
||||
.with_extensions(Arc::new(extensions.build()))
|
||||
.with_auth(codex_login::CodexAuth::from_external_chatgpt_tokens(
|
||||
"header.e30.initial",
|
||||
"account-a",
|
||||
/*chatgpt_plan_type*/ None,
|
||||
)?)
|
||||
.with_model_info_override("gpt-5.5", |model| {
|
||||
model.supports_search_tool = false;
|
||||
})
|
||||
.with_config(move |config| {
|
||||
config
|
||||
.features
|
||||
.enable(Feature::Plugins)
|
||||
.expect("enable plugins in fixture");
|
||||
config
|
||||
.features
|
||||
.enable(Feature::Apps)
|
||||
.expect("enable apps in fixture");
|
||||
config.chatgpt_base_url = apps_server.chatgpt_base_url;
|
||||
});
|
||||
let test = builder.build_with_auto_env(&server).await?;
|
||||
let calendar = ("calendar", "list_events", "Cloud Demo");
|
||||
let gmail = ("gmail", "search_email", "Gmail Demo");
|
||||
for (turn, reply, expected_plugin, new_account, tool_call) in [
|
||||
(
|
||||
"empty",
|
||||
Ok(PluginCatalog::default()),
|
||||
None,
|
||||
None,
|
||||
Some(calendar),
|
||||
),
|
||||
(
|
||||
"discovered",
|
||||
Ok(cloud_catalog()),
|
||||
Some("cloud:plugin_remote"),
|
||||
None,
|
||||
Some(calendar),
|
||||
),
|
||||
(
|
||||
"same-account-failure",
|
||||
Err("temporarily unavailable".to_string()),
|
||||
Some("cloud:plugin_remote"),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
("removed", Ok(PluginCatalog::default()), None, None, None),
|
||||
(
|
||||
"restored",
|
||||
Ok(cloud_catalog()),
|
||||
Some("cloud:plugin_remote"),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
(
|
||||
"auth-changed-failure",
|
||||
Err("unavailable after auth changed".to_string()),
|
||||
None,
|
||||
Some("account-b"),
|
||||
None,
|
||||
),
|
||||
(
|
||||
"new-account-discovered",
|
||||
Ok(gmail_catalog),
|
||||
Some("cloud:plugin_gmail"),
|
||||
None,
|
||||
Some(gmail),
|
||||
),
|
||||
(
|
||||
"in-flight-auth-change",
|
||||
Ok(cloud_catalog()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
] {
|
||||
let previous_discoveries = provider.calls.load(Ordering::SeqCst);
|
||||
if let Some(account) = new_account {
|
||||
// Do not manually refresh MCP: the auth transition must pick up this new inventory.
|
||||
*apps_tools
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) =
|
||||
vec![app_tool("gmail", "search_email")];
|
||||
let resources = provider
|
||||
.resources
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone()
|
||||
.context("cloud provider resource client missing")?;
|
||||
let previous_auth =
|
||||
resources.auth_cache_key_for_server(codex_mcp::CODEX_APPS_MCP_SERVER_NAME);
|
||||
codex_login::auth::login_with_chatgpt_auth_tokens(
|
||||
test.codex_home_path(),
|
||||
"header.e30.changed",
|
||||
account,
|
||||
/*chatgpt_plan_type*/ None,
|
||||
)?;
|
||||
test.thread_manager.auth_manager().reload().await;
|
||||
tokio::time::timeout(std::time::Duration::from_secs(10), async {
|
||||
loop {
|
||||
if resources.auth_cache_key_for_server(codex_mcp::CODEX_APPS_MCP_SERVER_NAME)
|
||||
!= previous_auth
|
||||
&& test
|
||||
.codex
|
||||
.thread_extension_data()
|
||||
.get::<SelectedPluginSnapshot>()
|
||||
.is_some_and(|snapshot| snapshot.plugins.is_empty())
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
assert_eq!(provider.calls.load(Ordering::SeqCst), previous_discoveries);
|
||||
}
|
||||
provider.set_reply(reply);
|
||||
let previous_calls = recorded_apps_tool_calls(&server).await.len();
|
||||
let call_id = format!("{turn}-call");
|
||||
let mut events = vec![ev_response_created("resp")];
|
||||
if let Some((connector, action, _)) = tool_call {
|
||||
events.push(ev_function_call_with_namespace(
|
||||
&call_id,
|
||||
&format!("mcp__codex_apps__{connector}"),
|
||||
&format!("_{action}"),
|
||||
"{}",
|
||||
));
|
||||
}
|
||||
events.push(ev_completed("resp"));
|
||||
let mut responses = vec![sse(events)];
|
||||
if tool_call.is_some() {
|
||||
responses.push(sse(vec![
|
||||
ev_response_created("after-tool"),
|
||||
ev_completed("after-tool"),
|
||||
]));
|
||||
}
|
||||
let response = mount_sse_sequence(&server, responses).await;
|
||||
*step_packages
|
||||
.0
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = None;
|
||||
if turn == "in-flight-auth-change" {
|
||||
provider.pause_discovery.store(true, Ordering::SeqCst);
|
||||
tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async {
|
||||
tokio::try_join!(
|
||||
test.submit_turn("What capabilities are available?"),
|
||||
async {
|
||||
provider.discovery_started.notified().await;
|
||||
let resources = provider
|
||||
.resources
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone()
|
||||
.context("cloud provider resource client missing")?;
|
||||
let previous_auth = resources
|
||||
.auth_cache_key_for_server(codex_mcp::CODEX_APPS_MCP_SERVER_NAME);
|
||||
// Rotate credentials without replacing the captured runtime/client.
|
||||
codex_login::auth::login_with_chatgpt_auth_tokens(
|
||||
test.codex_home_path(),
|
||||
"header.e30.rotated",
|
||||
"account-b",
|
||||
/*chatgpt_plan_type*/ None,
|
||||
)?;
|
||||
test.thread_manager.auth_manager().reload().await;
|
||||
while resources
|
||||
.auth_cache_key_for_server(codex_mcp::CODEX_APPS_MCP_SERVER_NAME)
|
||||
== previous_auth
|
||||
{
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
provider.resume_discovery.notify_one();
|
||||
anyhow::Ok(())
|
||||
}
|
||||
)
|
||||
})
|
||||
.await??;
|
||||
// The old request must not republish its catalog after credentials change.
|
||||
let state = test
|
||||
.codex
|
||||
.thread_extension_data()
|
||||
.get::<PluginsThreadState>()
|
||||
.context("plugin state missing")?;
|
||||
assert_eq!(state.cloud_catalog(), None);
|
||||
assert_eq!(
|
||||
state
|
||||
.contributor_state()
|
||||
.cloud_generation
|
||||
.as_ref()
|
||||
.and_then(|generation| generation.catalog.clone()),
|
||||
None,
|
||||
);
|
||||
} else {
|
||||
test.submit_turn("What capabilities are available?").await?;
|
||||
}
|
||||
let requests = response.requests();
|
||||
let request = &requests[0];
|
||||
if let Some((connector, action, plugin_name)) = tool_call {
|
||||
assert_eq!(requests.len(), 2);
|
||||
let calls = recorded_apps_tool_calls(&server).await;
|
||||
let output = requests[1].function_call_output(&call_id);
|
||||
assert_eq!(calls.len(), previous_calls + 1, "turn {turn}: {output}");
|
||||
let call = calls.last().context("expected an Apps tool call")?;
|
||||
assert_eq!(
|
||||
call["params"]["name"],
|
||||
json!(format!("{connector}_{action}"))
|
||||
);
|
||||
assert_eq!(call["params"]["arguments"], json!({}));
|
||||
let (other_namespace, other_tool) = if connector == "calendar" {
|
||||
("mcp__codex_apps__gmail", "_search_email")
|
||||
} else {
|
||||
("mcp__codex_apps__calendar", "_list_events")
|
||||
};
|
||||
assert!(request.tool_by_name(other_namespace, other_tool).is_none());
|
||||
let tool = request
|
||||
.tool_by_name(
|
||||
&format!("mcp__codex_apps__{connector}"),
|
||||
&format!("_{action}"),
|
||||
)
|
||||
.expect("the current account's tool should be advertised");
|
||||
assert_eq!(
|
||||
tool["description"]
|
||||
.as_str()
|
||||
.context("tool description missing")?
|
||||
.contains(plugin_name),
|
||||
expected_plugin.is_some(),
|
||||
"cloud discovery should attribute tools to the current account's plugin"
|
||||
);
|
||||
// The server echoes this URI in its result; it must reach the next model step.
|
||||
assert!(
|
||||
output["output"]
|
||||
.as_str()
|
||||
.context("tool output missing")?
|
||||
.contains(&format!(
|
||||
"connector://{connector}/tools/{connector}_{action}"
|
||||
)),
|
||||
"turn {turn}: {output}"
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
*step_packages
|
||||
.0
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner),
|
||||
Some(expected_plugin.into_iter().map(str::to_string).collect()),
|
||||
"turn {turn}"
|
||||
);
|
||||
assert_eq!(
|
||||
provider.calls.load(Ordering::SeqCst),
|
||||
previous_discoveries + 1,
|
||||
"turn {turn}"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -7,10 +7,14 @@ use codex_extension_api::McpServerContributor;
|
||||
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
|
||||
use codex_mcp::hosted_plugin_runtime_mcp_server_config;
|
||||
|
||||
mod cloud_plugin;
|
||||
#[cfg(test)]
|
||||
#[path = "event_stream_tests.rs"]
|
||||
mod event_stream_tests;
|
||||
mod plugin;
|
||||
mod plugin_contributor;
|
||||
mod plugin_contributor_state;
|
||||
mod plugin_providers;
|
||||
mod stream_manager;
|
||||
|
||||
pub use codex_core_plugins::PluginListQuery;
|
||||
@@ -18,6 +22,10 @@ pub use codex_core_plugins::PluginProvider;
|
||||
pub use codex_core_plugins::PluginProviderError;
|
||||
pub use codex_core_plugins::PluginProviderFuture;
|
||||
pub use codex_core_plugins::PluginProviderResult;
|
||||
pub use plugin_contributor::install_plugin_providers;
|
||||
pub use plugin_contributor::install_plugins;
|
||||
pub use plugin_contributor_state::PluginsThreadState;
|
||||
pub use plugin_providers::PluginProviders;
|
||||
pub use stream_manager::McpEventStreamManager;
|
||||
pub use stream_manager::McpEventStreamUpdate;
|
||||
|
||||
@@ -59,16 +67,6 @@ pub fn install(builder: &mut ExtensionRegistryBuilder<Config>) {
|
||||
builder.mcp_server_contributor(std::sync::Arc::new(HostedPluginRuntimeExtension));
|
||||
}
|
||||
|
||||
/// Installs discovery for MCP servers and apps declared by thread-selected plugins.
|
||||
pub fn install_plugins(
|
||||
builder: &mut ExtensionRegistryBuilder<Config>,
|
||||
environment_manager: std::sync::Arc<codex_exec_server::EnvironmentManager>,
|
||||
) {
|
||||
builder.mcp_server_contributor(std::sync::Arc::new(plugin::PluginContributor::new(
|
||||
environment_manager,
|
||||
)));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "lib_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
use codex_config::types::PluginMcpServerConfig;
|
||||
use codex_connectors_extension::PluginAppProvider;
|
||||
use codex_core::config::Config;
|
||||
use codex_core_plugins::ExecutorPluginProvider;
|
||||
use codex_core_plugins::loader::apply_configured_plugin_mcp_server_policies;
|
||||
use codex_core_plugins::loader::configured_plugin_mcp_server_policies;
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use codex_extension_api::ExtensionFuture;
|
||||
use codex_extension_api::McpServerContribution;
|
||||
use codex_extension_api::McpServerContributionContext;
|
||||
@@ -12,51 +10,19 @@ use codex_extension_api::McpServerContributor;
|
||||
use codex_features::Feature;
|
||||
use codex_protocol::capabilities::SelectedCapabilityRoot;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use self::provider::PluginMcpProvider;
|
||||
use crate::PluginsThreadState;
|
||||
use crate::cloud_plugin::CloudPluginMetadata;
|
||||
use crate::cloud_plugin::catalog_to_metadata;
|
||||
use crate::plugin_contributor::PluginContributor;
|
||||
use crate::plugin_contributor_state::CachedSelectedRoot;
|
||||
use crate::plugin_contributor_state::SelectedPluginMetadata;
|
||||
|
||||
mod discovery;
|
||||
mod provider;
|
||||
|
||||
/// Frozen MCP and app declarations for one selected package.
|
||||
///
|
||||
/// Each server config retains the stable logical environment ID. Reconnection may replace the
|
||||
/// concrete environment instance without changing that authority.
|
||||
#[derive(Clone)]
|
||||
struct SelectedPluginMetadata {
|
||||
plugin_id: String,
|
||||
plugin_display_name: String,
|
||||
servers: Vec<(String, codex_config::McpServerConfig)>,
|
||||
connector_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct PluginContributorState {
|
||||
cache: Mutex<Vec<CachedSelectedRoot>>,
|
||||
}
|
||||
|
||||
struct CachedSelectedRoot {
|
||||
root: SelectedCapabilityRoot,
|
||||
metadata: Option<SelectedPluginMetadata>,
|
||||
}
|
||||
|
||||
pub(crate) struct PluginContributor {
|
||||
plugin_provider: ExecutorPluginProvider,
|
||||
mcp_provider: PluginMcpProvider,
|
||||
app_provider: PluginAppProvider,
|
||||
}
|
||||
|
||||
impl PluginContributor {
|
||||
pub(crate) fn new(environment_manager: Arc<EnvironmentManager>) -> Self {
|
||||
Self {
|
||||
plugin_provider: ExecutorPluginProvider::new(Arc::clone(&environment_manager)),
|
||||
mcp_provider: PluginMcpProvider,
|
||||
app_provider: PluginAppProvider,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns metadata for one stable selected root.
|
||||
///
|
||||
/// Successful resolution, including a root that is not a plugin or declares no capabilities,
|
||||
@@ -65,20 +31,19 @@ impl PluginContributor {
|
||||
#[tracing::instrument(name = "mcp.plugin.metadata.load", skip_all)]
|
||||
async fn metadata_for_root(
|
||||
&self,
|
||||
state: &PluginContributorState,
|
||||
state: &PluginsThreadState,
|
||||
selected_root: &SelectedCapabilityRoot,
|
||||
) -> Option<SelectedPluginMetadata> {
|
||||
if let Some(cached) = state
|
||||
.cache
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.contributor_state()
|
||||
.executor_cache
|
||||
.iter()
|
||||
.find(|cached| cached.root == *selected_root)
|
||||
{
|
||||
return cached.metadata.clone();
|
||||
}
|
||||
|
||||
let plugin = match self.plugin_provider.resolve_bound(selected_root).await {
|
||||
let plugin = match self.providers.executor.resolve_bound(selected_root).await {
|
||||
Ok(plugin) => plugin,
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
@@ -95,8 +60,8 @@ impl PluginContributor {
|
||||
// executor-owned files. Read them together so a remote environment only
|
||||
// pays for the slower read instead of both reads back-to-back.
|
||||
let (servers, app_declarations) = tokio::join!(
|
||||
self.mcp_provider.load(&plugin),
|
||||
self.app_provider.load(&plugin)
|
||||
PluginMcpProvider.load(&plugin),
|
||||
PluginAppProvider.load(&plugin)
|
||||
);
|
||||
let servers = servers.unwrap_or_else(|err| {
|
||||
tracing::warn!(
|
||||
@@ -127,10 +92,8 @@ impl PluginContributor {
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
let mut cache = state
|
||||
.cache
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let mut state = state.contributor_state();
|
||||
let cache = &mut state.executor_cache;
|
||||
if let Some(cached) = cache.iter().find(|cached| cached.root == *selected_root) {
|
||||
return cached.metadata.clone();
|
||||
}
|
||||
@@ -155,9 +118,17 @@ impl McpServerContributor<Config> for PluginContributor {
|
||||
let Some(thread_store) = context.thread_store() else {
|
||||
return Vec::new();
|
||||
};
|
||||
let Some(selected_roots) = context.ready_selected_capability_roots() else {
|
||||
return Vec::new();
|
||||
};
|
||||
// Cloud projection is gated independently; executor projection remains unchanged.
|
||||
let cloud_plugins_enabled = self.providers.cloud.is_some()
|
||||
&& context.config().features.enabled(Feature::Plugins);
|
||||
let state = thread_store.get_or_init(PluginsThreadState::default);
|
||||
if !cloud_plugins_enabled || context.auth_changed() {
|
||||
// Clear stale cloud metadata before projecting a replacement runtime.
|
||||
state.contributor_state().cloud_generation = None;
|
||||
}
|
||||
let selected_roots = context
|
||||
.ready_selected_capability_roots()
|
||||
.unwrap_or_default();
|
||||
let plugin_policies =
|
||||
configured_plugin_mcp_server_policies(&context.config().config_layer_stack);
|
||||
let mut contributions = Vec::new();
|
||||
@@ -189,7 +160,6 @@ impl McpServerContributor<Config> for PluginContributor {
|
||||
));
|
||||
}
|
||||
} else {
|
||||
let state = thread_store.get_or_init(PluginContributorState::default);
|
||||
for (selection_order, selected_root) in selected_roots.iter().enumerate() {
|
||||
let Some(plugin) = self.metadata_for_root(&state, selected_root).await else {
|
||||
continue;
|
||||
@@ -204,6 +174,24 @@ impl McpServerContributor<Config> for PluginContributor {
|
||||
}
|
||||
}
|
||||
|
||||
// V1 has no remote plugin identities, so cloud packages are additive for now.
|
||||
// Clone the current snapshot after executor loading and release the lock before projection.
|
||||
if cloud_plugins_enabled && let Some(catalog) = state.cloud_catalog() {
|
||||
for CloudPluginMetadata {
|
||||
selection_order,
|
||||
selected_root_id,
|
||||
metadata,
|
||||
} in catalog_to_metadata(&catalog)
|
||||
{
|
||||
contributions.extend(project_metadata(
|
||||
context.config(),
|
||||
/*plugin_policy*/ None,
|
||||
selection_order,
|
||||
&selected_root_id,
|
||||
metadata,
|
||||
));
|
||||
}
|
||||
}
|
||||
contributions
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
//! Coordinates cloud turn refreshes and executor MCP contributions with thread-owned catalogs.
|
||||
//!
|
||||
//! One contributor is registered for both lifecycles. Provider instances are shared, while
|
||||
//! catalogs and executor metadata live together in thread extension data.
|
||||
|
||||
use crate::PluginProviders;
|
||||
use codex_core::config::Config;
|
||||
use codex_core_plugins::ExecutorPluginProvider;
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub(crate) struct PluginContributor {
|
||||
pub(crate) providers: PluginProviders,
|
||||
}
|
||||
|
||||
/// Installs one coordinator for cloud turn-start and executor MCP callbacks.
|
||||
/// The optional cloud provider never disables executor plugins, and Core's plugin policy applies
|
||||
/// to both paths.
|
||||
pub fn install_plugin_providers(
|
||||
builder: &mut ExtensionRegistryBuilder<Config>,
|
||||
providers: PluginProviders,
|
||||
) {
|
||||
let contributor = Arc::new(PluginContributor { providers });
|
||||
builder.thread_lifecycle_contributor(contributor.clone());
|
||||
builder.config_contributor(contributor.clone());
|
||||
builder.turn_lifecycle_contributor(contributor.clone());
|
||||
builder.mcp_server_contributor(contributor);
|
||||
}
|
||||
|
||||
/// Installs plugin support backed by selected environment roots.
|
||||
pub fn install_plugins(
|
||||
builder: &mut ExtensionRegistryBuilder<Config>,
|
||||
environment_manager: Arc<EnvironmentManager>,
|
||||
) {
|
||||
install_plugin_providers(
|
||||
builder,
|
||||
PluginProviders::new(Arc::new(ExecutorPluginProvider::new(environment_manager))),
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
//! Thread-owned cloud catalog and legacy selected-root cache, guarded by a short-lived mutex.
|
||||
|
||||
use codex_core_plugins::PluginCatalog;
|
||||
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
|
||||
use codex_mcp::McpResourceClient;
|
||||
use codex_mcp::McpResourceClientAuthKey;
|
||||
use codex_protocol::capabilities::SelectedCapabilityRoot;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::MutexGuard;
|
||||
|
||||
/// Plugin configuration and caches owned by one thread; mutable state stays extension-private.
|
||||
#[derive(Default)]
|
||||
pub struct PluginsThreadState {
|
||||
state: Mutex<PluginContributorState>,
|
||||
}
|
||||
|
||||
impl PluginsThreadState {
|
||||
/// Returns an owned cloud snapshot only while its auth generation is current.
|
||||
pub fn cloud_catalog(&self) -> Option<PluginCatalog> {
|
||||
self.contributor_state().cloud_catalog().cloned()
|
||||
}
|
||||
|
||||
pub(crate) fn contributor_state(&self) -> MutexGuard<'_, PluginContributorState> {
|
||||
self.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
}
|
||||
|
||||
/// `None` is unpublished; an empty catalog is a successful discovery result.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct PluginContributorState {
|
||||
pub(crate) plugins_enabled: bool,
|
||||
pub(crate) cloud_generation: Option<CloudPluginGeneration>,
|
||||
pub(crate) executor_cache: Vec<CachedSelectedRoot>,
|
||||
}
|
||||
|
||||
/// Keeps cloud plugin metadata within one published auth scope and Apps availability state.
|
||||
pub(crate) struct CloudPluginGeneration {
|
||||
pub(crate) auth_cache_key: Option<McpResourceClientAuthKey>,
|
||||
pub(crate) mcp_resources: Option<McpResourceClient>,
|
||||
pub(crate) catalog: Option<PluginCatalog>,
|
||||
}
|
||||
|
||||
impl CloudPluginGeneration {
|
||||
pub(crate) fn new(mcp_resources: Option<McpResourceClient>) -> Self {
|
||||
let auth_cache_key = mcp_resources
|
||||
.as_ref()
|
||||
.map(|client| client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME));
|
||||
Self {
|
||||
auth_cache_key,
|
||||
mcp_resources,
|
||||
catalog: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_current(&self) -> bool {
|
||||
self.auth_cache_key
|
||||
== self
|
||||
.mcp_resources
|
||||
.as_ref()
|
||||
.map(|client| client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME))
|
||||
}
|
||||
}
|
||||
|
||||
impl PluginContributorState {
|
||||
/// Returns the cloud catalog only while it belongs to the live auth generation.
|
||||
fn cloud_catalog(&self) -> Option<&PluginCatalog> {
|
||||
self.cloud_generation
|
||||
.as_ref()
|
||||
.filter(|generation| generation.is_current())
|
||||
.and_then(|generation| generation.catalog.as_ref())
|
||||
}
|
||||
|
||||
pub(crate) fn prepare_cloud_generation(
|
||||
&mut self,
|
||||
mcp_resources: Option<McpResourceClient>,
|
||||
) -> Option<McpResourceClientAuthKey> {
|
||||
let auth_cache_key = mcp_resources
|
||||
.as_ref()
|
||||
.map(|client| client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME));
|
||||
if self
|
||||
.cloud_generation
|
||||
.as_ref()
|
||||
.is_none_or(|generation| generation.auth_cache_key != auth_cache_key)
|
||||
{
|
||||
self.cloud_generation = Some(CloudPluginGeneration::new(mcp_resources));
|
||||
}
|
||||
auth_cache_key
|
||||
}
|
||||
|
||||
pub(crate) fn publish_cloud_catalog(
|
||||
&mut self,
|
||||
auth_cache_key: Option<&McpResourceClientAuthKey>,
|
||||
catalog: PluginCatalog,
|
||||
) -> bool {
|
||||
let Some(generation) = self
|
||||
.cloud_generation
|
||||
.as_mut()
|
||||
.filter(|generation| generation.auth_cache_key.as_ref() == auth_cache_key)
|
||||
.filter(|generation| generation.is_current())
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
generation.catalog = Some(catalog);
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct CachedSelectedRoot {
|
||||
pub(crate) root: SelectedCapabilityRoot,
|
||||
pub(crate) metadata: Option<SelectedPluginMetadata>,
|
||||
}
|
||||
|
||||
/// Frozen declarations retain the logical environment ID across executor reconnections.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct SelectedPluginMetadata {
|
||||
pub(crate) plugin_id: String,
|
||||
pub(crate) plugin_display_name: String,
|
||||
pub(crate) servers: Vec<(String, codex_config::McpServerConfig)>,
|
||||
pub(crate) connector_ids: Vec<String>,
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
//! Host-assembled plugin sources; discovery results belong to the active thread, not this registry.
|
||||
|
||||
use codex_core_plugins::ExecutorPluginProvider;
|
||||
use codex_core_plugins::PluginProvider;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Cloud discovery is optional and never replaces executor plugin support.
|
||||
pub struct PluginProviders {
|
||||
// CCA owns and supplies the cloud provider implementation.
|
||||
pub(crate) cloud: Option<Arc<dyn PluginProvider>>,
|
||||
// Keep concrete access to resolve_bound so loaders retain the exact executor filesystem.
|
||||
pub(crate) executor: Arc<ExecutorPluginProvider>,
|
||||
}
|
||||
|
||||
impl PluginProviders {
|
||||
pub fn new(executor: Arc<ExecutorPluginProvider>) -> Self {
|
||||
Self {
|
||||
cloud: None,
|
||||
executor,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_cloud_provider(mut self, provider: Arc<dyn PluginProvider>) -> Self {
|
||||
self.cloud = Some(provider);
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -1,8 +1,20 @@
|
||||
use anyhow::Context;
|
||||
use codex_config::AppToolApproval;
|
||||
use codex_config::McpServerToolConfig;
|
||||
use codex_config::test_support::CloudConfigBundleFixture;
|
||||
use codex_core::EnvironmentConfig;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config::ConfigBuilder;
|
||||
use codex_core::windows_sandbox::WindowsSandboxLevelExt;
|
||||
use codex_core_plugins::ExecutorPluginProvider;
|
||||
use codex_core_plugins::PluginCatalog;
|
||||
use codex_core_plugins::PluginCatalogEntry;
|
||||
use codex_core_plugins::PluginIdentity;
|
||||
use codex_core_plugins::PluginListQuery;
|
||||
use codex_core_plugins::PluginProvider;
|
||||
use codex_core_plugins::PluginProviderError;
|
||||
use codex_core_plugins::PluginProviderFuture;
|
||||
use codex_core_plugins::PluginSourceLocation;
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use codex_exec_server::ExecutorCapabilityDiscoveryCache;
|
||||
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
|
||||
@@ -12,12 +24,34 @@ use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_extension_api::McpServerContribution;
|
||||
use codex_extension_api::McpServerContributionContext;
|
||||
use codex_features::Feature;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_mcp_extension::PluginProviders;
|
||||
use codex_mcp_extension::PluginsThreadState;
|
||||
use codex_mcp_extension::install_plugin_providers;
|
||||
use codex_protocol::capabilities::CapabilityRootLocation;
|
||||
use codex_protocol::capabilities::SelectedCapabilityRoot;
|
||||
use codex_protocol::config_types::WindowsSandboxLevel;
|
||||
use codex_protocol::models::PermissionProfileSnapshot;
|
||||
use codex_utils_path_uri::PathUri;
|
||||
use core_test_support::apps_test_server::AppsTestServer;
|
||||
use core_test_support::apps_test_server::recorded_apps_tool_calls;
|
||||
use core_test_support::responses::ev_assistant_message;
|
||||
use core_test_support::responses::ev_completed;
|
||||
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_sequence;
|
||||
use core_test_support::responses::sse;
|
||||
use core_test_support::responses::start_mock_server;
|
||||
use core_test_support::skip_if_no_network;
|
||||
use core_test_support::test_codex::run_test_with_large_stack;
|
||||
use core_test_support::test_codex::test_codex;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
type TestResult = Result<(), Box<dyn std::error::Error>>;
|
||||
|
||||
@@ -405,7 +439,7 @@ async fn raw_selected_plugin_contributions(
|
||||
None
|
||||
};
|
||||
|
||||
Ok(registry.mcp_server_contributors()[0]
|
||||
let contributions = registry.mcp_server_contributors()[0]
|
||||
.contribute(McpServerContributionContext::for_step(
|
||||
config,
|
||||
&thread_init,
|
||||
@@ -414,5 +448,252 @@ async fn raw_selected_plugin_contributions(
|
||||
&selected_capability_roots,
|
||||
executor_capability_discovery.as_ref(),
|
||||
))
|
||||
.await)
|
||||
.await;
|
||||
Ok(contributions)
|
||||
}
|
||||
|
||||
struct CloudCatalogFixture {
|
||||
reply: Result<PluginCatalog, String>,
|
||||
calls: AtomicUsize,
|
||||
}
|
||||
|
||||
impl PluginProvider for CloudCatalogFixture {
|
||||
fn list(&self, _query: PluginListQuery) -> PluginProviderFuture<'_, PluginCatalog> {
|
||||
Box::pin(async {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
self.reply.clone().map_err(PluginProviderError::Message)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cloud_plugins_append_apps_without_changing_executor_tools() -> anyhow::Result<()> {
|
||||
run_test_with_large_stack(
|
||||
"cloud-plugin-projection",
|
||||
run_cloud_plugin_projection_scenarios,
|
||||
)
|
||||
}
|
||||
|
||||
async fn run_cloud_plugin_projection_scenarios() -> anyhow::Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
core_test_support::skip_if_remote!(Ok(()), "V1 plugin fixtures use host-local files");
|
||||
// Fresh sessions keep unavailable-cloud fallback independent of a previously cached catalog.
|
||||
for (case, installed, cloud_enabled, cloud_available, cloud_empty) in [
|
||||
("cloud_only", false, true, true, false),
|
||||
("cloud_and_executor", true, true, true, false),
|
||||
("unavailable_cloud", true, true, false, false),
|
||||
("rollout_disabled", true, false, true, false),
|
||||
("empty_cloud", true, true, true, true),
|
||||
] {
|
||||
let server = start_mock_server().await;
|
||||
let apps = AppsTestServer::mount(&server).await?;
|
||||
let plugin_server = start_mock_server().await;
|
||||
let plugin_mcp = AppsTestServer::mount_with_tools(
|
||||
&plugin_server,
|
||||
Arc::new(std::sync::Mutex::new(vec![json!({
|
||||
"name": "echo", "description": "Echo a note.",
|
||||
"annotations": {"readOnlyHint": true},
|
||||
"inputSchema": {"type": "object", "properties": {"title": {"type": "string"}}}
|
||||
})])),
|
||||
)
|
||||
.await?;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::body_partial_json(json!({"method": "tools/call"})))
|
||||
.respond_with(|request: &wiremock::Request| {
|
||||
let body: serde_json::Value = serde_json::from_slice(&request.body)
|
||||
.unwrap_or_else(|error| panic!("invalid MCP request JSON: {error}"));
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(json!({
|
||||
"jsonrpc": "2.0", "id": body["id"],
|
||||
"result": {"content": [], "structuredContent": {"echo": body["params"]["arguments"]["title"]}, "isError": false}
|
||||
}))
|
||||
})
|
||||
.with_priority(1)
|
||||
.mount(&plugin_server).await;
|
||||
let declaration =
|
||||
json!({"url": format!("{}/api/codex/ps/mcp", plugin_mcp.chatgpt_base_url)});
|
||||
let cloud = Arc::new(CloudCatalogFixture {
|
||||
reply: if cloud_empty {
|
||||
Ok(PluginCatalog::default())
|
||||
} else if cloud_available {
|
||||
Ok(PluginCatalog {
|
||||
entries: vec![PluginCatalogEntry {
|
||||
id: PluginIdentity::Remote {
|
||||
remote_plugin_id: "plugin_notes".into(),
|
||||
},
|
||||
display_name: "Cloud Notes".into(),
|
||||
version: Some("1".into()),
|
||||
mcp_servers: [("notes".into(), declaration.to_string())].into(),
|
||||
connector_ids: vec!["calendar".into()],
|
||||
locations: vec![PluginSourceLocation::Cloud {
|
||||
resource_uri: "plugin://plugin_notes".into(),
|
||||
bundle_uri: None,
|
||||
}],
|
||||
}],
|
||||
warnings: Vec::new(),
|
||||
})
|
||||
} else {
|
||||
Err("plugin-service unavailable".into())
|
||||
},
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let mut extensions = ExtensionRegistryBuilder::new();
|
||||
let mut providers = PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new(
|
||||
EnvironmentManager::default_for_tests(),
|
||||
))));
|
||||
if cloud_enabled {
|
||||
providers = providers.with_cloud_provider(cloud.clone());
|
||||
}
|
||||
install_plugin_providers(&mut extensions, providers);
|
||||
let mut builder = test_codex()
|
||||
.with_extensions(Arc::new(extensions.build()))
|
||||
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
|
||||
.with_model_info_override("gpt-5.5", |model| model.supports_search_tool = false)
|
||||
.with_config(move |config| {
|
||||
assert!(config.features.enable(Feature::Plugins).is_ok());
|
||||
assert!(config.features.enable(Feature::Apps).is_ok());
|
||||
assert!(
|
||||
config
|
||||
.features
|
||||
.enable(Feature::ExecutorCapabilityDiscovery)
|
||||
.is_ok()
|
||||
);
|
||||
config.chatgpt_base_url = apps.chatgpt_base_url;
|
||||
});
|
||||
let test = builder.build(&server).await?;
|
||||
if installed {
|
||||
let selection = test
|
||||
.codex
|
||||
.environment_selections()
|
||||
.await
|
||||
.into_iter()
|
||||
.next()
|
||||
.context("selected environment missing")?;
|
||||
let root = selection.cwd.join("notes")?;
|
||||
let root_path = root.to_abs_path()?;
|
||||
fs::create_dir_all(root_path.join(".codex-plugin"))?;
|
||||
fs::write(root_path.join(".codex-plugin/plugin.json"), json!({
|
||||
"name": "notes", "version": "1", "interface": {"displayName": "Installed Notes"},
|
||||
"mcpServers": "./.mcp.json"
|
||||
}).to_string())?;
|
||||
fs::write(
|
||||
root_path.join(".mcp.json"),
|
||||
json!({"mcpServers": {"notes": declaration}}).to_string(),
|
||||
)?;
|
||||
test.codex
|
||||
.environment_ready(
|
||||
&selection,
|
||||
EnvironmentConfig {
|
||||
allow_login_shell: false,
|
||||
workspace_roots: selection.workspace_roots.clone(),
|
||||
permission_profile: PermissionProfileSnapshot::legacy(
|
||||
test.config.permissions.permission_profile().clone(),
|
||||
),
|
||||
shell_environment_policy: Default::default(),
|
||||
windows_sandbox_level: WindowsSandboxLevel::from_config(&test.config),
|
||||
windows_sandbox_type: test.config.permissions.windows_sandbox_type,
|
||||
use_legacy_landlock: test.config.features.use_legacy_landlock(),
|
||||
exec_policy: None,
|
||||
mcp_policy: None,
|
||||
network_policy: None,
|
||||
selected_capability_roots: vec![SelectedCapabilityRoot {
|
||||
id: "notes@marketplace".into(),
|
||||
location: CapabilityRootLocation::Environment {
|
||||
environment_id: selection.environment_id.clone(),
|
||||
path: root,
|
||||
},
|
||||
}],
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
let response = mount_sse_sequence(
|
||||
&server,
|
||||
vec![
|
||||
sse(vec![
|
||||
ev_response_created("call"),
|
||||
ev_function_call_with_namespace(
|
||||
"notes-call",
|
||||
"mcp__notes",
|
||||
"echo",
|
||||
r#"{"title":"release check"}"#,
|
||||
),
|
||||
ev_completed("call"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_assistant_message("done", "Finished checking Notes."),
|
||||
ev_completed("done"),
|
||||
]),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
test.submit_turn("Use the Notes MCP to check the release.")
|
||||
.await?;
|
||||
let requests = response.requests();
|
||||
assert_eq!(requests.len(), 2, "{case}");
|
||||
let tool = requests[0].tool_by_name("mcp__notes", "echo");
|
||||
assert_eq!(tool.is_some(), installed, "{case}");
|
||||
let output = requests[1].function_call_output("notes-call");
|
||||
if installed {
|
||||
assert!(
|
||||
tool.context("Notes tool missing")?["description"]
|
||||
.as_str()
|
||||
.context("tool description missing")?
|
||||
.contains("Installed Notes")
|
||||
);
|
||||
assert!(
|
||||
output.to_string().contains("release check"),
|
||||
"{case}: {output}"
|
||||
);
|
||||
}
|
||||
let calls = recorded_apps_tool_calls(&plugin_server).await;
|
||||
assert_eq!(calls.len(), usize::from(installed), "{case}");
|
||||
if installed {
|
||||
assert_eq!(calls[0]["params"]["name"], "echo");
|
||||
assert_eq!(
|
||||
calls[0]["params"]["arguments"],
|
||||
json!({"title": "release check"})
|
||||
);
|
||||
}
|
||||
let initialized = plugin_server
|
||||
.received_requests()
|
||||
.await
|
||||
.context("MCP requests missing")?
|
||||
.iter()
|
||||
.any(|request| {
|
||||
serde_json::from_slice::<serde_json::Value>(&request.body)
|
||||
.is_ok_and(|body| body["method"] == "initialize")
|
||||
});
|
||||
assert_eq!(
|
||||
initialized, installed,
|
||||
"{case}: cloud-only declarations must not start MCPs"
|
||||
);
|
||||
assert_eq!(
|
||||
cloud.calls.load(Ordering::SeqCst),
|
||||
usize::from(cloud_enabled),
|
||||
"{case}"
|
||||
);
|
||||
let state = test
|
||||
.codex
|
||||
.thread_extension_data()
|
||||
.get::<PluginsThreadState>()
|
||||
.context("plugin state missing")?;
|
||||
assert_eq!(
|
||||
state.cloud_catalog().is_some(),
|
||||
cloud_enabled && cloud_available,
|
||||
"{case}"
|
||||
);
|
||||
let calendar = requests[0]
|
||||
.tool_by_name("mcp__codex_apps__calendar", "_list_events")
|
||||
.context("Calendar tool missing")?;
|
||||
assert_eq!(
|
||||
calendar["description"]
|
||||
.as_str()
|
||||
.context("Apps description missing")?
|
||||
.contains("Cloud Notes"),
|
||||
cloud_enabled && cloud_available && !cloud_empty,
|
||||
"{case}"
|
||||
);
|
||||
test.codex.shutdown_and_wait().await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user