Track used Codex app connectors

This commit is contained in:
John Gregg
2026-05-28 17:24:13 -07:00
parent bf72be5927
commit c14b7fab75
14 changed files with 273 additions and 11 deletions
+23 -2
View File
@@ -339,8 +339,14 @@ async fn handle_approved_mcp_tool_call(
};
let result = async {
let rewritten_arguments = rewrite?;
let request_meta =
build_mcp_tool_call_request_meta(turn_context, &server, call_id, metadata);
let used_connector_ids = sess.get_used_connector_ids().await;
let request_meta = build_mcp_tool_call_request_meta(
turn_context,
&server,
call_id,
metadata,
&used_connector_ids,
);
let result = execute_mcp_tool_call(
sess,
turn_context,
@@ -931,6 +937,9 @@ async fn maybe_track_codex_app_used(
let (connector_id, app_name) = metadata
.map(|metadata| (metadata.connector_id, metadata.app_name))
.unwrap_or((None, None));
if let Some(connector_id) = connector_id.as_deref() {
sess.record_used_connector_id(connector_id).await;
}
let invocation_type = if let Some(connector_id) = connector_id.as_deref() {
let mentioned_connector_ids = sess.get_connector_selection().await;
if mentioned_connector_ids.contains(connector_id) {
@@ -983,6 +992,7 @@ const MCP_TOOL_OPENAI_OUTPUT_TEMPLATE_META_KEY: &str = "openai/outputTemplate";
const MCP_TOOL_UI_RESOURCE_URI_META_KEY: &str = "ui/resourceUri";
const MCP_TOOL_PLUGIN_ID_META_KEY: &str = "plugin_id";
const MCP_TOOL_THREAD_ID_META_KEY: &str = "threadId";
const MCP_TOOL_USED_CONNECTOR_IDS_META_KEY: &str = "used_connector_ids";
async fn custom_mcp_tool_approval_mode(
sess: &Session,
@@ -1038,6 +1048,7 @@ fn build_mcp_tool_call_request_meta(
server: &str,
call_id: &str,
metadata: Option<&McpToolApprovalMetadata>,
used_connector_ids: &[String],
) -> Option<serde_json::Value> {
let mut request_meta = serde_json::Map::new();
@@ -1062,6 +1073,16 @@ fn build_mcp_tool_call_request_meta(
"call_id".to_string(),
serde_json::Value::String(call_id.to_string()),
);
codex_apps_meta.insert(
MCP_TOOL_USED_CONNECTOR_IDS_META_KEY.to_string(),
serde_json::Value::Array(
used_connector_ids
.iter()
.cloned()
.map(serde_json::Value::String)
.collect(),
),
);
request_meta.insert(
MCP_TOOL_CODEX_APPS_META_KEY.to_string(),
serde_json::Value::Object(codex_apps_meta),
+41 -1
View File
@@ -1034,6 +1034,7 @@ async fn mcp_tool_call_request_meta_includes_turn_metadata_for_custom_server() {
"custom_server",
"call-custom",
/*metadata*/ None,
&[],
)
.expect("custom servers should receive turn metadata");
let turn_metadata = meta
@@ -1076,6 +1077,7 @@ async fn mcp_tool_call_request_meta_includes_turn_started_at_unix_ms() {
"custom_server",
"call-custom",
/*metadata*/ None,
&[],
)
.expect("custom servers should receive turn metadata");
let turn_metadata = meta
@@ -1105,7 +1107,13 @@ async fn plugin_mcp_tool_call_request_meta_includes_plugin_id() {
metadata.plugin_id = Some("sample@test".to_string());
assert_eq!(
build_mcp_tool_call_request_meta(&turn_context, "sample", "call-plugin", Some(&metadata),),
build_mcp_tool_call_request_meta(
&turn_context,
"sample",
"call-plugin",
Some(&metadata),
&[],
),
Some(serde_json::json!({
crate::X_CODEX_TURN_METADATA_HEADER: expected_turn_metadata,
MCP_TOOL_PLUGIN_ID_META_KEY: "sample@test",
@@ -1182,6 +1190,7 @@ async fn codex_apps_tool_call_request_meta_includes_turn_metadata_and_codex_apps
CODEX_APPS_MCP_SERVER_NAME,
"call_abc123xyz789",
Some(&metadata),
&[],
),
Some(serde_json::json!({
crate::X_CODEX_TURN_METADATA_HEADER: expected_turn_metadata,
@@ -1190,6 +1199,7 @@ async fn codex_apps_tool_call_request_meta_includes_turn_metadata_and_codex_apps
"resource_uri": "connector://calendar/tools/calendar_create_event",
"contains_mcp_source": true,
"connector_id": "calendar",
"used_connector_ids": [],
},
}))
);
@@ -1209,16 +1219,46 @@ async fn codex_apps_tool_call_request_meta_includes_call_id_without_existing_cod
CODEX_APPS_MCP_SERVER_NAME,
"call_abc123xyz789",
/*metadata*/ None,
&[],
),
Some(serde_json::json!({
crate::X_CODEX_TURN_METADATA_HEADER: expected_turn_metadata,
MCP_TOOL_CODEX_APPS_META_KEY: {
"call_id": "call_abc123xyz789",
"used_connector_ids": [],
},
}))
);
}
#[tokio::test]
async fn codex_apps_tool_call_request_meta_includes_used_connector_ids() {
let (_, turn_context) = make_session_and_context().await;
let used_connector_ids = vec![
"connector_calendar".to_string(),
"connector_drive".to_string(),
];
let meta = build_mcp_tool_call_request_meta(
&turn_context,
CODEX_APPS_MCP_SERVER_NAME,
"call_abc123xyz789",
/*metadata*/ None,
&used_connector_ids,
)
.expect("codex apps tool call metadata");
assert_eq!(
meta.get(MCP_TOOL_CODEX_APPS_META_KEY)
.and_then(serde_json::Value::as_object)
.and_then(|meta| meta.get(MCP_TOOL_USED_CONNECTOR_IDS_META_KEY)),
Some(&serde_json::json!([
"connector_calendar",
"connector_drive"
]))
);
}
fn codex_apps_auth_failure_result() -> CallToolResult {
CallToolResult {
content: vec![serde_json::json!({
+32
View File
@@ -138,6 +138,7 @@ use codex_thread_store::LocalThreadStore;
use codex_thread_store::ReadThreadParams;
use codex_thread_store::ResumeThreadParams;
use codex_thread_store::ThreadEventPersistenceMode;
use codex_thread_store::ThreadMetadataPatch;
use codex_thread_store::ThreadPersistenceMetadata;
use codex_thread_store::ThreadStore;
use codex_utils_output_truncation::TruncationPolicy;
@@ -1167,6 +1168,37 @@ impl Session {
state.clear_connector_selection();
}
pub(crate) async fn record_used_connector_id(&self, connector_id: &str) {
let used_connector_ids = {
let mut state = self.state.lock().await;
state.record_used_connector_id(connector_id)
};
let Some(used_connector_ids) = used_connector_ids else {
return;
};
let Some(live_thread) = self.live_thread() else {
return;
};
if let Err(err) = live_thread
.update_metadata(
ThreadMetadataPatch {
used_connector_ids: Some(used_connector_ids),
..Default::default()
},
/*include_archived*/ true,
)
.await
{
warn!("failed to persist used connector metadata: {err}");
}
}
pub(crate) async fn get_used_connector_ids(&self) -> Vec<String> {
let state = self.state.lock().await;
state.get_used_connector_ids()
}
async fn record_initial_history(&self, conversation_history: InitialHistory) {
let turn_context = self.new_default_turn().await;
let is_subagent = {
+14 -1
View File
@@ -875,7 +875,20 @@ impl Session {
session_configuration.thread_name = thread_name.clone();
validate_config_lock_if_configured(&session_configuration).await?;
export_config_lock_if_configured(&session_configuration, thread_id).await?;
let state = SessionState::new(session_configuration.clone());
let initial_used_connector_ids = if let Some(state_db) = state_db_ctx.as_ref() {
match state_db.get_thread(thread_id).await {
Ok(Some(metadata)) => metadata.used_connector_ids,
Ok(None) => Vec::new(),
Err(err) => {
warn!("failed to load used connector metadata for {thread_id}: {err}");
Vec::new()
}
}
} else {
Vec::new()
};
let mut state = SessionState::new(session_configuration.clone());
state.set_used_connector_ids(initial_used_connector_ids);
let managed_network_requirements_configured = config
.config_layer_stack
.requirements_toml()
+32
View File
@@ -36,6 +36,7 @@ pub(crate) struct SessionState {
/// Startup prewarmed session prepared during session initialization.
pub(crate) startup_prewarm: Option<SessionStartupPrewarmHandle>,
pub(crate) active_connector_selection: HashSet<String>,
used_connector_ids: Vec<String>,
pub(crate) pending_session_start_sources: VecDeque<codex_hooks::SessionStartSource>,
granted_permissions: Option<AdditionalPermissionProfile>,
next_turn_is_first: bool,
@@ -56,6 +57,7 @@ impl SessionState {
auto_compact_window: AutoCompactWindow::new(),
startup_prewarm: None,
active_connector_selection: HashSet::new(),
used_connector_ids: Vec::new(),
pending_session_start_sources: VecDeque::new(),
granted_permissions: None,
next_turn_is_first: true,
@@ -222,6 +224,36 @@ impl SessionState {
self.active_connector_selection.clear();
}
pub(crate) fn set_used_connector_ids(&mut self, used_connector_ids: Vec<String>) {
self.used_connector_ids = Vec::new();
for connector_id in used_connector_ids {
if !self
.used_connector_ids
.iter()
.any(|used_connector_id| used_connector_id == &connector_id)
{
self.used_connector_ids.push(connector_id);
}
}
}
pub(crate) fn record_used_connector_id(&mut self, connector_id: &str) -> Option<Vec<String>> {
if self
.used_connector_ids
.iter()
.any(|used_connector_id| used_connector_id == connector_id)
{
return None;
}
self.used_connector_ids.push(connector_id.to_string());
Some(self.used_connector_ids.clone())
}
pub(crate) fn get_used_connector_ids(&self) -> Vec<String> {
self.used_connector_ids.clone()
}
pub(crate) fn queue_pending_session_start_source(
&mut self,
value: codex_hooks::SessionStartSource,
+20
View File
@@ -34,6 +34,26 @@ async fn clear_connector_selection_removes_entries() {
assert_eq!(state.get_connector_selection(), HashSet::new());
}
#[tokio::test]
async fn record_used_connector_id_preserves_first_use_order() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
assert_eq!(
state.record_used_connector_id("calendar"),
Some(vec!["calendar".to_string()])
);
assert_eq!(state.record_used_connector_id("calendar"), None);
assert_eq!(
state.record_used_connector_id("drive"),
Some(vec!["calendar".to_string(), "drive".to_string()])
);
assert_eq!(
state.get_used_connector_ids(),
vec!["calendar".to_string(), "drive".to_string()]
);
}
#[tokio::test]
async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() {
let session_configuration = make_session_configuration_for_tests().await;
@@ -0,0 +1 @@
ALTER TABLE threads ADD COLUMN used_connector_ids TEXT NOT NULL DEFAULT '[]';
+1
View File
@@ -493,6 +493,7 @@ mod tests {
git_sha: None,
git_branch: None,
git_origin_url: None,
used_connector_ids: Vec::new(),
}
}
@@ -107,6 +107,8 @@ pub struct ThreadMetadata {
pub git_branch: Option<String>,
/// The git origin URL, if known.
pub git_origin_url: Option<String>,
/// Connector IDs used by this thread, in first-use order.
pub used_connector_ids: Vec<String>,
}
/// Builder data required to construct [`ThreadMetadata`] without parsing filenames.
@@ -221,6 +223,7 @@ impl ThreadMetadataBuilder {
git_sha: self.git_sha.clone(),
git_branch: self.git_branch.clone(),
git_origin_url: self.git_origin_url.clone(),
used_connector_ids: Vec::new(),
}
}
}
@@ -311,6 +314,9 @@ impl ThreadMetadata {
if self.git_origin_url != other.git_origin_url {
diffs.push("git_origin_url");
}
if self.used_connector_ids != other.used_connector_ids {
diffs.push("used_connector_ids");
}
diffs
}
}
@@ -345,6 +351,7 @@ pub(crate) struct ThreadRow {
git_sha: Option<String>,
git_branch: Option<String>,
git_origin_url: Option<String>,
used_connector_ids: String,
}
impl ThreadRow {
@@ -374,6 +381,7 @@ impl ThreadRow {
git_sha: row.try_get("git_sha")?,
git_branch: row.try_get("git_branch")?,
git_origin_url: row.try_get("git_origin_url")?,
used_connector_ids: row.try_get("used_connector_ids")?,
})
}
}
@@ -407,11 +415,14 @@ impl TryFrom<ThreadRow> for ThreadMetadata {
git_sha,
git_branch,
git_origin_url,
used_connector_ids,
} = row;
let thread_source = thread_source
.map(|thread_source| thread_source.parse())
.transpose()
.map_err(anyhow::Error::msg)?;
let used_connector_ids = serde_json::from_str::<Vec<String>>(&used_connector_ids)
.map_err(|err| anyhow::anyhow!("invalid used_connector_ids JSON: {err}"))?;
Ok(Self {
id: ThreadId::try_from(id)?,
rollout_path: PathBuf::from(rollout_path),
@@ -438,6 +449,7 @@ impl TryFrom<ThreadRow> for ThreadMetadata {
git_sha,
git_branch,
git_origin_url,
used_connector_ids,
})
}
}
@@ -524,6 +536,7 @@ mod tests {
git_sha: None,
git_branch: None,
git_origin_url: None,
used_connector_ids: "[]".to_string(),
}
}
@@ -554,6 +567,7 @@ mod tests {
git_sha: None,
git_branch: None,
git_origin_url: None,
used_connector_ids: Vec::new(),
}
}
+4 -2
View File
@@ -190,7 +190,8 @@ SELECT
threads.archived_at,
threads.git_sha,
threads.git_branch,
threads.git_origin_url
threads.git_origin_url,
threads.used_connector_ids
FROM threads
"#,
);
@@ -560,7 +561,8 @@ SELECT
threads.archived_at,
threads.git_sha,
threads.git_branch,
threads.git_origin_url
threads.git_origin_url,
threads.used_connector_ids
FROM threads
WHERE threads.id = ? AND threads.memory_mode = 'enabled'
"#,
@@ -67,5 +67,6 @@ pub(super) fn test_thread_metadata(
git_sha: None,
git_branch: None,
git_origin_url: None,
used_connector_ids: Vec::new(),
}
}
+75 -5
View File
@@ -31,7 +31,8 @@ SELECT
threads.archived_at,
threads.git_sha,
threads.git_branch,
threads.git_origin_url
threads.git_origin_url,
threads.used_connector_ids
FROM threads
WHERE threads.id = ?
"#,
@@ -506,8 +507,9 @@ INSERT INTO threads (
git_sha,
git_branch,
git_origin_url,
used_connector_ids,
memory_mode
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO NOTHING
"#,
)
@@ -547,6 +549,7 @@ ON CONFLICT(id) DO NOTHING
.bind(metadata.git_sha.as_deref())
.bind(metadata.git_branch.as_deref())
.bind(metadata.git_origin_url.as_deref())
.bind(used_connector_ids_json(&metadata.used_connector_ids)?)
.bind("enabled")
.execute(self.pool.as_ref())
.await?;
@@ -568,6 +571,19 @@ ON CONFLICT(id) DO NOTHING
Ok(result.rows_affected() > 0)
}
pub async fn set_thread_used_connector_ids(
&self,
thread_id: ThreadId,
used_connector_ids: &[String],
) -> anyhow::Result<bool> {
let result = sqlx::query("UPDATE threads SET used_connector_ids = ? WHERE id = ?")
.bind(used_connector_ids_json(used_connector_ids)?)
.bind(thread_id.to_string())
.execute(self.pool.as_ref())
.await?;
Ok(result.rows_affected() > 0)
}
pub async fn update_thread_title(
&self,
thread_id: ThreadId,
@@ -712,8 +728,9 @@ INSERT INTO threads (
git_sha,
git_branch,
git_origin_url,
used_connector_ids,
memory_mode
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
rollout_path = excluded.rollout_path,
created_at = excluded.created_at,
@@ -740,7 +757,8 @@ ON CONFLICT(id) DO UPDATE SET
archived_at = excluded.archived_at,
git_sha = COALESCE(threads.git_sha, excluded.git_sha),
git_branch = COALESCE(threads.git_branch, excluded.git_branch),
git_origin_url = COALESCE(threads.git_origin_url, excluded.git_origin_url)
git_origin_url = COALESCE(threads.git_origin_url, excluded.git_origin_url),
used_connector_ids = threads.used_connector_ids
"#,
)
.bind(metadata.id.to_string())
@@ -779,6 +797,7 @@ ON CONFLICT(id) DO UPDATE SET
.bind(metadata.git_sha.as_deref())
.bind(metadata.git_branch.as_deref())
.bind(metadata.git_origin_url.as_deref())
.bind(used_connector_ids_json(&metadata.used_connector_ids)?)
.bind(creation_memory_mode.unwrap_or("enabled"))
.execute(self.pool.as_ref())
.await?;
@@ -942,11 +961,17 @@ SELECT
threads.archived_at,
threads.git_sha,
threads.git_branch,
threads.git_origin_url
threads.git_origin_url,
threads.used_connector_ids
"#,
);
}
fn used_connector_ids_json(used_connector_ids: &[String]) -> anyhow::Result<String> {
serde_json::to_string(used_connector_ids)
.map_err(|err| anyhow::anyhow!("failed to serialize used connector ids: {err}"))
}
pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option<String> {
items.iter().rev().find_map(|item| match item {
RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(),
@@ -1144,6 +1169,51 @@ mod tests {
assert_eq!(memory_mode, "disabled");
}
#[tokio::test]
async fn used_connector_ids_are_updated_explicitly_and_preserved_by_upsert() {
let codex_home = unique_temp_dir();
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string())
.await
.expect("state db should initialize");
let thread_id =
ThreadId::from_string("00000000-0000-0000-0000-000000000123").expect("valid thread id");
let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone());
runtime
.upsert_thread(&metadata)
.await
.expect("initial insert should succeed");
runtime
.set_thread_used_connector_ids(
thread_id,
&[
"connector_calendar".to_string(),
"connector_drive".to_string(),
],
)
.await
.expect("used connector update should succeed");
metadata.title = "updated title".to_string();
runtime
.upsert_thread(&metadata)
.await
.expect("upsert should succeed");
let persisted = runtime
.get_thread(thread_id)
.await
.expect("thread should load")
.expect("thread should exist");
assert_eq!(
persisted.used_connector_ids,
vec![
"connector_calendar".to_string(),
"connector_drive".to_string()
]
);
}
#[tokio::test]
async fn list_threads_updated_after_returns_oldest_changes_first() {
let codex_home = unique_temp_dir();
@@ -295,12 +295,21 @@ async fn apply_metadata_update(
metadata.git_branch = branch;
metadata.git_origin_url = origin_url;
}
if let Some(used_connector_ids) = patch.used_connector_ids {
metadata.used_connector_ids = used_connector_ids;
}
state_db
.upsert_thread(&metadata)
.await
.map_err(|err| ThreadStoreError::Internal {
message: format!("failed to update thread metadata for {thread_id}: {err}"),
})?;
state_db
.set_thread_used_connector_ids(thread_id, &metadata.used_connector_ids)
.await
.map_err(|err| ThreadStoreError::Internal {
message: format!("failed to update used connector ids for {thread_id}: {err}"),
})?;
if let Some(memory_mode) = patch.memory_mode {
state_db
.set_thread_memory_mode(thread_id, memory_mode_as_str(memory_mode))
+6
View File
@@ -529,6 +529,8 @@ pub struct ThreadMetadataPatch {
pub git_info: Option<GitInfoPatch>,
/// Thread memory behavior.
pub memory_mode: Option<MemoryMode>,
/// Connector IDs used by this thread, in first-use order.
pub used_connector_ids: Option<Vec<String>>,
}
impl ThreadMetadataPatch {
@@ -606,6 +608,9 @@ impl ThreadMetadataPatch {
if next.memory_mode.is_some() {
self.memory_mode = next.memory_mode;
}
if next.used_connector_ids.is_some() {
self.used_connector_ids = next.used_connector_ids;
}
}
pub fn is_empty(&self) -> bool {
@@ -631,6 +636,7 @@ impl ThreadMetadataPatch {
&& self.first_user_message.is_none()
&& self.git_info.is_none()
&& self.memory_mode.is_none()
&& self.used_connector_ids.is_none()
}
}