openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix skills catalog budgets and warning delivery

Browse Source
This commit is contained in:
jif-oai
2026-06-18 10:37:36 +02:00
parent 3fea72fda6
commit b3d7b2d45a
23 changed files with 309 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -3889,6 +3889,7 @@ dependencies = [
"codex-exec-server",
"codex-extension-api",
"codex-mcp",
"codex-otel",
"codex-protocol",
"codex-tools",
"codex-utils-absolute-path",
codex-rs/app-server/src/extensions.rs
+30 -9
View File
@@ -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
codex-rs/app-server/src/request_processors.rs
+1
View File
@@ -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;
codex-rs/app-server/src/request_processors/thread_lifecycle.rs
+8
View File
@@ -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,
codex-rs/app-server/src/thread_state.rs
+4
View File
@@ -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 {
codex-rs/core/src/session/mod.rs
+16 -4
View File
@@ -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());
}
codex-rs/core/src/session/turn.rs
+1
View File
@@ -652,6 +652,7 @@ async fn build_extension_turn_input_items(
.collect::<Vec<_>>();
let input = TurnInputContext {
thread_id: sess.thread_id(),
turn_id: turn_context.sub_id.to_string(),
user_input: user_input.to_vec(),
environments,
codex-rs/core/src/session/turn_context.rs
+1
View File
@@ -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,
codex-rs/ext/extension-api/examples/enabled_extensions.rs
+1
View File
@@ -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,
codex-rs/ext/extension-api/src/capabilities/events.rs
+6 -4
View File
@@ -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) {}
}
codex-rs/ext/extension-api/src/contributors.rs
+2
View File
@@ -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<Box<dyn Future<Output = T> + 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,
codex-rs/ext/extension-api/src/contributors/turn_input.rs
+3
View File
@@ -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.
codex-rs/ext/extension-api/tests/registry.rs
+10 -7
View File
@@ -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
codex-rs/ext/goal/src/events.rs
+12 -8
View File
@@ -22,13 +22,17 @@ impl GoalEventEmitter {
turn_id: Option<String>,
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,
}),
},
);
}
}
codex-rs/ext/goal/tests/goal_extension_backend.rs
+1 -1
View File
@@ -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);
}
}
codex-rs/ext/memories/src/tests.rs
+1
View File
@@ -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,
codex-rs/ext/skills/Cargo.toml
+1
View File
@@ -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 }
codex-rs/ext/skills/src/extension.rs
+20 -10
View File
@@ -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::<SessionTelemetry>();
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<C> SkillsExtension<C> {
.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 }),
},
);
}
}
codex-rs/ext/skills/src/render.rs
+72 -33
View File
@@ -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<i64>,
session_telemetry: Option<&SessionTelemetry>,
) -> Option<(AvailableSkillsInstructions, Option<String>)> {
let visible_entries = catalog
.entries
.iter()
.filter(|entry| entry.enabled && entry.prompt_visible)
.collect::<Vec<_>>();
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<String>,
omitted: usize,
}
impl BoundedSkillLines {
fn into_lines(mut self) -> Vec<String> {
if let Some(line) = omitted_skills_line(self.omitted) {
self.lines.push(line);
}
self.lines
}
}
fn bounded_skill_lines<'a>(
entries: impl IntoIterator<Item = &'a SkillCatalogEntry>,
) -> Vec<String> {
) -> 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<String> {
(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 {
codex-rs/ext/skills/tests/skills_extension.rs
+66 -102
View File
@@ -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<dyn std::error::Error>>;
@@ -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<Event>);
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,
codex-rs/mcp-server/src/codex_tool_runner.rs
+20 -15
View File
@@ -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)]
codex-rs/mcp-server/src/message_processor.rs
+16 -17
View File
@@ -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,
codex-rs/mcp-server/src/message_processor_tests.rs
+16 -8
View File
@@ -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",