From f2995fd4c0e211950f92cb6ebd5ad9db0c7ec136 Mon Sep 17 00:00:00 2001 From: Matthew Zeng Date: Fri, 30 Jan 2026 12:44:00 -0800 Subject: [PATCH] update --- codex-rs/core/src/codex.rs | 118 +++++++++++------- .../src/instructions/user_instructions.rs | 14 +++ codex-rs/core/src/mentions.rs | 82 +++++++++++- codex-rs/core/src/skills/injection.rs | 99 +++++++++++++++ codex-rs/core/src/skills/mod.rs | 1 + codex-rs/core/src/state/session.rs | 28 +++++ 6 files changed, 294 insertions(+), 48 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 1fb5572183..8a4282bf38 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -120,10 +120,10 @@ use crate::mcp::effective_mcp_servers; use crate::mcp::maybe_prompt_and_install_mcp_dependencies; use crate::mcp::with_codex_apps_mcp; use crate::mcp_connection_manager::McpConnectionManager; +use crate::mentions::CollectedToolMentions; use crate::mentions::build_connector_slug_counts; use crate::mentions::build_skill_name_counts; -use crate::mentions::collect_explicit_app_paths; -use crate::mentions::collect_tool_mentions_from_messages; +use crate::mentions::collect_app_ids_from_mentions; use crate::model_provider_info::CHAT_WIRE_API_DEPRECATION_SUMMARY; use crate::project_doc::get_user_instructions; use crate::protocol::AgentMessageContentDeltaEvent; @@ -170,9 +170,7 @@ use crate::skills::SkillsManager; use crate::skills::build_skill_injections; use crate::skills::collect_env_var_dependencies; use crate::skills::collect_explicit_skill_mentions; -use crate::skills::injection::ToolMentionKind; -use crate::skills::injection::app_id_from_path; -use crate::skills::injection::tool_kind_for_path; +use crate::skills::collect_explicit_skill_mentions_from_mentions; use crate::skills::resolve_skill_dependencies_for_turn; use crate::state::ActiveTurn; use crate::state::SessionServices; @@ -1755,6 +1753,7 @@ impl Session { ) { let mut state = self.state.lock().await; state.record_items(items.iter(), turn_context.truncation_policy); + state.update_mentions_from_items(items); } pub(crate) async fn record_model_warning(&self, message: impl Into, ctx: &TurnContext) { @@ -2021,6 +2020,10 @@ impl Session { // those spans, and `record_response_item_and_emit_turn_item` would drop them. self.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) .await; + { + let mut state = self.state.lock().await; + state.update_mentions_from_user_input(input); + } let turn_item = TurnItem::UserMessage(UserMessageItem::new(input)); self.emit_turn_item_started(turn_context, &turn_item).await; self.emit_turn_item_completed(turn_context, turn_item).await; @@ -3226,7 +3229,23 @@ pub(crate) async fn run_turn( } else { HashMap::new() }; - let mentioned_skills = skills_outcome.as_ref().map_or_else(Vec::new, |outcome| { + let (history_mentions, injected_skill_paths) = { + let state = sess.state.lock().await; + ( + state.tool_mentions.clone(), + state.injected_skill_paths.clone(), + ) + }; + let mentioned_skills_from_history = skills_outcome.as_ref().map_or_else(Vec::new, |outcome| { + collect_explicit_skill_mentions_from_mentions( + &history_mentions, + &outcome.skills, + &outcome.disabled_paths, + &skill_name_counts, + &connector_slug_counts, + ) + }); + let mentioned_skills_from_input = skills_outcome.as_ref().map_or_else(Vec::new, |outcome| { collect_explicit_skill_mentions( &input, &outcome.skills, @@ -3235,7 +3254,11 @@ pub(crate) async fn run_turn( &connector_slug_counts, ) }); - let explicit_app_paths = collect_explicit_app_paths(&input); + let mentioned_skills = merge_mentioned_skills( + mentioned_skills_from_history, + mentioned_skills_from_input, + &injected_skill_paths, + ); let config = turn_context.client.config(); if config @@ -3310,8 +3333,12 @@ pub(crate) async fn run_turn( }) .map(|user_message| user_message.message()) .collect::>(); + let tool_mentions = { + let state = sess.state.lock().await; + state.tool_mentions.clone() + }; let tool_selection = SamplingRequestToolSelection { - explicit_app_paths: &explicit_app_paths, + tool_mentions, skill_name_counts_lower: &skill_name_counts_lower, }; match run_sampling_request( @@ -3394,18 +3421,15 @@ async fn run_auto_compact(sess: &Arc, turn_context: &Arc) } } -fn filter_connectors_for_input( +fn filter_connectors_for_mentions( connectors: Vec, - input: &[ResponseItem], - explicit_app_paths: &[String], + mentions: &CollectedToolMentions, skill_name_counts_lower: &HashMap, ) -> Vec { - let user_messages = collect_user_messages(input); - if user_messages.is_empty() && explicit_app_paths.is_empty() { + if mentions.is_empty() { return Vec::new(); } - let mentions = collect_tool_mentions_from_messages(&user_messages); let mention_names_lower = mentions .plain_names .iter() @@ -3413,16 +3437,7 @@ fn filter_connectors_for_input( .collect::>(); let connector_slug_counts = build_connector_slug_counts(&connectors); - let mut allowed_connector_ids: HashSet = HashSet::new(); - for path in explicit_app_paths - .iter() - .chain(mentions.paths.iter()) - .filter(|path| tool_kind_for_path(path) == ToolMentionKind::App) - { - if let Some(connector_id) = app_id_from_path(path) { - allowed_connector_ids.insert(connector_id.to_string()); - } - } + let allowed_connector_ids = collect_app_ids_from_mentions(mentions); connectors .into_iter() @@ -3461,6 +3476,33 @@ fn connector_inserted_in_messages( connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&mention_slug) } +fn merge_mentioned_skills( + history: Vec, + current: Vec, + injected_skill_paths: &HashSet, +) -> Vec { + let mut merged = Vec::with_capacity(history.len() + current.len()); + let mut seen_paths: HashSet = HashSet::new(); + + for skill in history { + let path_str = skill.path.to_string_lossy(); + if injected_skill_paths.contains(path_str.as_ref()) { + continue; + } + if seen_paths.insert(skill.path.clone()) { + merged.push(skill); + } + } + + for skill in current { + if seen_paths.insert(skill.path.clone()) { + merged.push(skill); + } + } + + merged +} + fn filter_codex_apps_mcp_tools( mut mcp_tools: HashMap, connectors: &[connectors::AppInfo], @@ -3488,7 +3530,7 @@ fn codex_apps_connector_id(tool: &crate::mcp_connection_manager::ToolInfo) -> Op } struct SamplingRequestToolSelection<'a> { - explicit_app_paths: &'a [String], + tool_mentions: CollectedToolMentions, skill_name_counts_lower: &'a HashMap, } @@ -3519,10 +3561,9 @@ async fn run_sampling_request( .await?; let connectors_for_tools = if turn_context.client.config().features.enabled(Feature::Apps) { let connectors = connectors::accessible_connectors_from_mcp_tools(&mcp_tools); - Some(filter_connectors_for_input( + Some(filter_connectors_for_mentions( connectors, - &input, - tool_selection.explicit_app_paths, + &tool_selection.tool_mentions, tool_selection.skill_name_counts_lower, )) } else { @@ -3941,6 +3982,7 @@ mod tests { use codex_protocol::ThreadId; use codex_protocol::models::FunctionCallOutputPayload; + use crate::mentions::collect_tool_mentions_from_response_items; use crate::protocol::CompactedItem; use crate::protocol::CreditsSnapshot; use crate::protocol::InitialHistory; @@ -4080,15 +4122,11 @@ mod tests { make_connector("two", "Foo-Bar"), ]; let input = vec![user_message("use $foo-bar")]; - let explicit_app_paths = Vec::new(); + let mentions = collect_tool_mentions_from_response_items(&input); let skill_name_counts_lower = HashMap::new(); - let selected = filter_connectors_for_input( - connectors, - &input, - &explicit_app_paths, - &skill_name_counts_lower, - ); + let selected = + filter_connectors_for_mentions(connectors, &mentions, &skill_name_counts_lower); assert_eq!(selected, Vec::new()); } @@ -4097,15 +4135,11 @@ mod tests { fn filter_connectors_for_input_skips_when_skill_name_conflicts() { let connectors = vec![make_connector("one", "Todoist")]; let input = vec![user_message("use $todoist")]; - let explicit_app_paths = Vec::new(); + let mentions = collect_tool_mentions_from_response_items(&input); let skill_name_counts_lower = HashMap::from([("todoist".to_string(), 1)]); - let selected = filter_connectors_for_input( - connectors, - &input, - &explicit_app_paths, - &skill_name_counts_lower, - ); + let selected = + filter_connectors_for_mentions(connectors, &mentions, &skill_name_counts_lower); assert_eq!(selected, Vec::new()); } diff --git a/codex-rs/core/src/instructions/user_instructions.rs b/codex-rs/core/src/instructions/user_instructions.rs index 611bf4cbf8..7ac3563688 100644 --- a/codex-rs/core/src/instructions/user_instructions.rs +++ b/codex-rs/core/src/instructions/user_instructions.rs @@ -59,6 +59,20 @@ impl SkillInstructions { false } } + + pub fn extract_path(message: &[ContentItem]) -> Option { + let [ContentItem::InputText { text }] = message else { + return None; + }; + let path_start = text.find("")? + "".len(); + let path_end = text[path_start..].find("")? + path_start; + let path = text.get(path_start..path_end)?; + if path.is_empty() { + None + } else { + Some(path.to_string()) + } + } } impl From for ResponseItem { diff --git a/codex-rs/core/src/mentions.rs b/codex-rs/core/src/mentions.rs index 2f39c10d2f..c747085453 100644 --- a/codex-rs/core/src/mentions.rs +++ b/codex-rs/core/src/mentions.rs @@ -2,18 +2,42 @@ use std::collections::HashMap; use std::collections::HashSet; use std::path::PathBuf; +use codex_protocol::models::ResponseItem; use codex_protocol::user_input::UserInput; +use crate::compact::collect_user_messages; use crate::connectors; +use crate::instructions::SkillInstructions; use crate::skills::SkillMetadata; +use crate::skills::injection::ToolMentionKind; use crate::skills::injection::extract_tool_mentions; +use crate::skills::injection::tool_kind_for_path; +#[derive(Debug, Clone, Default)] pub(crate) struct CollectedToolMentions { pub(crate) plain_names: HashSet, pub(crate) paths: HashSet, } +impl CollectedToolMentions { + pub(crate) fn is_empty(&self) -> bool { + self.plain_names.is_empty() && self.paths.is_empty() + } + + pub(crate) fn extend_from(&mut self, other: &CollectedToolMentions) { + self.plain_names.extend(other.plain_names.iter().cloned()); + self.paths.extend(other.paths.iter().cloned()); + } +} + pub(crate) fn collect_tool_mentions_from_messages(messages: &[String]) -> CollectedToolMentions { + collect_tool_mentions_from_texts(messages.iter().map(String::as_str)) +} + +pub(crate) fn collect_tool_mentions_from_texts<'a, I>(messages: I) -> CollectedToolMentions +where + I: IntoIterator, +{ let mut plain_names = HashSet::new(); let mut paths = HashSet::new(); for message in messages { @@ -24,13 +48,59 @@ pub(crate) fn collect_tool_mentions_from_messages(messages: &[String]) -> Collec CollectedToolMentions { plain_names, paths } } -pub(crate) fn collect_explicit_app_paths(input: &[UserInput]) -> Vec { - input +pub(crate) fn collect_structured_tool_mentions_from_user_input( + input: &[UserInput], +) -> CollectedToolMentions { + let mut mentions = CollectedToolMentions::default(); + for item in input { + match item { + UserInput::Mention { name, path } => { + mentions.plain_names.insert(name.clone()); + mentions.paths.insert(path.clone()); + } + UserInput::Skill { name, path } => { + mentions.plain_names.insert(name.clone()); + mentions.paths.insert(path.to_string_lossy().into_owned()); + } + UserInput::Text { .. } | UserInput::Image { .. } | UserInput::LocalImage { .. } => {} + _ => {} + } + } + + mentions +} + +pub(crate) fn collect_tool_mentions_from_response_items( + items: &[ResponseItem], +) -> CollectedToolMentions { + let messages = collect_user_messages(items); + collect_tool_mentions_from_messages(&messages) +} + +pub(crate) fn collect_skill_instruction_paths(items: &[ResponseItem]) -> HashSet { + let mut paths = HashSet::new(); + for item in items { + let ResponseItem::Message { role, content, .. } = item else { + continue; + }; + if role != "user" { + continue; + } + let Some(path) = SkillInstructions::extract_path(content) else { + continue; + }; + paths.insert(path); + } + paths +} + +pub(crate) fn collect_app_ids_from_mentions(mentions: &CollectedToolMentions) -> HashSet { + mentions + .paths .iter() - .filter_map(|item| match item { - UserInput::Mention { path, .. } => Some(path.clone()), - _ => None, - }) + .filter(|path| tool_kind_for_path(path) == ToolMentionKind::App) + .filter_map(|path| crate::skills::injection::app_id_from_path(path)) + .map(str::to_string) .collect() } diff --git a/codex-rs/core/src/skills/injection.rs b/codex-rs/core/src/skills/injection.rs index 57de457e90..5e8e08f226 100644 --- a/codex-rs/core/src/skills/injection.rs +++ b/codex-rs/core/src/skills/injection.rs @@ -3,6 +3,7 @@ use std::collections::HashSet; use std::path::PathBuf; use crate::instructions::SkillInstructions; +use crate::mentions::CollectedToolMentions; use crate::skills::SkillMetadata; use codex_otel::OtelManager; use codex_protocol::models::ResponseItem; @@ -106,6 +107,34 @@ pub(crate) fn collect_explicit_skill_mentions( selected } +pub(crate) fn collect_explicit_skill_mentions_from_mentions( + mentions: &CollectedToolMentions, + skills: &[SkillMetadata], + disabled_paths: &HashSet, + skill_name_counts: &HashMap, + connector_slug_counts: &HashMap, +) -> Vec { + let selection_context = SkillSelectionContext { + skills, + disabled_paths, + skill_name_counts, + connector_slug_counts, + }; + let mut selected: Vec = Vec::new(); + let mut seen_names: HashSet = HashSet::new(); + let mut seen_paths: HashSet = HashSet::new(); + + select_skills_from_collected_mentions( + &selection_context, + mentions, + &mut seen_names, + &mut seen_paths, + &mut selected, + ); + + selected +} + struct SkillSelectionContext<'a> { skills: &'a [SkillMetadata], disabled_paths: &'a HashSet, @@ -308,6 +337,76 @@ fn select_skills_from_mentions( } } +fn select_skills_from_collected_mentions( + selection_context: &SkillSelectionContext<'_>, + mentions: &CollectedToolMentions, + seen_names: &mut HashSet, + seen_paths: &mut HashSet, + selected: &mut Vec, +) { + if mentions.is_empty() { + return; + } + + let mention_skill_paths: HashSet<&str> = mentions + .paths + .iter() + .filter(|path| { + !matches!( + tool_kind_for_path(path), + ToolMentionKind::App | ToolMentionKind::Mcp + ) + }) + .map(|path| normalize_skill_path(path)) + .collect(); + + for skill in selection_context.skills { + if selection_context.disabled_paths.contains(&skill.path) + || seen_paths.contains(&skill.path) + { + continue; + } + + let path_str = skill.path.to_string_lossy(); + if mention_skill_paths.contains(path_str.as_ref()) { + seen_paths.insert(skill.path.clone()); + seen_names.insert(skill.name.clone()); + selected.push(skill.clone()); + } + } + + for skill in selection_context.skills { + if selection_context.disabled_paths.contains(&skill.path) + || seen_paths.contains(&skill.path) + { + continue; + } + + if !mentions.plain_names.contains(skill.name.as_str()) { + continue; + } + + let skill_count = selection_context + .skill_name_counts + .get(skill.name.as_str()) + .copied() + .unwrap_or(0); + let connector_count = selection_context + .connector_slug_counts + .get(&skill.name.to_ascii_lowercase()) + .copied() + .unwrap_or(0); + if skill_count != 1 || connector_count != 0 { + continue; + } + + if seen_names.insert(skill.name.clone()) { + seen_paths.insert(skill.path.clone()); + selected.push(skill.clone()); + } + } +} + fn parse_linked_tool_mention<'a>( text: &'a str, text_bytes: &[u8], diff --git a/codex-rs/core/src/skills/mod.rs b/codex-rs/core/src/skills/mod.rs index 4d79bca258..f130c2a6f0 100644 --- a/codex-rs/core/src/skills/mod.rs +++ b/codex-rs/core/src/skills/mod.rs @@ -11,6 +11,7 @@ pub(crate) use env_var_dependencies::resolve_skill_dependencies_for_turn; pub(crate) use injection::SkillInjections; pub(crate) use injection::build_skill_injections; pub(crate) use injection::collect_explicit_skill_mentions; +pub(crate) use injection::collect_explicit_skill_mentions_from_mentions; pub use loader::load_skills; pub use manager::SkillsManager; pub use model::SkillError; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index deee0d0c7c..c7761b7db3 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -1,11 +1,16 @@ //! Session-wide mutable state. use codex_protocol::models::ResponseItem; +use codex_protocol::user_input::UserInput; use std::collections::HashMap; use std::collections::HashSet; use crate::codex::SessionConfiguration; use crate::context_manager::ContextManager; +use crate::mentions::CollectedToolMentions; +use crate::mentions::collect_skill_instruction_paths; +use crate::mentions::collect_structured_tool_mentions_from_user_input; +use crate::mentions::collect_tool_mentions_from_response_items; use crate::protocol::RateLimitSnapshot; use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; @@ -19,6 +24,8 @@ pub(crate) struct SessionState { pub(crate) server_reasoning_included: bool, pub(crate) dependency_env: HashMap, pub(crate) mcp_dependency_prompted: HashSet, + pub(crate) tool_mentions: CollectedToolMentions, + pub(crate) injected_skill_paths: HashSet, /// Whether the session's initial context has been seeded into history. /// /// TODO(owen): This is a temporary solution to avoid updating a thread's updated_at @@ -37,6 +44,8 @@ impl SessionState { server_reasoning_included: false, dependency_env: HashMap::new(), mcp_dependency_prompted: HashSet::new(), + tool_mentions: CollectedToolMentions::default(), + injected_skill_paths: HashSet::new(), initial_context_seeded: false, } } @@ -56,6 +65,25 @@ impl SessionState { pub(crate) fn replace_history(&mut self, items: Vec) { self.history.replace(items); + self.rebuild_mentions_from_history(); + } + + pub(crate) fn update_mentions_from_items(&mut self, items: &[ResponseItem]) { + let mentions = collect_tool_mentions_from_response_items(items); + self.tool_mentions.extend_from(&mentions); + let skill_paths = collect_skill_instruction_paths(items); + self.injected_skill_paths.extend(skill_paths); + } + + pub(crate) fn update_mentions_from_user_input(&mut self, input: &[UserInput]) { + let mentions = collect_structured_tool_mentions_from_user_input(input); + self.tool_mentions.extend_from(&mentions); + } + + fn rebuild_mentions_from_history(&mut self) { + let items = self.history.raw_items(); + self.tool_mentions = collect_tool_mentions_from_response_items(items); + self.injected_skill_paths = collect_skill_instruction_paths(items); } pub(crate) fn set_token_info(&mut self, info: Option) {