From df5f12285454ce0a6e78a1ce084b57b6860add41 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Thu, 18 Jun 2026 11:39:02 -0700 Subject: [PATCH 1/5] Add Config for Time Reminders (varlatency 1/n) (#28822) ## Summary Example: > [features.current_time_reminder] enabled = true reminder_interval_model_requests = 1 clock_source = "system" ## Testing - `just test -p codex-core varlatency` - `just test -p codex-core lock_contains_prompts_and_materializes_features` - `just fix -p codex-core -p codex-config -p codex-features` --- codex-rs/config/src/schema.rs | 9 ++++ codex-rs/core/config.schema.json | 40 ++++++++++++++ codex-rs/core/src/config/config_tests.rs | 61 ++++++++++++++++++++++ codex-rs/core/src/config/mod.rs | 58 ++++++++++++++++++++ codex-rs/core/src/session/config_lock.rs | 20 +++++++ codex-rs/features/src/feature_configs.rs | 30 +++++++++++ codex-rs/features/src/lib.rs | 22 ++++++++ codex-rs/thread-manager-sample/src/main.rs | 1 + 8 files changed, 241 insertions(+) diff --git a/codex-rs/config/src/schema.rs b/codex-rs/config/src/schema.rs index 6d3f901bf4..fad074e9c3 100644 --- a/codex-rs/config/src/schema.rs +++ b/codex-rs/config/src/schema.rs @@ -53,6 +53,15 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema { ); continue; } + if feature.id == codex_features::Feature::CurrentTimeReminder { + validation.properties.insert( + feature.key.to_string(), + schema_gen.subschema_for::>(), + ); + continue; + } if feature.id == codex_features::Feature::AppsMcpPathOverride { validation.properties.insert( feature.key.to_string(), diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 852f6de39a..7faf569f7d 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -476,6 +476,9 @@ "connectors": { "type": "boolean" }, + "current_time_reminder": { + "$ref": "#/definitions/FeatureToml_for_CurrentTimeReminderConfigToml" + }, "default_mode_request_user_input": { "type": "boolean" }, @@ -799,6 +802,30 @@ }, "type": "object" }, + "CurrentTimeReminderConfigToml": { + "additionalProperties": false, + "properties": { + "clock_source": { + "$ref": "#/definitions/CurrentTimeSource" + }, + "enabled": { + "type": "boolean" + }, + "reminder_interval_model_requests": { + "format": "uint64", + "minimum": 1.0, + "type": "integer" + } + }, + "type": "object" + }, + "CurrentTimeSource": { + "enum": [ + "system", + "external" + ], + "type": "string" + }, "DebugConfigLockToml": { "additionalProperties": false, "properties": { @@ -891,6 +918,16 @@ } ] }, + "FeatureToml_for_CurrentTimeReminderConfigToml": { + "anyOf": [ + { + "type": "boolean" + }, + { + "$ref": "#/definitions/CurrentTimeReminderConfigToml" + } + ] + }, "FeatureToml_for_MultiAgentV2ConfigToml": { "anyOf": [ { @@ -4679,6 +4716,9 @@ "connectors": { "type": "boolean" }, + "current_time_reminder": { + "$ref": "#/definitions/FeatureToml_for_CurrentTimeReminderConfigToml" + }, "default_mode_request_user_input": { "type": "boolean" }, diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 2244e3e1b7..b13c9575f1 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -522,6 +522,67 @@ async fn load_config_rejects_enabled_rollout_budget_without_limit() -> std::io:: Ok(()) } +#[tokio::test] +async fn load_config_resolves_current_time_reminder() -> std::io::Result<()> { + for (config_toml, expected) in [ + ( + r#" +[features] +current_time_reminder = true +"#, + CurrentTimeReminderConfig::default(), + ), + ( + r#" +[features.current_time_reminder] +enabled = true +reminder_interval_model_requests = 4 +clock_source = "external" +"#, + CurrentTimeReminderConfig { + reminder_interval_model_requests: 4, + clock_source: CurrentTimeSource::External, + }, + ), + ] { + let config = load_current_time_reminder_config(config_toml).await?; + assert!(config.features.enabled(Feature::CurrentTimeReminder)); + assert_eq!(config.current_time_reminder, Some(expected)); + } + Ok(()) +} + +#[tokio::test] +async fn load_config_rejects_zero_current_time_reminder_interval() -> std::io::Result<()> { + let error = load_current_time_reminder_config( + r#" +[features.current_time_reminder] +enabled = true +reminder_interval_model_requests = 0 +"#, + ) + .await + .expect_err("zero reminder interval should be rejected"); + + assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); + assert_eq!( + error.to_string(), + "features.current_time_reminder.reminder_interval_model_requests must be positive" + ); + Ok(()) +} + +async fn load_current_time_reminder_config(config_toml: &str) -> std::io::Result { + let codex_home = tempdir()?; + let config_toml = toml::from_str(config_toml).expect("TOML should deserialize"); + Config::load_from_base_config_with_overrides( + config_toml, + ConfigOverrides::default(), + codex_home.abs(), + ) + .await +} + #[test] fn rejects_provider_auth_with_env_key() { let err = toml::from_str::( diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 0427613361..a1854067f5 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -58,6 +58,8 @@ use codex_core_plugins::PluginsConfigInput; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::LOCAL_FS; use codex_features::CodeModeConfigToml; +use codex_features::CurrentTimeReminderConfigToml; +use codex_features::CurrentTimeSource; use codex_features::Feature; use codex_features::FeatureConfigSource; use codex_features::FeatureOverrides; @@ -1022,6 +1024,8 @@ pub struct Config { /// Shared token budget for the root thread and its sub-agents. pub rollout_budget: Option, + /// Current-time reminder configuration, when enabled. + pub current_time_reminder: Option, /// Centralized feature flags; source of truth for feature gating. pub features: ManagedFeatures, @@ -1075,6 +1079,21 @@ pub struct RolloutBudgetConfig { pub prefill_token_weight: f64, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct CurrentTimeReminderConfig { + pub reminder_interval_model_requests: u64, + pub clock_source: CurrentTimeSource, +} + +impl Default for CurrentTimeReminderConfig { + fn default() -> Self { + Self { + reminder_interval_model_requests: 1, + clock_source: CurrentTimeSource::System, + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct MultiAgentV2Config { pub max_concurrent_threads_per_session: usize, @@ -2539,6 +2558,34 @@ fn resolve_rollout_budget_config( })) } +fn resolve_current_time_reminder_config( + config_toml: &ConfigToml, + features: &ManagedFeatures, +) -> std::io::Result> { + if !features.enabled(Feature::CurrentTimeReminder) { + return Ok(None); + } + + let base = current_time_reminder_toml_config(config_toml.features.as_ref()); + let default = CurrentTimeReminderConfig::default(); + let reminder_interval_model_requests = base + .and_then(|config| config.reminder_interval_model_requests) + .unwrap_or(default.reminder_interval_model_requests); + if reminder_interval_model_requests == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "features.current_time_reminder.reminder_interval_model_requests must be positive", + )); + } + + Ok(Some(CurrentTimeReminderConfig { + reminder_interval_model_requests, + clock_source: base + .and_then(|config| config.clock_source) + .unwrap_or(default.clock_source), + })) +} + fn resolve_terminal_resize_reflow_config(config_toml: &ConfigToml) -> TerminalResizeReflowConfig { let Some(tui) = config_toml.tui.as_ref() else { return TerminalResizeReflowConfig::default(); @@ -2578,6 +2625,15 @@ fn multi_agent_v2_toml_config(features: Option<&FeaturesToml>) -> Option<&MultiA } } +fn current_time_reminder_toml_config( + features: Option<&FeaturesToml>, +) -> Option<&CurrentTimeReminderConfigToml> { + match features?.current_time_reminder.as_ref()? { + FeatureToml::Enabled(_) => None, + FeatureToml::Config(config) => Some(config), + } +} + fn network_proxy_toml_config(features: Option<&FeaturesToml>) -> Option<&NetworkProxyConfigToml> { match features?.network_proxy.as_ref()? { FeatureToml::Enabled(_) => None, @@ -3217,6 +3273,7 @@ impl Config { let code_mode = resolve_code_mode_config(&cfg); let multi_agent_v2 = resolve_multi_agent_v2_config(&cfg); let rollout_budget = resolve_rollout_budget_config(&cfg, &features)?; + let current_time_reminder = resolve_current_time_reminder_config(&cfg, &features)?; let terminal_resize_reflow = resolve_terminal_resize_reflow_config(&cfg); let agent_roles = @@ -3758,6 +3815,7 @@ impl Config { ghost_snapshot, multi_agent_v2, rollout_budget, + current_time_reminder, features, suppress_unstable_features_warning: cfg .suppress_unstable_features_warning diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index b5f1689ad2..61c39a5dd3 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -2,6 +2,7 @@ use anyhow::Context; use codex_config::config_toml::ConfigLockfileToml; use codex_config::config_toml::ConfigToml; use codex_config::types::MemoriesToml; +use codex_features::CurrentTimeReminderConfigToml; use codex_features::Feature; use codex_features::FeatureToml; use codex_features::FeaturesToml; @@ -155,6 +156,12 @@ fn save_config_resolved_fields( rollout_budget.enabled = Some(config.features.enabled(Feature::RolloutBudget)); features.rollout_budget = Some(FeatureToml::Config(rollout_budget)); } + if let Some(current_time_reminder) = config.current_time_reminder.as_ref() { + let mut current_time_reminder: CurrentTimeReminderConfigToml = + resolved_config_to_toml(current_time_reminder, "features.current_time_reminder")?; + current_time_reminder.enabled = Some(config.features.enabled(Feature::CurrentTimeReminder)); + features.current_time_reminder = Some(FeatureToml::Config(current_time_reminder)); + } lock_config.memories = Some(resolved_config_to_toml::( &config.memories, "memories", @@ -227,6 +234,11 @@ mod tests { .features .enable(Feature::RolloutBudget) .expect("rollout_budget should be enableable in tests"); + config.current_time_reminder = Some(crate::config::CurrentTimeReminderConfig::default()); + config + .features + .enable(Feature::CurrentTimeReminder) + .expect("current_time_reminder should be enableable in tests"); sc.original_config_do_not_use = Arc::new(config); sc.base_instructions = "resolved instructions".to_string(); sc.developer_instructions = Some("resolved developer instructions".to_string()); @@ -302,6 +314,14 @@ mod tests { prefill_token_weight: Some(0.25), })) ); + assert_eq!( + features.current_time_reminder, + Some(FeatureToml::Config(CurrentTimeReminderConfigToml { + enabled: Some(true), + reminder_interval_model_requests: Some(1), + clock_source: Some(codex_features::CurrentTimeSource::System), + })) + ); assert_eq!(lockfile.version, crate::config_lock::CONFIG_LOCK_VERSION); } diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index ad4b6da63c..dc78d8cd43 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -102,6 +102,36 @@ impl FeatureConfig for RolloutBudgetConfigToml { } } +#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum CurrentTimeSource { + #[default] + System, + External, +} + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct CurrentTimeReminderConfigToml { + #[serde(skip_serializing_if = "Option::is_none")] + pub enabled: Option, + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(range(min = 1))] + pub reminder_interval_model_requests: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub clock_source: Option, +} + +impl FeatureConfig for CurrentTimeReminderConfigToml { + fn enabled(&self) -> Option { + self.enabled + } + + fn set_enabled(&mut self, enabled: bool) { + self.enabled = Some(enabled); + } +} + #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] #[serde(deny_unknown_fields)] pub(crate) struct RemovedAppsMcpPathOverrideConfigToml { diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 731ea20a10..58f291407a 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -17,6 +17,8 @@ use toml::Table; mod feature_configs; mod legacy; pub use feature_configs::CodeModeConfigToml; +pub use feature_configs::CurrentTimeReminderConfigToml; +pub use feature_configs::CurrentTimeSource; pub use feature_configs::MultiAgentV2ConfigToml; pub use feature_configs::NetworkProxyConfigToml; pub use feature_configs::NetworkProxyDomainPermissionToml; @@ -206,6 +208,8 @@ pub enum Feature { TokenBudget, /// Track and report a shared token budget across a session's agent threads. RolloutBudget, + /// Add current-time reminders to model-visible context. + CurrentTimeReminder, /// Expose an input-interruptible sleep tool. SleepTool, /// Route MCP tool approval prompts through the MCP elicitation request path. @@ -623,6 +627,8 @@ pub struct FeaturesToml { pub multi_agent_v2: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub rollout_budget: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub current_time_reminder: Option>, #[serde(default, rename = "apps_mcp_path_override", skip_serializing)] #[schemars(skip)] removed_apps_mcp_path_override: Option>, @@ -658,6 +664,13 @@ impl FeaturesToml { if let Some(enabled) = self.rollout_budget.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::RolloutBudget.key().to_string(), enabled); } + if let Some(enabled) = self + .current_time_reminder + .as_ref() + .and_then(FeatureToml::enabled) + { + entries.insert(Feature::CurrentTimeReminder.key().to_string(), enabled); + } if let Some(enabled) = self.network_proxy.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::NetworkProxy.key().to_string(), enabled); } @@ -670,6 +683,7 @@ impl FeaturesToml { code_mode, multi_agent_v2, rollout_budget, + current_time_reminder, removed_apps_mcp_path_override: _, network_proxy, entries, @@ -685,6 +699,8 @@ impl FeaturesToml { materialize_resolved_feature_enabled(multi_agent_v2, enabled); } else if spec.id == Feature::RolloutBudget { materialize_resolved_feature_enabled(rollout_budget, enabled); + } else if spec.id == Feature::CurrentTimeReminder { + materialize_resolved_feature_enabled(current_time_reminder, enabled); } else if spec.id == Feature::NetworkProxy { materialize_resolved_feature_enabled(network_proxy, enabled); } else { @@ -1184,6 +1200,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::CurrentTimeReminder, + key: "current_time_reminder", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::SleepTool, key: "sleep_tool", diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index ca11ff9072..1534ba86a0 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -278,6 +278,7 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R ghost_snapshot: GhostSnapshotConfig::default(), multi_agent_v2: MultiAgentV2Config::default(), rollout_budget: None, + current_time_reminder: None, features: Default::default(), suppress_unstable_features_warning: false, active_project: ProjectConfig { trust_level: None }, From 32a696dbacaa1383745455ea2a77d5477891ed0b Mon Sep 17 00:00:00 2001 From: rka-oai Date: Thu, 18 Jun 2026 11:52:19 -0700 Subject: [PATCH 2/5] [codex] rollout budget implementation (varlength 2/N) (#28494) ## Stack Depends on #28746. This PR implements shared rollout-budget accounting and model-visible reminders using the configuration defined in #28746. # Description / Main changes to Core: `AgentControl` will now be the area where "rollout level" features & accounting will have to live. It is incorrectly named for this responsibility, but I think it can hold all the necessary shared state & features (rollout token budget, mutliple thread interruption responsibilitym etc) In this PR, we have one "token ledger" that each thread will subtract from when sampling. The "charge" will occur when response.completed() is done and the calculation will be done on the responses api usage carrier. The calculation will weigh sampling and pre-fill tokens as specified. Every time the budget crosses the configured reminder threshold, a developer message is appended before the thread's next request This remaining budget will _always_ be restated/reminded after a compaction event. Expiration and fan-out interruption will be in the stacked follow-up (and also live in Agent Control). ## Reminders "You have weighted {session_tokens_left} tokens left in the shared session token budget." The first request in each thread context receives the current remainder. Later reminders are emitted after aggregate weighted usage crosses a configured interval. If several intervals are crossed before a thread sends another request, Core inserts one reminder with the latest remainder. Compaction response usage is charged before the next context starts. The next reminder is appended after the compaction summary, leaving the initial context content stable. ## Tests Integration coverage verifies: - weighted output and non-cached input accounting - initial and periodic reminders - shared accounting between a root and sub-agent - post-compaction remainder and message placement Local checks: - `just fmt` - `just test -p codex-core rollout_budget` - `git diff --check` The full workspace test suite was not run locally. --- codex-rs/core/src/agent/control.rs | 19 +- codex-rs/core/src/compact_remote_v2.rs | 1 + codex-rs/core/src/context/mod.rs | 2 + codex-rs/core/src/context/rollout_budget.rs | 27 ++ codex-rs/core/src/event_mapping.rs | 1 + codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/rollout_budget.rs | 108 +++++++ codex-rs/core/src/session/handlers.rs | 4 + codex-rs/core/src/session/mod.rs | 2 + codex-rs/core/src/session/rollout_budget.rs | 30 ++ codex-rs/core/src/session/turn.rs | 9 +- codex-rs/core/src/thread_manager.rs | 21 +- codex-rs/core/tests/suite/mod.rs | 1 + codex-rs/core/tests/suite/rollout_budget.rs | 304 ++++++++++++++++++++ 14 files changed, 521 insertions(+), 9 deletions(-) create mode 100644 codex-rs/core/src/context/rollout_budget.rs create mode 100644 codex-rs/core/src/rollout_budget.rs create mode 100644 codex-rs/core/src/session/rollout_budget.rs create mode 100644 codex-rs/core/tests/suite/rollout_budget.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 01f0e1415d..410c97cbd5 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -6,7 +6,9 @@ use crate::agent::role::resolve_role_config; use crate::agent::status::is_final; use crate::codex_thread::ThreadConfigSnapshot; use crate::config::Config; +use crate::config::RolloutBudgetConfig; use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::rollout_budget::RolloutBudget; use crate::session::emit_subagent_session_started; use crate::session_prefix::format_inter_agent_completion_message; use crate::session_prefix::format_subagent_context_line; @@ -100,15 +102,24 @@ pub(crate) struct AgentControl { state: Arc, v2_residency: Arc, agent_execution_limiter: Arc, + /// Session-scoped state shared by the root thread and every cloned sub-agent control handle. + rollout_budget: Arc, } impl AgentControl { /// Construct a new `AgentControl` that can spawn/message agents via the given manager state. - pub(crate) fn new(manager: Weak) -> Self { - Self { + pub(crate) fn new( + manager: Weak, + rollout_budget: Option, + ) -> Self { + let control = Self { manager, ..Default::default() + }; + if let Some(rollout_budget) = rollout_budget { + control.rollout_budget.configure(rollout_budget); } + control } pub(crate) fn with_session_id(mut self, session_id: SessionId, max_threads: usize) -> Self { @@ -121,6 +132,10 @@ impl AgentControl { self.session_id } + pub(crate) fn rollout_budget(&self) -> &RolloutBudget { + self.rollout_budget.as_ref() + } + /// Send rich user input items to an existing agent thread. pub(crate) async fn send_input( &self, diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 6c3a939024..f9c2ddc710 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -284,6 +284,7 @@ async fn run_remote_compact_task_inner_impl( token_usage, } = compaction_output_result?; if let Some(token_usage) = token_usage { + sess.record_rollout_budget_usage(&token_usage); analytics_details.active_context_tokens_before = Some(token_usage.input_tokens); analytics_details.compaction_summary_tokens = Some(token_usage.output_tokens); analytics_details.cached_input_tokens = Some(token_usage.cached_input_tokens); diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index a1f204e475..7fc1375781 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -24,6 +24,7 @@ mod realtime_end_instructions; mod realtime_start_instructions; mod realtime_start_with_instructions; mod recommended_plugins_instructions; +mod rollout_budget; mod subagent_notification; mod token_budget_context; mod turn_aborted; @@ -64,6 +65,7 @@ pub(crate) use realtime_end_instructions::RealtimeEndInstructions; pub(crate) use realtime_start_instructions::RealtimeStartInstructions; pub(crate) use realtime_start_with_instructions::RealtimeStartWithInstructions; pub(crate) use recommended_plugins_instructions::RecommendedPluginsInstructions; +pub(crate) use rollout_budget::RolloutBudgetContext; pub(crate) use subagent_notification::SubagentNotification; pub(crate) use token_budget_context::TokenBudgetContext; pub(crate) use token_budget_context::TokenBudgetRemainingContext; diff --git a/codex-rs/core/src/context/rollout_budget.rs b/codex-rs/core/src/context/rollout_budget.rs new file mode 100644 index 0000000000..33ed724b8e --- /dev/null +++ b/codex-rs/core/src/context/rollout_budget.rs @@ -0,0 +1,27 @@ +use super::ContextualUserFragment; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RolloutBudgetContext { + pub(crate) remaining_tokens: i64, +} + +impl ContextualUserFragment for RolloutBudgetContext { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("\n", "\n") + } + + fn body(&self) -> String { + format!( + "You have {} weighted tokens left in the shared session token budget.", + self.remaining_tokens + ) + } +} diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index 9e96ebff52..73193984b3 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -33,6 +33,7 @@ const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[ SKILLS_INSTRUCTIONS_OPEN_TAG, "", "", + "", ]; pub(crate) fn is_contextual_user_message_content(message: &[ContentItem]) -> bool { diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 0641d7226e..0e193baa1d 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -131,6 +131,7 @@ pub use agents_md::DEFAULT_AGENTS_MD_FILENAME; pub use agents_md::LOCAL_AGENTS_MD_FILENAME; pub use agents_md::LoadedAgentsMd; mod rollout; +mod rollout_budget; pub(crate) mod safety; mod session_rollout_init_error; pub mod shell; diff --git a/codex-rs/core/src/rollout_budget.rs b/codex-rs/core/src/rollout_budget.rs new file mode 100644 index 0000000000..accabac761 --- /dev/null +++ b/codex-rs/core/src/rollout_budget.rs @@ -0,0 +1,108 @@ +use crate::config::RolloutBudgetConfig; +use codex_protocol::ThreadId; +use codex_protocol::protocol::TokenUsage; +use std::collections::HashMap; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::OnceLock; + +pub(crate) struct RolloutBudgetReminder { + pub(crate) remaining_tokens: i64, + reminder_index: i64, +} + +/// Shared accounting and reminder state for one root-thread session tree. +#[derive(Default)] +pub(crate) struct RolloutBudget { + state: OnceLock>, +} + +struct RolloutBudgetState { + config: RolloutBudgetConfig, + weighted_tokens_used: f64, + /// Last reminder delivered to each thread, so every thread observes crossed thresholds. + deliveries: HashMap, +} + +struct ThreadBudgetDelivery { + window_id: String, + reminder_index: i64, +} + +impl RolloutBudget { + pub(crate) fn configure(&self, config: RolloutBudgetConfig) { + self.state.get_or_init(|| { + Mutex::new(RolloutBudgetState { + config, + weighted_tokens_used: 0.0, + deliveries: HashMap::new(), + }) + }); + } + + pub(crate) fn record_usage(&self, usage: &TokenUsage) { + let Some(mut state) = self.lock() else { + return; + }; + state.weighted_tokens_used += usage.output_tokens.max(0) as f64 + * state.config.sampling_token_weight + + usage.non_cached_input() as f64 * state.config.prefill_token_weight; + } + + pub(crate) fn pending_reminder( + &self, + thread_id: ThreadId, + window_id: &str, + ) -> Option { + let state = self.lock()?; + let reminder_index = (state.weighted_tokens_used + / state.config.reminder_interval_tokens as f64) + .floor() as i64; + if state.deliveries.get(&thread_id).is_some_and(|delivery| { + delivery.window_id.as_str() == window_id && delivery.reminder_index >= reminder_index + }) { + return None; + } + Some(RolloutBudgetReminder { + remaining_tokens: (state.config.limit_tokens as f64 - state.weighted_tokens_used) + .max(0.0) + .floor() as i64, + reminder_index, + }) + } + + pub(crate) fn mark_reminder_delivered( + &self, + thread_id: ThreadId, + window_id: &str, + reminder: RolloutBudgetReminder, + ) { + // Mark delivery only after history insertion; cancellation before then should retry it. + let Some(mut state) = self.lock() else { + return; + }; + state.deliveries.insert( + thread_id, + ThreadBudgetDelivery { + window_id: window_id.to_string(), + reminder_index: reminder.reminder_index, + }, + ); + } + + /// Forces the next sampling request for `thread_id` to restate the current remainder. + pub(crate) fn rearm_reminder(&self, thread_id: ThreadId) { + let Some(mut state) = self.lock() else { + return; + }; + state.deliveries.remove(&thread_id); + } + + fn lock(&self) -> Option> { + self.state.get().map(|state| { + state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + }) + } +} diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 4ed79d04e3..5299747f3d 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -525,6 +525,10 @@ pub async fn thread_rollback(sess: &Arc, sub_id: String, num_turns: u32 .collect::>(); sess.apply_rollout_reconstruction(turn_context.as_ref(), replay_items.as_slice()) .await; + sess.services + .agent_control + .rollout_budget() + .rearm_reminder(sess.thread_id()); sess.recompute_token_usage(turn_context.as_ref()).await; sess.persist_rollout_items(&[RolloutItem::EventMsg(rollback_msg.clone())]) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 84e60a3d2d..5aa56b68f8 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -211,6 +211,7 @@ mod input_queue; mod mcp; mod multi_agents; mod review; +mod rollout_budget; mod rollout_reconstruction; #[allow(clippy::module_inception)] pub(crate) mod session; @@ -3360,6 +3361,7 @@ impl Session { } state.token_info() }; + self.record_rollout_budget_usage(token_usage); if let Some(token_info) = token_info.as_ref() { for contributor in self.services.extensions.token_usage_contributors() { contributor diff --git a/codex-rs/core/src/session/rollout_budget.rs b/codex-rs/core/src/session/rollout_budget.rs new file mode 100644 index 0000000000..4627d95ee7 --- /dev/null +++ b/codex-rs/core/src/session/rollout_budget.rs @@ -0,0 +1,30 @@ +use super::session::Session; +use super::turn_context::TurnContext; +use crate::context::ContextualUserFragment; +use codex_protocol::protocol::TokenUsage; + +pub(super) async fn maybe_record_reminder( + sess: &Session, + turn_context: &TurnContext, + window_id: &str, +) { + let budget = sess.services.agent_control.rollout_budget(); + let Some(reminder) = budget.pending_reminder(sess.thread_id(), window_id) else { + return; + }; + let response_item = ContextualUserFragment::into(crate::context::RolloutBudgetContext { + remaining_tokens: reminder.remaining_tokens, + }); + sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) + .await; + budget.mark_reminder_delivered(sess.thread_id(), window_id, reminder); +} + +impl Session { + pub(crate) fn record_rollout_budget_usage(&self, usage: &TokenUsage) { + self.services + .agent_control + .rollout_budget() + .record_usage(usage); + } +} diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 257e741203..657e00ff35 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -218,6 +218,14 @@ pub(crate) async fn run_turn( break; } + let window_id = sess.current_window_id().await; + super::rollout_budget::maybe_record_reminder( + sess.as_ref(), + turn_context.as_ref(), + &window_id, + ) + .await; + // Construct the input that we will send to the model. let sampling_request_input: Vec = async { sess.clone_history() @@ -227,7 +235,6 @@ pub(crate) async fn run_turn( .instrument(trace_span!("run_turn.prepare_sampling_request_input")) .await; - let window_id = sess.current_window_id().await; let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( sess.installation_id.clone(), window_id, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 63335c0129..b452dc45ab 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -621,6 +621,7 @@ impl ThreadManager { options: StartThreadOptions, forked_from_thread_id: Option, ) -> CodexResult { + let agent_control = self.agent_control_for_config(&options.config); let (resumed_session_source, resumed_thread_source) = options .initial_history .get_resumed_session_sources() @@ -631,7 +632,7 @@ impl ThreadManager { options.config, options.initial_history, Arc::clone(&self.state.auth_manager), - self.agent_control(), + agent_control, session_source, /*parent_thread_id*/ None, forked_from_thread_id, @@ -710,6 +711,7 @@ impl ThreadManager { auth_manager: Arc, parent_trace: Option, ) -> CodexResult { + let agent_control = self.agent_control_for_config(&config); let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), &config.cwd, @@ -721,7 +723,7 @@ impl ThreadManager { config, initial_history, auth_manager, - self.agent_control(), + agent_control, session_source, /*parent_thread_id*/ None, /*forked_from_thread_id*/ None, @@ -743,6 +745,7 @@ impl ThreadManager { config: Config, user_shell_override: crate::shell::Shell, ) -> CodexResult { + let agent_control = self.agent_control_for_config(&config); let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), &config.cwd, @@ -751,7 +754,7 @@ impl ThreadManager { config, InitialHistory::New, Arc::clone(&self.state.auth_manager), - self.agent_control(), + agent_control, /*parent_thread_id*/ None, /*forked_from_thread_id*/ None, /*thread_source*/ None, @@ -772,6 +775,7 @@ impl ThreadManager { auth_manager: Arc, user_shell_override: crate::shell::Shell, ) -> CodexResult { + let agent_control = self.agent_control_for_config(&config); let initial_history = self.initial_history_from_rollout_path(rollout_path).await?; let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), @@ -784,7 +788,7 @@ impl ThreadManager { config, initial_history, auth_manager, - self.agent_control(), + agent_control, session_source, /*parent_thread_id*/ None, /*forked_from_thread_id*/ None, @@ -952,11 +956,12 @@ impl ThreadManager { self.state.environment_manager.as_ref(), &config.cwd, ); + let agent_control = self.agent_control_for_config(&config); Box::pin(self.state.spawn_thread( config, history, Arc::clone(&self.state.auth_manager), - self.agent_control(), + agent_control, /*parent_thread_id*/ None, forked_from_thread_id, thread_source, @@ -971,7 +976,11 @@ impl ThreadManager { } pub(crate) fn agent_control(&self) -> AgentControl { - AgentControl::new(Arc::downgrade(&self.state)) + AgentControl::new(Arc::downgrade(&self.state), /*rollout_budget*/ None) + } + + fn agent_control_for_config(&self, config: &Config) -> AgentControl { + AgentControl::new(Arc::downgrade(&self.state), config.rollout_budget) } #[cfg(test)] diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 275458ffdc..0fdd958b39 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -100,6 +100,7 @@ mod resume; mod resume_warning; mod review; mod rmcp_client; +mod rollout_budget; mod rollout_list_find; mod safety_check_downgrade; mod search_tool; diff --git a/codex-rs/core/tests/suite/rollout_budget.rs b/codex-rs/core/tests/suite/rollout_budget.rs new file mode 100644 index 0000000000..c2a044eb52 --- /dev/null +++ b/codex-rs/core/tests/suite/rollout_budget.rs @@ -0,0 +1,304 @@ +use anyhow::Result; +use codex_core::config::RolloutBudgetConfig; +use codex_features::Feature; +use codex_model_provider_info::built_in_model_providers; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use core_test_support::responses::ResponsesRequest; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_completed_with_tokens; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once_match; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::time::Duration; +use tokio::time::timeout; + +const ROLLOUT_BUDGET: RolloutBudgetConfig = RolloutBudgetConfig { + limit_tokens: 100, + reminder_interval_tokens: 25, + sampling_token_weight: 1.0, + prefill_token_weight: 1.0, +}; + +fn rollout_budget_texts(request: &ResponsesRequest) -> Vec { + request + .message_input_texts("developer") + .into_iter() + .filter(|text| text.starts_with("")) + .collect() +} + +fn rollout_budget_message(remaining_tokens: i64) -> String { + format!( + "\nYou have {remaining_tokens} weighted tokens left in the shared session token budget.\n" + ) +} + +fn wire_request_contains(request: &wiremock::Request, text: &str) -> bool { + std::str::from_utf8(&request.body).is_ok_and(|body| body.contains(text)) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn adds_weighted_initial_and_periodic_reminders() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + json!({ + "type": "response.completed", + "response": { + "id": "resp-1", + "usage": { + "input_tokens": 60, + "input_tokens_details": { "cached_tokens": 40 }, + "output_tokens": 15, + "output_tokens_details": null, + "total_tokens": 75 + } + } + }), + ]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.rollout_budget = Some(RolloutBudgetConfig { + sampling_token_weight: 2.0, + prefill_token_weight: 0.5, + ..ROLLOUT_BUDGET + }); + }) + .build(&server) + .await?; + + test.submit_turn("first turn").await?; + test.submit_turn("second turn").await?; + + let requests = responses.requests(); + assert_eq!( + rollout_budget_texts(&requests[0]), + vec![rollout_budget_message(/*remaining_tokens*/ 100)] + ); + assert_eq!( + rollout_budget_texts(&requests[1]), + vec![ + rollout_budget_message(/*remaining_tokens*/ 100), + rollout_budget_message(/*remaining_tokens*/ 60), + ] + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { + skip_if_no_network!(Ok(())); + + const ROOT_PROMPT: &str = "spawn a budget worker"; + const CHILD_PROMPT: &str = "consume child budget"; + const FOLLOW_UP_PROMPT: &str = "report the shared budget"; + const SPAWN_CALL_ID: &str = "spawn-budget-worker"; + + let server = start_mock_server().await; + let spawn_args = json!({ + "message": CHILD_PROMPT, + "task_name": "budget_worker", + }) + .to_string(); + mount_sse_once_match( + &server, + |request: &wiremock::Request| wire_request_contains(request, ROOT_PROMPT), + sse(vec![ + ev_response_created("root-1"), + ev_function_call(SPAWN_CALL_ID, "spawn_agent", &spawn_args), + ev_completed_with_tokens("root-1", /*total_tokens*/ 10), + ]), + ) + .await; + mount_sse_once_match( + &server, + |request: &wiremock::Request| wire_request_contains(request, "\"type\":\"agent_message\""), + sse(vec![ + ev_response_created("child-1"), + ev_completed_with_tokens("child-1", /*total_tokens*/ 30), + ]), + ) + .await; + mount_sse_once_match( + &server, + |request: &wiremock::Request| { + wire_request_contains(request, SPAWN_CALL_ID) + && !wire_request_contains(request, "\"type\":\"agent_message\"") + }, + sse(vec![ + ev_response_created("root-2"), + ev_completed_with_tokens("root-2", /*total_tokens*/ 10), + ]), + ) + .await; + let follow_up = mount_sse_once_match( + &server, + |request: &wiremock::Request| wire_request_contains(request, FOLLOW_UP_PROMPT), + sse(vec![ev_response_created("root-3"), ev_completed("root-3")]), + ) + .await; + + let test = test_codex() + .with_config(|config| { + config + .features + .enable(Feature::Collab) + .expect("test config should allow multi-agent tools"); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow multi-agent v2"); + config.rollout_budget = Some(ROLLOUT_BUDGET); + }) + .build(&server) + .await?; + + let mut created_threads = test.thread_manager.subscribe_thread_created(); + test.submit_turn(ROOT_PROMPT).await?; + let child_thread_id = timeout(Duration::from_secs(10), created_threads.recv()).await??; + let child_thread = test.thread_manager.get_thread(child_thread_id).await?; + wait_for_event(child_thread.as_ref(), |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_turn(FOLLOW_UP_PROMPT).await?; + + let request = follow_up.single_request(); + assert_eq!( + rollout_budget_texts(&request).last(), + Some(&rollout_budget_message(/*remaining_tokens*/ 50)) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn restates_the_current_remainder_after_compaction() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 20), + ]), + sse(vec![ + ev_response_created("resp-compact"), + ev_assistant_message("msg-compact", "compact summary"), + ev_completed_with_tokens("resp-compact", /*total_tokens*/ 10), + ]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone(); + model_provider.name = "OpenAI-compatible test provider".to_string(); + model_provider.base_url = Some(format!("{}/v1", server.uri())); + model_provider.supports_websockets = false; + let test = test_codex() + .with_config(move |config| { + config.model_provider = model_provider; + config.rollout_budget = Some(RolloutBudgetConfig { + reminder_interval_tokens: 50, + ..ROLLOUT_BUDGET + }); + }) + .build(&server) + .await?; + + test.submit_turn("first turn").await?; + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_turn("second turn").await?; + + let requests = responses.requests(); + assert_eq!( + rollout_budget_texts(&requests[2]), + vec![rollout_budget_message(/*remaining_tokens*/ 70)], + "a new context window should restate the current remainder" + ); + let request_body = requests[2].body_json().to_string(); + let summary_position = request_body + .find("compact summary") + .expect("post-compaction request should contain the summary"); + let reminder_position = request_body + .find("You have 70 weighted tokens left in the shared session token budget.") + .expect("post-compaction request should contain the current remainder"); + assert!( + summary_position < reminder_position, + "the current remainder should follow the compaction summary" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn restates_the_current_remainder_after_rollback() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 30), + ]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.rollout_budget = Some(RolloutBudgetConfig { + reminder_interval_tokens: 50, + ..ROLLOUT_BUDGET + }); + }) + .build(&server) + .await?; + + test.submit_turn("rolled-back turn").await?; + test.codex + .submit(Op::ThreadRollback { num_turns: 1 }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::ThreadRolledBack(_)) + }) + .await; + test.submit_turn("turn after rollback").await?; + + let requests = responses.requests(); + assert_eq!( + rollout_budget_texts(&requests[1]), + vec![rollout_budget_message(/*remaining_tokens*/ 70)], + "rollback should rearm the current budget reminder without refunding usage" + ); + + Ok(()) +} From 21a599fa56472a7cea8132c5a47a4374d4d5aa17 Mon Sep 17 00:00:00 2001 From: Gabriel Peal Date: Thu, 18 Jun 2026 11:54:49 -0700 Subject: [PATCH 3/5] Support `openai/form` extended form elicitations (#27500) # Summary Allow App Server clients to opt into `openai/form` MCP elicitations. --- .../analytics/src/analytics_client_tests.rs | 4 + codex-rs/app-server-client/src/lib.rs | 32 +- codex-rs/app-server-client/src/remote.rs | 4 +- .../schema/json/ClientRequest.json | 4 + .../McpServerElicitationRequestParams.json | 21 + .../schema/json/ServerRequest.json | 21 + .../codex_app_server_protocol.schemas.json | 25 + .../codex_app_server_protocol.v2.schemas.json | 4 + .../schema/json/v1/InitializeParams.json | 4 + .../typescript/InitializeCapabilities.ts | 4 + .../v2/McpServerElicitationRequestParams.ts | 2 +- .../src/protocol/common.rs | 8 +- .../app-server-protocol/src/protocol/v1.rs | 3 + .../src/protocol/v2/mcp.rs | 18 + .../src/protocol/v2/tests.rs | 34 + codex-rs/app-server-test-client/src/lib.rs | 1 + codex-rs/app-server/README.md | 12 + codex-rs/app-server/src/message_processor.rs | 17 + .../initialize_processor.rs | 19 +- .../request_processors/thread_processor.rs | 14 + .../src/request_processors/turn_processor.rs | 12 + .../app-server/tests/suite/v2/attestation.rs | 1 + .../suite/v2/connection_handling_websocket.rs | 2 +- .../tests/suite/v2/experimental_api.rs | 8 + .../app-server/tests/suite/v2/initialize.rs | 1 + .../tests/suite/v2/mcp_server_elicitation.rs | 813 +++++++++++++----- .../tests/suite/v2/thread_status.rs | 1 + codex-rs/codex-mcp/src/connection_manager.rs | 2 + .../codex-mcp/src/connection_manager_tests.rs | 41 +- codex-rs/codex-mcp/src/elicitation.rs | 53 +- codex-rs/codex-mcp/src/mcp/mod.rs | 2 + codex-rs/codex-mcp/src/rmcp_client.rs | 58 +- codex-rs/core/src/codex_delegate.rs | 4 + codex-rs/core/src/codex_thread.rs | 7 + codex-rs/core/src/connectors.rs | 1 + codex-rs/core/src/mcp_tool_call_tests.rs | 1 + codex-rs/core/src/session/mcp.rs | 63 +- codex-rs/core/src/session/mcp_tests.rs | 48 +- codex-rs/core/src/session/mod.rs | 3 + codex-rs/core/src/session/session.rs | 7 + codex-rs/core/src/session/tests.rs | 5 + .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/test_support.rs | 9 +- codex-rs/core/src/thread_manager.rs | 33 +- codex-rs/core/src/thread_manager_tests.rs | 14 + .../src/tools/handlers/multi_agents_tests.rs | 1 + codex-rs/core/tests/common/test_codex.rs | 31 +- codex-rs/core/tests/suite/agents_md.rs | 2 + .../core/tests/suite/compact_resume_fork.rs | 1 + codex-rs/core/tests/suite/fork_thread.rs | 1 + .../core/tests/suite/realtime_conversation.rs | 1 + codex-rs/core/tests/suite/resume_warning.rs | 1 + codex-rs/core/tests/suite/rmcp_client.rs | 119 ++- .../tests/suite/subagent_notifications.rs | 1 + .../tests/suite/unstable_features_warning.rs | 2 + codex-rs/exec/src/lib.rs | 1 + codex-rs/memories/write/src/runtime.rs | 1 + codex-rs/protocol/src/approvals.rs | 13 +- .../rmcp-client/src/bin/test_stdio_server.rs | 47 + .../src/elicitation_client_service.rs | 115 ++- .../rmcp-client/src/logging_client_handler.rs | 3 +- codex-rs/rmcp-client/src/rmcp_client.rs | 20 +- codex-rs/tui/src/app/thread_routing.rs | 5 +- codex-rs/tui/src/chatwidget/tool_requests.rs | 3 +- codex-rs/tui/src/lib.rs | 2 + codex-rs/tui/src/onboarding/auth.rs | 1 + 67 files changed, 1486 insertions(+), 328 deletions(-) diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 55467c06b6..fcdd97af8e 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -774,6 +774,7 @@ fn sample_initialize_fact(connection_id: u64) -> AnalyticsFact { experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), }, product_client_id: DEFAULT_ORIGINATOR.to_string(), @@ -1669,6 +1670,7 @@ async fn initialize_caches_client_and_thread_lifecycle_publishes_once_initialize experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), }, product_client_id: DEFAULT_ORIGINATOR.to_string(), @@ -1818,6 +1820,7 @@ async fn compaction_event_ingests_custom_fact() { experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), }, product_client_id: DEFAULT_ORIGINATOR.to_string(), @@ -1946,6 +1949,7 @@ async fn guardian_review_event_ingests_custom_fact_with_optional_target_item() { experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), }, product_client_id: DEFAULT_ORIGINATOR.to_string(), diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index bfcf5522d8..b7a958eff1 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -350,6 +350,8 @@ pub struct InProcessClientStartArgs { pub client_version: String, /// Whether experimental APIs are requested at initialize time. pub experimental_api: bool, + /// Whether MCP servers may send `openai/form` elicitation requests. + pub mcp_server_openai_form_elicitation: bool, /// Notification methods this client opts out of receiving. pub opt_out_notification_methods: Vec, /// Queue capacity for command/event channels (clamped to at least 1). @@ -374,6 +376,7 @@ impl InProcessClientStartArgs { } else { Some(self.opt_out_notification_methods.clone()) }, + mcp_server_openai_form_elicitation: self.mcp_server_openai_form_elicitation, }; InitializeParams { @@ -1044,6 +1047,7 @@ mod tests { client_name: "codex-app-server-client-test".to_string(), client_version: "0.0.0-test".to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity, }) @@ -1237,11 +1241,25 @@ mod tests { client_name: "codex-app-server-client-test".to_string(), client_version: "0.0.0-test".to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: 8, } } + #[test] + fn remote_initialize_params_forward_openai_form_capability() { + let mut args = test_remote_connect_args("ws://localhost/rpc".to_string()); + args.mcp_server_openai_form_elicitation = true; + + assert!( + args.initialize_params() + .capabilities + .expect("initialize capabilities") + .mcp_server_openai_form_elicitation + ); + } + #[tokio::test] async fn typed_request_roundtrip_works() { let client = start_test_client(SessionSource::Exec).await; @@ -1512,6 +1530,7 @@ mod tests { client_name: "codex-app-server-client-test".to_string(), client_version: "0.0.0-test".to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: 8, }) @@ -1600,6 +1619,7 @@ mod tests { client_name: "codex-app-server-client-test".to_string(), client_version: "0.0.0-test".to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: 8, }) @@ -1619,6 +1639,7 @@ mod tests { client_name: "codex-app-server-client-test".to_string(), client_version: "0.0.0-test".to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: 8, }) @@ -2189,7 +2210,7 @@ mod tests { } #[tokio::test] - async fn runtime_start_args_forward_environment_manager() { + async fn runtime_start_args_forward_environment_manager_and_openai_form_capability() { let config = Arc::new(build_test_config().await); let environment_manager = Arc::new( EnvironmentManager::create_for_tests( @@ -2222,12 +2243,20 @@ mod tests { client_name: "codex-app-server-client-test".to_string(), client_version: "0.0.0-test".to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: true, opt_out_notification_methods: Vec::new(), channel_capacity: DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, } .into_runtime_start_args(); assert_eq!(runtime_args.config, config); + assert!( + runtime_args + .initialize + .capabilities + .expect("initialize capabilities") + .mcp_server_openai_form_elicitation + ); assert!(Arc::ptr_eq( &runtime_args.environment_manager, &environment_manager @@ -2263,6 +2292,7 @@ mod tests { client_name: "codex-app-server-client-test".to_string(), client_version: "0.0.0-test".to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, } diff --git a/codex-rs/app-server-client/src/remote.rs b/codex-rs/app-server-client/src/remote.rs index e1c9f16c4a..5575fe4ad4 100644 --- a/codex-rs/app-server-client/src/remote.rs +++ b/codex-rs/app-server-client/src/remote.rs @@ -86,11 +86,12 @@ pub struct RemoteAppServerConnectArgs { pub client_name: String, pub client_version: String, pub experimental_api: bool, + pub mcp_server_openai_form_elicitation: bool, pub opt_out_notification_methods: Vec, pub channel_capacity: usize, } impl RemoteAppServerConnectArgs { - fn initialize_params(&self) -> InitializeParams { + pub(crate) fn initialize_params(&self) -> InitializeParams { let capabilities = InitializeCapabilities { experimental_api: self.experimental_api, request_attestation: false, @@ -99,6 +100,7 @@ impl RemoteAppServerConnectArgs { } else { Some(self.opt_out_notification_methods.clone()) }, + mcp_server_openai_form_elicitation: self.mcp_server_openai_form_elicitation, }; InitializeParams { diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 5ebb861de0..e183c88b08 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -1272,6 +1272,10 @@ "description": "Opt into receiving experimental API methods and fields.", "type": "boolean" }, + "mcpServerOpenaiFormElicitation": { + "description": "Allow downstream MCP servers to request OpenAI extended form elicitations.", + "type": "boolean" + }, "optOutNotificationMethods": { "description": "Exact notification method names that should be suppressed for this connection (for example `thread/started`).", "items": { diff --git a/codex-rs/app-server-protocol/schema/json/McpServerElicitationRequestParams.json b/codex-rs/app-server-protocol/schema/json/McpServerElicitationRequestParams.json index aa7fa817ae..3fc6971330 100644 --- a/codex-rs/app-server-protocol/schema/json/McpServerElicitationRequestParams.json +++ b/codex-rs/app-server-protocol/schema/json/McpServerElicitationRequestParams.json @@ -557,6 +557,27 @@ ], "type": "object" }, + { + "properties": { + "_meta": true, + "message": { + "type": "string" + }, + "mode": { + "enum": [ + "openai/form" + ], + "type": "string" + }, + "requestedSchema": true + }, + "required": [ + "message", + "mode", + "requestedSchema" + ], + "type": "object" + }, { "properties": { "_meta": true, diff --git a/codex-rs/app-server-protocol/schema/json/ServerRequest.json b/codex-rs/app-server-protocol/schema/json/ServerRequest.json index 31ca31de2c..c6ac80dee5 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ServerRequest.json @@ -1390,6 +1390,27 @@ ], "type": "object" }, + { + "properties": { + "_meta": true, + "message": { + "type": "string" + }, + "mode": { + "enum": [ + "openai/form" + ], + "type": "string" + }, + "requestedSchema": true + }, + "required": [ + "message", + "mode", + "requestedSchema" + ], + "type": "object" + }, { "properties": { "_meta": true, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index c75ae416fa..8ef80b7b54 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -2900,6 +2900,10 @@ "description": "Opt into receiving experimental API methods and fields.", "type": "boolean" }, + "mcpServerOpenaiFormElicitation": { + "description": "Allow downstream MCP servers to request OpenAI extended form elicitations.", + "type": "boolean" + }, "optOutNotificationMethods": { "description": "Exact notification method names that should be suppressed for this connection (for example `thread/started`).", "items": { @@ -3655,6 +3659,27 @@ ], "type": "object" }, + { + "properties": { + "_meta": true, + "message": { + "type": "string" + }, + "mode": { + "enum": [ + "openai/form" + ], + "type": "string" + }, + "requestedSchema": true + }, + "required": [ + "message", + "mode", + "requestedSchema" + ], + "type": "object" + }, { "properties": { "_meta": true, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index a4a0267dc5..00d7a4bbdf 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -7440,6 +7440,10 @@ "description": "Opt into receiving experimental API methods and fields.", "type": "boolean" }, + "mcpServerOpenaiFormElicitation": { + "description": "Allow downstream MCP servers to request OpenAI extended form elicitations.", + "type": "boolean" + }, "optOutNotificationMethods": { "description": "Exact notification method names that should be suppressed for this connection (for example `thread/started`).", "items": { diff --git a/codex-rs/app-server-protocol/schema/json/v1/InitializeParams.json b/codex-rs/app-server-protocol/schema/json/v1/InitializeParams.json index af5c509249..75f0860ddd 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/InitializeParams.json +++ b/codex-rs/app-server-protocol/schema/json/v1/InitializeParams.json @@ -30,6 +30,10 @@ "description": "Opt into receiving experimental API methods and fields.", "type": "boolean" }, + "mcpServerOpenaiFormElicitation": { + "description": "Allow downstream MCP servers to request OpenAI extended form elicitations.", + "type": "boolean" + }, "optOutNotificationMethods": { "description": "Exact notification method names that should be suppressed for this connection (for example `thread/started`).", "items": { diff --git a/codex-rs/app-server-protocol/schema/typescript/InitializeCapabilities.ts b/codex-rs/app-server-protocol/schema/typescript/InitializeCapabilities.ts index c5043e3b64..dcc4dffb0b 100644 --- a/codex-rs/app-server-protocol/schema/typescript/InitializeCapabilities.ts +++ b/codex-rs/app-server-protocol/schema/typescript/InitializeCapabilities.ts @@ -14,6 +14,10 @@ experimentalApi: boolean, * Opt into `attestation/generate` requests for upstream `x-oai-attestation`. */ requestAttestation: boolean, +/** + * Allow downstream MCP servers to request OpenAI extended form elicitations. + */ +mcpServerOpenaiFormElicitation?: boolean, /** * Exact notification method names that should be suppressed for this * connection (for example `thread/started`). diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerElicitationRequestParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerElicitationRequestParams.ts index 90d60f77c7..a4f1e732c6 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerElicitationRequestParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerElicitationRequestParams.ts @@ -13,4 +13,4 @@ export type McpServerElicitationRequestParams = { threadId: string, * context is app-server correlation rather than part of the protocol identity of the * elicitation itself. */ -turnId: string | null, serverName: string, } & ({ "mode": "form", _meta: JsonValue | null, message: string, requestedSchema: McpElicitationSchema, } | { "mode": "url", _meta: JsonValue | null, message: string, url: string, elicitationId: string, }); +turnId: string | null, serverName: string, } & ({ "mode": "form", _meta: JsonValue | null, message: string, requestedSchema: McpElicitationSchema, } | { "mode": "openai/form", _meta: JsonValue | null, message: string, requestedSchema: JsonValue, } | { "mode": "url", _meta: JsonValue | null, message: string, url: string, elicitationId: string, }); diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index edeabfc7ce..c1fac138a8 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -2150,7 +2150,7 @@ mod tests { } #[test] - fn serialize_initialize_with_opt_out_notification_methods() -> Result<()> { + fn serialize_initialize_capabilities() -> Result<()> { let request = ClientRequest::Initialize { request_id: RequestId::Integer(42), params: v1::InitializeParams { @@ -2162,6 +2162,7 @@ mod tests { capabilities: Some(v1::InitializeCapabilities { experimental_api: true, request_attestation: true, + mcp_server_openai_form_elicitation: true, opt_out_notification_methods: Some(vec![ "thread/started".to_string(), "item/agentMessage/delta".to_string(), @@ -2183,6 +2184,7 @@ mod tests { "capabilities": { "experimentalApi": true, "requestAttestation": true, + "mcpServerOpenaiFormElicitation": true, "optOutNotificationMethods": [ "thread/started", "item/agentMessage/delta" @@ -2196,7 +2198,7 @@ mod tests { } #[test] - fn deserialize_initialize_with_opt_out_notification_methods() -> Result<()> { + fn deserialize_initialize_capabilities() -> Result<()> { let request: ClientRequest = serde_json::from_value(json!({ "method": "initialize", "id": 42, @@ -2209,6 +2211,7 @@ mod tests { "capabilities": { "experimentalApi": true, "requestAttestation": true, + "mcpServerOpenaiFormElicitation": true, "optOutNotificationMethods": [ "thread/started", "item/agentMessage/delta" @@ -2230,6 +2233,7 @@ mod tests { capabilities: Some(v1::InitializeCapabilities { experimental_api: true, request_attestation: true, + mcp_server_openai_form_elicitation: true, opt_out_notification_methods: Some(vec![ "thread/started".to_string(), "item/agentMessage/delta".to_string(), diff --git a/codex-rs/app-server-protocol/src/protocol/v1.rs b/codex-rs/app-server-protocol/src/protocol/v1.rs index f83674d4c3..dccea51a36 100644 --- a/codex-rs/app-server-protocol/src/protocol/v1.rs +++ b/codex-rs/app-server-protocol/src/protocol/v1.rs @@ -50,6 +50,9 @@ pub struct InitializeCapabilities { /// Opt into `attestation/generate` requests for upstream `x-oai-attestation`. #[serde(default)] pub request_attestation: bool, + /// Allow downstream MCP servers to request OpenAI extended form elicitations. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub mcp_server_openai_form_elicitation: bool, /// Exact notification method names that should be suppressed for this /// connection (for example `thread/started`). #[ts(optional = nullable)] diff --git a/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs b/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs index 37197a2237..cdfeb524f5 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs @@ -632,6 +632,15 @@ pub enum McpServerElicitationRequest { message: String, requested_schema: McpElicitationSchema, }, + #[serde(rename = "openai/form", rename_all = "camelCase")] + #[ts(rename = "openai/form", rename_all = "camelCase")] + OpenAiForm { + #[serde(rename = "_meta")] + #[ts(rename = "_meta")] + meta: Option, + message: String, + requested_schema: JsonValue, + }, #[serde(rename_all = "camelCase")] #[ts(rename_all = "camelCase")] Url { @@ -658,6 +667,15 @@ impl TryFrom for McpServerElicitationRequest { message, requested_schema: serde_json::from_value(requested_schema)?, }), + CoreElicitationRequest::OpenAiForm { + meta, + message, + requested_schema, + } => Ok(Self::OpenAiForm { + meta, + message, + requested_schema, + }), CoreElicitationRequest::Url { meta, message, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index b973b8b4ae..bb66aa632e 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -1901,6 +1901,40 @@ fn mcp_server_elicitation_request_from_core_form_request() { ); } +#[test] +fn mcp_server_elicitation_request_from_core_openai_form_request() { + let requested_schema = json!({ + "type": "object", + "properties": { + "template": { + "type": "openai/imagePicker", + "title": "Template", + "items": [{ + "id": "monthly-review", + "title": "Monthly review", + "image": "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4=", + }], + }, + }, + "required": ["template"], + }); + let request = McpServerElicitationRequest::try_from(CoreElicitationRequest::OpenAiForm { + meta: None, + message: "Choose a report".to_string(), + requested_schema: requested_schema.clone(), + }) + .expect("OpenAI form request should convert"); + + assert_eq!( + request, + McpServerElicitationRequest::OpenAiForm { + meta: None, + message: "Choose a report".to_string(), + requested_schema, + } + ); +} + #[test] fn mcp_elicitation_schema_matches_mcp_2025_11_25_primitives() { let schema: McpElicitationSchema = serde_json::from_value(json!({ diff --git a/codex-rs/app-server-test-client/src/lib.rs b/codex-rs/app-server-test-client/src/lib.rs index 0be6e68bb2..856d75585e 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -1668,6 +1668,7 @@ impl CodexClient { .map(|method| (*method).to_string()) .collect(), ), + mcp_server_openai_form_elicitation: false, }), }, }; diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index a6c5ddab7d..183c7c6007 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -86,6 +86,13 @@ Clients must send a single `initialize` request per transport connection before `initialize.params.capabilities` also supports per-connection notification opt-out via `optOutNotificationMethods`, which is a list of exact method names to suppress for that connection. Matching is exact (no wildcards/prefixes). Unknown method names are accepted and ignored. +Clients that handle OpenAI extended MCP forms, including a fallback for +unsupported field types, set +`initialize.params.capabilities.mcpServerOpenaiFormElicitation` to `true`. +App-server then advertises the downstream `openai/form` MCP extension for +threads started, resumed, or forked by that connection. Clients that cannot +handle the request envelope omit the field or set it to `false`. + Applications building on top of `codex app-server` should identify themselves via the `clientInfo` parameter. **Important**: `clientInfo.name` is used to identify the client for the OpenAI Compliance Logs Platform. If @@ -1470,12 +1477,17 @@ Order of messages: 1. `mcpServer/elicitation/request` (request) — includes `threadId`, nullable `turnId`, `serverName`, and either: - a form request: `{ "mode": "form", "message": "...", "requestedSchema": { ... } }` + - an OpenAI extended form request: `{ "mode": "openai/form", "message": "...", "requestedSchema": { ... } }` - a URL request: `{ "mode": "url", "message": "...", "url": "...", "elicitationId": "..." }` 2. Client response — `{ "action": "accept", "content": ... }`, `{ "action": "decline", "content": null }`, or `{ "action": "cancel", "content": null }`. 3. `serverRequest/resolved` — `{ threadId, requestId }` confirms the pending request has been resolved or cleared, including lifecycle cleanup on turn start/complete/interrupt. `turnId` is best-effort. When the elicitation is correlated with an active turn, the request includes that turn id; otherwise it is `null`. +For `openai/form`, app-server forwards `requestedSchema` as opaque JSON. The +client owns validation and rendering of supported field types and must return a +valid `decline` or `cancel` response when it cannot render a form. + For MCP tool approval elicitations, form request `meta` includes `codex_approval_kind: "mcp_tool_call"` and may include `persist: "session"`, `persist: "always"`, or `persist: ["session", "always"]` to advertise whether diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 3544592b49..9e9437bd34 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -223,6 +223,7 @@ pub(crate) struct InitializedConnectionSessionState { pub(crate) app_server_client_name: String, pub(crate) client_version: String, pub(crate) request_attestation: bool, + pub(crate) supports_openai_form_elicitation: bool, } impl Default for ConnectionSessionState { @@ -274,6 +275,12 @@ impl ConnectionSessionState { .is_some_and(|session| session.request_attestation) } + pub(crate) fn supports_openai_form_elicitation(&self) -> bool { + self.initialized + .get() + .is_some_and(|session| session.supports_openai_form_elicitation) + } + pub(crate) fn initialize(&self, session: InitializedConnectionSessionState) -> Result<(), ()> { self.initialized.set(session).map_err(|_| ()) } @@ -884,6 +891,7 @@ impl MessageProcessor { let serialization_scope = codex_request.serialization_scope(); let app_server_client_name = session.app_server_client_name().map(str::to_string); let client_version = session.client_version().map(str::to_string); + let supports_openai_form_elicitation = session.supports_openai_form_elicitation(); let error_request_id = connection_request_id.clone(); let rpc_gate = Arc::clone(&session.rpc_gate); let processor = Arc::clone(self); @@ -899,6 +907,7 @@ impl MessageProcessor { request_context, app_server_client_name, client_version, + supports_openai_form_elicitation, ) .await; if let Err(error) = result { @@ -928,6 +937,7 @@ impl MessageProcessor { request_context: RequestContext, app_server_client_name: Option, client_version: Option, + supports_openai_form_elicitation: bool, ) -> Result<(), JSONRPCErrorError> { let connection_id = connection_request_id.connection_id; let request_id = ConnectionRequestId { @@ -1080,6 +1090,7 @@ impl MessageProcessor { params, app_server_client_name.clone(), client_version.clone(), + supports_openai_form_elicitation, request_context, ) .await @@ -1096,6 +1107,8 @@ impl MessageProcessor { params, app_server_client_name.clone(), client_version.clone(), + /*supports_openai_form_elicitation*/ + supports_openai_form_elicitation, ) .await } @@ -1106,6 +1119,8 @@ impl MessageProcessor { params, app_server_client_name.clone(), client_version.clone(), + /*supports_openai_form_elicitation*/ + supports_openai_form_elicitation, ) .await } @@ -1305,6 +1320,8 @@ impl MessageProcessor { params, app_server_client_name.clone(), client_version.clone(), + /*supports_openai_form_elicitation*/ + supports_openai_form_elicitation, ) .await } diff --git a/codex-rs/app-server/src/request_processors/initialize_processor.rs b/codex-rs/app-server/src/request_processors/initialize_processor.rs index a40007db11..cfdad27f50 100644 --- a/codex-rs/app-server/src/request_processors/initialize_processor.rs +++ b/codex-rs/app-server/src/request_processors/initialize_processor.rs @@ -67,17 +67,13 @@ impl InitializeRequestProcessor { // experimental API). Proposed direction is instance-global first-write-wins // with initialize-time mismatch rejection. let analytics_initialize_params = params.clone(); - let (experimental_api_enabled, request_attestation, opt_out_notification_methods) = - match params.capabilities { - Some(capabilities) => ( - capabilities.experimental_api, - capabilities.request_attestation, - capabilities - .opt_out_notification_methods - .unwrap_or_default(), - ), - None => (false, false, Vec::new()), - }; + let capabilities = params.capabilities.unwrap_or_default(); + let experimental_api_enabled = capabilities.experimental_api; + let request_attestation = capabilities.request_attestation; + let supports_openai_form_elicitation = capabilities.mcp_server_openai_form_elicitation; + let opt_out_notification_methods = capabilities + .opt_out_notification_methods + .unwrap_or_default(); let ClientInfo { name, title: _title, @@ -101,6 +97,7 @@ impl InitializeRequestProcessor { app_server_client_name: name.clone(), client_version: version, request_attestation, + supports_openai_form_elicitation, }) .is_err() { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 91005d6f71..63883cbdfe 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -417,6 +417,7 @@ impl ThreadRequestProcessor { params: ThreadStartParams, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, request_context: RequestContext, ) -> Result, JSONRPCErrorError> { self.thread_start_inner( @@ -424,6 +425,7 @@ impl ThreadRequestProcessor { params, app_server_client_name, app_server_client_version, + supports_openai_form_elicitation, request_context, ) .await @@ -446,12 +448,14 @@ impl ThreadRequestProcessor { params: ThreadResumeParams, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, ) -> Result, JSONRPCErrorError> { self.thread_resume_inner( request_id, params, app_server_client_name, app_server_client_version, + supports_openai_form_elicitation, ) .await .map(|()| None) @@ -463,12 +467,14 @@ impl ThreadRequestProcessor { params: ThreadForkParams, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, ) -> Result, JSONRPCErrorError> { self.thread_fork_inner( request_id, params, app_server_client_name, app_server_client_version, + supports_openai_form_elicitation, ) .await .map(|()| None) @@ -874,6 +880,7 @@ impl ThreadRequestProcessor { params: ThreadStartParams, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, request_context: RequestContext, ) -> Result<(), JSONRPCErrorError> { let ThreadStartParams { @@ -945,6 +952,7 @@ impl ThreadRequestProcessor { request_id, app_server_client_name, app_server_client_version, + supports_openai_form_elicitation, config, typesafe_overrides, dynamic_tools, @@ -1018,6 +1026,7 @@ impl ThreadRequestProcessor { request_id: ConnectionRequestId, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, config_overrides: Option>, typesafe_overrides: ConfigOverrides, dynamic_tools: Option>, @@ -1146,6 +1155,7 @@ impl ThreadRequestProcessor { parent_trace: request_trace, environments, thread_extension_init, + supports_openai_form_elicitation, }) .instrument(tracing::info_span!( "app_server.thread_start.create_thread", @@ -2506,6 +2516,7 @@ impl ThreadRequestProcessor { params: ThreadResumeParams, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, ) -> Result<(), JSONRPCErrorError> { if let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) && self @@ -2650,6 +2661,7 @@ impl ThreadRequestProcessor { thread_history, self.auth_manager.clone(), self.request_trace_context(&request_id).await, + supports_openai_form_elicitation, ) .await { @@ -3267,6 +3279,7 @@ impl ThreadRequestProcessor { params: ThreadForkParams, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, ) -> Result<(), JSONRPCErrorError> { let ThreadForkParams { thread_id, @@ -3376,6 +3389,7 @@ impl ThreadRequestProcessor { }), thread_source.map(Into::into), self.request_trace_context(&request_id).await, + supports_openai_form_elicitation, ) .await .map_err(|err| match err { diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index fbfcb797b6..bd7aaf1ba0 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -101,12 +101,14 @@ impl TurnRequestProcessor { params: TurnStartParams, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, ) -> Result, JSONRPCErrorError> { self.turn_start_inner( request_id, params, app_server_client_name, app_server_client_version, + /*supports_openai_form_elicitation*/ supports_openai_form_elicitation, ) .await .map(|response| Some(response.into())) @@ -384,6 +386,7 @@ impl TurnRequestProcessor { params: TurnStartParams, app_server_client_name: Option, app_server_client_version: Option, + supports_openai_form_elicitation: bool, ) -> Result { let (thread_id, thread) = self.load_thread(¶ms.thread_id) @@ -410,6 +413,14 @@ impl TurnRequestProcessor { .inspect_err(|error| { self.track_error_response(&request_id, error, /*error_type*/ None); })?; + thread + .set_openai_form_elicitation_support(supports_openai_form_elicitation) + .await + .map_err(|err| { + internal_error(format!( + "failed to update OpenAI form elicitation support: {err}" + )) + })?; let environment_selections = resolve_turn_environment_selections(self.thread_manager.as_ref(), params.environments)?; @@ -1164,6 +1175,7 @@ impl TurnRequestProcessor { }), /*thread_source*/ None, self.request_trace_context(request_id).await, + /*supports_openai_form_elicitation*/ false, ) .await .map_err(|err| { diff --git a/codex-rs/app-server/tests/suite/v2/attestation.rs b/codex-rs/app-server/tests/suite/v2/attestation.rs index d007614730..ea567a2d43 100644 --- a/codex-rs/app-server/tests/suite/v2/attestation.rs +++ b/codex-rs/app-server/tests/suite/v2/attestation.rs @@ -81,6 +81,7 @@ async fn attestation_generate_round_trip_adds_header_to_responses_websocket_hand experimental_api: true, request_attestation: true, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ), ) diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index 0e00bc6309..648d60f3b5 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs @@ -706,7 +706,7 @@ pub(super) async fn send_request( send_jsonrpc(stream, message).await } -async fn send_jsonrpc(stream: &mut WsClient, message: JSONRPCMessage) -> Result<()> { +pub(super) async fn send_jsonrpc(stream: &mut WsClient, message: JSONRPCMessage) -> Result<()> { let payload = serde_json::to_string(&message)?; stream .send(WebSocketMessage::Text(payload.into())) diff --git a/codex-rs/app-server/tests/suite/v2/experimental_api.rs b/codex-rs/app-server/tests/suite/v2/experimental_api.rs index 3b4f6e949a..62ad4c0677 100644 --- a/codex-rs/app-server/tests/suite/v2/experimental_api.rs +++ b/codex-rs/app-server/tests/suite/v2/experimental_api.rs @@ -39,6 +39,7 @@ async fn mock_experimental_method_requires_experimental_api_capability() -> Resu experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ) .await?; @@ -70,6 +71,7 @@ async fn realtime_conversation_start_requires_experimental_api_capability() -> R experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ) .await?; @@ -116,6 +118,7 @@ async fn thread_memory_mode_set_requires_experimental_api_capability() -> Result experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ) .await?; @@ -150,6 +153,7 @@ async fn thread_settings_update_requires_experimental_api_capability() -> Result experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ) .await?; @@ -184,6 +188,7 @@ async fn realtime_webrtc_start_requires_experimental_api_capability() -> Result< experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ) .await?; @@ -234,6 +239,7 @@ async fn thread_start_mock_field_requires_experimental_api_capability() -> Resul experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ) .await?; @@ -272,6 +278,7 @@ async fn thread_start_without_dynamic_tools_allows_without_experimental_api_capa experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ) .await?; @@ -309,6 +316,7 @@ async fn thread_start_granular_approval_policy_requires_experimental_api_capabil experimental_api: false, request_attestation: false, opt_out_notification_methods: None, + mcp_server_openai_form_elicitation: false, }), ) .await?; diff --git a/codex-rs/app-server/tests/suite/v2/initialize.rs b/codex-rs/app-server/tests/suite/v2/initialize.rs index 0cf8504d3b..6fb7aa3bc3 100644 --- a/codex-rs/app-server/tests/suite/v2/initialize.rs +++ b/codex-rs/app-server/tests/suite/v2/initialize.rs @@ -214,6 +214,7 @@ async fn initialize_opt_out_notification_methods_filters_notifications() -> Resu experimental_api: true, request_attestation: false, opt_out_notification_methods: Some(vec!["thread/started".to_string()]), + mcp_server_openai_form_elicitation: false, }), ), ) diff --git a/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs b/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs index b4bd1d377a..113f58f336 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs @@ -14,6 +14,9 @@ use axum::http::StatusCode; use axum::http::Uri; use axum::http::header::AUTHORIZATION; use axum::routing::get; +use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::InitializeCapabilities; +use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::McpElicitationSchema; @@ -24,6 +27,7 @@ use codex_app_server_protocol::McpServerElicitationRequestResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ServerRequestResolvedNotification; +use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnCompletedNotification; @@ -34,6 +38,7 @@ use codex_app_server_protocol::UserInput as V2UserInput; use codex_config::types::AuthCredentialsStoreMode; use core_test_support::assert_regex_match; use core_test_support::responses; +use core_test_support::responses::ResponseMock; use pretty_assertions::assert_eq; use rmcp::handler::server::ServerHandler; use rmcp::model::BooleanSchema; @@ -41,14 +46,18 @@ use rmcp::model::CallToolRequestParams; use rmcp::model::CallToolResult; use rmcp::model::Content; use rmcp::model::CreateElicitationRequestParams; +use rmcp::model::CustomRequest; use rmcp::model::ElicitationAction; use rmcp::model::ElicitationSchema; +use rmcp::model::InitializeRequestParams; +use rmcp::model::InitializeResult; use rmcp::model::JsonObject; use rmcp::model::ListToolsResult; use rmcp::model::Meta; use rmcp::model::PrimitiveSchema; use rmcp::model::ServerCapabilities; use rmcp::model::ServerInfo; +use rmcp::model::ServerRequest as McpServerRequest; use rmcp::model::Tool; use rmcp::model::ToolAnnotations; use rmcp::service::RequestContext; @@ -63,6 +72,15 @@ use tokio::net::TcpListener; use tokio::task::JoinHandle; use tokio::time::timeout; +use super::connection_handling_websocket::WsClient; +use super::connection_handling_websocket::connect_websocket; +use super::connection_handling_websocket::read_jsonrpc_message; +use super::connection_handling_websocket::read_notification_for_method; +use super::connection_handling_websocket::read_response_for_id; +use super::connection_handling_websocket::send_jsonrpc; +use super::connection_handling_websocket::send_request; +use super::connection_handling_websocket::spawn_websocket_server; + const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); const CONNECTOR_ID: &str = "calendar"; const CONNECTOR_NAME: &str = "Calendar"; @@ -71,9 +89,278 @@ const CALLABLE_TOOL_NAME: &str = "_confirm_action"; const TOOL_NAME: &str = "calendar_confirm_action"; const TOOL_CALL_ID: &str = "call-calendar-confirm"; const ELICITATION_MESSAGE: &str = "Allow this request?"; +const OPENAI_FORM_MESSAGE: &str = "Select a template"; +const IMAGE_DATA_URL: &str = + "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4="; + +#[derive(Clone, Copy)] +enum ElicitationScenario { + StandardForm, + OpenAiForm, +} #[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn mcp_server_elicitation_round_trip() -> Result<()> { +async fn mcp_server_form_elicitation_round_trip() -> Result<()> { + let mut fixture = ElicitationRoundTripFixture::start(ElicitationScenario::StandardForm).await?; + let (request_id, params) = fixture.read_elicitation().await?; + let requested_schema: McpElicitationSchema = serde_json::from_value(serde_json::to_value( + ElicitationSchema::builder() + .required_property("confirmed", PrimitiveSchema::Boolean(BooleanSchema::new())) + .build() + .map_err(anyhow::Error::msg)?, + )?)?; + assert_eq!( + params, + McpServerElicitationRequestParams { + thread_id: fixture.thread_id.clone(), + turn_id: Some(fixture.turn_id.clone()), + server_name: "codex_apps".to_string(), + request: McpServerElicitationRequest::Form { + meta: None, + message: ELICITATION_MESSAGE.to_string(), + requested_schema, + }, + } + ); + + fixture + .accept(request_id.clone(), json!({ "confirmed": true })) + .await?; + fixture.finish(request_id, "accepted").await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn mcp_server_openai_form_elicitation_round_trip() -> Result<()> { + let mut fixture = ElicitationRoundTripFixture::start(ElicitationScenario::OpenAiForm).await?; + let (request_id, params) = fixture.read_elicitation().await?; + assert_eq!( + params, + McpServerElicitationRequestParams { + thread_id: fixture.thread_id.clone(), + turn_id: Some(fixture.turn_id.clone()), + server_name: "codex_apps".to_string(), + request: McpServerElicitationRequest::OpenAiForm { + meta: None, + message: OPENAI_FORM_MESSAGE.to_string(), + requested_schema: json!({ + "type": "object", + "properties": { + "template": { + "type": "openai/imagePicker", + "title": "Template", + "items": [{ + "id": "monthly-review", + "title": "Monthly review", + "image": IMAGE_DATA_URL, + }], + }, + }, + "required": ["template"], + }), + }, + } + ); + + fixture + .accept(request_id.clone(), json!({ "template": "monthly-review" })) + .await?; + fixture.finish(request_id, "accepted monthly-review").await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn openai_form_capability_follows_the_turn_starting_connection() -> Result<()> { + let (responses_server, response_mock, apps_server_url, apps_server_handle) = + start_elicitation_services(ElicitationScenario::OpenAiForm).await?; + let codex_home = TempDir::new()?; + write_config_toml(codex_home.path(), &responses_server.uri(), &apps_server_url)?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("chatgpt-token") + .account_id("account-123") + .chatgpt_user_id("user-123") + .chatgpt_account_id("account-123"), + AuthCredentialsStoreMode::File, + )?; + + let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?; + let mut supported_client = connect_websocket(bind_addr).await?; + initialize_websocket_client( + &mut supported_client, + /*id*/ 1, + "supported-client", + /*supports_openai_form_elicitation*/ true, + ) + .await?; + + send_request( + &mut supported_client, + "thread/start", + /*id*/ 2, + Some(serde_json::to_value(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + })?), + ) + .await?; + let ThreadStartResponse { thread, .. } = + to_response(read_response_for_id(&mut supported_client, /*id*/ 2).await?)?; + + send_request( + &mut supported_client, + "turn/start", + /*id*/ 3, + Some(serde_json::to_value(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "Warm up connectors.".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model".to_string()), + ..Default::default() + })?), + ) + .await?; + let _: TurnStartResponse = + to_response(read_response_for_id(&mut supported_client, /*id*/ 3).await?)?; + let _: TurnCompletedNotification = serde_json::from_value( + read_notification_for_method(&mut supported_client, "turn/completed") + .await? + .params + .expect("turn/completed params"), + )?; + + let mut unsupported_client = connect_websocket(bind_addr).await?; + initialize_websocket_client( + &mut unsupported_client, + /*id*/ 4, + "unsupported-client", + /*supports_openai_form_elicitation*/ false, + ) + .await?; + send_request( + &mut unsupported_client, + "thread/resume", + /*id*/ 5, + Some(serde_json::to_value(ThreadResumeParams { + thread_id: thread.id.clone(), + ..Default::default() + })?), + ) + .await?; + let _ = read_response_for_id(&mut unsupported_client, /*id*/ 5).await?; + + send_request( + &mut supported_client, + "turn/start", + /*id*/ 6, + Some(serde_json::to_value(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "Use [$calendar](app://calendar) to run the calendar tool.".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model".to_string()), + ..Default::default() + })?), + ) + .await?; + let TurnStartResponse { turn } = + to_response(read_response_for_id(&mut supported_client, /*id*/ 6).await?)?; + + let (request_id, params) = loop { + let JSONRPCMessage::Request(request) = read_jsonrpc_message(&mut supported_client).await? + else { + continue; + }; + let request: ServerRequest = serde_json::from_value(serde_json::to_value(request)?)?; + let ServerRequest::McpServerElicitationRequest { request_id, params } = request else { + continue; + }; + break (request_id, params); + }; + assert_eq!( + params.request, + McpServerElicitationRequest::OpenAiForm { + meta: None, + message: OPENAI_FORM_MESSAGE.to_string(), + requested_schema: json!({ + "type": "object", + "properties": { + "template": { + "type": "openai/imagePicker", + "title": "Template", + "items": [{ + "id": "monthly-review", + "title": "Monthly review", + "image": IMAGE_DATA_URL, + }], + }, + }, + "required": ["template"], + }), + } + ); + send_jsonrpc( + &mut supported_client, + JSONRPCMessage::Response(JSONRPCResponse { + id: request_id, + result: serde_json::to_value(McpServerElicitationRequestResponse { + action: McpServerElicitationAction::Accept, + content: Some(json!({ "template": "monthly-review" })), + meta: None, + })?, + }), + ) + .await?; + + let completed: TurnCompletedNotification = serde_json::from_value( + read_notification_for_method(&mut supported_client, "turn/completed") + .await? + .params + .expect("turn/completed params"), + )?; + assert_eq!(completed.thread_id, thread.id); + assert_eq!(completed.turn.id, turn.id); + assert_eq!(completed.turn.status, TurnStatus::Completed); + assert_eq!(response_mock.requests().len(), 3); + + process.kill().await?; + apps_server_handle.abort(); + let _ = apps_server_handle.await; + Ok(()) +} + +async fn initialize_websocket_client( + client: &mut WsClient, + id: i64, + name: &str, + supports_openai_form_elicitation: bool, +) -> Result<()> { + send_request( + client, + "initialize", + id, + Some(serde_json::to_value(InitializeParams { + client_info: ClientInfo { + name: name.to_string(), + title: None, + version: "0.1.0".to_string(), + }, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + mcp_server_openai_form_elicitation: supports_openai_form_elicitation, + ..Default::default() + }), + })?), + ) + .await?; + let _ = read_response_for_id(client, id).await?; + Ok(()) +} + +async fn start_elicitation_services( + scenario: ElicitationScenario, +) -> Result<(wiremock::MockServer, ResponseMock, String, JoinHandle<()>)> { let responses_server = responses::start_mock_server().await; let tool_call_arguments = serde_json::to_string(&json!({}))?; let response_mock = responses::mount_sse_sequence( @@ -102,201 +389,223 @@ async fn mcp_server_elicitation_round_trip() -> Result<()> { ], ) .await; + let (apps_server_url, apps_server_handle) = start_apps_server(scenario).await?; + Ok(( + responses_server, + response_mock, + apps_server_url, + apps_server_handle, + )) +} - let (apps_server_url, apps_server_handle) = start_apps_server().await?; +struct ElicitationRoundTripFixture { + mcp: TestAppServer, + response_mock: ResponseMock, + _responses_server: wiremock::MockServer, + thread_id: String, + turn_id: String, + apps_server_handle: JoinHandle<()>, +} - let codex_home = TempDir::new()?; - write_config_toml(codex_home.path(), &responses_server.uri(), &apps_server_url)?; - write_chatgpt_auth( - codex_home.path(), - ChatGptAuthFixture::new("chatgpt-token") - .account_id("account-123") - .chatgpt_user_id("user-123") - .chatgpt_account_id("account-123"), - AuthCredentialsStoreMode::File, - )?; +impl ElicitationRoundTripFixture { + async fn start(scenario: ElicitationScenario) -> Result { + let (responses_server, response_mock, apps_server_url, apps_server_handle) = + start_elicitation_services(scenario).await?; + let codex_home = TempDir::new()?; + write_config_toml(codex_home.path(), &responses_server.uri(), &apps_server_url)?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("chatgpt-token") + .account_id("account-123") + .chatgpt_user_id("user-123") + .chatgpt_account_id("account-123"), + AuthCredentialsStoreMode::File, + )?; - let mut mcp = TestAppServer::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.initialize_with_capabilities( + ClientInfo { + name: "codex-app-server-tests".to_string(), + title: None, + version: "0.1.0".to_string(), + }, + Some(InitializeCapabilities { + experimental_api: true, + mcp_server_openai_form_elicitation: true, + ..Default::default() + }), + ), + ) + .await??; - let thread_start_id = mcp - .send_thread_start_request(ThreadStartParams { - model: Some("mock-model".to_string()), - ..Default::default() + let thread_start_id = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; + + let warmup_turn_start_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Warm up connectors.".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let warmup_turn_start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(warmup_turn_start_id)), + ) + .await??; + let _: TurnStartResponse = to_response(warmup_turn_start_resp)?; + let warmup_completed = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + let warmup_completed: TurnCompletedNotification = serde_json::from_value( + warmup_completed + .params + .clone() + .expect("warmup turn/completed params"), + )?; + assert_eq!(warmup_completed.thread_id, thread.id); + assert_eq!(warmup_completed.turn.status, TurnStatus::Completed); + + let turn_start_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Use [$calendar](app://calendar) to run the calendar tool.".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let turn_start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), + ) + .await??; + let TurnStartResponse { turn } = to_response(turn_start_resp)?; + + Ok(Self { + mcp, + response_mock, + _responses_server: responses_server, + thread_id: thread.id, + turn_id: turn.id, + apps_server_handle, }) - .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; - - let warmup_turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Warm up connectors.".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model".to_string()), - ..Default::default() - }) - .await?; - let warmup_turn_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(warmup_turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response(warmup_turn_start_resp)?; - - let warmup_completed = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), - ) - .await??; - let warmup_completed: TurnCompletedNotification = serde_json::from_value( - warmup_completed - .params - .clone() - .expect("warmup turn/completed params"), - )?; - assert_eq!(warmup_completed.thread_id, thread.id); - assert_eq!(warmup_completed.turn.status, TurnStatus::Completed); - - let turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Use [$calendar](app://calendar) to run the calendar tool.".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model".to_string()), - ..Default::default() - }) - .await?; - let turn_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response(turn_start_resp)?; - - let server_req = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_request_message(), - ) - .await??; - let ServerRequest::McpServerElicitationRequest { request_id, params } = server_req else { - panic!("expected McpServerElicitationRequest request, got: {server_req:?}"); - }; - let requested_schema: McpElicitationSchema = serde_json::from_value(serde_json::to_value( - ElicitationSchema::builder() - .required_property("confirmed", PrimitiveSchema::Boolean(BooleanSchema::new())) - .build() - .map_err(anyhow::Error::msg)?, - )?)?; - - assert_eq!( - params, - McpServerElicitationRequestParams { - thread_id: thread.id.clone(), - turn_id: Some(turn.id.clone()), - server_name: "codex_apps".to_string(), - request: McpServerElicitationRequest::Form { - meta: None, - message: ELICITATION_MESSAGE.to_string(), - requested_schema, - }, - } - ); - - let resolved_request_id = request_id.clone(); - mcp.send_response( - request_id, - serde_json::to_value(McpServerElicitationRequestResponse { - action: McpServerElicitationAction::Accept, - content: Some(json!({ - "confirmed": true, - })), - meta: None, - })?, - ) - .await?; - - let mut saw_resolved = false; - loop { - let message = timeout(DEFAULT_READ_TIMEOUT, mcp.read_next_message()).await??; - let JSONRPCMessage::Notification(notification) = message else { - continue; - }; - - match notification.method.as_str() { - "serverRequest/resolved" => { - let resolved: ServerRequestResolvedNotification = serde_json::from_value( - notification - .params - .clone() - .expect("serverRequest/resolved params"), - )?; - assert_eq!( - resolved, - ServerRequestResolvedNotification { - thread_id: thread.id.clone(), - request_id: resolved_request_id.clone(), - } - ); - saw_resolved = true; - } - "turn/completed" => { - let completed: TurnCompletedNotification = serde_json::from_value( - notification.params.clone().expect("turn/completed params"), - )?; - assert!(saw_resolved, "serverRequest/resolved should arrive first"); - assert_eq!(completed.thread_id, thread.id); - assert_eq!(completed.turn.id, turn.id); - assert_eq!(completed.turn.status, TurnStatus::Completed); - break; - } - _ => {} - } } - let requests = response_mock.requests(); - assert_eq!(requests.len(), 3); - let function_call_output = requests[2].function_call_output(TOOL_CALL_ID); - assert_eq!( - function_call_output.get("type"), - Some(&Value::String("function_call_output".to_string())) - ); - assert_eq!( - function_call_output.get("call_id"), - Some(&Value::String(TOOL_CALL_ID.to_string())) - ); - let output = function_call_output - .get("output") - .and_then(Value::as_str) - .expect("function_call_output output should be a JSON string"); - let payload = assert_regex_match( - r#"(?s)^Wall time: [0-9]+(?:\.[0-9]+)? seconds\nOutput:\n(.*)$"#, - output, - ) - .get(1) - .expect("wall-time wrapped output should include payload") - .as_str(); - assert_eq!( - serde_json::from_str::(payload)?, - json!([{ - "type": "text", - "text": "accepted" - }]) - ); + async fn read_elicitation(&mut self) -> Result<(RequestId, McpServerElicitationRequestParams)> { + let request = timeout( + DEFAULT_READ_TIMEOUT, + self.mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::McpServerElicitationRequest { request_id, params } = request else { + panic!("expected McpServerElicitationRequest request, got: {request:?}"); + }; + Ok((request_id, params)) + } - apps_server_handle.abort(); - let _ = apps_server_handle.await; - Ok(()) + async fn accept(&mut self, request_id: RequestId, content: Value) -> Result<()> { + self.mcp + .send_response( + request_id, + serde_json::to_value(McpServerElicitationRequestResponse { + action: McpServerElicitationAction::Accept, + content: Some(content), + meta: None, + })?, + ) + .await + } + + async fn finish(mut self, request_id: RequestId, expected_text: &str) -> Result<()> { + let mut resolved = false; + loop { + let message = timeout(DEFAULT_READ_TIMEOUT, self.mcp.read_next_message()).await??; + let JSONRPCMessage::Notification(notification) = message else { + continue; + }; + match notification.method.as_str() { + "serverRequest/resolved" => { + let notification: ServerRequestResolvedNotification = serde_json::from_value( + notification + .params + .clone() + .expect("serverRequest/resolved params"), + )?; + assert_eq!(notification.thread_id, self.thread_id); + assert_eq!(notification.request_id, request_id); + resolved = true; + } + "turn/completed" => { + let notification: TurnCompletedNotification = serde_json::from_value( + notification.params.clone().expect("turn/completed params"), + )?; + assert!( + resolved, + "server request should resolve before turn completion" + ); + assert_eq!(notification.thread_id, self.thread_id); + assert_eq!(notification.turn.id, self.turn_id); + assert_eq!(notification.turn.status, TurnStatus::Completed); + break; + } + _ => {} + } + } + + let requests = self.response_mock.requests(); + assert_eq!(requests.len(), 3); + let function_call_output = requests[2].function_call_output(TOOL_CALL_ID); + assert_eq!( + function_call_output.get("type"), + Some(&Value::String("function_call_output".to_string())) + ); + assert_eq!( + function_call_output.get("call_id"), + Some(&Value::String(TOOL_CALL_ID.to_string())) + ); + let output = function_call_output + .get("output") + .and_then(Value::as_str) + .expect("function_call_output output should be a JSON string"); + let payload = assert_regex_match( + r#"(?s)^Wall time: [0-9]+(?:\.[0-9]+)? seconds\nOutput:\n(.*)$"#, + output, + ) + .get(1) + .expect("wall-time wrapped output should include payload") + .as_str(); + assert_eq!( + serde_json::from_str::(payload)?, + json!([{ "type": "text", "text": expected_text }]) + ); + + self.apps_server_handle.abort(); + let _ = self.apps_server_handle.await; + Ok(()) + } } #[derive(Clone)] @@ -305,10 +614,33 @@ struct AppsServerState { expected_account_id: String, } -#[derive(Clone, Default)] -struct ElicitationAppsMcpServer; +#[derive(Clone)] +struct ElicitationAppsMcpServer { + scenario: ElicitationScenario, +} impl ServerHandler for ElicitationAppsMcpServer { + async fn initialize( + &self, + request: InitializeRequestParams, + context: RequestContext, + ) -> Result { + if matches!(self.scenario, ElicitationScenario::OpenAiForm) { + assert_eq!( + request + .capabilities + .extensions + .as_ref() + .and_then(|extensions| extensions.get("openai/form")) + .cloned() + .map(Value::Object), + Some(json!({})) + ); + } + context.peer.set_peer_info(request); + Ok(self.get_info()) + } + fn get_info(&self) -> ServerInfo { ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) .with_protocol_version(rmcp::model::ProtocolVersion::V_2025_06_18) @@ -351,40 +683,91 @@ impl ServerHandler for ElicitationAppsMcpServer { _request: CallToolRequestParams, context: RequestContext, ) -> Result { - let requested_schema = ElicitationSchema::builder() - .required_property("confirmed", PrimitiveSchema::Boolean(BooleanSchema::new())) - .build() - .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?; - - let result = context - .peer - .create_elicitation(CreateElicitationRequestParams::FormElicitationParams { - meta: None, - message: ELICITATION_MESSAGE.to_string(), - requested_schema, - }) - .await - .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?; - - let output = match result.action { - ElicitationAction::Accept => { + match self.scenario { + ElicitationScenario::StandardForm => { + let requested_schema = ElicitationSchema::builder() + .required_property("confirmed", PrimitiveSchema::Boolean(BooleanSchema::new())) + .build() + .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?; + let result = context + .peer + .create_elicitation(CreateElicitationRequestParams::FormElicitationParams { + meta: None, + message: ELICITATION_MESSAGE.to_string(), + requested_schema, + }) + .await + .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?; assert_eq!( result.content, Some(json!({ "confirmed": true, })) ); - "accepted" + let output = match result.action { + ElicitationAction::Accept => "accepted", + ElicitationAction::Decline => "declined", + ElicitationAction::Cancel => "cancelled", + }; + Ok(CallToolResult::success(vec![Content::text(output)])) } - ElicitationAction::Decline => "declined", - ElicitationAction::Cancel => "cancelled", - }; - - Ok(CallToolResult::success(vec![Content::text(output)])) + ElicitationScenario::OpenAiForm => { + let result = context + .peer + .send_request(McpServerRequest::CustomRequest(CustomRequest::new( + "openai/form", + Some(json!({ + "message": OPENAI_FORM_MESSAGE, + "requestedSchema": { + "type": "object", + "properties": { + "template": { + "type": "openai/imagePicker", + "title": "Template", + "items": [{ + "id": "monthly-review", + "title": "Monthly review", + "image": IMAGE_DATA_URL, + }], + }, + }, + "required": ["template"], + }, + })), + ))) + .await + .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?; + let result = match result { + rmcp::model::ClientResult::CustomResult(result) => result.0, + rmcp::model::ClientResult::CreateElicitationResult(result) => { + serde_json::to_value(result) + .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))? + } + result => { + return Err(rmcp::ErrorData::internal_error( + format!("unexpected OpenAI form response: {result:?}"), + None, + )); + } + }; + assert_eq!( + result, + json!({ + "action": "accept", + "content": { + "template": "monthly-review", + }, + }) + ); + Ok(CallToolResult::success(vec![Content::text( + "accepted monthly-review", + )])) + } + } } } -async fn start_apps_server() -> Result<(String, JoinHandle<()>)> { +async fn start_apps_server(scenario: ElicitationScenario) -> Result<(String, JoinHandle<()>)> { let state = Arc::new(AppsServerState { expected_bearer: "Bearer chatgpt-token".to_string(), expected_account_id: "account-123".to_string(), @@ -394,7 +777,7 @@ async fn start_apps_server() -> Result<(String, JoinHandle<()>)> { let addr = listener.local_addr()?; let mcp_service = StreamableHttpService::new( - move || Ok(ElicitationAppsMcpServer), + move || Ok(ElicitationAppsMcpServer { scenario }), Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default(), ); diff --git a/codex-rs/app-server/tests/suite/v2/thread_status.rs b/codex-rs/app-server/tests/suite/v2/thread_status.rs index b8349a77f1..e922bd13c1 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_status.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_status.rs @@ -148,6 +148,7 @@ async fn thread_status_changed_can_be_opted_out() -> Result<()> { experimental_api: true, request_attestation: false, opt_out_notification_methods: Some(vec!["thread/status/changed".to_string()]), + mcp_server_openai_form_elicitation: false, }), ), ) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index e50532714a..e9c835e863 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -133,6 +133,7 @@ impl McpConnectionManager { host_owned_codex_apps_enabled: bool, prefix_mcp_tool_names: bool, client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, tool_plugin_provenance: ToolPluginProvenance, auth: Option<&CodexAuth>, elicitation_reviewer: Option, @@ -209,6 +210,7 @@ impl McpConnectionManager { runtime_context.clone(), runtime_auth_provider, client_elicitation_capability.clone(), + supports_openai_form_elicitation, ); clients.insert(server_name.clone(), async_managed_client.clone()); let tx_event = tx_event.clone(); diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 20bb8d43cc..c8d8f79081 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -251,13 +251,15 @@ async fn disabled_permissions_auto_accept_elicitation_with_empty_form_schema() { let response = sender( NumberOrString::Number(1), - CreateElicitationRequestParams::FormElicitationParams { - meta: None, - message: "Confirm?".to_string(), - requested_schema: rmcp::model::ElicitationSchema::builder() - .build() - .expect("schema should build"), - }, + codex_rmcp_client::Elicitation::Mcp( + CreateElicitationRequestParams::FormElicitationParams { + meta: None, + message: "Confirm?".to_string(), + requested_schema: rmcp::model::ElicitationSchema::builder() + .build() + .expect("schema should build"), + }, + ), ) .await .expect("elicitation should auto accept"); @@ -284,17 +286,19 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel let response = sender( NumberOrString::Number(1), - CreateElicitationRequestParams::FormElicitationParams { - meta: None, - message: "What should I say?".to_string(), - requested_schema: rmcp::model::ElicitationSchema::builder() - .required_property( - "message", - rmcp::model::PrimitiveSchema::String(rmcp::model::StringSchema::new()), - ) - .build() - .expect("schema should build"), - }, + codex_rmcp_client::Elicitation::Mcp( + CreateElicitationRequestParams::FormElicitationParams { + meta: None, + message: "What should I say?".to_string(), + requested_schema: rmcp::model::ElicitationSchema::builder() + .required_property( + "message", + rmcp::model::PrimitiveSchema::String(rmcp::model::StringSchema::new()), + ) + .build() + .expect("schema should build"), + }, + ), ) .await .expect("elicitation should auto decline"); @@ -1265,6 +1269,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { /*host_owned_codex_apps_enabled*/ false, /*prefix_mcp_tool_names*/ true, ElicitationCapability::default(), + /*supports_openai_form_elicitation*/ false, ToolPluginProvenance::default(), /*auth*/ None, /*elicitation_reviewer*/ None, diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index a51cd7c623..abd5a46899 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -22,11 +22,11 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; +use codex_rmcp_client::Elicitation; use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::SendElicitation; use futures::future::BoxFuture; use futures::future::FutureExt; -use rmcp::model::CreateElicitationRequestParams; use rmcp::model::ElicitationAction; use rmcp::model::RequestId; use tokio::sync::Mutex; @@ -36,7 +36,7 @@ use tokio::sync::oneshot; pub struct ElicitationReviewRequest { pub server_name: String, pub request_id: RequestId, - pub elicitation: CreateElicitationRequestParams, + pub elicitation: Elicitation, } pub trait ElicitationReviewer: Send + Sync { @@ -172,11 +172,13 @@ impl ElicitationRequestManager { } let request = match elicitation { - CreateElicitationRequestParams::FormElicitationParams { - meta, - message, - requested_schema, - } => ElicitationRequest::Form { + Elicitation::Mcp( + rmcp::model::CreateElicitationRequestParams::FormElicitationParams { + meta, + message, + requested_schema, + }, + ) => ElicitationRequest::Form { meta: meta .map(serde_json::to_value) .transpose() @@ -185,12 +187,14 @@ impl ElicitationRequestManager { requested_schema: serde_json::to_value(requested_schema) .context("failed to serialize MCP elicitation schema")?, }, - CreateElicitationRequestParams::UrlElicitationParams { - meta, - message, - url, - elicitation_id, - } => ElicitationRequest::Url { + Elicitation::Mcp( + rmcp::model::CreateElicitationRequestParams::UrlElicitationParams { + meta, + message, + url, + elicitation_id, + }, + ) => ElicitationRequest::Url { meta: meta .map(serde_json::to_value) .transpose() @@ -199,6 +203,15 @@ impl ElicitationRequestManager { url, elicitation_id, }, + Elicitation::OpenAiForm { + meta, + message, + requested_schema, + } => ElicitationRequest::OpenAiForm { + meta, + message, + requested_schema, + }, }; let (tx, rx) = oneshot::channel(); { @@ -243,14 +256,18 @@ pub(crate) fn elicitation_is_rejected_by_policy(approval_policy: AskForApproval) type ResponderMap = HashMap<(String, RequestId), oneshot::Sender>; -fn can_auto_accept_elicitation(elicitation: &CreateElicitationRequestParams) -> bool { +fn can_auto_accept_elicitation(elicitation: &Elicitation) -> bool { match elicitation { - CreateElicitationRequestParams::FormElicitationParams { - requested_schema, .. - } => { + Elicitation::Mcp(rmcp::model::CreateElicitationRequestParams::FormElicitationParams { + requested_schema, + .. + }) => { // Auto-accept confirm/approval elicitations without schema requirements. requested_schema.properties.is_empty() } - CreateElicitationRequestParams::UrlElicitationParams { .. } => false, + Elicitation::Mcp(rmcp::model::CreateElicitationRequestParams::UrlElicitationParams { + .. + }) + | Elicitation::OpenAiForm { .. } => false, } } diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 7a585cb5d8..d87602438f 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -305,6 +305,7 @@ pub async fn read_mcp_resource( host_owned_codex_apps_enabled, config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), + /*supports_openai_form_elicitation*/ false, tool_plugin_provenance(config), auth, /*elicitation_reviewer*/ None, @@ -379,6 +380,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( host_owned_codex_apps_enabled, config.prefix_mcp_tool_names, config.client_elicitation_capability.clone(), + /*supports_openai_form_elicitation*/ false, tool_plugin_provenance, auth, /*elicitation_reviewer*/ None, diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 3ba1adcc94..aa56650e8a 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -7,6 +7,7 @@ //! [`crate::connection_manager`]. use std::borrow::Cow; +use std::collections::BTreeMap; use std::collections::HashMap; use std::env; use std::ffi::OsString; @@ -62,6 +63,7 @@ use rmcp::model::ClientCapabilities; use rmcp::model::ElicitationCapability; use rmcp::model::Implementation; use rmcp::model::InitializeRequestParams; +use rmcp::model::JsonObject; use rmcp::model::ProtocolVersion; use rmcp::model::Tool as RmcpTool; use tokio_util::sync::CancellationToken; @@ -70,6 +72,7 @@ use tracing::warn; /// MCP server capability indicating that Codex should include [`SandboxState`] /// in tool-call request `_meta` under this key. pub const MCP_SANDBOX_STATE_META_CAPABILITY: &str = "codex/sandbox-state-meta"; +pub const OPENAI_FORM_CAPABILITY: &str = "openai/form"; pub(crate) const MCP_TOOLS_LIST_DURATION_METRIC: &str = "codex.mcp.tools.list.duration_ms"; pub(crate) const MCP_TOOLS_FETCH_UNCACHED_DURATION_METRIC: &str = @@ -151,6 +154,7 @@ impl AsyncManagedClient { runtime_context: McpRuntimeContext, runtime_auth_provider: Option, client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, ) -> Self { let tool_filter = server .configured_config() @@ -204,6 +208,7 @@ impl AsyncManagedClient { elicitation_requests, codex_apps_tools_cache_context, client_elicitation_capability, + supports_openai_form_elicitation, }, ) .await @@ -483,14 +488,12 @@ async fn start_server_task( elicitation_requests, codex_apps_tools_cache_context, client_elicitation_capability, + supports_openai_form_elicitation, } = params; - let mut capabilities = ClientCapabilities::default(); - capabilities.elicitation = Some(client_elicitation_capability); - let params = InitializeRequestParams::new( - capabilities, - Implementation::new("codex-mcp-client", env!("CARGO_PKG_VERSION")).with_title("Codex"), - ) - .with_protocol_version(ProtocolVersion::V_2025_06_18); + let params = mcp_initialize_request_params( + client_elicitation_capability, + supports_openai_form_elicitation, + ); let send_elicitation = elicitation_requests.make_sender(server_name.clone(), tx_event); @@ -550,6 +553,25 @@ async fn start_server_task( Ok(managed) } +fn mcp_initialize_request_params( + client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, +) -> InitializeRequestParams { + let mut capabilities = ClientCapabilities::default(); + capabilities.elicitation = Some(client_elicitation_capability); + if supports_openai_form_elicitation { + capabilities.extensions = Some(BTreeMap::from([( + OPENAI_FORM_CAPABILITY.to_string(), + JsonObject::new(), + )])); + } + InitializeRequestParams::new( + capabilities, + Implementation::new("codex-mcp-client", env!("CARGO_PKG_VERSION")).with_title("Codex"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18) +} + fn mcp_server_info_from_implementation(server_info: Implementation) -> McpServerInfo { McpServerInfo { name: server_info.name, @@ -574,6 +596,7 @@ struct StartServerTaskParams { elicitation_requests: ElicitationRequestManager, codex_apps_tools_cache_context: Option, client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, } async fn make_rmcp_client( @@ -668,6 +691,27 @@ mod tests { use rmcp::model::JsonObject; use rmcp::model::Meta; + #[test] + fn mcp_initialize_advertises_openai_form_only_when_supported() { + let unsupported = mcp_initialize_request_params( + ElicitationCapability::default(), + /*supports_openai_form_elicitation*/ false, + ); + assert_eq!(unsupported.capabilities.extensions, None); + + let supported = mcp_initialize_request_params( + ElicitationCapability::default(), + /*supports_openai_form_elicitation*/ true, + ); + assert_eq!( + supported.capabilities.extensions, + Some(BTreeMap::from([( + OPENAI_FORM_CAPABILITY.to_string(), + JsonObject::new(), + )])) + ); + } + fn tool_with_connector_meta() -> RmcpTool { RmcpTool::new( "capture_file_upload", diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index b3f0d6c43c..b7b1fe7580 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -113,6 +113,10 @@ pub(crate) async fn run_codex_thread_interactive( parent_trace: None, environment_selections: parent_ctx.environments.to_selections(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), + supports_openai_form_elicitation: parent_session + .services + .supports_openai_form_elicitation + .load(std::sync::atomic::Ordering::Relaxed), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 1df188fc31..8c44857297 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -335,6 +335,13 @@ impl CodexThread { .await } + pub async fn set_openai_form_elicitation_support(&self, supported: bool) -> anyhow::Result<()> { + self.codex + .session + .set_openai_form_elicitation_support(supported) + .await + } + /// Preview persistent thread settings overrides without committing them. pub async fn preview_thread_settings_overrides( &self, diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 496f14e86a..6d73015595 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -291,6 +291,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( host_owned_codex_apps_enabled, mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, + /*supports_openai_form_elicitation*/ false, ToolPluginProvenance::default(), auth.as_ref(), /*elicitation_reviewer*/ None, diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 4249c0d4b0..24abfd9e09 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1314,6 +1314,7 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: /*host_owned_codex_apps_enabled*/ true, turn_context.config.prefix_mcp_tool_names(), rmcp::model::ElicitationCapability::default(), + /*supports_openai_form_elicitation*/ false, codex_mcp::ToolPluginProvenance::default(), auth.as_ref(), /*elicitation_reviewer*/ None, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index d1b8982223..d03576aadc 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -16,7 +16,7 @@ use codex_protocol::mcp_approval_meta::TOOL_DESCRIPTION_KEY as MCP_ELICITATION_T use codex_protocol::mcp_approval_meta::TOOL_NAME_KEY as MCP_ELICITATION_TOOL_NAME_KEY; use codex_protocol::mcp_approval_meta::TOOL_PARAMS_KEY as MCP_ELICITATION_TOOL_PARAMS_KEY; use codex_protocol::mcp_approval_meta::TOOL_TITLE_KEY as MCP_ELICITATION_TOOL_TITLE_KEY; -use rmcp::model::CreateElicitationRequestParams; +use codex_rmcp_client::Elicitation; use rmcp::model::ElicitationAction; use rmcp::model::Meta; use serde_json::Map; @@ -143,6 +143,15 @@ impl Session { requested_schema, } } + McpServerElicitationRequest::OpenAiForm { + meta, + message, + requested_schema, + } => codex_protocol::approvals::ElicitationRequest::OpenAiForm { + meta, + message, + requested_schema, + }, McpServerElicitationRequest::Url { meta, message, @@ -353,6 +362,9 @@ impl Session { host_owned_codex_apps_enabled, mcp_config.prefix_mcp_tool_names, mcp_config.client_elicitation_capability, + self.services + .supports_openai_form_elicitation + .load(std::sync::atomic::Ordering::Relaxed), tool_plugin_provenance, auth.as_ref(), elicitation_reviewer, @@ -419,6 +431,34 @@ impl Session { .await; } + pub(crate) async fn set_openai_form_elicitation_support( + &self, + supported: bool, + ) -> anyhow::Result<()> { + if self + .services + .supports_openai_form_elicitation + .load(std::sync::atomic::Ordering::Relaxed) + == supported + { + return Ok(()); + } + + let config = self.get_config().await; + let refresh_config = McpServerRefreshConfig { + mcp_servers: serde_json::to_value(config.mcp_servers.get())?, + mcp_oauth_credentials_store_mode: serde_json::to_value( + config.mcp_oauth_credentials_store_mode, + )?, + auth_keyring_backend_kind: serde_json::to_value(config.auth_keyring_backend_kind())?, + }; + self.services + .supports_openai_form_elicitation + .store(supported, std::sync::atomic::Ordering::Relaxed); + *self.pending_mcp_server_refresh_config.lock().await = Some(refresh_config); + Ok(()) + } + pub(crate) async fn refresh_mcp_servers_now( &self, turn_context: &TurnContext, @@ -510,12 +550,15 @@ fn guardian_elicitation_review_request( request: &ElicitationReviewRequest, ) -> GuardianElicitationReview { let (meta, requested_schema) = match &request.elicitation { - CreateElicitationRequestParams::FormElicitationParams { + Elicitation::Mcp(rmcp::model::CreateElicitationRequestParams::FormElicitationParams { meta, requested_schema, .. - } => (meta, Some(requested_schema)), - CreateElicitationRequestParams::UrlElicitationParams { meta, .. } => { + }) => (meta, Some(requested_schema)), + Elicitation::Mcp(rmcp::model::CreateElicitationRequestParams::UrlElicitationParams { + meta, + .. + }) => { return if meta_requests_approval_request(meta) { GuardianElicitationReview::Decline( "guardian MCP elicitation review only supports form elicitations", @@ -524,6 +567,7 @@ fn guardian_elicitation_review_request( GuardianElicitationReview::NotRequested }; } + Elicitation::OpenAiForm { .. } => return GuardianElicitationReview::NotRequested, }; let Some(meta) = meta.as_ref().map(|meta| &meta.0) else { @@ -585,13 +629,10 @@ fn guardian_elicitation_review_request( )) } -fn elicitation_connector_id(elicitation: &CreateElicitationRequestParams) -> Option<&str> { - match elicitation { - CreateElicitationRequestParams::FormElicitationParams { meta, .. } - | CreateElicitationRequestParams::UrlElicitationParams { meta, .. } => meta - .as_ref() - .and_then(|meta| metadata_str(&meta.0, MCP_ELICITATION_CONNECTOR_ID_KEY)), - } +fn elicitation_connector_id(elicitation: &Elicitation) -> Option<&str> { + elicitation + .meta() + .and_then(|meta| metadata_str(meta, MCP_ELICITATION_CONNECTOR_ID_KEY)) } fn meta_requests_approval_request(meta: &Option) -> bool { diff --git a/codex-rs/core/src/session/mcp_tests.rs b/codex-rs/core/src/session/mcp_tests.rs index 0e3664019b..e5d0184bb7 100644 --- a/codex-rs/core/src/session/mcp_tests.rs +++ b/codex-rs/core/src/session/mcp_tests.rs @@ -30,13 +30,15 @@ fn form_request(meta: Option) -> ElicitationReviewRequest { ElicitationReviewRequest { server_name: "browser-use".to_string(), request_id: rmcp::model::NumberOrString::Number(7), - elicitation: CreateElicitationRequestParams::FormElicitationParams { - meta, - message: "Allow origin?".to_string(), - requested_schema: ElicitationSchema::builder() - .build() - .expect("schema should build"), - }, + elicitation: Elicitation::Mcp( + rmcp::model::CreateElicitationRequestParams::FormElicitationParams { + meta, + message: "Allow origin?".to_string(), + requested_schema: ElicitationSchema::builder() + .build() + .expect("schema should build"), + }, + ), } } @@ -171,12 +173,14 @@ fn guardian_elicitation_review_request_declines_unsupported_opt_in_shapes() { let url_request = ElicitationReviewRequest { server_name: "browser-use".to_string(), request_id: rmcp::model::NumberOrString::Number(8), - elicitation: CreateElicitationRequestParams::UrlElicitationParams { - meta: guardian_meta(Some(json!({}))), - message: "Open URL".to_string(), - url: "https://example.com".to_string(), - elicitation_id: "elicit-1".to_string(), - }, + elicitation: Elicitation::Mcp( + rmcp::model::CreateElicitationRequestParams::UrlElicitationParams { + meta: guardian_meta(Some(json!({}))), + message: "Open URL".to_string(), + url: "https://example.com".to_string(), + elicitation_id: "elicit-1".to_string(), + }, + ), }; assert!(matches!( guardian_elicitation_review_request(&url_request), @@ -186,14 +190,16 @@ fn guardian_elicitation_review_request_declines_unsupported_opt_in_shapes() { let non_empty_schema_request = ElicitationReviewRequest { server_name: "browser-use".to_string(), request_id: rmcp::model::NumberOrString::Number(9), - elicitation: CreateElicitationRequestParams::FormElicitationParams { - meta: guardian_meta(Some(json!({}))), - message: "Allow origin?".to_string(), - requested_schema: ElicitationSchema::builder() - .required_property("confirmed", PrimitiveSchema::Boolean(BooleanSchema::new())) - .build() - .expect("schema should build"), - }, + elicitation: Elicitation::Mcp( + rmcp::model::CreateElicitationRequestParams::FormElicitationParams { + meta: guardian_meta(Some(json!({}))), + message: "Allow origin?".to_string(), + requested_schema: ElicitationSchema::builder() + .required_property("confirmed", PrimitiveSchema::Boolean(BooleanSchema::new())) + .build() + .expect("schema should build"), + }, + ), }; assert!(matches!( guardian_elicitation_review_request(&non_empty_schema_request), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5aa56b68f8..f92bd63680 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -435,6 +435,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) parent_trace: Option, pub(crate) environment_selections: Vec, pub(crate) thread_extension_init: ExtensionDataInit, + pub(crate) supports_openai_form_elicitation: bool, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, @@ -517,6 +518,7 @@ impl Codex { parent_trace: _, environment_selections, thread_extension_init, + supports_openai_form_elicitation, analytics_events_client, thread_store, attestation_provider, @@ -659,6 +661,7 @@ impl Codex { mcp_manager.clone(), extensions, thread_extension_init, + supports_openai_form_elicitation, agent_control, environment_manager, inherited_environments, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index c4bc064a47..2b61468c15 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -484,6 +484,7 @@ impl Session { mcp_manager: Arc, extensions: Arc>, thread_extension_init: ExtensionDataInit, + supports_openai_form_elicitation: bool, agent_control: AgentControl, environment_manager: Arc, inherited_environments: Option, @@ -1009,6 +1010,9 @@ impl Session { session_extension_data, thread_extension_data, mcp_thread_init, + supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new( + supports_openai_form_elicitation, + ), agent_control, network_proxy: arc_swap::ArcSwapOption::from(network_proxy.map(Arc::new)), network_proxy_audit_metadata, @@ -1149,6 +1153,9 @@ impl Session { host_owned_codex_apps_enabled, config.prefix_mcp_tool_names(), client_elicitation_capability, + sess.services + .supports_openai_form_elicitation + .load(std::sync::atomic::Ordering::Relaxed), tool_plugin_provenance, auth, Some(sess.mcp_elicitation_reviewer()), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 63e8783cc6..b13cd029ce 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4856,6 +4856,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, AgentControl::default(), environment_manager, /*inherited_environments*/ None, @@ -5018,6 +5019,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), + supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new(false), agent_control, network_proxy: arc_swap::ArcSwapOption::from(None), network_proxy_audit_metadata: crate::config::NetworkProxyAuditMetadata::default(), @@ -5203,6 +5205,7 @@ async fn make_session_with_config_and_rx( mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, AgentControl::default(), environment_manager, /*inherited_environments*/ None, @@ -5307,6 +5310,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, agent_control, environment_manager, /*inherited_environments*/ None, @@ -7061,6 +7065,7 @@ where ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), + supports_openai_form_elicitation: std::sync::atomic::AtomicBool::new(false), agent_control, network_proxy: arc_swap::ArcSwapOption::from(None), network_proxy_audit_metadata: crate::config::NetworkProxyAuditMetadata::default(), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index b9727a0401..0bed032351 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -732,6 +732,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { parent_trace: None, environment_selections: Vec::new(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), + supports_openai_form_elicitation: false, analytics_events_client: None, thread_store, attestation_provider: None, diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 9724d8b3ad..914fd4c32d 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use crate::SkillsService; use crate::agent::AgentControl; @@ -67,6 +68,7 @@ pub(crate) struct SessionServices { pub(crate) extensions: Arc>, pub(crate) session_extension_data: ExtensionData, pub(crate) thread_extension_data: ExtensionData, + pub(crate) supports_openai_form_elicitation: AtomicBool, pub(crate) mcp_thread_init: ExtensionDataInit, pub(crate) agent_control: AgentControl, pub(crate) network_proxy: ArcSwapOption, diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index 60e80ca290..1db714fada 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -112,9 +112,14 @@ pub async fn start_thread_with_user_shell_override( thread_manager: &ThreadManager, config: Config, user_shell_override: crate::shell::Shell, + supports_openai_form_elicitation: bool, ) -> codex_protocol::error::Result { thread_manager - .start_thread_with_user_shell_override_for_tests(config, user_shell_override) + .start_thread_with_user_shell_override_for_tests( + config, + user_shell_override, + supports_openai_form_elicitation, + ) .await } @@ -124,6 +129,7 @@ pub async fn resume_thread_from_rollout_with_user_shell_override( rollout_path: PathBuf, auth_manager: Arc, user_shell_override: crate::shell::Shell, + supports_openai_form_elicitation: bool, ) -> codex_protocol::error::Result { thread_manager .resume_thread_from_rollout_with_user_shell_override_for_tests( @@ -131,6 +137,7 @@ pub async fn resume_thread_from_rollout_with_user_shell_override( rollout_path, auth_manager, user_shell_override, + supports_openai_form_elicitation, ) .await } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b452dc45ab..77a5520026 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -186,6 +186,7 @@ pub struct StartThreadOptions { pub parent_trace: Option, pub environments: Vec, pub thread_extension_init: ExtensionDataInit, + pub supports_openai_form_elicitation: bool, } pub(crate) struct ResumeThreadWithHistoryOptions { @@ -604,6 +605,7 @@ impl ThreadManager { parent_trace: None, environments, thread_extension_init: ExtensionDataInit::default(), + supports_openai_form_elicitation: false, })) .await } @@ -644,6 +646,7 @@ impl ThreadManager { options.parent_trace, options.environments, options.thread_extension_init, + options.supports_openai_form_elicitation, /*user_shell_override*/ None, )) .await @@ -692,6 +695,7 @@ impl ThreadManager { rollout_path: PathBuf, auth_manager: Arc, parent_trace: Option, + supports_openai_form_elicitation: bool, ) -> CodexResult { let initial_history = self.initial_history_from_rollout_path(rollout_path).await?; Box::pin(self.resume_thread_with_history( @@ -699,6 +703,7 @@ impl ThreadManager { initial_history, auth_manager, parent_trace, + supports_openai_form_elicitation, )) .await } @@ -710,6 +715,7 @@ impl ThreadManager { initial_history: InitialHistory, auth_manager: Arc, parent_trace: Option, + supports_openai_form_elicitation: bool, ) -> CodexResult { let agent_control = self.agent_control_for_config(&config); let environments = default_thread_environment_selections( @@ -735,6 +741,7 @@ impl ThreadManager { parent_trace, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, /*user_shell_override*/ None, )) .await @@ -744,6 +751,7 @@ impl ThreadManager { &self, config: Config, user_shell_override: crate::shell::Shell, + supports_openai_form_elicitation: bool, ) -> CodexResult { let agent_control = self.agent_control_for_config(&config); let environments = default_thread_environment_selections( @@ -763,6 +771,7 @@ impl ThreadManager { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, /*user_shell_override*/ Some(user_shell_override), )) .await @@ -774,6 +783,7 @@ impl ThreadManager { rollout_path: PathBuf, auth_manager: Arc, user_shell_override: crate::shell::Shell, + supports_openai_form_elicitation: bool, ) -> CodexResult { let agent_control = self.agent_control_for_config(&config); let initial_history = self.initial_history_from_rollout_path(rollout_path).await?; @@ -800,6 +810,7 @@ impl ThreadManager { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, /*user_shell_override*/ Some(user_shell_override), )) .await @@ -880,8 +891,15 @@ impl ThreadManager { { let snapshot = snapshot.into(); let history = self.initial_history_from_rollout_path(path).await?; - self.fork_thread_from_history(snapshot, config, history, thread_source, parent_trace) - .await + self.fork_thread_from_history( + snapshot, + config, + history, + thread_source, + parent_trace, + /*supports_openai_form_elicitation*/ false, + ) + .await } async fn initial_history_from_rollout_path( @@ -910,6 +928,7 @@ impl ThreadManager { history: InitialHistory, thread_source: Option, parent_trace: Option, + supports_openai_form_elicitation: bool, ) -> CodexResult where S: Into, @@ -920,6 +939,7 @@ impl ThreadManager { history, thread_source, parent_trace, + supports_openai_form_elicitation, ) .await } @@ -931,6 +951,7 @@ impl ThreadManager { history: InitialHistory, thread_source: Option, parent_trace: Option, + supports_openai_form_elicitation: bool, ) -> CodexResult { // `forked_from_id()` describes this history's existing lineage. When // forking a resumed thread, the child copies the resumed thread itself. @@ -970,6 +991,7 @@ impl ThreadManager { parent_trace, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, /*user_shell_override*/ None, )) .await @@ -1230,6 +1252,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, )) .await @@ -1267,6 +1290,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, )) .await @@ -1305,6 +1329,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, )) .await @@ -1326,6 +1351,7 @@ impl ThreadManagerState { parent_trace: Option, environments: Vec, thread_extension_init: ExtensionDataInit, + supports_openai_form_elicitation: bool, user_shell_override: Option, ) -> CodexResult { Box::pin(self.spawn_thread_with_source( @@ -1344,6 +1370,7 @@ impl ThreadManagerState { parent_trace, environments, thread_extension_init, + supports_openai_form_elicitation, user_shell_override, )) .await @@ -1367,6 +1394,7 @@ impl ThreadManagerState { parent_trace: Option, environments: Vec, thread_extension_init: ExtensionDataInit, + supports_openai_form_elicitation: bool, user_shell_override: Option, ) -> CodexResult { let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); @@ -1434,6 +1462,7 @@ impl ThreadManagerState { parent_trace, environment_selections: environments, thread_extension_init, + supports_openai_form_elicitation, analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index be952bbcbc..ee010b5e1d 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -325,6 +325,7 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { parent_trace: None, environments: Vec::new(), thread_extension_init: Default::default(), + supports_openai_form_elicitation: false, }) .await .expect("internal thread should start"); @@ -463,6 +464,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() parent_trace: None, environments: Vec::new(), thread_extension_init: selected_root_init("selected-a", "env-a"), + supports_openai_form_elicitation: false, }) .await .expect("start first thread"); @@ -477,6 +479,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() parent_trace: None, environments: Vec::new(), thread_extension_init: selected_root_init("selected-b", "env-b"), + supports_openai_form_elicitation: false, }) .await .expect("start second thread"); @@ -567,6 +570,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { parent_trace: None, environments: environments.clone(), thread_extension_init: Default::default(), + supports_openai_form_elicitation: false, }) .await .expect("start source thread"); @@ -593,6 +597,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { rollout_path.clone(), auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("resume source thread"); @@ -729,6 +734,7 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { rollout_path, auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("resume active source thread"); @@ -792,6 +798,7 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { rollout_path, auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("resume stopped source thread"); @@ -842,6 +849,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { parent_trace: None, environments: Vec::new(), thread_extension_init: Default::default(), + supports_openai_form_elicitation: false, }) .await .expect("start source thread"); @@ -868,6 +876,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { rollout_path, auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("resume source thread"); @@ -947,6 +956,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { }), auth_manager.clone(), /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("seed rollout path in store"); @@ -963,6 +973,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { rollout_path.clone(), auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("resume from rollout path"); @@ -1254,6 +1265,7 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor ]), auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("create source thread from completed history"); @@ -1368,6 +1380,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { ]), auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("create source thread from explicit partial history"); @@ -1458,6 +1471,7 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ ]), auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("create source thread from partial history"); diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 557164f4de..ed15a7483a 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -2832,6 +2832,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { })]), AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")), /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("start thread"); diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index daf6f79b5c..b10e26bf0c 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -15,6 +15,7 @@ use anyhow::Result; use anyhow::anyhow; use codex_config::CloudConfigBundleLoader; use codex_core::CodexThread; +use codex_core::StartThreadOptions; use codex_core::ThreadManager; use codex_core::config::Config; use codex_core::resolve_installation_id; @@ -39,6 +40,7 @@ use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::Op; use codex_protocol::protocol::RealtimeConversationVersion as RealtimeWsVersion; use codex_protocol::protocol::SandboxPolicy; @@ -258,6 +260,7 @@ pub struct TestCodexBuilder { exec_server_url: Option, extensions: Arc>, user_instructions_provider: Option>, + supports_openai_form_elicitation: bool, } impl TestCodexBuilder { @@ -354,6 +357,11 @@ impl TestCodexBuilder { self } + pub fn with_openai_form_elicitation(mut self) -> Self { + self.supports_openai_form_elicitation = true; + self + } + pub fn with_windows_cmd_shell(self) -> Self { if cfg!(windows) { self.with_user_shell(get_shell_by_model_provided_path(&PathBuf::from("cmd.exe"))) @@ -574,6 +582,7 @@ impl TestCodexBuilder { path, auth_manager, user_shell_override, + self.supports_openai_form_elicitation, ), ) .await? @@ -585,6 +594,7 @@ impl TestCodexBuilder { path, auth_manager, /*parent_trace*/ None, + self.supports_openai_form_elicitation, )) .await? } @@ -594,11 +604,29 @@ impl TestCodexBuilder { thread_manager.as_ref(), config.clone(), user_shell_override, + self.supports_openai_form_elicitation, ), ) .await? } - (None, None) => Box::pin(thread_manager.start_thread(config.clone())).await?, + (None, None) => { + let environments = thread_manager.default_environment_selections(&config.cwd); + Box::pin( + thread_manager.start_thread_with_options(StartThreadOptions { + config: config.clone(), + initial_history: InitialHistory::New, + session_source: None, + thread_source: None, + dynamic_tools: Vec::new(), + metrics_service_name: None, + parent_trace: None, + environments, + thread_extension_init: Default::default(), + supports_openai_form_elicitation: self.supports_openai_form_elicitation, + }), + ) + .await? + } }; Ok(TestCodex { @@ -1143,6 +1171,7 @@ pub fn test_codex() -> TestCodexBuilder { exec_server_url: None, extensions: empty_extension_registry(), user_instructions_provider: None, + supports_openai_form_elicitation: false, } } diff --git a/codex-rs/core/tests/suite/agents_md.rs b/codex-rs/core/tests/suite/agents_md.rs index a4ab404065..5b79e9f1e6 100644 --- a/codex-rs/core/tests/suite/agents_md.rs +++ b/codex-rs/core/tests/suite/agents_md.rs @@ -448,6 +448,7 @@ async fn loads_user_instructions_without_a_primary_environment() -> Result<()> { parent_trace: None, environments: Vec::new(), thread_extension_init: Default::default(), + supports_openai_form_elicitation: false, }) .await?; assert_eq!(provider.load_count(), 2); @@ -664,6 +665,7 @@ async fn multi_environment_thread_loads_every_project_and_keeps_creation_snapsho }, ], thread_extension_init: Default::default(), + supports_openai_form_elicitation: false, }) .await?; assert_eq!(provider.load_count(), 2); diff --git a/codex-rs/core/tests/suite/compact_resume_fork.rs b/codex-rs/core/tests/suite/compact_resume_fork.rs index 909dfb8c82..3270a922b0 100644 --- a/codex-rs/core/tests/suite/compact_resume_fork.rs +++ b/codex-rs/core/tests/suite/compact_resume_fork.rs @@ -830,6 +830,7 @@ async fn resume_conversation( path, auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, )) .await .expect("resume conversation") diff --git a/codex-rs/core/tests/suite/fork_thread.rs b/codex-rs/core/tests/suite/fork_thread.rs index a1fdf79a8e..d51ba7d2f6 100644 --- a/codex-rs/core/tests/suite/fork_thread.rs +++ b/codex-rs/core/tests/suite/fork_thread.rs @@ -201,6 +201,7 @@ async fn fork_thread_from_history_does_not_require_source_rollout_path() { }), /*thread_source*/ None, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("fork from stored history"); diff --git a/codex-rs/core/tests/suite/realtime_conversation.rs b/codex-rs/core/tests/suite/realtime_conversation.rs index dccf460f0b..0788709952 100644 --- a/codex-rs/core/tests/suite/realtime_conversation.rs +++ b/codex-rs/core/tests/suite/realtime_conversation.rs @@ -2267,6 +2267,7 @@ async fn conversation_startup_context_current_thread_selects_many_turns_by_budge InitialHistory::Forked(history), auth_manager_from_auth(CodexAuth::from_api_key("dummy")), /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await?; let codex = resumed_thread.thread; diff --git a/codex-rs/core/tests/suite/resume_warning.rs b/codex-rs/core/tests/suite/resume_warning.rs index 546219ce90..0b9ff204d4 100644 --- a/codex-rs/core/tests/suite/resume_warning.rs +++ b/codex-rs/core/tests/suite/resume_warning.rs @@ -110,6 +110,7 @@ async fn emits_warning_when_resumed_model_differs() { initial_history, auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("resume conversation"); diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 656ca69317..2c64b5fba2 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -50,6 +50,7 @@ use core_test_support::assert_regex_match; use core_test_support::responses; use core_test_support::responses::mount_models_once; use core_test_support::responses::mount_sse_once; +use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::skip_if_wine_exec; use core_test_support::stdio_server_bin; @@ -339,13 +340,23 @@ async fn call_cwd_tool( fixture: &TestCodex, server_name: &str, call_id: &str, +) -> anyhow::Result { + call_structured_tool(server, fixture, server_name, "cwd", call_id).await +} + +async fn call_structured_tool( + server: &MockServer, + fixture: &TestCodex, + server_name: &str, + tool_name: &str, + call_id: &str, ) -> anyhow::Result { let namespace = format!("mcp__{server_name}"); mount_sse_once( server, responses::sse(vec![ responses::ev_response_created("resp-1"), - responses::ev_function_call_with_namespace(call_id, &namespace, "cwd", r#"{}"#), + responses::ev_function_call_with_namespace(call_id, &namespace, tool_name, r#"{}"#), responses::ev_completed("resp-1"), ]), ) @@ -353,7 +364,7 @@ async fn call_cwd_tool( mount_sse_once( server, responses::sse(vec![ - responses::ev_assistant_message("msg-1", "rmcp cwd tool completed successfully."), + responses::ev_assistant_message("msg-1", "rmcp tool completed successfully."), responses::ev_completed("resp-2"), ]), ) @@ -361,7 +372,7 @@ async fn call_cwd_tool( fixture .codex - .submit(read_only_user_turn(fixture, "call the rmcp cwd tool")) + .submit(read_only_user_turn(fixture, "call the requested rmcp tool")) .await?; wait_for_event(&fixture.codex, |ev| { @@ -378,7 +389,7 @@ async fn call_cwd_tool( let structured_content = end .result .as_ref() - .expect("rmcp cwd tool should return success") + .expect("rmcp tool should return success") .structured_content .as_ref() .expect("structured content") @@ -388,6 +399,106 @@ async fn call_cwd_tool( Ok(structured_content) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn openai_form_capability_is_advertised_to_mcp_servers() -> anyhow::Result<()> { + assert_openai_form_capability_advertisement(/*expected*/ true).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn openai_form_capability_is_not_advertised_by_default() -> anyhow::Result<()> { + assert_openai_form_capability_advertisement(/*expected*/ false).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn openai_form_capability_updates_for_loaded_thread() -> anyhow::Result<()> { + skip_if_wine_exec!( + Ok(()), + "requires a Windows test_stdio_server in the Wine-exec environment" + ); + + let server = start_mock_server().await; + let server_name = "capabilities"; + let command = stdio_server_bin()?; + let fixture = test_codex() + .with_config(move |config| { + insert_mcp_server( + config, + server_name, + stdio_transport(command, /*env*/ None, Vec::new()), + TestMcpServerOptions::default(), + ); + }) + .build(&server) + .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + + let unsupported = call_structured_tool( + &server, + &fixture, + server_name, + "client_capabilities", + "call-client-capabilities-unsupported", + ) + .await?; + assert_eq!( + unsupported, + json!({ "supportsOpenaiFormElicitation": false }) + ); + + fixture + .codex + .set_openai_form_elicitation_support(/*supported*/ true) + .await?; + let supported = call_structured_tool( + &server, + &fixture, + server_name, + "client_capabilities", + "call-client-capabilities-supported", + ) + .await?; + assert_eq!(supported, json!({ "supportsOpenaiFormElicitation": true })); + Ok(()) +} + +async fn assert_openai_form_capability_advertisement(expected: bool) -> anyhow::Result<()> { + skip_if_wine_exec!( + Ok(()), + "requires a Windows test_stdio_server in the Wine-exec environment" + ); + + let server = start_mock_server().await; + let server_name = "capabilities"; + let command = stdio_server_bin()?; + let mut builder = test_codex().with_config(move |config| { + insert_mcp_server( + config, + server_name, + stdio_transport(command, /*env*/ None, Vec::new()), + TestMcpServerOptions::default(), + ); + }); + if expected { + builder = builder.with_openai_form_elicitation(); + } + let fixture = builder.build(&server).await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + + let structured = call_structured_tool( + &server, + &fixture, + server_name, + "client_capabilities", + "call-client-capabilities", + ) + .await?; + assert_eq!( + structured, + json!({ "supportsOpenaiFormElicitation": expected }) + ); + Ok(()) +} + fn assert_cwd_tool_output(structured: &Value, expected_cwd: &Path) { let actual_cwd = structured .get("cwd") diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 2501aa0dde..a31d7e3616 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -754,6 +754,7 @@ async fn subagent_stop_replaces_stop_and_skips_internal_subagents() -> Result<() parent_trace: None, environments: Vec::new(), thread_extension_init: Default::default(), + supports_openai_form_elicitation: false, }) .await?; diff --git a/codex-rs/core/tests/suite/unstable_features_warning.rs b/codex-rs/core/tests/suite/unstable_features_warning.rs index e66c674f92..783acc638b 100644 --- a/codex-rs/core/tests/suite/unstable_features_warning.rs +++ b/codex-rs/core/tests/suite/unstable_features_warning.rs @@ -47,6 +47,7 @@ async fn emits_warning_when_unstable_features_enabled_via_config() { InitialHistory::New, auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("spawn conversation"); @@ -93,6 +94,7 @@ async fn suppresses_warning_when_configured() { InitialHistory::New, auth_manager, /*parent_trace*/ None, + /*supports_openai_form_elicitation*/ false, ) .await .expect("spawn conversation"); diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 05bd5cc1df..a1f4b95b50 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -556,6 +556,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result client_name: "codex_exec".to_string(), client_version: env!("CARGO_PKG_VERSION").to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, }; diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 8dae6ca617..1d8fcb567b 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -313,6 +313,7 @@ impl MemoryStartupContext { parent_trace: None, environments, thread_extension_init: Default::default(), + supports_openai_form_elicitation: false, }) .await?; diff --git a/codex-rs/protocol/src/approvals.rs b/codex-rs/protocol/src/approvals.rs index 314f6bd923..5cf4f27f18 100644 --- a/codex-rs/protocol/src/approvals.rs +++ b/codex-rs/protocol/src/approvals.rs @@ -342,6 +342,15 @@ pub enum ElicitationRequest { message: String, requested_schema: JsonValue, }, + #[serde(rename = "openai/form")] + #[ts(rename = "openai/form")] + OpenAiForm { + #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] + #[ts(optional, rename = "_meta")] + meta: Option, + message: String, + requested_schema: JsonValue, + }, Url { #[serde(rename = "_meta", default, skip_serializing_if = "Option::is_none")] #[ts(optional, rename = "_meta")] @@ -355,7 +364,9 @@ pub enum ElicitationRequest { impl ElicitationRequest { pub fn message(&self) -> &str { match self { - Self::Form { message, .. } | Self::Url { message, .. } => message, + Self::Form { message, .. } + | Self::OpenAiForm { message, .. } + | Self::Url { message, .. } => message, } } } diff --git a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs index 4407a27eeb..9ff36e349d 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -4,6 +4,8 @@ use std::collections::HashMap; use std::collections::hash_map::Entry; use std::sync::Arc; use std::sync::OnceLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use rmcp::ErrorData as McpError; @@ -11,6 +13,8 @@ use rmcp::ServiceExt; use rmcp::handler::server::ServerHandler; use rmcp::model::CallToolRequestParams; use rmcp::model::CallToolResult; +use rmcp::model::InitializeRequestParams; +use rmcp::model::InitializeResult; use rmcp::model::JsonObject; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; @@ -38,6 +42,7 @@ struct TestToolServer { tools: Arc>, resources: Arc>, resource_templates: Arc>, + supports_openai_form_elicitation: Arc, } const MEMO_URI: &str = "memo://codex/example-note"; @@ -68,6 +73,7 @@ impl TestToolServer { let tools = vec![ Self::echo_tool(), Self::echo_dash_tool(), + Self::client_capabilities_tool(), Self::cwd_tool(), Self::sync_tool(), Self::sync_readonly_tool(), @@ -81,6 +87,7 @@ impl TestToolServer { tools: Arc::new(tools), resources: Arc::new(resources), resource_templates: Arc::new(resource_templates), + supports_openai_form_elicitation: Arc::new(AtomicBool::new(false)), } } @@ -166,6 +173,24 @@ impl TestToolServer { tool } + fn client_capabilities_tool() -> Tool { + #[expect(clippy::expect_used)] + let schema: JsonObject = serde_json::from_value(serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": false + })) + .expect("client capabilities tool schema should deserialize"); + + let mut tool = Tool::new( + Cow::Borrowed("client_capabilities"), + Cow::Borrowed("Return capabilities advertised by the MCP client."), + Arc::new(schema), + ); + tool.annotations = Some(ToolAnnotations::new().read_only(true)); + tool + } + fn sync_tool() -> Tool { #[expect(clippy::expect_used)] let schema: JsonObject = serde_json::from_value(json!({ @@ -396,6 +421,23 @@ struct ImageScenarioArgs { } impl ServerHandler for TestToolServer { + async fn initialize( + &self, + request: InitializeRequestParams, + context: rmcp::service::RequestContext, + ) -> Result { + self.supports_openai_form_elicitation.store( + request + .capabilities + .extensions + .as_ref() + .is_some_and(|extensions| extensions.contains_key("openai/form")), + Ordering::Relaxed, + ); + context.peer.set_peer_info(request); + Ok(self.get_info()) + } + fn get_info(&self) -> ServerInfo { let mut capabilities = ServerCapabilities::builder() .enable_tools() @@ -481,6 +523,11 @@ impl ServerHandler for TestToolServer { context: rmcp::service::RequestContext, ) -> Result { match request.name.as_ref() { + "client_capabilities" => Ok(Self::structured_result(json!({ + "supportsOpenaiFormElicitation": self + .supports_openai_form_elicitation + .load(Ordering::Relaxed), + }))), "sandbox_meta" => Ok(Self::structured_result(serde_json::Value::Object( context.meta.0, ))), diff --git a/codex-rs/rmcp-client/src/elicitation_client_service.rs b/codex-rs/rmcp-client/src/elicitation_client_service.rs index 49f11f0a76..2227ee4d63 100644 --- a/codex-rs/rmcp-client/src/elicitation_client_service.rs +++ b/codex-rs/rmcp-client/src/elicitation_client_service.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use rmcp::RoleClient; use rmcp::model::ClientInfo; use rmcp::model::ClientResult; +use rmcp::model::CustomRequest; use rmcp::model::CustomResult; use rmcp::model::ElicitationAction; use rmcp::model::Meta; @@ -12,7 +13,9 @@ use rmcp::model::ServerRequest; use rmcp::service::NotificationContext; use rmcp::service::RequestContext; use rmcp::service::Service; +use serde::Deserialize; use serde::Serialize; +use serde_json::Map; use serde_json::Value; use crate::logging_client_handler::LoggingClientHandler; @@ -22,10 +25,21 @@ use crate::rmcp_client::ElicitationResponse; use crate::rmcp_client::SendElicitation; const MCP_PROGRESS_TOKEN_META_KEY: &str = "progressToken"; +const OPENAI_FORM_METHOD: &str = "openai/form"; + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct OpenAiFormRequestParams { + #[serde(rename = "_meta")] + meta: Option, + message: String, + requested_schema: Value, +} #[derive(Clone)] pub(crate) struct ElicitationClientService { handler: LoggingClientHandler, + supports_openai_form: bool, send_elicitation: Arc, pause_state: ElicitationPauseState, } @@ -36,12 +50,18 @@ impl ElicitationClientService { send_elicitation: SendElicitation, pause_state: ElicitationPauseState, ) -> Self { + let supports_openai_form = client_info + .capabilities + .extensions + .as_ref() + .is_some_and(|extensions| extensions.contains_key(OPENAI_FORM_METHOD)); let send_elicitation = Arc::new(send_elicitation); Self { handler: LoggingClientHandler::new( client_info, clone_send_elicitation(Arc::clone(&send_elicitation)), ), + supports_openai_form, send_elicitation, pause_state, } @@ -73,11 +93,23 @@ impl Service for ElicitationClientService { ) -> Result { match request { ServerRequest::CreateElicitationRequest(request) => { - let response = self.create_elicitation(request.params, context).await?; + let response = self + .create_elicitation(Elicitation::Mcp(request.params), context) + .await?; // RMCP's typed CreateElicitationResult does not model result-level `_meta`. let result = elicitation_response_result(response)?; Ok(ClientResult::CustomResult(result)) } + ServerRequest::CustomRequest(request) + if request.method == OPENAI_FORM_METHOD && self.supports_openai_form => + { + let response = self + .create_elicitation(openai_form_elicitation(request)?, context) + .await?; + Ok(ClientResult::CustomResult(elicitation_response_result( + response, + )?)) + } request => { >::handle_request( &self.handler, @@ -107,6 +139,18 @@ impl Service for ElicitationClientService { } } +fn openai_form_elicitation(request: CustomRequest) -> Result { + let params = request + .params_as::() + .map_err(|err| rmcp::ErrorData::invalid_params(err.to_string(), None))? + .ok_or_else(|| rmcp::ErrorData::invalid_params("missing params", None))?; + Ok(Elicitation::OpenAiForm { + meta: params.meta, + message: params.message, + requested_schema: params.requested_schema, + }) +} + fn restore_context_meta(mut request: Elicitation, mut context_meta: Meta) -> Elicitation { // RMCP lifts JSON-RPC `_meta` into RequestContext before invoking services. context_meta.remove(MCP_PROGRESS_TOKEN_META_KEY); @@ -114,10 +158,20 @@ fn restore_context_meta(mut request: Elicitation, mut context_meta: Meta) -> Eli return request; } - request - .meta_mut() - .get_or_insert_with(Meta::new) - .extend(context_meta); + match &mut request { + Elicitation::Mcp(request) => request + .meta_mut() + .get_or_insert_with(Meta::new) + .extend(context_meta), + Elicitation::OpenAiForm { meta, .. } => { + let meta = meta + .get_or_insert_with(|| Value::Object(Map::new())) + .as_object_mut(); + if let Some(meta) = meta { + meta.extend(context_meta.0); + } + } + } request } @@ -165,7 +219,7 @@ mod tests { #[test] fn restore_context_meta_adds_elicitation_meta_and_removes_progress_token() { let request = restore_context_meta( - form_request(/*meta*/ None), + Elicitation::Mcp(form_request(/*meta*/ None)), meta(json!({ "progressToken": "progress-token", "persist": ["session", "always"], @@ -174,9 +228,54 @@ mod tests { assert_eq!( request, - form_request(Some(meta(json!({ + Elicitation::Mcp(form_request(Some(meta(json!({ "persist": ["session", "always"], - })))) + }))))) + ); + } + + #[test] + fn parses_openai_form_custom_requests() { + let elicitation = openai_form_elicitation(CustomRequest::new( + OPENAI_FORM_METHOD, + Some(json!({ + "message": "Select a template", + "requestedSchema": { + "type": "object", + "properties": { + "template": { + "type": "openai/imagePicker", + "items": [{ + "id": "monthly-review", + "title": "Monthly review", + "image": "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4=" + }] + } + } + } + })), + )) + .expect("valid openai/form request"); + + assert_eq!( + elicitation, + Elicitation::OpenAiForm { + meta: None, + message: "Select a template".to_string(), + requested_schema: json!({ + "type": "object", + "properties": { + "template": { + "type": "openai/imagePicker", + "items": [{ + "id": "monthly-review", + "title": "Monthly review", + "image": "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4=" + }] + } + } + }), + } ); } diff --git a/codex-rs/rmcp-client/src/logging_client_handler.rs b/codex-rs/rmcp-client/src/logging_client_handler.rs index 0c3da0fe2c..e575966cff 100644 --- a/codex-rs/rmcp-client/src/logging_client_handler.rs +++ b/codex-rs/rmcp-client/src/logging_client_handler.rs @@ -17,6 +17,7 @@ use tracing::error; use tracing::info; use tracing::warn; +use crate::rmcp_client::Elicitation; use crate::rmcp_client::SendElicitation; #[derive(Clone)] @@ -40,7 +41,7 @@ impl ClientHandler for LoggingClientHandler { request: CreateElicitationRequestParams, context: RequestContext, ) -> Result { - (self.send_elicitation)(context.id, request) + (self.send_elicitation)(context.id, Elicitation::Mcp(request)) .await .map(Into::into) .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None)) diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index e9b0148622..4a9b89de3c 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -40,6 +40,7 @@ use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; +use rmcp::model::RequestParamsMeta; use rmcp::model::ServerResult; use rmcp::model::Tool; use rmcp::service::RoleClient; @@ -251,7 +252,24 @@ fn remaining_operation_timeout( } } -pub type Elicitation = CreateElicitationRequestParams; +#[derive(Debug, Clone, PartialEq)] +pub enum Elicitation { + Mcp(CreateElicitationRequestParams), + OpenAiForm { + meta: Option, + message: String, + requested_schema: serde_json::Value, + }, +} + +impl Elicitation { + pub fn meta(&self) -> Option<&serde_json::Map> { + match self { + Self::Mcp(request) => request.meta().map(|meta| &meta.0), + Self::OpenAiForm { meta, .. } => meta.as_ref().and_then(serde_json::Value::as_object), + } + } +} #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index f436ac51b3..76cf56ef06 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -292,7 +292,10 @@ impl App { message: message.clone(), }, )), - codex_app_server_protocol::McpServerElicitationRequest::Url { .. } => { + codex_app_server_protocol::McpServerElicitationRequest::OpenAiForm { + .. + } + | codex_app_server_protocol::McpServerElicitationRequest::Url { .. } => { self.app_event_tx.resolve_elicitation( thread_id, params.server_name.clone(), diff --git a/codex-rs/tui/src/chatwidget/tool_requests.rs b/codex-rs/tui/src/chatwidget/tool_requests.rs index 99b0c65207..f3ff2391f6 100644 --- a/codex-rs/tui/src/chatwidget/tool_requests.rs +++ b/codex-rs/tui/src/chatwidget/tool_requests.rs @@ -376,7 +376,8 @@ impl ChatWidget { self.bottom_pane .push_approval_request(request, &self.config.features); } - McpServerElicitationRequest::Url { .. } => { + McpServerElicitationRequest::OpenAiForm { .. } + | McpServerElicitationRequest::Url { .. } => { self.app_event_tx.resolve_elicitation( thread_id, params.server_name, diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index b5540f878b..55cf893076 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -399,6 +399,7 @@ async fn connect_remote_app_server( client_name: "codex-tui".to_string(), client_version: env!("CARGO_PKG_VERSION").to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, }) @@ -562,6 +563,7 @@ where client_name: "codex-tui".to_string(), client_version: env!("CARGO_PKG_VERSION").to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, }) diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index fceb7fb059..0170ae49d4 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -1055,6 +1055,7 @@ mod tests { client_name: "test".to_string(), client_version: "test".to_string(), experimental_api: true, + mcp_server_openai_form_elicitation: false, opt_out_notification_methods: Vec::new(), channel_capacity: DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, }) From 01a2df29479e8e982a5b6a9097d21814ad24000f Mon Sep 17 00:00:00 2001 From: Tom Date: Thu, 18 Jun 2026 12:13:31 -0700 Subject: [PATCH 4/5] [codex] Make thread store turn filter optional (#28949) Make `ListItemsParams::turn_id` optional so callers can list persisted items across an entire thread or narrow the result to one turn. This aligns the thread-store API and documentation with thread-wide item listing while preserving the optional turn-filter behavior for implementations. --- codex-rs/thread-store/src/in_memory.rs | 2 +- codex-rs/thread-store/src/store.rs | 2 +- codex-rs/thread-store/src/types.rs | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index 7eda54a9c0..a60409b740 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -79,7 +79,7 @@ mod tests { let items_err = store .list_items(ListItemsParams { thread_id, - turn_id: "turn_1".to_string(), + turn_id: None, include_archived: true, cursor: None, page_size: 10, diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index ef02e85d50..65ad1ce8a8 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -102,7 +102,7 @@ pub trait ThreadStore: Any + Send + Sync { }) } - /// Lists persisted items within a stored turn. + /// Lists persisted items within a stored thread, optionally filtered to a turn. fn list_items(&self, _params: ListItemsParams) -> ThreadStoreFuture<'_, ItemPage> { Box::pin(async { Err(ThreadStoreError::Unsupported { diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index b969c6b957..2e3c60723f 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -334,13 +334,13 @@ pub struct TurnPage { pub backwards_cursor: Option, } -/// Parameters for listing persisted items within a single turn. +/// Parameters for listing persisted items within a thread. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct ListItemsParams { /// Thread id to read. pub thread_id: ThreadId, - /// Turn id to hydrate. - pub turn_id: String, + /// Optional turn id to filter by. When omitted, returns items across the thread. + pub turn_id: Option, /// Whether archived threads are eligible. pub include_archived: bool, /// Opaque cursor returned by a previous list call. @@ -361,7 +361,7 @@ pub struct StoredThreadItem { pub materialized_thread_item_json: Vec, } -/// A page of persisted items within a turn. +/// A page of persisted items within a thread, optionally filtered to a turn. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct ItemPage { /// Items returned for this page. From 752ed90d78ad9444dd449fbc8e71ad2121865be0 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Thu, 18 Jun 2026 12:18:42 -0700 Subject: [PATCH 5/5] current time reminders impl for system clock (varlatency 2/n) (#28824) Stacked on #28822. ## Summary - add a host-injectable current-time provider with a built-in system implementation - record UTC developer reminders in history immediately before due model requests - keep cadence state per session and force a refresh after compaction This does NOT include the app server client <-> server clock logic. This PR is only for the reminder message & system clock that will be used in prod. ## Testing - `just test -p codex-core varlatency_` - `just clippy -p codex-core -p codex-app-server -p codex-mcp-server -p codex-thread-manager-sample` - `just fmt` --- codex-rs/app-server/src/mcp_refresh.rs | 1 + codex-rs/app-server/src/message_processor.rs | 1 + codex-rs/core/src/codex_delegate.rs | 1 + .../core/src/context/current_time_reminder.rs | 35 +++ codex-rs/core/src/context/mod.rs | 2 + codex-rs/core/src/current_time.rs | 41 +++ codex-rs/core/src/lib.rs | 3 + codex-rs/core/src/prompt_debug.rs | 1 + codex-rs/core/src/session/mod.rs | 5 + codex-rs/core/src/session/session.rs | 6 + codex-rs/core/src/session/tests.rs | 5 + .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/session/time_reminder.rs | 61 ++++ codex-rs/core/src/session/turn.rs | 70 +++-- codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/state/session.rs | 3 + codex-rs/core/src/thread_manager.rs | 6 + codex-rs/core/src/thread_manager_tests.rs | 11 + .../src/tools/handlers/multi_agents_tests.rs | 1 + codex-rs/core/tests/common/test_codex.rs | 9 + codex-rs/core/tests/suite/client.rs | 1 + .../core/tests/suite/current_time_reminder.rs | 274 ++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 1 + codex-rs/mcp-server/src/message_processor.rs | 1 + codex-rs/thread-manager-sample/src/main.rs | 1 + 25 files changed, 516 insertions(+), 27 deletions(-) create mode 100644 codex-rs/core/src/context/current_time_reminder.rs create mode 100644 codex-rs/core/src/current_time.rs create mode 100644 codex-rs/core/src/session/time_reminder.rs create mode 100644 codex-rs/core/tests/suite/current_time_reminder.rs diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index a5659174f4..84a6b8e83b 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -250,6 +250,7 @@ mod tests { Some(state_db.clone()), "11111111-1111-4111-8111-111111111111".to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ) }); thread_manager.start_thread(good_config).await?; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 9e9437bd34..0f9f6d3321 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -379,6 +379,7 @@ impl MessageProcessor { outgoing.clone(), thread_state_manager.clone(), )), + /*external_time_provider*/ None, ) }); let models_manager = thread_manager.get_models_manager(); diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index b7b1fe7580..dab902ec37 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -120,6 +120,7 @@ pub(crate) async fn run_codex_thread_interactive( analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), + external_time_provider: Some(Arc::clone(&parent_session.services.time_provider)), inherited_multi_agent_version: Some(MultiAgentVersion::Disabled), })) .or_cancel(&cancel_token) diff --git a/codex-rs/core/src/context/current_time_reminder.rs b/codex-rs/core/src/context/current_time_reminder.rs new file mode 100644 index 0000000000..0e3fe0b1e8 --- /dev/null +++ b/codex-rs/core/src/context/current_time_reminder.rs @@ -0,0 +1,35 @@ +use chrono::DateTime; +use chrono::Utc; + +use super::ContextualUserFragment; + +pub(crate) struct CurrentTimeReminder { + current_time: DateTime, +} + +impl CurrentTimeReminder { + pub(crate) fn new(current_time: DateTime) -> Self { + Self { current_time } + } +} + +impl ContextualUserFragment for CurrentTimeReminder { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } + + fn body(&self) -> String { + format!( + "It is {}.", + self.current_time.format("%Y-%m-%d %H:%M:%S UTC") + ) + } +} diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 7fc1375781..3c2f487231 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -6,6 +6,7 @@ mod available_plugins_instructions; mod available_skills_instructions; mod collaboration_mode_instructions; mod contextual_user_message; +mod current_time_reminder; mod environment_context; mod guardian_followup_review_reminder; mod hook_additional_context; @@ -44,6 +45,7 @@ pub(crate) use codex_core_skills::SkillInstructions; pub(crate) use collaboration_mode_instructions::CollaborationModeInstructions; pub(crate) use contextual_user_message::is_contextual_user_fragment; pub(crate) use contextual_user_message::parse_visible_hook_prompt_message; +pub(crate) use current_time_reminder::CurrentTimeReminder; pub(crate) use environment_context::EnvironmentContext; pub(crate) use guardian_followup_review_reminder::GuardianFollowupReviewReminder; pub(crate) use hook_additional_context::HookAdditionalContext; diff --git a/codex-rs/core/src/current_time.rs b/codex-rs/core/src/current_time.rs new file mode 100644 index 0000000000..7740552e1e --- /dev/null +++ b/codex-rs/core/src/current_time.rs @@ -0,0 +1,41 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use anyhow::Result; +use anyhow::anyhow; +use chrono::DateTime; +use chrono::Utc; +use codex_features::CurrentTimeSource; +use codex_protocol::ThreadId; + +use crate::config::CurrentTimeReminderConfig; + +pub type TimeFuture<'a> = Pin>> + Send + 'a>>; + +/// Host integration boundary for obtaining the current time. +pub trait TimeProvider: Send + Sync { + fn current_time(&self, thread_id: ThreadId) -> TimeFuture<'_>; +} + +pub(crate) struct SystemTimeProvider; + +impl TimeProvider for SystemTimeProvider { + fn current_time(&self, _thread_id: ThreadId) -> TimeFuture<'_> { + Box::pin(async { Ok(Utc::now()) }) + } +} + +pub(crate) fn resolve_time_provider( + config: Option<&CurrentTimeReminderConfig>, + external_provider: Option>, +) -> Result> { + match config.map(|config| config.clock_source).unwrap_or_default() { + CurrentTimeSource::System => Ok(Arc::new(SystemTimeProvider)), + CurrentTimeSource::External => external_provider.ok_or_else(|| { + anyhow!( + "features.current_time_reminder.clock_source is external, but no external current-time provider is available" + ) + }), + } +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 0e193baa1d..121c8d2be2 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -37,6 +37,7 @@ pub mod config; pub mod connectors; pub mod context; mod context_manager; +mod current_time; mod environment_selection; pub mod exec; pub mod exec_env; @@ -186,6 +187,8 @@ pub use client_common::ResponseEvent; pub use client_common::ResponseStream; pub use codex_prompts::REVIEW_PROMPT; pub use compact::content_items_to_text; +pub use current_time::TimeFuture; +pub use current_time::TimeProvider; pub use event_mapping::parse_turn_item; pub use exec_policy::ExecPolicyError; pub use exec_policy::check_execpolicy_for_warnings; diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 0da42d8a4e..bd9289b770 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -60,6 +60,7 @@ pub async fn build_prompt_input( state_db.clone(), installation_id, /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let thread = thread_manager.start_thread(config).await?; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index f92bd63680..b55c4b11d7 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -31,6 +31,7 @@ use crate::context::NetworkRuleSaved; use crate::context::PermissionsInstructions; use crate::context::PersonalitySpecInstructions; use crate::context::RecommendedPluginsInstructions; +use crate::current_time::TimeProvider; use crate::default_skill_metadata_budget; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::exec_policy::ExecPolicyManager; @@ -215,6 +216,7 @@ mod rollout_budget; mod rollout_reconstruction; #[allow(clippy::module_inception)] pub(crate) mod session; +pub(crate) mod time_reminder; mod token_budget; pub(crate) mod turn; pub(crate) mod turn_context; @@ -439,6 +441,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, + pub(crate) external_time_provider: Option>, pub(crate) inherited_multi_agent_version: Option, } @@ -522,6 +525,7 @@ impl Codex { analytics_events_client, thread_store, attestation_provider, + external_time_provider, inherited_multi_agent_version, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); @@ -669,6 +673,7 @@ impl Codex { thread_store, parent_rollout_thread_trace, attestation_provider, + external_time_provider, multi_agent_version, )) .await diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 2b61468c15..af2d1abeb8 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -492,6 +492,7 @@ impl Session { thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, + external_time_provider: Option>, multi_agent_version: Option, ) -> anyhow::Result> { debug!( @@ -516,6 +517,10 @@ impl Session { } InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id, }; + let time_provider = crate::current_time::resolve_time_provider( + config.current_time_reminder.as_ref(), + external_time_provider, + )?; let mcp_thread_init = thread_extension_init.clone(); let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( thread_id.to_string(), @@ -1022,6 +1027,7 @@ impl Session { live_thread: live_thread_init.as_ref().cloned(), thread_store: Arc::clone(&thread_store), attestation_provider: attestation_provider.clone(), + time_provider, model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), thread_id, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index b13cd029ce..0b1cc045c4 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4867,6 +4867,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + /*external_time_provider*/ None, Some(config.multi_agent_version_from_features()), ) .await; @@ -5032,6 +5033,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*state_db*/ None, )), attestation_provider: None, + time_provider: Arc::new(crate::current_time::SystemTimeProvider), model_client: ModelClient::new( Some(auth_manager.clone()), thread_id, @@ -5216,6 +5218,7 @@ async fn make_session_with_config_and_rx( )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + /*external_time_provider*/ None, Some(config.multi_agent_version_from_features()), ) .await?; @@ -5328,6 +5331,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + /*external_time_provider*/ None, Some(config.multi_agent_version_from_features()), ) .await?; @@ -7078,6 +7082,7 @@ where state_db, )), attestation_provider: None, + time_provider: Arc::new(crate::current_time::SystemTimeProvider), model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), thread_id, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 0bed032351..5de331ee1f 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -736,6 +736,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { analytics_events_client: None, thread_store, attestation_provider: None, + external_time_provider: None, inherited_multi_agent_version: None, }) .await diff --git a/codex-rs/core/src/session/time_reminder.rs b/codex-rs/core/src/session/time_reminder.rs new file mode 100644 index 0000000000..058132a895 --- /dev/null +++ b/codex-rs/core/src/session/time_reminder.rs @@ -0,0 +1,61 @@ +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; + +use super::session::Session; +use super::turn_context::TurnContext; +use crate::context::ContextualUserFragment; + +#[derive(Default)] +pub(crate) struct CurrentTimeReminderState { + model_requests_since_delivery: u64, + last_window_id: Option, +} + +impl CurrentTimeReminderState { + fn take_reminder_due(&mut self, window_id: &str, interval: u64) -> bool { + self.model_requests_since_delivery = self.model_requests_since_delivery.saturating_add(1); + let reminder_is_due = self.last_window_id.as_deref() != Some(window_id) + || self.model_requests_since_delivery >= interval; + + if reminder_is_due { + self.model_requests_since_delivery = 0; + self.last_window_id = Some(window_id.to_string()); + } + + reminder_is_due + } +} + +pub(super) async fn maybe_record_current_time_reminder( + sess: &Session, + turn_context: &TurnContext, + window_id: &str, +) -> CodexResult<()> { + let Some(config) = turn_context.config.current_time_reminder else { + return Ok(()); + }; + + let reminder_is_due = { + let mut state = sess.state.lock().await; + state + .current_time_reminder + .take_reminder_due(window_id, config.reminder_interval_model_requests) + }; + if !reminder_is_due { + return Ok(()); + } + + let current_time = sess + .services + .time_provider + .current_time(sess.thread_id) + .await + .map_err(|err| CodexErr::Fatal(format!("failed to read current time: {err:#}")))?; + + let response_item = + ContextualUserFragment::into(crate::context::CurrentTimeReminder::new(current_time)); + sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) + .await; + + Ok(()) +} diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 657e00ff35..25e4e1a6af 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -226,34 +226,50 @@ pub(crate) async fn run_turn( ) .await; - // Construct the input that we will send to the model. - let sampling_request_input: Vec = async { - sess.clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities) - } - .instrument(trace_span!("run_turn.prepare_sampling_request_input")) - .await; + let sampling_request_result: CodexResult<_> = async { + super::time_reminder::maybe_record_current_time_reminder( + sess.as_ref(), + turn_context.as_ref(), + &window_id, + ) + .await?; - let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( - sess.installation_id.clone(), - window_id, - CodexResponsesRequestKind::Turn, - ); - let tokens_before_sampling = sess.get_total_token_usage().await; - match run_sampling_request( - Arc::clone(&sess), - Arc::clone(&turn_context), - Arc::clone(&turn_extension_data), - Arc::clone(&turn_diff_tracker), - &mut client_session, - &responses_metadata, - sampling_request_input, - cancellation_token.child_token(), - ) - .await - { - Ok((sampling_request_output, sampling_request_input)) => { + // Construct the input that we will send to the model. + let sampling_request_input: Vec = async { + sess.clone_history() + .await + .for_prompt(&turn_context.model_info.input_modalities) + } + .instrument(trace_span!("run_turn.prepare_sampling_request_input")) + .await; + + let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( + sess.installation_id.clone(), + window_id, + CodexResponsesRequestKind::Turn, + ); + let tokens_before_sampling = sess.get_total_token_usage().await; + let (sampling_request_output, sampling_request_input) = run_sampling_request( + Arc::clone(&sess), + Arc::clone(&turn_context), + Arc::clone(&turn_extension_data), + Arc::clone(&turn_diff_tracker), + &mut client_session, + &responses_metadata, + sampling_request_input, + cancellation_token.child_token(), + ) + .await?; + + Ok(( + tokens_before_sampling, + sampling_request_output, + sampling_request_input, + )) + } + .await; + match sampling_request_result { + Ok((tokens_before_sampling, sampling_request_output, sampling_request_input)) => { let SamplingRequestResult { needs_follow_up: model_needs_follow_up, last_agent_message: sampling_request_last_agent_message, diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 914fd4c32d..853b96e921 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -8,6 +8,7 @@ use crate::attestation::AttestationProvider; use crate::client::ModelClient; use crate::config::NetworkProxyAuditMetadata; use crate::config::StartedNetworkProxy; +use crate::current_time::TimeProvider; use crate::environment_selection::ThreadEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; @@ -79,6 +80,7 @@ pub(crate) struct SessionServices { pub(crate) live_thread: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, + pub(crate) time_provider: Arc, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 269d3e0f60..2bd9da5ce3 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -13,6 +13,7 @@ use super::auto_compact_window::AutoCompactWindowSnapshot; use crate::context_manager::ContextManager; use crate::session::PreviousTurnSettings; use crate::session::session::SessionConfiguration; +use crate::session::time_reminder::CurrentTimeReminderState; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::TokenUsage; @@ -36,6 +37,7 @@ pub(crate) struct SessionState { auto_compact_window: AutoCompactWindow, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, + pub(crate) current_time_reminder: CurrentTimeReminderState, pub(crate) active_connector_selection: HashSet, pub(crate) pending_session_start_sources: VecDeque, granted_permissions_by_environment_id: HashMap, @@ -56,6 +58,7 @@ impl SessionState { previous_turn_settings: None, auto_compact_window: AutoCompactWindow::new(), startup_prewarm: None, + current_time_reminder: CurrentTimeReminderState::default(), active_connector_selection: HashSet::new(), pending_session_start_sources: VecDeque::new(), granted_permissions_by_environment_id: HashMap::new(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 77a5520026..814c8fe764 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -4,6 +4,7 @@ use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; use crate::config::Config; use crate::config::ThreadStoreConfig; +use crate::current_time::TimeProvider; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::environment_selection::default_thread_environment_selections; use crate::mcp::McpManager; @@ -215,6 +216,7 @@ pub(crate) struct ThreadManagerState { user_instructions_provider: Arc, thread_store: Arc, attestation_provider: Option>, + external_time_provider: Option>, session_source: SessionSource, installation_id: String, analytics_events_client: Option, @@ -269,6 +271,7 @@ impl ThreadManager { state_db: Option, installation_id: String, attestation_provider: Option>, + external_time_provider: Option>, ) -> Self { let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); @@ -300,6 +303,7 @@ impl ThreadManager { user_instructions_provider, thread_store, attestation_provider, + external_time_provider, auth_manager, session_source, installation_id, @@ -405,6 +409,7 @@ impl ThreadManager { ), thread_store, attestation_provider: None, + external_time_provider: None, auth_manager, session_source: SessionSource::Exec, installation_id, @@ -1466,6 +1471,7 @@ impl ThreadManagerState { analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), + external_time_provider: self.external_time_provider.clone(), inherited_multi_agent_version: multi_agent_version, })) .await?; diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index ee010b5e1d..bfb7865bb2 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -440,6 +440,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let selected_root_init = |id: &str, environment_id: &str| { let mut init = codex_extension_api::ExtensionDataInit::new(); @@ -548,6 +549,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let selected_cwd = AbsolutePathBuf::try_from(config.cwd.as_path().join("selected")).expect("absolute path"); @@ -671,6 +673,7 @@ async fn explicit_installation_id_skips_codex_home_file() { state_db.clone(), installation_id.clone(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let thread = manager @@ -711,6 +714,7 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let source = manager @@ -770,6 +774,7 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let source = manager @@ -836,6 +841,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { state_db.clone(), TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let source = manager @@ -929,6 +935,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { state_db, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let source = manager @@ -1033,6 +1040,7 @@ async fn new_uses_active_provider_for_model_refresh() { /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let _ = manager.list_models(RefreshStrategy::Online).await; @@ -1254,6 +1262,7 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor state_db.clone(), TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let source = manager @@ -1362,6 +1371,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { state_db.clone(), TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let source = manager @@ -1460,6 +1470,7 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ state_db.clone(), TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let source = manager diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index ed15a7483a..eefea1eb8f 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -4279,6 +4279,7 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr state_db.clone(), "11111111-1111-4111-8111-111111111111".to_string(), /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let parent = manager diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index b10e26bf0c..42b9bf87ca 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -17,6 +17,7 @@ use codex_config::CloudConfigBundleLoader; use codex_core::CodexThread; use codex_core::StartThreadOptions; use codex_core::ThreadManager; +use codex_core::TimeProvider; use codex_core::config::Config; use codex_core::resolve_installation_id; use codex_core::shell::Shell; @@ -261,6 +262,7 @@ pub struct TestCodexBuilder { extensions: Arc>, user_instructions_provider: Option>, supports_openai_form_elicitation: bool, + external_time_provider: Option>, } impl TestCodexBuilder { @@ -362,6 +364,11 @@ impl TestCodexBuilder { self } + pub fn with_external_time_provider(mut self, provider: Arc) -> Self { + self.external_time_provider = Some(provider); + self + } + pub fn with_windows_cmd_shell(self) -> Self { if cfg!(windows) { self.with_user_shell(get_shell_by_model_provided_path(&PathBuf::from("cmd.exe"))) @@ -568,6 +575,7 @@ impl TestCodexBuilder { state_db.clone(), installation_id, /*attestation_provider*/ None, + /*external_time_provider*/ self.external_time_provider.clone(), ); let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); @@ -1172,6 +1180,7 @@ pub fn test_codex() -> TestCodexBuilder { extensions: empty_extension_registry(), user_instructions_provider: None, supports_openai_form_elicitation: false, + external_time_provider: None, } } diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 54e24d49cb..e329618828 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1259,6 +1259,7 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { /*state_db*/ None, installation_id, /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let NewThread { thread: codex, .. } = thread_manager .start_thread(config.clone()) diff --git a/codex-rs/core/tests/suite/current_time_reminder.rs b/codex-rs/core/tests/suite/current_time_reminder.rs new file mode 100644 index 0000000000..5b5970a0ee --- /dev/null +++ b/codex-rs/core/tests/suite/current_time_reminder.rs @@ -0,0 +1,274 @@ +use std::sync::Arc; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use anyhow::Result; +use anyhow::anyhow; +use chrono::DateTime; +use chrono::Utc; +use codex_core::TimeFuture; +use codex_core::TimeProvider; +use codex_core::config::CurrentTimeReminderConfig; +use codex_features::CurrentTimeSource; +use codex_features::Feature; +use codex_model_provider_info::built_in_model_providers; +use codex_protocol::ThreadId; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::CodexErrorInfo; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use core_test_support::assert_regex_match; +use core_test_support::responses::ResponsesRequest; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::json; + +const FIRST_REMINDER: &str = "It is 2026-06-17 17:34:15 UTC."; +const SECOND_REMINDER: &str = "It is 2026-06-17 17:35:15 UTC."; +const FIRST_TIME_UNIX_SECONDS: i64 = 1_781_717_655; + +struct TestTimeProvider(AtomicI64); + +impl Default for TestTimeProvider { + fn default() -> Self { + Self(AtomicI64::new(FIRST_TIME_UNIX_SECONDS)) + } +} + +impl TimeProvider for TestTimeProvider { + fn current_time(&self, _thread_id: ThreadId) -> TimeFuture<'_> { + let timestamp = self.0.fetch_add(60, Ordering::Relaxed); + Box::pin(async move { + Ok(DateTime::::from_timestamp(timestamp, 0) + .expect("test timestamp should be valid")) + }) + } +} + +struct FailingTimeProvider; + +impl TimeProvider for FailingTimeProvider { + fn current_time(&self, _thread_id: ThreadId) -> TimeFuture<'_> { + Box::pin(async { Err(anyhow!("test clock unavailable")) }) + } +} + +fn current_time_reminders(request: &ResponsesRequest) -> Vec { + request + .message_input_texts("developer") + .into_iter() + .filter(|text| text.starts_with("It is ")) + .collect() +} + +fn enable_current_time_reminder( + config: &mut codex_core::config::Config, + interval: u64, + clock_source: CurrentTimeSource, +) { + config + .features + .enable(Feature::CurrentTimeReminder) + .expect("test config should allow current-time reminders"); + config.current_time_reminder = Some(CurrentTimeReminderConfig { + reminder_interval_model_requests: interval, + clock_source, + }); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn current_time_reminders_follow_request_interval_and_persist_in_history() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let tool_args = json!({ + "command": "echo current time", + "timeout_ms": 1_000, + }); + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + "current-time-tool-call", + "shell_command", + &serde_json::to_string(&tool_args)?, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ]), + sse(vec![ev_response_created("resp-3"), ev_completed("resp-3")]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + enable_current_time_reminder(config, /*interval*/ 2, CurrentTimeSource::External) + }) + .with_external_time_provider(Arc::new(TestTimeProvider::default())) + .build(&server) + .await?; + + test.submit_turn_with_permission_profile("first turn", PermissionProfile::Disabled) + .await?; + test.submit_turn("second turn").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + assert_eq!(current_time_reminders(&requests[0]), vec![FIRST_REMINDER]); + assert_eq!(current_time_reminders(&requests[1]), vec![FIRST_REMINDER]); + assert_eq!( + current_time_reminders(&requests[2]), + vec![FIRST_REMINDER, SECOND_REMINDER] + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn system_time_source_adds_current_time_reminder() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let test = test_codex() + .with_config(|config| { + enable_current_time_reminder(config, /*interval*/ 1, CurrentTimeSource::System) + }) + .build(&server) + .await?; + + test.submit_turn("what time is it?").await?; + + let reminders = current_time_reminders(&responses.single_request()); + assert_eq!(reminders.len(), 1); + assert_regex_match( + r"^It is \d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} UTC\.$", + &reminders[0], + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn current_time_reminder_is_refreshed_after_compaction() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ + ev_response_created("resp-compact"), + ev_assistant_message("msg-compact", "compact summary"), + ev_completed("resp-compact"), + ]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone(); + model_provider.name = "OpenAI-compatible test provider".to_string(); + model_provider.base_url = Some(format!("{}/v1", server.uri())); + model_provider.supports_websockets = false; + let test = test_codex() + .with_config(move |config| { + config.model_provider = model_provider; + enable_current_time_reminder(config, /*interval*/ 50, CurrentTimeSource::External); + }) + .with_external_time_provider(Arc::new(TestTimeProvider::default())) + .build(&server) + .await?; + + test.submit_turn("before compact").await?; + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_turn("after compact").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + assert_eq!( + current_time_reminders(&requests[2]), + vec![SECOND_REMINDER], + "a new context window should force a fresh reminder before the next model request" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn time_provider_failure_stops_before_inference() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_once( + &server, + sse(vec![ + ev_response_created("unused-response"), + ev_completed("unused-response"), + ]), + ) + .await; + let test = test_codex() + .with_config(|config| { + enable_current_time_reminder(config, /*interval*/ 1, CurrentTimeSource::External) + }) + .with_external_time_provider(Arc::new(FailingTimeProvider)) + .build(&server) + .await?; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "fail before inference".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + + let EventMsg::Error(error) = + wait_for_event(&test.codex, |event| matches!(event, EventMsg::Error(_))).await + else { + unreachable!(); + }; + assert_eq!( + error.message, + "Fatal error: failed to read current time: test clock unavailable" + ); + assert_eq!(error.codex_error_info, Some(CodexErrorInfo::Other)); + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert!(responses.requests().is_empty()); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 0fdd958b39..1e200ab486 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -48,6 +48,7 @@ mod compact; mod compact_remote; mod compact_remote_parity; mod compact_resume_fork; +mod current_time_reminder; mod deprecation_notice; mod exec; mod exec_policy; diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index ce19575a9c..5f36091fee 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -78,6 +78,7 @@ impl MessageProcessor { state_db.clone(), installation_id, /*attestation_provider*/ None, + /*external_time_provider*/ None, )); Self { outgoing, diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 1534ba86a0..26ba71ae46 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -136,6 +136,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { state_db, installation_id, /*attestation_provider*/ None, + /*external_time_provider*/ None, ); let NewThread {