diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 3cd7fc809..dbfa92ff4 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1411,6 +1411,27 @@ impl Session { state.clear_mcp_tool_selection(); } + // Merges connector IDs into the session-level explicit connector selection. + pub(crate) async fn merge_connector_selection( + &self, + connector_ids: HashSet, + ) -> HashSet { + let mut state = self.state.lock().await; + state.merge_connector_selection(connector_ids) + } + + // Returns the connector IDs currently selected for this session. + pub(crate) async fn get_connector_selection(&self) -> HashSet { + let state = self.state.lock().await; + state.get_connector_selection() + } + + // Clears connector IDs that were accumulated for explicit selection. + pub(crate) async fn clear_connector_selection(&self) { + let mut state = self.state.lock().await; + state.clear_connector_selection(); + } + async fn record_initial_history(&self, conversation_history: InitialHistory) { let turn_context = self.new_default_turn().await; match conversation_history { @@ -4056,7 +4077,7 @@ pub(crate) async fn run_turn( .await, ); - let connector_slug_counts = if turn_context.config.features.enabled(Feature::Apps) { + let available_connectors = if turn_context.config.features.enabled(Feature::Apps) { let mcp_tools = match sess .services .mcp_connection_manager @@ -4069,11 +4090,16 @@ pub(crate) async fn run_turn( Ok(mcp_tools) => mcp_tools, Err(_) => return None, }; - let connectors = connectors::accessible_connectors_from_mcp_tools(&mcp_tools); - build_connector_slug_counts(&connectors) + connectors::accessible_connectors_from_mcp_tools(&mcp_tools) } else { - HashMap::new() + Vec::new() }; + let connector_slug_counts = build_connector_slug_counts(&available_connectors); + let skill_name_counts_lower = skills_outcome + .as_ref() + .map_or_else(HashMap::new, |outcome| { + build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1 + }); let mentioned_skills = skills_outcome.as_ref().map_or_else(Vec::new, |outcome| { collect_explicit_skill_mentions( &input, @@ -4082,7 +4108,6 @@ pub(crate) async fn run_turn( &connector_slug_counts, ) }); - let explicitly_enabled_connectors = collect_explicit_app_ids(&input); let config = turn_context.config.clone(); if config .features @@ -4119,6 +4144,15 @@ pub(crate) async fn run_turn( .await; } + let mut explicitly_enabled_connectors = collect_explicit_app_ids(&input); + explicitly_enabled_connectors.extend(collect_explicit_app_ids_from_skill_items( + &skill_items, + &available_connectors, + &skill_name_counts_lower, + )); + sess.merge_connector_selection(explicitly_enabled_connectors.clone()) + .await; + let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input.clone()); let response_item: ResponseItem = initial_input_for_turn.clone().into(); sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item) @@ -4349,6 +4383,57 @@ async fn run_auto_compact(sess: &Arc, turn_context: &Arc) Ok(()) } +fn collect_explicit_app_ids_from_skill_items( + skill_items: &[ResponseItem], + connectors: &[connectors::AppInfo], + skill_name_counts_lower: &HashMap, +) -> HashSet { + if skill_items.is_empty() || connectors.is_empty() { + return HashSet::new(); + } + + let skill_messages = skill_items + .iter() + .filter_map(|item| match item { + ResponseItem::Message { content, .. } => { + content.iter().find_map(|content_item| match content_item { + ContentItem::InputText { text } => Some(text.clone()), + _ => None, + }) + } + _ => None, + }) + .collect::>(); + if skill_messages.is_empty() { + return HashSet::new(); + } + + let mentions = collect_tool_mentions_from_messages(&skill_messages); + let mention_names_lower = mentions + .plain_names + .iter() + .map(|name| name.to_ascii_lowercase()) + .collect::>(); + let mut connector_ids = mentions + .paths + .iter() + .filter(|path| tool_kind_for_path(path) == ToolMentionKind::App) + .filter_map(|path| app_id_from_path(path).map(str::to_string)) + .collect::>(); + + let connector_slug_counts = build_connector_slug_counts(connectors); + for connector in connectors { + let slug = connectors::connector_mention_slug(connector); + let connector_count = connector_slug_counts.get(&slug).copied().unwrap_or(0); + let skill_count = skill_name_counts_lower.get(&slug).copied().unwrap_or(0); + if connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&slug) { + connector_ids.insert(connector.id.clone()); + } + } + + connector_ids +} + fn filter_connectors_for_input( connectors: Vec, input: &[ResponseItem], @@ -4594,6 +4679,9 @@ async fn built_tools( .or_cancel(cancellation_token) .await?; + let mut effective_explicitly_enabled_connectors = explicitly_enabled_connectors.clone(); + effective_explicitly_enabled_connectors.extend(sess.get_connector_selection().await); + let connectors_for_tools = if turn_context.config.features.enabled(Feature::Apps) { let skill_name_counts_lower = skills_outcome.map_or_else(HashMap::new, |outcome| { build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1 @@ -4602,7 +4690,7 @@ async fn built_tools( Some(filter_connectors_for_input( connectors, input, - explicitly_enabled_connectors, + &effective_explicitly_enabled_connectors, &skill_name_counts_lower, )) } else { @@ -5442,6 +5530,18 @@ mod tests { } } + fn skill_message(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + } + } + fn make_connector(id: &str, name: &str) -> AppInfo { AppInfo { id: id.to_string(), @@ -5574,6 +5674,49 @@ mod tests { assert_eq!(selected, Vec::new()); } + #[test] + fn collect_explicit_app_ids_from_skill_items_includes_linked_mentions() { + let connectors = vec![make_connector("calendar", "Calendar")]; + let skill_items = vec![skill_message( + "\ndemo\n/tmp/skills/demo/SKILL.md\nuse [$calendar](app://calendar)\n", + )]; + + let connector_ids = + collect_explicit_app_ids_from_skill_items(&skill_items, &connectors, &HashMap::new()); + + assert_eq!(connector_ids, HashSet::from(["calendar".to_string()])); + } + + #[test] + fn collect_explicit_app_ids_from_skill_items_resolves_unambiguous_plain_mentions() { + let connectors = vec![make_connector("calendar", "Calendar")]; + let skill_items = vec![skill_message( + "\ndemo\n/tmp/skills/demo/SKILL.md\nuse $calendar\n", + )]; + + let connector_ids = + collect_explicit_app_ids_from_skill_items(&skill_items, &connectors, &HashMap::new()); + + assert_eq!(connector_ids, HashSet::from(["calendar".to_string()])); + } + + #[test] + fn collect_explicit_app_ids_from_skill_items_skips_plain_mentions_with_skill_conflicts() { + let connectors = vec![make_connector("calendar", "Calendar")]; + let skill_items = vec![skill_message( + "\ndemo\n/tmp/skills/demo/SKILL.md\nuse $calendar\n", + )]; + let skill_name_counts_lower = HashMap::from([("calendar".to_string(), 1)]); + + let connector_ids = collect_explicit_app_ids_from_skill_items( + &skill_items, + &connectors, + &skill_name_counts_lower, + ); + + assert_eq!(connector_ids, HashSet::::new()); + } + #[test] fn search_tool_selection_keeps_codex_apps_tools_without_mentions() { let selected_tool_names = vec![ diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index a2df7c239..70fa78a28 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -30,6 +30,7 @@ pub(crate) struct SessionState { /// Startup regular task pre-created during session initialization. pub(crate) startup_regular_task: Option, pub(crate) active_mcp_tool_selection: Option>, + pub(crate) active_connector_selection: HashSet, } impl SessionState { @@ -47,6 +48,7 @@ impl SessionState { previous_model: None, startup_regular_task: None, active_mcp_tool_selection: None, + active_connector_selection: HashSet::new(), } } @@ -175,6 +177,25 @@ impl SessionState { pub(crate) fn clear_mcp_tool_selection(&mut self) { self.active_mcp_tool_selection = None; } + + // Adds connector IDs to the active set and returns the merged selection. + pub(crate) fn merge_connector_selection(&mut self, connector_ids: I) -> HashSet + where + I: IntoIterator, + { + self.active_connector_selection.extend(connector_ids); + self.active_connector_selection.clone() + } + + // Returns the current connector selection tracked on session state. + pub(crate) fn get_connector_selection(&self) -> HashSet { + self.active_connector_selection.clone() + } + + // Removes all currently tracked connector selections. + pub(crate) fn clear_connector_selection(&mut self) { + self.active_connector_selection.clear(); + } } // Sometimes new snapshots don't include credits or plan information. @@ -272,6 +293,35 @@ mod tests { assert_eq!(state.get_mcp_tool_selection(), None); } + #[tokio::test] + // Verifies connector merging deduplicates repeated IDs. + async fn merge_connector_selection_deduplicates_entries() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + let merged = state.merge_connector_selection([ + "calendar".to_string(), + "calendar".to_string(), + "drive".to_string(), + ]); + + assert_eq!( + merged, + HashSet::from(["calendar".to_string(), "drive".to_string()]) + ); + } + + #[tokio::test] + // Verifies clearing connector selection removes all saved IDs. + async fn clear_connector_selection_removes_entries() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + state.merge_connector_selection(["calendar".to_string()]); + + state.clear_connector_selection(); + + assert_eq!(state.get_connector_selection(), HashSet::new()); + } + #[tokio::test] async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { let session_configuration = make_session_configuration_for_tests().await; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index cf2462b83..e05912f5f 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -121,6 +121,7 @@ impl Session { ) { self.abort_all_tasks(TurnAbortReason::Replaced).await; self.clear_mcp_tool_selection().await; + self.clear_connector_selection().await; self.seed_initial_context_if_needed(turn_context.as_ref()) .await;