Resolve templated connector suggestions

This commit is contained in:
Noah MacCallum
2026-05-29 14:14:56 -07:00
parent bcd8d541f5
commit 25d081d10b
2 changed files with 253 additions and 19 deletions
+97 -17
View File
@@ -29,6 +29,7 @@ use codex_config::types::AppToolApproval;
use codex_config::types::AppsConfigToml;
use codex_config::types::ToolSuggestDiscoverableType;
use codex_core_plugins::PluginsManager;
use codex_core_plugins::remote::RemotePluginServiceConfig;
use codex_features::Feature;
use codex_login::AuthManager;
use codex_login::CodexAuth;
@@ -42,6 +43,7 @@ use codex_mcp::codex_apps_tools_cache_key;
use codex_mcp::compute_auth_statuses;
use codex_mcp::host_owned_codex_apps_enabled;
use codex_mcp::with_codex_apps_mcp;
use codex_plugin::AppConnectorId;
const CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS: Duration = Duration::from_secs(30);
@@ -115,9 +117,14 @@ pub(crate) async fn list_tool_suggest_discoverable_tools_with_auth(
accessible_connectors: &[AppInfo],
loaded_plugin_app_connector_ids: &[String],
) -> anyhow::Result<Vec<DiscoverableTool>> {
let connector_ids = tool_suggest_connector_ids(config).await;
let connector_ids = tool_suggest_connector_ids(config, auth).await;
let directory_connectors = codex_connectors::merge::merge_plugin_connectors(
cached_directory_connectors_for_tool_suggest_with_auth(config, auth).await,
resolve_template_directory_connectors(
config,
auth,
cached_directory_connectors_for_tool_suggest_with_auth(config, auth).await,
)
.await,
connector_ids.iter().cloned(),
);
let discoverable_connectors =
@@ -406,35 +413,108 @@ fn write_cached_accessible_connectors(
});
}
async fn tool_suggest_connector_ids(config: &Config) -> HashSet<String> {
async fn tool_suggest_connector_ids(config: &Config, auth: Option<&CodexAuth>) -> HashSet<String> {
let plugins_input = config.plugins_config_input();
let mut connector_ids = PluginsManager::new(config.codex_home.to_path_buf())
let connector_ids = PluginsManager::new(config.codex_home.to_path_buf())
.plugins_for_config(&plugins_input)
.await
.capability_summaries()
.iter()
.flat_map(|plugin| plugin.app_connector_ids.iter())
.map(|connector_id| connector_id.0.clone())
.collect::<HashSet<_>>();
connector_ids.extend(
config
.tool_suggest
.discoverables
.iter()
.filter(|discoverable| discoverable.kind == ToolSuggestDiscoverableType::Connector)
.map(|discoverable| discoverable.id.clone()),
);
.cloned()
.chain(
config
.tool_suggest
.discoverables
.iter()
.filter(|discoverable| discoverable.kind == ToolSuggestDiscoverableType::Connector)
.map(|discoverable| AppConnectorId(discoverable.id.clone())),
)
.collect::<Vec<_>>();
let remote_plugin_service_config = RemotePluginServiceConfig {
chatgpt_base_url: config.chatgpt_base_url.clone(),
};
let mut connector_ids = codex_core_plugins::remote::resolve_remote_plugin_app_ids(
&remote_plugin_service_config,
auth,
&connector_ids,
)
.await
.into_iter()
.map(|connector_id| connector_id.0)
.collect::<HashSet<_>>();
let disabled_connector_ids = config
.tool_suggest
.disabled_tools
.iter()
.filter(|disabled_tool| disabled_tool.kind == ToolSuggestDiscoverableType::Connector)
.map(|disabled_tool| disabled_tool.id.as_str())
.collect::<HashSet<_>>();
connector_ids.retain(|connector_id| !disabled_connector_ids.contains(connector_id.as_str()));
.map(|disabled_tool| AppConnectorId(disabled_tool.id.clone()))
.collect::<Vec<_>>();
let mut disabled_connector_ids = codex_core_plugins::remote::resolve_remote_plugin_app_ids(
&remote_plugin_service_config,
auth,
&disabled_connector_ids,
)
.await
.into_iter()
.map(|connector_id| connector_id.0)
.collect::<HashSet<_>>();
disabled_connector_ids.extend(
config
.tool_suggest
.disabled_tools
.iter()
.filter(|disabled_tool| disabled_tool.kind == ToolSuggestDiscoverableType::Connector)
.map(|disabled_tool| disabled_tool.id.clone()),
);
connector_ids.retain(|connector_id| !disabled_connector_ids.contains(connector_id));
connector_ids
}
async fn resolve_template_directory_connectors(
config: &Config,
auth: Option<&CodexAuth>,
connectors: Vec<AppInfo>,
) -> Vec<AppInfo> {
let remote_plugin_service_config = RemotePluginServiceConfig {
chatgpt_base_url: config.chatgpt_base_url.clone(),
};
let mut resolved_connectors = Vec::new();
let mut seen_connector_ids = HashSet::new();
for connector in connectors {
let resolved_connector_ids = codex_core_plugins::remote::resolve_remote_plugin_app_ids(
&remote_plugin_service_config,
auth,
&[AppConnectorId(connector.id.clone())],
)
.await;
if resolved_connector_ids.is_empty() {
continue;
}
for connector_id in resolved_connector_ids {
if !seen_connector_ids.insert(connector_id.clone()) {
continue;
}
if connector_id.0 == connector.id {
resolved_connectors.push(connector.clone());
continue;
}
let mut resolved_connector = connector.clone();
resolved_connector.id = connector_id.0;
resolved_connector.install_url =
Some(codex_connectors::metadata::connector_install_url(
&resolved_connector.name,
&resolved_connector.id,
));
resolved_connectors.push(resolved_connector);
}
}
resolved_connectors
}
async fn cached_directory_connectors_for_tool_suggest_with_auth(
config: &Config,
auth: Option<&CodexAuth>,
+156 -2
View File
@@ -29,6 +29,12 @@ use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use tempfile::tempdir;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
fn annotations(destructive_hint: Option<bool>, open_world_hint: Option<bool>) -> ToolAnnotations {
ToolAnnotations::from_raw(
@@ -1180,7 +1186,7 @@ discoverables = [
.expect("config should load");
assert_eq!(
tool_suggest_connector_ids(&config).await,
tool_suggest_connector_ids(&config, None).await,
HashSet::from(["connector_2128aebfecb84f64a069897515042a44".to_string()])
);
}
@@ -1209,11 +1215,159 @@ disabled_tools = [
.expect("config should load");
assert_eq!(
tool_suggest_connector_ids(&config).await,
tool_suggest_connector_ids(&config, None).await,
HashSet::from(["connector_gmail".to_string()])
);
}
#[tokio::test]
async fn tool_suggest_connector_ids_resolve_template_connectors() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path(
"/backend-api/ps/connectors/by_template_id/templated_apps_Databricks",
))
.and(header("authorization", "Bearer Access Token"))
.and(header("chatgpt-account-id", "account_id"))
.respond_with(
ResponseTemplate::new(200)
.set_body_string(r#"{"connector_ids":["asdk_app_databricks_workspace"]}"#),
)
.mount(&server)
.await;
let codex_home = tempdir().expect("tempdir should succeed");
std::fs::write(
codex_home.path().join(CONFIG_TOML_FILE),
r#"
[tool_suggest]
discoverables = [
{ type = "connector", id = "templated_apps_Databricks" }
]
"#,
)
.expect("write config");
let mut config = ConfigBuilder::default()
.codex_home(codex_home.path().to_path_buf())
.build()
.await
.expect("config should load");
config.chatgpt_base_url = format!("{}/backend-api", server.uri());
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
assert_eq!(
tool_suggest_connector_ids(&config, Some(&auth)).await,
HashSet::from(["asdk_app_databricks_workspace".to_string()])
);
}
#[tokio::test]
async fn tool_suggest_resolves_template_connectors_before_returning_install_entries() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path(
"/backend-api/ps/connectors/by_template_id/templated_apps_Databricks",
))
.and(header("authorization", "Bearer Access Token"))
.and(header("chatgpt-account-id", "account_id"))
.respond_with(
ResponseTemplate::new(200)
.set_body_string(r#"{"connector_ids":["asdk_app_databricks_workspace"]}"#),
)
.mount(&server)
.await;
let codex_home = tempdir().expect("tempdir should succeed");
std::fs::write(
codex_home.path().join(CONFIG_TOML_FILE),
r#"
[features]
apps = true
[tool_suggest]
discoverables = [
{ type = "connector", id = "templated_apps_Databricks" }
]
"#,
)
.expect("write config");
let mut config = ConfigBuilder::default()
.codex_home(codex_home.path().to_path_buf())
.build()
.await
.expect("config should load");
config.chatgpt_base_url = format!("{}/backend-api", server.uri());
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let discoverable_tools =
list_tool_suggest_discoverable_tools_with_auth(&config, Some(&auth), &[], &[])
.await
.expect("discoverable tools should load");
assert_eq!(
discoverable_tools,
vec![DiscoverableTool::from(plugin_connector_to_app_info(
"asdk_app_databricks_workspace".to_string(),
))]
);
}
#[tokio::test]
async fn template_directory_connector_metadata_is_carried_to_resolved_connector() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path(
"/backend-api/ps/connectors/by_template_id/templated_apps_Databricks",
))
.and(header("authorization", "Bearer Access Token"))
.and(header("chatgpt-account-id", "account_id"))
.respond_with(
ResponseTemplate::new(200)
.set_body_string(r#"{"connector_ids":["asdk_app_databricks_workspace"]}"#),
)
.mount(&server)
.await;
let codex_home = tempdir().expect("tempdir should succeed");
let mut config = ConfigBuilder::default()
.codex_home(codex_home.path().to_path_buf())
.build()
.await
.expect("config should load");
config.chatgpt_base_url = format!("{}/backend-api", server.uri());
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let resolved = resolve_template_directory_connectors(
&config,
Some(&auth),
vec![AppInfo {
id: "templated_apps_Databricks".to_string(),
name: "Databricks".to_string(),
description: Some("Query Databricks".to_string()),
install_url: Some(connector_install_url(
"Databricks",
"templated_apps_Databricks",
)),
..app("templated_apps_Databricks")
}],
)
.await;
assert_eq!(
resolved,
vec![AppInfo {
id: "asdk_app_databricks_workspace".to_string(),
name: "Databricks".to_string(),
description: Some("Query Databricks".to_string()),
install_url: Some(connector_install_url(
"Databricks",
"asdk_app_databricks_workspace",
)),
..app("asdk_app_databricks_workspace")
}]
);
}
#[tokio::test]
async fn tool_suggest_uses_connector_id_fallback_when_directory_cache_is_empty() {
let codex_home = tempdir().expect("tempdir should succeed");