mirror of
https://github.com/openai/codex.git
synced 2026-10-28 18:58:46 +00:00
Propagate Apps tool refreshes to existing threads (#43900)
## Why Refreshing installed Apps without a thread should update the tools available to existing threads on their next turn. ## What changed - Publish live tool catalogs to clients with matching transport, auth, protocol, and listing settings within the same account and home directory. Ignore the thread attribution header when matching scopes. - Adopt updated tools before catalog reads and new calls, preserving running calls and rejecting calls prepared against an outdated catalog revision. - Keep the newest successful fetch per scope so older refreshes cannot overwrite newer tools, and exclude disk snapshots from live updates. ## Testing Add regression coverage for refresh propagation to an existing thread without another tools listing, scope and account isolation, out-of-order refreshes, late client startup, and running versus stale prepared calls. GitOrigin-RevId: 7a5ee34e23742ce374c6647dc8928b76ea622448
This commit is contained in:
@@ -103,6 +103,74 @@ async fn installed_apps_force_refresh_only_refreshes_tools_snapshot() -> Result<
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||||
|
async fn installed_apps_threadless_refresh_updates_existing_thread_tools() -> Result<()> {
|
||||||
|
let fixture = InstalledAppsFixture::start().await?;
|
||||||
|
fixture.set_tools(vec![connector_tool("alpha", "Alpha")?]);
|
||||||
|
let responses_server = responses::start_mock_server().await;
|
||||||
|
let codex_home = configured_codex_home(fixture.base_url())?;
|
||||||
|
MockResponsesConfig::new(&responses_server.uri())
|
||||||
|
.with_root_config(&format!(
|
||||||
|
"chatgpt_base_url = {:?}\nmcp_oauth_credentials_store = \"file\"",
|
||||||
|
fixture.base_url(),
|
||||||
|
))
|
||||||
|
.enable_feature(Feature::Apps)
|
||||||
|
.disable_feature(Feature::CodeMode)
|
||||||
|
.disable_feature(Feature::CodeModeOnly)
|
||||||
|
.disable_feature(Feature::ToolSearch)
|
||||||
|
.write(codex_home.path())?;
|
||||||
|
let mut app_server = start_app_server(codex_home.path()).await?;
|
||||||
|
let ThreadStartResponse { thread, .. } = app_server
|
||||||
|
.start_thread(ThreadStartParams::default())
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
for (index, (connector_id, connector_name)) in [("alpha", "Alpha"), ("beta", "Beta")]
|
||||||
|
.into_iter()
|
||||||
|
.enumerate()
|
||||||
|
{
|
||||||
|
if index > 0 {
|
||||||
|
fixture.set_tools(vec![connector_tool(connector_id, connector_name)?]);
|
||||||
|
send_installed_request(&mut app_server, /*force_refresh*/ true).await?;
|
||||||
|
}
|
||||||
|
let response_id = format!("response-{connector_id}");
|
||||||
|
let response = responses::mount_sse_once(
|
||||||
|
&responses_server,
|
||||||
|
responses::sse(vec![
|
||||||
|
responses::ev_response_created(&response_id),
|
||||||
|
responses::ev_assistant_message("message", "done"),
|
||||||
|
responses::ev_completed(&response_id),
|
||||||
|
]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
app_server
|
||||||
|
.start_turn_and_wait_for_completion(TurnStartParams {
|
||||||
|
thread_id: thread.id.clone(),
|
||||||
|
input: vec![UserInput::Text {
|
||||||
|
text: "Show the available apps".to_string(),
|
||||||
|
text_elements: Vec::new(),
|
||||||
|
}],
|
||||||
|
..Default::default()
|
||||||
|
})
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let body = response.single_request().body_json();
|
||||||
|
for candidate in ["alpha", "beta"] {
|
||||||
|
assert_eq!(
|
||||||
|
responses::namespace_child_tool(
|
||||||
|
&body,
|
||||||
|
&format!("mcp__codex_apps__{candidate}"),
|
||||||
|
&format!("connector_{candidate}"),
|
||||||
|
)
|
||||||
|
.is_some(),
|
||||||
|
candidate == connector_id,
|
||||||
|
"the existing thread should expose the refreshed Apps catalog: {body}",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
assert_eq!(fixture.list_tools_calls(), index + 1);
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn installed_apps_global_disable_retains_tool_derived_identities() -> Result<()> {
|
async fn installed_apps_global_disable_retains_tool_derived_identities() -> Result<()> {
|
||||||
let fixture = InstalledAppsFixture::start().await?;
|
let fixture = InstalledAppsFixture::start().await?;
|
||||||
|
|||||||
@@ -77,7 +77,10 @@ async fn test_step(
|
|||||||
.await
|
.await
|
||||||
.expect("create in-process MCP client"),
|
.expect("create in-process MCP client"),
|
||||||
);
|
);
|
||||||
let tool_catalog = Arc::new(ClientToolCatalog::new(vec![tool.clone()]));
|
let tool_catalog = Arc::new(ClientToolCatalog::new(
|
||||||
|
vec![tool.clone()],
|
||||||
|
/*updates*/ None,
|
||||||
|
));
|
||||||
let managed_client = Arc::new(ManagedClient {
|
let managed_client = Arc::new(ManagedClient {
|
||||||
_auth_change_notifications: None,
|
_auth_change_notifications: None,
|
||||||
client: Arc::clone(&client),
|
client: Arc::clone(&client),
|
||||||
|
|||||||
@@ -4,17 +4,23 @@
|
|||||||
//! catalog. Snapshot reads and revision-checked calls keep the locks private;
|
//! catalog. Snapshot reads and revision-checked calls keep the locks private;
|
||||||
//! successful refreshes publish after calls using the current catalog finish.
|
//! successful refreshes publish after calls using the current catalog finish.
|
||||||
//! Refresh results retain raw tools and eligibility from the same published runtime.
|
//! Refresh results retain raw tools and eligibility from the same published runtime.
|
||||||
|
//! Optional live updates are adopted before reading or starting another call.
|
||||||
|
|
||||||
use std::collections::HashSet;
|
use std::collections::HashSet;
|
||||||
use std::future::Future;
|
use std::future::Future;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
|
use codex_connectors::ConnectorRuntimeSnapshot;
|
||||||
use tokio::sync::Mutex;
|
use tokio::sync::Mutex;
|
||||||
use tokio::sync::RwLock;
|
use tokio::sync::RwLock;
|
||||||
|
use tokio::sync::RwLockReadGuard;
|
||||||
|
use tokio::sync::watch;
|
||||||
|
|
||||||
use crate::tools::ToolInfo;
|
use crate::tools::ToolInfo;
|
||||||
|
|
||||||
|
type ToolCatalogUpdates = watch::Receiver<Option<Arc<ConnectorRuntimeSnapshot<ToolInfo>>>>;
|
||||||
|
|
||||||
/// The exact Apps catalog returned by an awaited refresh of one published runtime.
|
/// The exact Apps catalog returned by an awaited refresh of one published runtime.
|
||||||
pub struct CodexAppsToolSnapshot {
|
pub struct CodexAppsToolSnapshot {
|
||||||
/// Raw installed tools, including tools hidden or disabled for the model.
|
/// Raw installed tools, including tools hidden or disabled for the model.
|
||||||
@@ -31,24 +37,65 @@ pub(crate) struct ClientToolCatalog {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) struct ToolCatalogSnapshot {
|
pub(crate) struct ToolCatalogSnapshot {
|
||||||
/// Zero is the startup catalog; successful explicit refreshes advance it.
|
/// Advances on explicit refresh or adoption of changed live tools.
|
||||||
pub(crate) revision: u64,
|
pub(crate) revision: u64,
|
||||||
pub(crate) tools: Vec<ToolInfo>,
|
pub(crate) tools: Vec<ToolInfo>,
|
||||||
|
updates: Option<ToolCatalogUpdates>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ClientToolCatalog {
|
impl ClientToolCatalog {
|
||||||
pub(crate) fn new(tools: Vec<ToolInfo>) -> Self {
|
pub(crate) fn new(tools: Vec<ToolInfo>, mut updates: Option<ToolCatalogUpdates>) -> Self {
|
||||||
|
let tools = updates
|
||||||
|
.as_mut()
|
||||||
|
.and_then(|updates| {
|
||||||
|
updates
|
||||||
|
.borrow_and_update()
|
||||||
|
.as_ref()
|
||||||
|
.map(|snapshot| snapshot.tools().to_vec())
|
||||||
|
})
|
||||||
|
.unwrap_or(tools);
|
||||||
Self {
|
Self {
|
||||||
current: RwLock::new(ToolCatalogSnapshot { revision: 0, tools }),
|
current: RwLock::new(ToolCatalogSnapshot {
|
||||||
|
revision: 0,
|
||||||
|
tools,
|
||||||
|
updates,
|
||||||
|
}),
|
||||||
refresh_lock: Mutex::new(()),
|
refresh_lock: Mutex::new(()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn read<R>(&self, read: impl FnOnce(&ToolCatalogSnapshot) -> R) -> R {
|
pub(crate) async fn read<R>(&self, read: impl FnOnce(&ToolCatalogSnapshot) -> R) -> R {
|
||||||
let current = self.current.read().await;
|
let current = self.read_current().await;
|
||||||
read(¤t)
|
read(¤t)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn read_current(&self) -> RwLockReadGuard<'_, ToolCatalogSnapshot> {
|
||||||
|
loop {
|
||||||
|
{
|
||||||
|
let current = self.current.read().await;
|
||||||
|
if !current
|
||||||
|
.updates
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|updates| updates.has_changed().unwrap_or(false))
|
||||||
|
{
|
||||||
|
return current;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut current = self.current.write().await;
|
||||||
|
if let Some(updates) = current.updates.as_mut()
|
||||||
|
&& updates.has_changed().unwrap_or(false)
|
||||||
|
{
|
||||||
|
let snapshot = updates.borrow_and_update().clone();
|
||||||
|
if let Some(snapshot) = snapshot
|
||||||
|
&& current.tools != snapshot.tools()
|
||||||
|
{
|
||||||
|
current.tools = snapshot.tools().to_vec();
|
||||||
|
current.revision += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Serialize fetching and publication, leaving the current catalog usable during the fetch.
|
/// Serialize fetching and publication, leaving the current catalog usable during the fetch.
|
||||||
/// The publication callback runs alongside the exact-client update under the write lock.
|
/// The publication callback runs alongside the exact-client update under the write lock.
|
||||||
#[expect(
|
#[expect(
|
||||||
@@ -65,7 +112,16 @@ impl ClientToolCatalog {
|
|||||||
let (tools, context) = fetch().await?;
|
let (tools, context) = fetch().await?;
|
||||||
let mut current = self.current.write().await;
|
let mut current = self.current.write().await;
|
||||||
let result = publish(&tools, context);
|
let result = publish(&tools, context);
|
||||||
current.tools = tools;
|
current.tools = current
|
||||||
|
.updates
|
||||||
|
.as_mut()
|
||||||
|
.and_then(|updates| {
|
||||||
|
updates
|
||||||
|
.borrow_and_update()
|
||||||
|
.as_ref()
|
||||||
|
.map(|snapshot| snapshot.tools().to_vec())
|
||||||
|
})
|
||||||
|
.unwrap_or(tools);
|
||||||
current.revision += 1;
|
current.revision += 1;
|
||||||
Ok(result)
|
Ok(result)
|
||||||
}
|
}
|
||||||
@@ -84,7 +140,7 @@ impl ClientToolCatalog {
|
|||||||
F: FnOnce() -> Fut,
|
F: FnOnce() -> Fut,
|
||||||
Fut: Future<Output = R>,
|
Fut: Future<Output = R>,
|
||||||
{
|
{
|
||||||
let current = self.current.read().await;
|
let current = self.read_current().await;
|
||||||
if current.revision != expected_revision {
|
if current.revision != expected_revision {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -336,8 +336,28 @@ impl McpConnectionSet {
|
|||||||
let shares_codex_apps_tools_cache = is_host_owned_codex_apps
|
let shares_codex_apps_tools_cache = is_host_owned_codex_apps
|
||||||
&& should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token);
|
&& should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token);
|
||||||
let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| {
|
let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| {
|
||||||
|
// Tools/list has no thread selection or UI capabilities. Only equivalent
|
||||||
|
// transport/auth and listing settings may share executable Apps tools.
|
||||||
|
let mut transport = configured_config.transport.clone();
|
||||||
|
if let McpServerTransportConfig::StreamableHttp {
|
||||||
|
http_headers: Some(headers),
|
||||||
|
..
|
||||||
|
} = &mut transport
|
||||||
|
{
|
||||||
|
// mcp_server_config_for_url in codex-rs/codex-mcp/src/mcp/mod.rs
|
||||||
|
// adds thread attribution that threadless discovery does not carry.
|
||||||
|
headers.retain(|name, _| !name.eq_ignore_ascii_case("originator"));
|
||||||
|
}
|
||||||
|
let mut scope = serde_json::json!([
|
||||||
|
transport,
|
||||||
|
&configured_config.auth,
|
||||||
|
protocol_mode.preferred_protocol_version().as_str(),
|
||||||
|
catalog_item_limit,
|
||||||
|
]);
|
||||||
|
scope.sort_all_objects();
|
||||||
codex_apps_tools_cache
|
codex_apps_tools_cache
|
||||||
.context(codex_home.clone(), codex_apps_tools_cache_key.clone())
|
.context(codex_home.clone(), codex_apps_tools_cache_key.clone())
|
||||||
|
.with_live_scope(scope.to_string())
|
||||||
});
|
});
|
||||||
// The reserved Codex Apps registration follows the shared
|
// The reserved Codex Apps registration follows the shared
|
||||||
// AuthManager across refreshes. In the hosted-plugin path, this
|
// AuthManager across refreshes. In the hosted-plugin path, this
|
||||||
|
|||||||
@@ -446,7 +446,7 @@ impl McpConnectionSet {
|
|||||||
.client()
|
.client()
|
||||||
.await
|
.await
|
||||||
.context("failed to get client")?;
|
.context("failed to get client")?;
|
||||||
let (client_tools, tools, list_start) = managed_client
|
let (tools, list_start) = managed_client
|
||||||
.tool_catalog
|
.tool_catalog
|
||||||
.refresh(
|
.refresh(
|
||||||
|| async {
|
|| async {
|
||||||
@@ -474,8 +474,8 @@ impl McpConnectionSet {
|
|||||||
Ok((client_tools, (fetch_ticket, list_start)))
|
Ok((client_tools, (fetch_ticket, list_start)))
|
||||||
},
|
},
|
||||||
|client_tools, (fetch_ticket, list_start)| {
|
|client_tools, (fetch_ticket, list_start)| {
|
||||||
// Discovery may accept another client's newer fetch. The catalog
|
// Discovery can accept another scope's winner; executable catalogs
|
||||||
// retains the tools fetched through this exact connection.
|
// receive only the latest successful fetch from their own scope.
|
||||||
let tools = match (
|
let tools = match (
|
||||||
managed_client.codex_apps_tools_cache_context.as_ref(),
|
managed_client.codex_apps_tools_cache_context.as_ref(),
|
||||||
fetch_ticket,
|
fetch_ticket,
|
||||||
@@ -489,10 +489,14 @@ impl McpConnectionSet {
|
|||||||
(None, None) => client_tools.to_vec(),
|
(None, None) => client_tools.to_vec(),
|
||||||
_ => unreachable!("Codex Apps fetch ticket requires cache context"),
|
_ => unreachable!("Codex Apps fetch ticket requires cache context"),
|
||||||
};
|
};
|
||||||
(client_tools.to_vec(), tools, list_start)
|
(tools, list_start)
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
|
let client_tools = managed_client
|
||||||
|
.tool_catalog
|
||||||
|
.read(|catalog| catalog.tools.clone())
|
||||||
|
.await;
|
||||||
emit_duration(
|
emit_duration(
|
||||||
MCP_TOOLS_LIST_DURATION_METRIC,
|
MCP_TOOLS_LIST_DURATION_METRIC,
|
||||||
list_start.elapsed(),
|
list_start.elapsed(),
|
||||||
|
|||||||
@@ -458,7 +458,7 @@ async fn create_test_managed_client(tools: Vec<ToolInfo>) -> ManagedClient {
|
|||||||
.expect("create in-process RMCP client"),
|
.expect("create in-process RMCP client"),
|
||||||
),
|
),
|
||||||
server_info: create_test_server_info("Ready"),
|
server_info: create_test_server_info("Ready"),
|
||||||
tool_catalog: Arc::new(ClientToolCatalog::new(tools)),
|
tool_catalog: Arc::new(ClientToolCatalog::new(tools, /*updates*/ None)),
|
||||||
tool_timeout: None,
|
tool_timeout: None,
|
||||||
server_instructions: None,
|
server_instructions: None,
|
||||||
server_supports_sandbox_state_meta_capability: false,
|
server_supports_sandbox_state_meta_capability: false,
|
||||||
@@ -758,7 +758,7 @@ pub(crate) async fn create_test_manager_with_ready_apps_client(
|
|||||||
_auth_change_notifications: None,
|
_auth_change_notifications: None,
|
||||||
client,
|
client,
|
||||||
server_info: create_test_server_info("Codex Apps"),
|
server_info: create_test_server_info("Codex Apps"),
|
||||||
tool_catalog: Arc::new(ClientToolCatalog::new(vec![tool])),
|
tool_catalog: Arc::new(ClientToolCatalog::new(vec![tool], /*updates*/ None)),
|
||||||
tool_timeout: Some(Duration::from_secs(5)),
|
tool_timeout: Some(Duration::from_secs(5)),
|
||||||
server_instructions: None,
|
server_instructions: None,
|
||||||
server_supports_sandbox_state_meta_capability: false,
|
server_supports_sandbox_state_meta_capability: false,
|
||||||
@@ -4941,6 +4941,90 @@ async fn reconcile_reusable_server_with_mcp_config(
|
|||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn apps_catalog_broadcast_preserves_running_calls_and_rejects_stale_calls()
|
||||||
|
-> anyhow::Result<()> {
|
||||||
|
let codex_home = tempdir()?;
|
||||||
|
let context = create_codex_apps_tools_cache_context(
|
||||||
|
codex_home.path().to_path_buf(),
|
||||||
|
/*account_id*/ None,
|
||||||
|
/*chatgpt_user_id*/ None,
|
||||||
|
)
|
||||||
|
.with_live_scope("apps".to_string());
|
||||||
|
let original = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "original")];
|
||||||
|
let updated = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "updated")];
|
||||||
|
let catalog = ClientToolCatalog::new(original.clone(), context.subscribe());
|
||||||
|
|
||||||
|
store_current_tools(&context, original.clone());
|
||||||
|
assert_eq!(
|
||||||
|
catalog.read(|catalog| catalog.revision).await,
|
||||||
|
0,
|
||||||
|
"an unchanged startup broadcast must not invalidate prepared calls"
|
||||||
|
);
|
||||||
|
|
||||||
|
let (release, released) = tokio::sync::oneshot::channel::<()>();
|
||||||
|
let running = catalog.run_with_revision(
|
||||||
|
/*expected_revision*/ 0,
|
||||||
|
|| async { released.await.unwrap() },
|
||||||
|
);
|
||||||
|
tokio::pin!(running);
|
||||||
|
assert!(futures::poll!(&mut running).is_pending());
|
||||||
|
store_current_tools(&context, updated.clone());
|
||||||
|
let stale = catalog.run_with_revision(
|
||||||
|
/*expected_revision*/ 0,
|
||||||
|
|| async { panic!("stale preparation") },
|
||||||
|
);
|
||||||
|
tokio::pin!(stale);
|
||||||
|
assert!(
|
||||||
|
futures::poll!(&mut stale).is_pending(),
|
||||||
|
"publication does not wait, but adoption must wait for the running call"
|
||||||
|
);
|
||||||
|
release.send(()).unwrap();
|
||||||
|
assert_eq!(running.await, Some(()));
|
||||||
|
assert_eq!(stale.await, None::<()>);
|
||||||
|
assert_eq!(
|
||||||
|
catalog
|
||||||
|
.read(|catalog| (catalog.revision, catalog.tools.clone()))
|
||||||
|
.await,
|
||||||
|
(1, updated.clone())
|
||||||
|
);
|
||||||
|
|
||||||
|
// A client whose startup finishes late adopts the already-published result.
|
||||||
|
let late = ClientToolCatalog::new(original, context.subscribe());
|
||||||
|
assert_eq!(late.read(|catalog| catalog.tools.clone()).await, updated);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn apps_catalog_broadcast_survives_an_older_local_refresh() -> anyhow::Result<()> {
|
||||||
|
let codex_home = tempdir()?;
|
||||||
|
let context = create_codex_apps_tools_cache_context(
|
||||||
|
codex_home.path().to_path_buf(),
|
||||||
|
/*account_id*/ None,
|
||||||
|
/*chatgpt_user_id*/ None,
|
||||||
|
)
|
||||||
|
.with_live_scope("apps".to_string());
|
||||||
|
let catalog = ClientToolCatalog::new(Vec::new(), context.subscribe());
|
||||||
|
let older_ticket = context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh);
|
||||||
|
let newer = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer")];
|
||||||
|
store_current_tools(&context, newer.clone());
|
||||||
|
assert_eq!(catalog.read(|catalog| catalog.tools.clone()).await, newer);
|
||||||
|
catalog
|
||||||
|
.refresh(
|
||||||
|
|| async { Ok((Vec::new(), older_ticket)) },
|
||||||
|
|tools, ticket| {
|
||||||
|
context.publish_if_newest_accepted(
|
||||||
|
ticket,
|
||||||
|
&create_test_server_info("Apps"),
|
||||||
|
tools.to_vec(),
|
||||||
|
)
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
assert_eq!(catalog.read(|catalog| catalog.tools.clone()).await, newer);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn refreshed_catalog_follows_reused_client_without_mutating_old_bindings()
|
async fn refreshed_catalog_follows_reused_client_without_mutating_old_bindings()
|
||||||
-> anyhow::Result<()> {
|
-> anyhow::Result<()> {
|
||||||
@@ -5083,7 +5167,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a
|
|||||||
_auth_change_notifications: None,
|
_auth_change_notifications: None,
|
||||||
client,
|
client,
|
||||||
server_info: create_test_server_info("Mutable tools"),
|
server_info: create_test_server_info("Mutable tools"),
|
||||||
tool_catalog: Arc::new(ClientToolCatalog::new(initial_tools)),
|
tool_catalog: Arc::new(ClientToolCatalog::new(initial_tools, /*updates*/ None)),
|
||||||
tool_timeout: None,
|
tool_timeout: None,
|
||||||
server_instructions: initialize.instructions,
|
server_instructions: initialize.instructions,
|
||||||
server_supports_sandbox_state_meta_capability: false,
|
server_supports_sandbox_state_meta_capability: false,
|
||||||
|
|||||||
@@ -1005,7 +1005,12 @@ async fn start_server_task(
|
|||||||
_auth_change_notifications: auth_change_notifications,
|
_auth_change_notifications: auth_change_notifications,
|
||||||
client: Arc::clone(&client),
|
client: Arc::clone(&client),
|
||||||
server_info,
|
server_info,
|
||||||
tool_catalog: Arc::new(ClientToolCatalog::new(client_tools)),
|
tool_catalog: Arc::new(ClientToolCatalog::new(
|
||||||
|
client_tools,
|
||||||
|
codex_apps_tools_cache_context
|
||||||
|
.as_ref()
|
||||||
|
.and_then(ConnectorRuntimeContext::subscribe),
|
||||||
|
)),
|
||||||
tool_timeout: None,
|
tool_timeout: None,
|
||||||
server_instructions: initialize_result.instructions,
|
server_instructions: initialize_result.instructions,
|
||||||
server_supports_sandbox_state_meta_capability,
|
server_supports_sandbox_state_meta_capability,
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ use crate::mcp::sanitize_responses_api_tool_name;
|
|||||||
|
|
||||||
const LEGACY_MCP_TOOL_NAME_PREFIX: &str = "mcp__";
|
const LEGACY_MCP_TOOL_NAME_PREFIX: &str = "mcp__";
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||||
pub struct ToolInfo {
|
pub struct ToolInfo {
|
||||||
/// Raw MCP server name used for routing the tool call.
|
/// Raw MCP server name used for routing the tool call.
|
||||||
pub server_name: String,
|
pub server_name: String,
|
||||||
|
|||||||
@@ -4,6 +4,8 @@
|
|||||||
//! workspace. Disk is best-effort cold-start persistence; a context reads it
|
//! workspace. Disk is best-effort cold-start persistence; a context reads it
|
||||||
//! once when created and never rereads it. Full connector metadata is
|
//! once when created and never rereads it. Full connector metadata is
|
||||||
//! owned by the connector metadata store, not by this module.
|
//! owned by the connector metadata store, not by this module.
|
||||||
|
//! Live catalog subscriptions publish only successful fetches from a matching
|
||||||
|
//! discovery scope, never disk snapshots or another scope's discovery winner.
|
||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
@@ -22,6 +24,7 @@ use codex_protocol::mcp::McpServerInfo;
|
|||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
use serde::de::DeserializeOwned;
|
use serde::de::DeserializeOwned;
|
||||||
|
use tokio::sync::watch;
|
||||||
|
|
||||||
use self::persistence::load_cached_codex_apps_server_info;
|
use self::persistence::load_cached_codex_apps_server_info;
|
||||||
use self::persistence::load_cached_connector_runtime_for_identity;
|
use self::persistence::load_cached_connector_runtime_for_identity;
|
||||||
@@ -31,6 +34,8 @@ use self::persistence::tools_cache_path;
|
|||||||
|
|
||||||
const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_publish.duration_ms";
|
const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_publish.duration_ms";
|
||||||
|
|
||||||
|
type LiveCatalog<T> = watch::Sender<Option<Arc<ConnectorRuntimeSnapshot<T>>>>;
|
||||||
|
|
||||||
/// Values stored in the connector runtime's persisted tool snapshot.
|
/// Values stored in the connector runtime's persisted tool snapshot.
|
||||||
///
|
///
|
||||||
/// The runtime uses the connector-owned Codex Apps cache layout for every
|
/// The runtime uses the connector-owned Codex Apps cache layout for every
|
||||||
@@ -93,6 +98,7 @@ pub fn connector_runtime_cache_path(codex_home: &Path, auth: Option<&CodexAuth>)
|
|||||||
pub struct ConnectorRuntimeSnapshot<T> {
|
pub struct ConnectorRuntimeSnapshot<T> {
|
||||||
tools: Vec<T>,
|
tools: Vec<T>,
|
||||||
refreshed_at: SystemTime,
|
refreshed_at: SystemTime,
|
||||||
|
generation: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<T> ConnectorRuntimeSnapshot<T> {
|
impl<T> ConnectorRuntimeSnapshot<T> {
|
||||||
@@ -166,24 +172,47 @@ impl<T: ConnectorRuntimePayload> ConnectorRuntimeManager<T> {
|
|||||||
.entry(identity.clone())
|
.entry(identity.clone())
|
||||||
.or_insert_with(|| Arc::new(ConnectorRuntimeEntry::new(identity, self.disk_cache)))
|
.or_insert_with(|| Arc::new(ConnectorRuntimeEntry::new(identity, self.disk_cache)))
|
||||||
.clone();
|
.clone();
|
||||||
ConnectorRuntimeContext { entry }
|
ConnectorRuntimeContext {
|
||||||
|
entry,
|
||||||
|
live_catalog: None,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Handle to one shared account/workspace connector runtime.
|
/// Handle to one shared account/workspace connector runtime.
|
||||||
pub struct ConnectorRuntimeContext<T: ConnectorRuntimePayload> {
|
pub struct ConnectorRuntimeContext<T: ConnectorRuntimePayload> {
|
||||||
entry: Arc<ConnectorRuntimeEntry<T>>,
|
entry: Arc<ConnectorRuntimeEntry<T>>,
|
||||||
|
live_catalog: Option<LiveCatalog<T>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<T: ConnectorRuntimePayload> Clone for ConnectorRuntimeContext<T> {
|
impl<T: ConnectorRuntimePayload> Clone for ConnectorRuntimeContext<T> {
|
||||||
fn clone(&self) -> Self {
|
fn clone(&self) -> Self {
|
||||||
Self {
|
Self {
|
||||||
entry: Arc::clone(&self.entry),
|
entry: Arc::clone(&self.entry),
|
||||||
|
live_catalog: self.live_catalog.clone(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<T: ConnectorRuntimePayload> ConnectorRuntimeContext<T> {
|
impl<T: ConnectorRuntimePayload> ConnectorRuntimeContext<T> {
|
||||||
|
/// Groups executable catalogs by equivalent discovery inputs within this account/home.
|
||||||
|
/// The caller must include the endpoint, auth configuration and listing protocol in `scope`.
|
||||||
|
/// Account-wide discovery reads are unaffected.
|
||||||
|
pub fn with_live_scope(mut self, scope: String) -> Self {
|
||||||
|
self.live_catalog = Some(
|
||||||
|
lock_unpoisoned(&self.entry.live_catalogs)
|
||||||
|
.entry(scope)
|
||||||
|
.or_insert_with(|| watch::channel(None).0)
|
||||||
|
.clone(),
|
||||||
|
);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Subscribes to accepted live tools without refetching or waiting for other clients.
|
||||||
|
pub fn subscribe(&self) -> Option<watch::Receiver<Option<Arc<ConnectorRuntimeSnapshot<T>>>>> {
|
||||||
|
self.live_catalog.as_ref().map(watch::Sender::subscribe)
|
||||||
|
}
|
||||||
|
|
||||||
pub fn current_snapshot(&self) -> Option<Arc<ConnectorRuntimeSnapshot<T>>> {
|
pub fn current_snapshot(&self) -> Option<Arc<ConnectorRuntimeSnapshot<T>>> {
|
||||||
self.entry.current_snapshot.load_full()
|
self.entry.current_snapshot.load_full()
|
||||||
}
|
}
|
||||||
@@ -254,6 +283,23 @@ impl<T: ConnectorRuntimePayload> ConnectorRuntimeContext<T> {
|
|||||||
) -> Arc<ConnectorRuntimeSnapshot<T>> {
|
) -> Arc<ConnectorRuntimeSnapshot<T>> {
|
||||||
let publish_start = Instant::now();
|
let publish_start = Instant::now();
|
||||||
let mut last_accepted_generation = lock_unpoisoned(&self.entry.last_accepted_generation);
|
let mut last_accepted_generation = lock_unpoisoned(&self.entry.last_accepted_generation);
|
||||||
|
let snapshot = Arc::new(ConnectorRuntimeSnapshot {
|
||||||
|
tools,
|
||||||
|
refreshed_at: SystemTime::now(),
|
||||||
|
generation: ticket.generation,
|
||||||
|
});
|
||||||
|
if let Some(live_catalog) = &self.live_catalog {
|
||||||
|
live_catalog.send_if_modified(|current| {
|
||||||
|
if current
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|current| current.generation >= ticket.generation)
|
||||||
|
{
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
*current = Some(Arc::clone(&snapshot));
|
||||||
|
true
|
||||||
|
});
|
||||||
|
}
|
||||||
if ticket.generation <= *last_accepted_generation
|
if ticket.generation <= *last_accepted_generation
|
||||||
&& let Some(snapshot) = self.current_snapshot()
|
&& let Some(snapshot) = self.current_snapshot()
|
||||||
{
|
{
|
||||||
@@ -266,11 +312,6 @@ impl<T: ConnectorRuntimePayload> ConnectorRuntimeContext<T> {
|
|||||||
return snapshot;
|
return snapshot;
|
||||||
}
|
}
|
||||||
|
|
||||||
let snapshot = Arc::new(ConnectorRuntimeSnapshot {
|
|
||||||
tools,
|
|
||||||
refreshed_at: SystemTime::now(),
|
|
||||||
});
|
|
||||||
|
|
||||||
*last_accepted_generation = ticket.generation;
|
*last_accepted_generation = ticket.generation;
|
||||||
self.entry
|
self.entry
|
||||||
.current_snapshot
|
.current_snapshot
|
||||||
@@ -326,6 +367,7 @@ struct ConnectorRuntimeEntry<T: ConnectorRuntimePayload> {
|
|||||||
current_snapshot: ArcSwapOption<ConnectorRuntimeSnapshot<T>>,
|
current_snapshot: ArcSwapOption<ConnectorRuntimeSnapshot<T>>,
|
||||||
next_fetch_generation: AtomicU64,
|
next_fetch_generation: AtomicU64,
|
||||||
last_accepted_generation: Mutex<u64>,
|
last_accepted_generation: Mutex<u64>,
|
||||||
|
live_catalogs: Mutex<HashMap<String, LiveCatalog<T>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<T: ConnectorRuntimePayload> ConnectorRuntimeEntry<T> {
|
impl<T: ConnectorRuntimePayload> ConnectorRuntimeEntry<T> {
|
||||||
@@ -342,6 +384,7 @@ impl<T: ConnectorRuntimePayload> ConnectorRuntimeEntry<T> {
|
|||||||
current_snapshot: ArcSwapOption::from(current_snapshot),
|
current_snapshot: ArcSwapOption::from(current_snapshot),
|
||||||
next_fetch_generation: AtomicU64::new(0),
|
next_fetch_generation: AtomicU64::new(0),
|
||||||
last_accepted_generation: Mutex::new(0),
|
last_accepted_generation: Mutex::new(0),
|
||||||
|
live_catalogs: Mutex::new(HashMap::new()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -64,6 +64,7 @@ pub(crate) fn load_cached_connector_runtime_for_identity<T: ConnectorRuntimePayl
|
|||||||
ConnectorRuntimeSnapshot {
|
ConnectorRuntimeSnapshot {
|
||||||
tools: cache.tools,
|
tools: cache.tools,
|
||||||
refreshed_at: modified_at,
|
refreshed_at: modified_at,
|
||||||
|
generation: 0,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -233,6 +234,7 @@ pub(crate) fn write_cached_codex_apps_tools_for_test<T>(
|
|||||||
let snapshot = ConnectorRuntimeSnapshot {
|
let snapshot = ConnectorRuntimeSnapshot {
|
||||||
tools: tools.to_vec(),
|
tools: tools.to_vec(),
|
||||||
refreshed_at: SystemTime::now(),
|
refreshed_at: SystemTime::now(),
|
||||||
|
generation: 0,
|
||||||
};
|
};
|
||||||
cache_context
|
cache_context
|
||||||
.entry
|
.entry
|
||||||
@@ -263,6 +265,7 @@ where
|
|||||||
let snapshot = ConnectorRuntimeSnapshot {
|
let snapshot = ConnectorRuntimeSnapshot {
|
||||||
tools: tools.to_vec(),
|
tools: tools.to_vec(),
|
||||||
refreshed_at: SystemTime::now(),
|
refreshed_at: SystemTime::now(),
|
||||||
|
generation: 0,
|
||||||
};
|
};
|
||||||
write_cached_connector_runtime(cache_context, &snapshot)
|
write_cached_connector_runtime(cache_context, &snapshot)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -492,6 +492,47 @@ fn codex_apps_tools_cache_scopes_non_utf8_home_disk_paths() {
|
|||||||
assert_ne!(cache_paths[0], cache_paths[1]);
|
assert_ne!(cache_paths[0], cache_paths[1]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn live_catalogs_isolate_scopes_and_reject_older_refreshes() {
|
||||||
|
let codex_home = tempdir().expect("tempdir");
|
||||||
|
let manager = ConnectorRuntimeManager::<TestTool>::new_without_cache();
|
||||||
|
let context = manager.context(
|
||||||
|
codex_home.path().to_path_buf(),
|
||||||
|
ConnectorRuntimeContextKey::personal(
|
||||||
|
Some("account".to_string()),
|
||||||
|
/*chatgpt_user_id*/ None,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
let scope_a = context.clone().with_live_scope("endpoint-a".to_string());
|
||||||
|
let scope_b = context.with_live_scope("endpoint-b".to_string());
|
||||||
|
let updates_a = scope_a.subscribe().expect("scope A subscription");
|
||||||
|
let updates_b = scope_b.subscribe().expect("scope B subscription");
|
||||||
|
let server_info = create_test_server_info("Codex Apps");
|
||||||
|
let older_a = scope_a.begin_fetch(ConnectorRuntimeFetchSource::Startup);
|
||||||
|
let newer_a = scope_a.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh);
|
||||||
|
let newest_b = scope_b.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh);
|
||||||
|
let tools_a = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer-a")];
|
||||||
|
let tools_b = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newest-b")];
|
||||||
|
|
||||||
|
scope_b.publish_if_newest_accepted(newest_b, &server_info, tools_b.clone());
|
||||||
|
assert!(updates_a.borrow().is_none());
|
||||||
|
let discovery = scope_a.publish_if_newest_accepted(newer_a, &server_info, tools_a.clone());
|
||||||
|
assert_eq!(discovery, tools_b);
|
||||||
|
scope_a.publish_if_newest_accepted(
|
||||||
|
older_a,
|
||||||
|
&server_info,
|
||||||
|
vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "stale-a")],
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
updates_a.borrow().as_ref().expect("scope A tools").tools(),
|
||||||
|
&tools_a
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
updates_b.borrow().as_ref().expect("scope B tools").tools(),
|
||||||
|
&tools_b
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn contexts_for_different_identities_keep_isolated_snapshots() {
|
fn contexts_for_different_identities_keep_isolated_snapshots() {
|
||||||
let codex_home = tempdir().expect("tempdir");
|
let codex_home = tempdir().expect("tempdir");
|
||||||
@@ -504,6 +545,8 @@ fn contexts_for_different_identities_keep_isolated_snapshots() {
|
|||||||
is_workspace_account: false,
|
is_workspace_account: false,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
let context_a = context_a.with_live_scope("same-endpoint".to_string());
|
||||||
|
let updates_a = context_a.subscribe().expect("account A subscription");
|
||||||
let tools_a = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "tool-a")];
|
let tools_a = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "tool-a")];
|
||||||
let snapshot_a = context_a.publish_runtime_if_newest_accepted(
|
let snapshot_a = context_a.publish_runtime_if_newest_accepted(
|
||||||
context_a.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh),
|
context_a.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh),
|
||||||
@@ -519,6 +562,8 @@ fn contexts_for_different_identities_keep_isolated_snapshots() {
|
|||||||
is_workspace_account: false,
|
is_workspace_account: false,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
let context_b = context_b.with_live_scope("same-endpoint".to_string());
|
||||||
|
let updates_b = context_b.subscribe().expect("account B subscription");
|
||||||
let same_context_a = manager.context(
|
let same_context_a = manager.context(
|
||||||
codex_home.path().to_path_buf(),
|
codex_home.path().to_path_buf(),
|
||||||
ConnectorRuntimeContextKey {
|
ConnectorRuntimeContextKey {
|
||||||
@@ -527,6 +572,7 @@ fn contexts_for_different_identities_keep_isolated_snapshots() {
|
|||||||
is_workspace_account: false,
|
is_workspace_account: false,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
let same_context_a = same_context_a.with_live_scope("same-endpoint".to_string());
|
||||||
|
|
||||||
assert!(Arc::ptr_eq(
|
assert!(Arc::ptr_eq(
|
||||||
&snapshot_a,
|
&snapshot_a,
|
||||||
@@ -566,6 +612,22 @@ fn contexts_for_different_identities_keep_isolated_snapshots() {
|
|||||||
&snapshot_b,
|
&snapshot_b,
|
||||||
&context_b.current_snapshot().expect("context B snapshot")
|
&context_b.current_snapshot().expect("context B snapshot")
|
||||||
));
|
));
|
||||||
|
assert_eq!(
|
||||||
|
updates_a
|
||||||
|
.borrow()
|
||||||
|
.as_ref()
|
||||||
|
.expect("account A tools")
|
||||||
|
.tools(),
|
||||||
|
&newer_tools_a
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
updates_b
|
||||||
|
.borrow()
|
||||||
|
.as_ref()
|
||||||
|
.expect("account B tools")
|
||||||
|
.tools(),
|
||||||
|
&tools_b
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
Reference in New Issue
Block a user