diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 6fd863daf5..010be9282d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3889,6 +3889,7 @@ dependencies = [ "codex-exec-server", "codex-extension-api", "codex-mcp", + "codex-otel", "codex-protocol", "codex-tools", "codex-utils-absolute-path", diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index c91d5d38ea..b78bd6daea 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -113,7 +113,7 @@ struct AppServerExtensionEventSink { } impl ExtensionEventSink for AppServerExtensionEventSink { - fn emit(&self, event: Event) { + fn emit(&self, thread_id: ThreadId, event: Event) { match event.msg { EventMsg::ThreadGoalUpdated(thread_goal_event) => { let thread_id = thread_goal_event.thread_id; @@ -148,10 +148,20 @@ impl ExtensionEventSink for AppServerExtensionEventSink { }); } EventMsg::Warning(warning_event) => { - let Ok(thread_id) = ThreadId::from_string(&event.id) else { - tracing::warn!(event_id = %event.id, "dropping extension warning with invalid thread id"); - return; - }; + if let Some(listener_command_tx) = self + .thread_state_manager + .current_listener_command_tx(thread_id) + { + let command = ThreadListenerCommand::EmitWarning { + message: warning_event.message.clone(), + }; + if listener_command_tx.send(command).is_ok() { + return; + } + tracing::warn!( + "failed to enqueue extension warning for {thread_id}: listener command channel is closed" + ); + } let outgoing = Arc::clone(&self.outgoing); tokio::spawn(async move { outgoing @@ -165,7 +175,7 @@ impl ExtensionEventSink for AppServerExtensionEventSink { }); } msg => { - tracing::debug!(event_id = %event.id, ?msg, "dropping unsupported extension event"); + tracing::debug!(%thread_id, event_id = %event.id, ?msg, "dropping unsupported extension event"); } } } @@ -203,7 +213,7 @@ mod tests { use super::*; #[tokio::test] - async fn app_server_event_sink_uses_listener_fifo_for_goal_updates_and_clears() { + async fn app_server_event_sink_uses_listener_fifo_for_extension_events() { let (outgoing_tx, _outgoing_rx) = mpsc::channel(4); let outgoing = Arc::new(OutgoingMessageSender::new( outgoing_tx, @@ -216,14 +226,23 @@ mod tests { let sink = app_server_extension_event_sink(outgoing, thread_state_manager); for turn_id in ["turn-1", "turn-2"] { - sink.emit(thread_goal_updated_event(thread_id, turn_id)); + sink.emit(thread_id, thread_goal_updated_event(thread_id, turn_id)); } + sink.emit( + thread_id, + Event { + id: "turn-2".to_string(), + msg: EventMsg::Warning(codex_protocol::protocol::WarningEvent { + message: "skill warning".to_string(), + }), + }, + ); listener_command_tx .send(ThreadListenerCommand::EmitThreadGoalCleared) .expect("listener command channel should be open"); let mut observed = Vec::new(); - for _ in 0..3 { + for _ in 0..4 { let command = timeout(Duration::from_secs(1), listener_command_rx.recv()) .await .expect("timed out waiting for listener command") @@ -235,6 +254,7 @@ mod tests { ThreadListenerCommand::EmitThreadGoalCleared => { observed.push("cleared".to_string()) } + ThreadListenerCommand::EmitWarning { message } => observed.push(message), _ => panic!("unexpected listener command"), } } @@ -243,6 +263,7 @@ mod tests { vec![ "turn-1".to_string(), "turn-2".to_string(), + "skill warning".to_string(), "cleared".to_string() ], observed diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 8c202a806d..fc32e9f079 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -276,6 +276,7 @@ use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::TurnSteerParams; use codex_app_server_protocol::TurnSteerResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_app_server_protocol::WarningNotification; use codex_app_server_protocol::WindowsSandboxReadiness; use codex_app_server_protocol::WindowsSandboxReadinessResponse; use codex_app_server_protocol::WindowsSandboxSetupCompletedNotification; diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index fd9e93e189..55bf23d247 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -505,6 +505,14 @@ pub(super) async fn handle_thread_listener_command( ThreadListenerCommand::EmitThreadGoalSnapshot { state_db } => { send_thread_goal_snapshot_notification(outgoing, conversation_id, &state_db).await; } + ThreadListenerCommand::EmitWarning { message } => { + outgoing + .send_server_notification(ServerNotification::Warning(WarningNotification { + thread_id: Some(conversation_id.to_string()), + message, + })) + .await; + } ThreadListenerCommand::ResolveServerRequest { request_id, completion_tx, diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 6d2b48a4c8..97e86b6755 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -56,6 +56,10 @@ pub(crate) enum ThreadListenerCommand { EmitThreadGoalSnapshot { state_db: StateDbHandle, }, + // EmitWarning is used to order extension warnings with core turn events. + EmitWarning { + message: String, + }, // ResolveServerRequest is used to notify the client that the request has been resolved. // It is executed in the thread listener's context to ensure that the resolved notification is ordered with regard to the request itself. ResolveServerRequest { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 00757d5210..381c721878 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2967,9 +2967,14 @@ impl Session { } } let context_contributors = self.services.extensions.context_contributors().to_vec(); + let mut extension_developer_policy = Vec::new(); + let mut extension_developer_capabilities = Vec::new(); + let mut extension_contextual_user = Vec::new(); + let mut extension_separate_developer = Vec::new(); for contributor in context_contributors { for fragment in contributor .contribute(ContextContributionContext { + thread_id: self.thread_id(), session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, turn_store: turn_context.extension_data.as_ref(), @@ -2978,18 +2983,22 @@ impl Session { .await { match fragment.slot() { - PromptSlot::DeveloperPolicy | PromptSlot::DeveloperCapabilities => { - developer_sections.push(fragment.text().to_string()); + PromptSlot::DeveloperPolicy => { + extension_developer_policy.push(fragment.text().to_string()); + } + PromptSlot::DeveloperCapabilities => { + extension_developer_capabilities.push(fragment.text().to_string()); } PromptSlot::ContextualUser => { - contextual_user_sections.push(fragment.text().to_string()); + extension_contextual_user.push(fragment.text().to_string()); } PromptSlot::SeparateDeveloper => { - separate_developer_sections.push(fragment.text().to_string()); + extension_separate_developer.push(fragment.text().to_string()); } } } } + developer_sections.extend(extension_developer_capabilities); let loaded_plugins = self .services .plugins_manager @@ -3023,6 +3032,9 @@ impl Session { { developer_sections.push(plugin_instructions.render()); } + developer_sections.extend(extension_developer_policy); + contextual_user_sections.extend(extension_contextual_user); + separate_developer_sections.extend(extension_separate_developer); if let Some(user_instructions) = turn_context.user_instructions.as_deref() { contextual_user_sections.push(user_instructions.to_string()); } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 257e741203..22987498c5 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -652,6 +652,7 @@ async fn build_extension_turn_input_items( .collect::>(); let input = TurnInputContext { + thread_id: sess.thread_id(), turn_id: turn_context.sub_id.to_string(), user_input: user_input.to_vec(), environments, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 91880bdefa..2a75989aa9 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -536,6 +536,7 @@ impl Session { )); let (current_date, timezone) = local_time_context(); let extension_data = Arc::new(codex_extension_api::ExtensionData::new(sub_id.clone())); + extension_data.insert(session_telemetry_for_context.clone()); extension_data.insert(skills_snapshot.clone()); TurnContext { sub_id, diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions.rs b/codex-rs/ext/extension-api/examples/enabled_extensions.rs index 130884dcde..37dbffaa64 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions.rs @@ -79,6 +79,7 @@ async fn contribute_prompt( fragments.extend( contributor .contribute(ContextContributionContext { + thread_id: codex_protocol::ThreadId::new(), session_store, thread_store, turn_store: &turn_store, diff --git a/codex-rs/ext/extension-api/src/capabilities/events.rs b/codex-rs/ext/extension-api/src/capabilities/events.rs index d8fda1aea7..3e799acaab 100644 --- a/codex-rs/ext/extension-api/src/capabilities/events.rs +++ b/codex-rs/ext/extension-api/src/capabilities/events.rs @@ -1,13 +1,15 @@ +use codex_protocol::ThreadId; use codex_protocol::protocol::Event; /// Host-provided fire-and-forget sink for extension-generated events. /// /// Extensions construct protocol events with the correlation id appropriate for -/// the callback they are handling, then leave persistence, ordering, transport -/// fanout, and logging decisions to the host. +/// the callback they are handling and identify the owning thread explicitly, +/// then leave persistence, ordering, transport fanout, and logging decisions to +/// the host. pub trait ExtensionEventSink: Send + Sync { /// Queue one protocol event for host-owned delivery. - fn emit(&self, event: Event); + fn emit(&self, thread_id: ThreadId, event: Event); } /// Event sink used when the host does not expose extension event emission. @@ -15,5 +17,5 @@ pub trait ExtensionEventSink: Send + Sync { pub struct NoopExtensionEventSink; impl ExtensionEventSink for NoopExtensionEventSink { - fn emit(&self, _event: Event) {} + fn emit(&self, _thread_id: ThreadId, _event: Event) {} } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index bafb73074e..a57951a253 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -3,6 +3,7 @@ use std::pin::Pin; use std::sync::Arc; use codex_context_fragments::ContextualUserFragment; +use codex_protocol::ThreadId; use codex_protocol::items::TurnItem; use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::TokenUsageInfo; @@ -44,6 +45,7 @@ pub type ExtensionFuture<'a, T> = Pin + Send + 'a>>; /// Host context available while extensions contribute initial model input. #[derive(Clone, Copy)] pub struct ContextContributionContext<'a> { + pub thread_id: ThreadId, pub session_store: &'a ExtensionData, pub thread_store: &'a ExtensionData, pub turn_store: &'a ExtensionData, diff --git a/codex-rs/ext/extension-api/src/contributors/turn_input.rs b/codex-rs/ext/extension-api/src/contributors/turn_input.rs index 9e4723440d..481a3a9eef 100644 --- a/codex-rs/ext/extension-api/src/contributors/turn_input.rs +++ b/codex-rs/ext/extension-api/src/contributors/turn_input.rs @@ -1,5 +1,6 @@ use std::path::PathBuf; +use codex_protocol::ThreadId; use codex_protocol::user_input::UserInput; /// Host-owned turn environment summary visible to turn-input contributors. @@ -16,6 +17,8 @@ pub struct TurnInputEnvironment { /// Turn facts supplied before the host records turn-local model input items. #[derive(Debug, Clone)] pub struct TurnInputContext { + /// Thread that owns this turn. + pub thread_id: ThreadId, /// Stable host-owned turn identifier. pub turn_id: String, /// User input submitted for this turn. diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index c4cc6fcf41..8e5bf49466 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/codex-rs/ext/extension-api/tests/registry.rs @@ -201,6 +201,7 @@ async fn contributors_preserve_registration_order() { fragments.extend( contributor .contribute(ContextContributionContext { + thread_id: codex_protocol::ThreadId::new(), session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, @@ -323,7 +324,7 @@ struct RecordingEventSink { } impl ExtensionEventSink for RecordingEventSink { - fn emit(&self, event: Event) { + fn emit(&self, _thread_id: codex_protocol::ThreadId, event: Event) { let EventMsg::Warning(warning) = event.msg else { panic!("test sink only accepts warning events"); }; @@ -339,13 +340,15 @@ fn custom_event_sink_survives_registry_build() { let sink = Arc::new(RecordingEventSink::default()); let builder = ExtensionRegistryBuilder::<()>::with_event_sink(sink.clone()); - builder - .event_sink() - .emit(warning_event("builder", "before")); + builder.event_sink().emit( + codex_protocol::ThreadId::new(), + warning_event("builder", "before"), + ); let registry = builder.build(); - registry - .event_sink() - .emit(warning_event("registry", "after")); + registry.event_sink().emit( + codex_protocol::ThreadId::new(), + warning_event("registry", "after"), + ); assert_eq!( sink.events diff --git a/codex-rs/ext/goal/src/events.rs b/codex-rs/ext/goal/src/events.rs index ab9eda4056..faa50eacfe 100644 --- a/codex-rs/ext/goal/src/events.rs +++ b/codex-rs/ext/goal/src/events.rs @@ -22,13 +22,17 @@ impl GoalEventEmitter { turn_id: Option, goal: ThreadGoal, ) { - self.sink.emit(Event { - id: event_id.into(), - msg: EventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent { - thread_id: goal.thread_id, - turn_id, - goal, - }), - }); + let thread_id = goal.thread_id; + self.sink.emit( + thread_id, + Event { + id: event_id.into(), + msg: EventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent { + thread_id, + turn_id, + goal, + }), + }, + ); } } diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 32043c9808..5ec329b484 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1407,7 +1407,7 @@ impl RecordingEventSink { } impl ExtensionEventSink for RecordingEventSink { - fn emit(&self, event: Event) { + fn emit(&self, _thread_id: ThreadId, event: Event) { self.events().push(event); } } diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index 8b359284e5..2f451d6703 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -185,6 +185,7 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() { let fragments = extension .contribute(ContextContributionContext { + thread_id: Default::default(), session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, diff --git a/codex-rs/ext/skills/Cargo.toml b/codex-rs/ext/skills/Cargo.toml index afd5111e95..61089be1c7 100644 --- a/codex-rs/ext/skills/Cargo.toml +++ b/codex-rs/ext/skills/Cargo.toml @@ -18,6 +18,7 @@ codex-core-skills = { workspace = true } codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-mcp = { workspace = true } +codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index d16f5ccba7..3e3cf98ea0 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -20,6 +20,7 @@ use codex_extension_api::ToolExecutor; use codex_extension_api::TurnInputContext; use codex_extension_api::TurnInputContributor; use codex_mcp::McpResourceClient; +use codex_otel::SessionTelemetry; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; @@ -134,17 +135,23 @@ where ) .await; for warning in &catalog.warnings { - self.emit_warning(thread_store.level_id(), warning.clone()); + self.emit_warning( + context.thread_id, + context.turn_store.level_id(), + warning.clone(), + ); } + let session_telemetry = context.turn_store.get::(); let Some((fragment, warning)) = available_skills_fragment( &catalog, host_snapshot.as_deref(), context.model_context_window, + session_telemetry.as_deref(), ) else { return Vec::new(); }; if let Some(warning) = warning { - self.emit_warning(thread_store.level_id(), warning); + self.emit_warning(context.thread_id, context.turn_store.level_id(), warning); } vec![PromptFragment::developer_capability(fragment.render())] }) @@ -206,7 +213,7 @@ where }; let catalog = self.list_skills(query, &thread_state).await; for warning in &catalog.warnings { - self.emit_warning(&input.turn_id, warning.clone()); + self.emit_warning(input.thread_id, &input.turn_id, warning.clone()); } let selected_entries = collect_explicit_skill_mentions(&input.user_input, &catalog); @@ -227,7 +234,7 @@ where "Skill `{}` exceeded the main prompt context limit and was truncated.", entry.name ); - self.emit_warning(&input.turn_id, warning.clone()); + self.emit_warning(input.thread_id, &input.turn_id, warning.clone()); warnings.push(warning); } let fragment = SkillInstructions { @@ -247,7 +254,7 @@ where } Err(message) => { let warning = format!("Failed to load skill `{}`: {message}", entry.name); - self.emit_warning(&input.turn_id, warning.clone()); + self.emit_warning(input.thread_id, &input.turn_id, warning.clone()); warnings.push(warning); } } @@ -332,11 +339,14 @@ impl SkillsExtension { .map_err(|err| err.message) } - fn emit_warning(&self, turn_id: &str, message: String) { - self.event_sink.emit(Event { - id: turn_id.to_string(), - msg: EventMsg::Warning(WarningEvent { message }), - }); + fn emit_warning(&self, thread_id: codex_protocol::ThreadId, turn_id: &str, message: String) { + self.event_sink.emit( + thread_id, + Event { + id: turn_id.to_string(), + msg: EventMsg::Warning(WarningEvent { message }), + }, + ); } } diff --git a/codex-rs/ext/skills/src/render.rs b/codex-rs/ext/skills/src/render.rs index 94787570fc..1718a4d8ca 100644 --- a/codex-rs/ext/skills/src/render.rs +++ b/codex-rs/ext/skills/src/render.rs @@ -9,8 +9,12 @@ use codex_core_skills::build_available_skills; use codex_core_skills::default_skill_metadata_budget; use codex_core_skills::render::SkillCatalogMode; use codex_core_skills::render::SkillRenderSideEffects; +use codex_extension_api::ContextualUserFragment; +use codex_otel::SessionTelemetry; +use codex_utils_string::approx_token_count; const MAX_AVAILABLE_SKILLS_BYTES: usize = 8_000; +const MAX_AVAILABLE_SKILLS_TOKENS: usize = 10_000; const MAX_MAIN_PROMPT_BYTES: usize = 8_000; pub(crate) const MAX_SKILL_NAME_BYTES: usize = 256; pub(crate) const MAX_SKILL_PATH_BYTES: usize = 1_024; @@ -19,16 +23,29 @@ pub(crate) fn available_skills_fragment( catalog: &SkillCatalog, host_snapshot: Option<&HostSkillsSnapshot>, model_context_window: Option, + session_telemetry: Option<&SessionTelemetry>, ) -> Option<(AvailableSkillsInstructions, Option)> { let visible_entries = catalog .entries .iter() .filter(|entry| entry.enabled && entry.prompt_visible) .collect::>(); + let side_effects = session_telemetry + .map_or(SkillRenderSideEffects::None, |session_telemetry| { + SkillRenderSideEffects::ThreadStart { session_telemetry } + }); + let host_available = host_snapshot.and_then(|host_snapshot| { + build_available_skills( + host_snapshot.outcome(), + default_skill_metadata_budget(model_context_window), + side_effects, + ) + }); + if visible_entries .iter() .any(|entry| entry.authority.kind == SkillSourceKind::Host) - && let Some(host_snapshot) = host_snapshot + && let Some(available) = host_available { let external_lines = bounded_skill_lines( visible_entries @@ -36,30 +53,8 @@ pub(crate) fn available_skills_fragment( .copied() .filter(|entry| entry.authority.kind != SkillSourceKind::Host), ); - if let Some(available) = build_available_skills( - host_snapshot.outcome(), - default_skill_metadata_budget(model_context_window), - SkillRenderSideEffects::None, - ) { - let warning = available.warning_message.clone(); - return Some(( - AvailableSkillsInstructions::from_available_skills_with_additional_lines( - available, - external_lines, - ), - warning, - )); - } - if !external_lines.is_empty() { - return Some(( - AvailableSkillsInstructions::from_skill_lines( - external_lines, - SkillCatalogMode::Mixed, - ), - None, - )); - } - return None; + let warning = available.warning_message.clone(); + return Some((bounded_mixed_fragment(available, external_lines), warning)); } let mode = if visible_entries @@ -70,7 +65,7 @@ pub(crate) fn available_skills_fragment( } else { SkillCatalogMode::HostOnly }; - let skill_lines = bounded_skill_lines(visible_entries); + let skill_lines = bounded_skill_lines(visible_entries).into_lines(); (!skill_lines.is_empty()).then(|| { ( AvailableSkillsInstructions::from_skill_lines(skill_lines, mode), @@ -79,9 +74,23 @@ pub(crate) fn available_skills_fragment( }) } +struct BoundedSkillLines { + lines: Vec, + omitted: usize, +} + +impl BoundedSkillLines { + fn into_lines(mut self) -> Vec { + if let Some(line) = omitted_skills_line(self.omitted) { + self.lines.push(line); + } + self.lines + } +} + fn bounded_skill_lines<'a>( entries: impl IntoIterator, -) -> Vec { +) -> BoundedSkillLines { let mut total_bytes = 0usize; let mut omitted = 0usize; let mut skill_lines = Vec::new(); @@ -101,13 +110,43 @@ fn bounded_skill_lines<'a>( skill_lines.push(line); } - if omitted > 0 { - let skill_word = if omitted == 1 { "skill" } else { "skills" }; - skill_lines.push(format!( - "- {omitted} additional {skill_word} omitted from this bounded skills list." - )); + BoundedSkillLines { + lines: skill_lines, + omitted, } - skill_lines +} + +fn bounded_mixed_fragment( + available: codex_core_skills::AvailableSkills, + mut external: BoundedSkillLines, +) -> AvailableSkillsInstructions { + loop { + let mut additional_lines = external.lines.clone(); + if let Some(line) = omitted_skills_line(external.omitted) { + additional_lines.push(line); + } + let fragment = AvailableSkillsInstructions::from_available_skills_with_additional_lines( + available.clone(), + additional_lines, + ); + if approx_token_count(&fragment.render()) <= MAX_AVAILABLE_SKILLS_TOKENS { + return fragment; + } + if external.lines.pop().is_none() { + return AvailableSkillsInstructions::from_available_skills_with_additional_lines( + available, + Vec::new(), + ); + } + external.omitted = external.omitted.saturating_add(1); + } +} + +fn omitted_skills_line(omitted: usize) -> Option { + (omitted > 0).then(|| { + let skill_word = if omitted == 1 { "skill" } else { "skills" }; + format!("- {omitted} additional {skill_word} omitted from this bounded skills list.") + }) } fn render_skill_line(entry: &SkillCatalogEntry, description: &str) -> String { diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index 65fd40803d..b3347afa1f 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -15,6 +15,7 @@ use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadStartInput; use codex_extension_api::TurnInputContext; +use codex_protocol::ThreadId; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Event; @@ -43,6 +44,7 @@ use codex_skills_extension::provider::SkillProviderFuture; use codex_skills_extension::provider::SkillReadRequest; use codex_skills_extension::provider::SkillSearchRequest; use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_string::approx_token_count; use pretty_assertions::assert_eq; type TestResult = Result<(), Box>; @@ -114,20 +116,32 @@ async fn installed_extension_catalog_uses_host_service_snapshot() -> TestResult } #[tokio::test] -async fn mixed_catalog_preserves_host_budget_and_external_locator() -> TestResult { - let host_path = - AbsolutePathBuf::try_from(std::env::temp_dir().join("host-demo").join("SKILL.md"))?; - let mut outcome = SkillLoadOutcome::default(); - outcome.skills.push(SkillMetadata { - name: "host-demo".to_string(), - description: "x".repeat(5_000), - short_description: None, - interface: None, - dependencies: None, - policy: None, - path_to_skills_md: host_path, - scope: SkillScope::User, - plugin_id: None, +async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_input() -> TestResult +{ + let read_requests = Arc::new(Mutex::new(Vec::new())); + let mut executor_entries = vec![test_entry( + SkillSourceKind::Executor, + "env-1", + "executor/lint-fix", + "lint-fix/SKILL.md", + )]; + executor_entries.extend((0..20).map(|index| { + test_entry( + SkillSourceKind::Executor, + "env-1", + &format!("executor/extra-{index}"), + &format!("extra-{index}/SKILL.md"), + ) + .with_short_description(Some("x".repeat(1_000))) + })); + let executor_provider = Arc::new(StaticSkillProvider { + catalog: SkillCatalog { + entries: executor_entries, + warnings: Vec::new(), + }, + read_requests: Arc::clone(&read_requests), + list_calls: None, + fail_first_list: false, }); let providers = SkillProviders::new() .with_host_provider(Arc::new(StaticSkillProvider { @@ -144,25 +158,13 @@ async fn mixed_catalog_preserves_host_budget_and_external_locator() -> TestResul list_calls: None, fail_first_list: false, })) - .with_executor_provider(Arc::new(StaticSkillProvider { - catalog: SkillCatalog { - entries: vec![test_entry( - SkillSourceKind::Executor, - "env-1", - "executor/lint-fix", - "lint-fix/SKILL.md", - )], - warnings: Vec::new(), - }, - read_requests: Arc::new(Mutex::new(Vec::new())), - list_calls: None, - fail_first_list: false, - })); + .with_executor_provider(executor_provider); let (event_tx, event_rx) = std::sync::mpsc::channel(); let mut builder = ExtensionRegistryBuilder::with_event_sink(Arc::new(ChannelEventSink(event_tx))); install_with_providers(&mut builder, providers, skills_extension_config); let registry = builder.build(); + let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); thread_store.insert(vec![SelectedCapabilityRoot { @@ -184,82 +186,32 @@ async fn mixed_catalog_preserves_host_budget_and_external_locator() -> TestResul thread_store: &thread_store, }) .await; - let turn_store = ExtensionData::new("turn-1"); - turn_store.insert(HostSkillsSnapshot::new(Arc::new(outcome))); - let fragments = registry.context_contributors()[0] - .contribute(ContextContributionContext { - session_store: &session_store, - thread_store: &thread_store, - turn_store: &turn_store, - model_context_window: Some(10_000), - }) - .await; - - assert_eq!(1, fragments.len()); - let text = fragments[0].text(); - assert!(text.contains("host-demo")); - assert!(text.contains("environment resource: skill://executor/lint-fix/SKILL.md")); - let EventMsg::Warning(warning) = event_rx.try_recv()?.msg else { - panic!("expected host budget warning"); - }; - assert!(warning.message.contains("2% skills context budget")); - - Ok(()) -} - -#[tokio::test] -async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_input() -> TestResult -{ - let read_requests = Arc::new(Mutex::new(Vec::new())); - let executor_provider = Arc::new(StaticSkillProvider { - catalog: SkillCatalog { - entries: vec![test_entry( - SkillSourceKind::Executor, - "env-1", - "executor/lint-fix", - "lint-fix/SKILL.md", - )], - warnings: Vec::new(), - }, - read_requests: Arc::clone(&read_requests), - list_calls: None, - fail_first_list: false, + let initial_turn_store = ExtensionData::new("initial-turn"); + let mut outcome = SkillLoadOutcome::default(); + outcome.skills.push(SkillMetadata { + name: "host-demo".to_string(), + description: "x".repeat(40_000), + short_description: None, + interface: None, + dependencies: None, + policy: None, + path_to_skills_md: AbsolutePathBuf::try_from( + std::env::temp_dir().join("host-demo").join("SKILL.md"), + )?, + scope: SkillScope::User, + plugin_id: None, }); - let providers = SkillProviders::new().with_executor_provider(executor_provider); - let mut builder = ExtensionRegistryBuilder::new(); - install_with_providers(&mut builder, providers, skills_extension_config); - let registry = builder.build(); - - let session_store = ExtensionData::new("session"); - let thread_store = ExtensionData::new("thread"); - thread_store.insert(vec![SelectedCapabilityRoot { - id: "lint-fix".to_string(), - location: CapabilityRootLocation::Environment { - environment_id: "env-1".to_string(), - path: "/skills/lint-fix".to_string(), - }, - }]); - let session_source = SessionSource::Cli; - let config = default_config(); - registry.thread_lifecycle_contributors()[0] - .on_thread_start(ThreadStartInput { - config: &config, - session_source: &session_source, - persistent_thread_state_available: true, - environments: &[], + initial_turn_store.insert(HostSkillsSnapshot::new(Arc::new(outcome))); + let prompt_fragments = registry.context_contributors()[0] + .contribute(ContextContributionContext { + thread_id: ThreadId::default(), session_store: &session_store, thread_store: &thread_store, + turn_store: &initial_turn_store, + model_context_window: Some(400_000), }) .await; - - let prompt_fragments = registry.context_contributors()[0] - .contribute(prompt_context( - &session_store, - &thread_store, - &ExtensionData::new("initial-turn"), - )) - .await; assert_eq!(1, prompt_fragments.len()); assert!( prompt_fragments[0] @@ -267,16 +219,23 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in .starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) ); assert!(prompt_fragments[0].text().contains("lint-fix")); + assert!(prompt_fragments[0].text().contains("host-demo")); assert!( prompt_fragments[0] .text() .contains("(environment resource: skill://executor/lint-fix/SKILL.md)") ); + assert!(approx_token_count(prompt_fragments[0].text()) <= 10_000); + let EventMsg::Warning(warning) = event_rx.try_recv()?.1.msg else { + panic!("expected host budget warning"); + }; + assert!(warning.message.contains("2% skills context budget")); let turn_store = ExtensionData::new("turn-1"); let fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { + thread_id: ThreadId::default(), turn_id: "turn-1".to_string(), user_input: vec![UserInput::Text { text: "$lint-fix please".to_string(), @@ -312,6 +271,7 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in let next_fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { + thread_id: ThreadId::default(), turn_id: "turn-2".to_string(), user_input: vec![UserInput::Text { text: "no skill this time".to_string(), @@ -376,7 +336,7 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult { )) .await; assert!(initial_fragments.is_empty()); - let EventMsg::Warning(warning) = event_rx.try_recv()?.msg else { + let EventMsg::Warning(warning) = event_rx.try_recv()?.1.msg else { panic!("expected warning event"); }; assert_eq!( @@ -388,6 +348,7 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult { let fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { + thread_id: ThreadId::default(), turn_id: turn_id.to_string(), user_input: vec![UserInput::Text { text: "$first".to_string(), @@ -467,6 +428,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te let fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { + thread_id: ThreadId::default(), turn_id: "turn-1".to_string(), user_input: vec![UserInput::Mention { name: "lint-fix".to_string(), @@ -553,6 +515,7 @@ async fn prompt_hidden_skill_can_still_be_invoked() -> TestResult { let fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { + thread_id: ThreadId::default(), turn_id: "turn-1".to_string(), user_input: vec![UserInput::Text { text: "$hidden-skill".to_string(), @@ -588,11 +551,11 @@ struct StaticSkillProvider { fail_first_list: bool, } -struct ChannelEventSink(std::sync::mpsc::Sender); +struct ChannelEventSink(std::sync::mpsc::Sender<(ThreadId, Event)>); impl ExtensionEventSink for ChannelEventSink { - fn emit(&self, event: Event) { - let _ = self.0.send(event); + fn emit(&self, thread_id: ThreadId, event: Event) { + let _ = self.0.send((thread_id, event)); } } @@ -700,6 +663,7 @@ fn prompt_context<'a>( turn_store: &'a ExtensionData, ) -> ContextContributionContext<'a> { ContextContributionContext { + thread_id: ThreadId::default(), session_store, thread_store, turn_store, diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index d605b640ca..1044a3dd93 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -4,6 +4,7 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex; use crate::exec_approval::handle_exec_approval_request; use crate::outgoing_message::OutgoingMessageSender; @@ -27,7 +28,6 @@ use rmcp::model::CallToolResult; use rmcp::model::Content; use rmcp::model::RequestId; use serde_json::json; -use tokio::sync::Mutex; /// To adhere to MCP `tools/call` response format, include the Codex /// `threadId` in the `structured_content` field of the response. @@ -98,7 +98,7 @@ pub async fn run_codex_tool_session( let sub_id = id.to_string(); running_requests_id_to_codex_uuid .lock() - .await + .unwrap_or_else(std::sync::PoisonError::into_inner) .insert(id.clone(), thread_id); let submission = Submission { id: sub_id.clone(), @@ -126,7 +126,10 @@ pub async fn run_codex_tool_session( ); outgoing.send_response(id.clone(), result).await; // unregister the id so we don't keep it in the map - running_requests_id_to_codex_uuid.lock().await.remove(&id); + running_requests_id_to_codex_uuid + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&id); return; } @@ -150,10 +153,11 @@ pub async fn run_codex_tool_session_reply( ) { running_requests_id_to_codex_uuid .lock() - .await + .unwrap_or_else(std::sync::PoisonError::into_inner) .insert(request_id.clone(), thread_id); - if let Err(e) = thread - .submit(Op::UserInput { + let submission = Submission { + id: request_id.to_string(), + op: Op::UserInput { items: vec![UserInput::Text { text: prompt, // MCP tool prompts are plain text with no UI element ranges. @@ -163,9 +167,11 @@ pub async fn run_codex_tool_session_reply( responsesapi_client_metadata: None, additional_context: Default::default(), thread_settings: Default::default(), - }) - .await - { + }, + client_user_message_id: None, + trace: None, + }; + if let Err(e) = thread.submit_with_id(submission).await { tracing::error!("Failed to submit user input: {e}"); let result = create_call_tool_result_with_thread_id( thread_id, @@ -176,7 +182,7 @@ pub async fn run_codex_tool_session_reply( // unregister the id so we don't keep it in the map running_requests_id_to_codex_uuid .lock() - .await + .unwrap_or_else(std::sync::PoisonError::into_inner) .remove(&request_id); return; } @@ -310,11 +316,6 @@ async fn run_codex_tool_session_inner( thread_id, text, /*is_error*/ None, ); outgoing.send_response(request_id.clone(), result).await; - // unregister the id so we don't keep it in the map - running_requests_id_to_codex_uuid - .lock() - .await - .remove(&request_id); break; } EventMsg::SessionConfigured(_) => { @@ -409,6 +410,10 @@ async fn run_codex_tool_session_inner( } } } + running_requests_id_to_codex_uuid + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&request_id); } #[cfg(test)] diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 073fe1bf73..2099bd6aad 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex; use codex_arg0::Arg0DispatchPaths; use codex_core::StateDbHandle; @@ -32,7 +33,6 @@ use rmcp::model::JsonRpcResponse; use rmcp::model::RequestId; use rmcp::model::ServerCapabilities; use serde_json::json; -use tokio::sync::Mutex; use tokio::task; use crate::codex_tool_config::CodexToolCallParam; @@ -527,7 +527,10 @@ impl MessageProcessor { // Obtain the thread id while holding the first lock, then release. let thread_id = { - let map_guard = self.running_requests_id_to_codex_uuid.lock().await; + let map_guard = self + .running_requests_id_to_codex_uuid + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); match map_guard.get(&request_id) { Some(id) => *id, None => { @@ -563,7 +566,7 @@ impl MessageProcessor { // unregister the id so we don't keep it in the map self.running_requests_id_to_codex_uuid .lock() - .await + .unwrap_or_else(std::sync::PoisonError::into_inner) .remove(&request_id); } @@ -604,22 +607,18 @@ struct McpServerExtensionEventSink { } impl ExtensionEventSink for McpServerExtensionEventSink { - fn emit(&self, event: Event) { - let Ok(thread_id) = ThreadId::from_string(&event.id) else { - tracing::warn!(event_id = %event.id, "dropping extension event with invalid thread id"); - return; - }; + fn emit(&self, thread_id: ThreadId, event: Event) { + let request_id = self + .running_requests + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .find_map(|(request_id, running_thread_id)| { + (*running_thread_id == thread_id && request_id.to_string() == event.id) + .then(|| request_id.clone()) + }); let outgoing = Arc::clone(&self.outgoing); - let running_requests = Arc::clone(&self.running_requests); tokio::spawn(async move { - let request_id = - running_requests - .lock() - .await - .iter() - .find_map(|(request_id, running_thread_id)| { - (*running_thread_id == thread_id).then(|| request_id.clone()) - }); outgoing .send_event_as_notification( &event, diff --git a/codex-rs/mcp-server/src/message_processor_tests.rs b/codex-rs/mcp-server/src/message_processor_tests.rs index 8ceaf013a2..f7fee88160 100644 --- a/codex-rs/mcp-server/src/message_processor_tests.rs +++ b/codex-rs/mcp-server/src/message_processor_tests.rs @@ -1,3 +1,4 @@ +use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::WarningEvent; use pretty_assertions::assert_eq; @@ -14,18 +15,25 @@ async fn extension_event_sink_forwards_warning_with_request_and_thread_metadata( let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx)); let thread_id = ThreadId::new(); let request_id = RequestId::String("request-1".into()); - let running_requests = Arc::new(Mutex::new(HashMap::from([(request_id.clone(), thread_id)]))); + let stale_request_id = RequestId::String("stale-request".into()); + let running_requests = Arc::new(Mutex::new(HashMap::from([ + (stale_request_id, thread_id), + (request_id.clone(), thread_id), + ]))); let sink = McpServerExtensionEventSink { outgoing, running_requests, }; - sink.emit(Event { - id: thread_id.to_string(), - msg: EventMsg::Warning(WarningEvent { - message: "skill warning".to_string(), - }), - }); + sink.emit( + thread_id, + Event { + id: request_id.to_string(), + msg: EventMsg::Warning(WarningEvent { + message: "skill warning".to_string(), + }), + }, + ); let OutgoingMessage::Notification(OutgoingNotification { method, params }) = outgoing_rx.recv().await.expect("notification") @@ -40,7 +48,7 @@ async fn extension_event_sink_forwards_warning_with_request_and_thread_metadata( "requestId": request_id, "threadId": thread_id, }, - "id": thread_id.to_string(), + "id": request_id.to_string(), "msg": { "type": "warning", "message": "skill warning",