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/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 3544592b49..0f9f6d3321 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(|_| ()) } @@ -372,6 +379,7 @@ impl MessageProcessor { outgoing.clone(), thread_state_manager.clone(), )), + /*external_time_provider*/ None, ) }); let models_manager = thread_manager.get_models_manager(); @@ -884,6 +892,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 +908,7 @@ impl MessageProcessor { request_context, app_server_client_name, client_version, + supports_openai_form_elicitation, ) .await; if let Err(error) = result { @@ -928,6 +938,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 +1091,7 @@ impl MessageProcessor { params, app_server_client_name.clone(), client_version.clone(), + supports_openai_form_elicitation, request_context, ) .await @@ -1096,6 +1108,8 @@ impl MessageProcessor { params, app_server_client_name.clone(), client_version.clone(), + /*supports_openai_form_elicitation*/ + supports_openai_form_elicitation, ) .await } @@ -1106,6 +1120,8 @@ impl MessageProcessor { params, app_server_client_name.clone(), client_version.clone(), + /*supports_openai_form_elicitation*/ + supports_openai_form_elicitation, ) .await } @@ -1305,6 +1321,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/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/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/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index b3f0d6c43c..dab902ec37 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -113,9 +113,14 @@ 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(), + 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/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/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/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/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/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 a1f204e475..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; @@ -24,6 +25,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; @@ -43,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; @@ -64,6 +67,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/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/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..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; @@ -131,6 +132,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; @@ -185,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/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/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/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/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/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/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 84e60a3d2d..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; @@ -211,9 +212,11 @@ mod input_queue; mod mcp; mod multi_agents; mod review; +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; @@ -434,9 +437,11 @@ 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>, + pub(crate) external_time_provider: Option>, pub(crate) inherited_multi_agent_version: Option, } @@ -516,9 +521,11 @@ impl Codex { parent_trace: _, environment_selections, thread_extension_init, + supports_openai_form_elicitation, 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); @@ -658,6 +665,7 @@ impl Codex { mcp_manager.clone(), extensions, thread_extension_init, + supports_openai_form_elicitation, agent_control, environment_manager, inherited_environments, @@ -665,6 +673,7 @@ impl Codex { thread_store, parent_rollout_thread_trace, attestation_provider, + external_time_provider, multi_agent_version, )) .await @@ -3360,6 +3369,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/session.rs b/codex-rs/core/src/session/session.rs index c4bc064a47..af2d1abeb8 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, @@ -491,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!( @@ -515,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(), @@ -1009,6 +1015,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, @@ -1018,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, @@ -1149,6 +1159,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..0b1cc045c4 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, @@ -4866,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; @@ -5018,6 +5020,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(), @@ -5030,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, @@ -5203,6 +5207,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, @@ -5213,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?; @@ -5307,6 +5313,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, @@ -5324,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?; @@ -7061,6 +7069,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(), @@ -7073,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 b9727a0401..5de331ee1f 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -732,9 +732,11 @@ 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, + 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 257e741203..25e4e1a6af 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -218,35 +218,58 @@ pub(crate) async fn run_turn( break; } - // 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")) + let window_id = sess.current_window_id().await; + super::rollout_budget::maybe_record_reminder( + sess.as_ref(), + turn_context.as_ref(), + &window_id, + ) .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, - 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)) => { + let sampling_request_result: CodexResult<_> = async { + super::time_reminder::maybe_record_current_time_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() + .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 9724d8b3ad..853b96e921 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; @@ -7,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; @@ -67,6 +69,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, @@ -77,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/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 63335c0129..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; @@ -186,6 +187,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 { @@ -214,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, @@ -268,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(); @@ -299,6 +303,7 @@ impl ThreadManager { user_instructions_provider, thread_store, attestation_provider, + external_time_provider, auth_manager, session_source, installation_id, @@ -404,6 +409,7 @@ impl ThreadManager { ), thread_store, attestation_provider: None, + external_time_provider: None, auth_manager, session_source: SessionSource::Exec, installation_id, @@ -604,6 +610,7 @@ impl ThreadManager { parent_trace: None, environments, thread_extension_init: ExtensionDataInit::default(), + supports_openai_form_elicitation: false, })) .await } @@ -621,6 +628,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 +639,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, @@ -643,6 +651,7 @@ impl ThreadManager { options.parent_trace, options.environments, options.thread_extension_init, + options.supports_openai_form_elicitation, /*user_shell_override*/ None, )) .await @@ -691,6 +700,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( @@ -698,6 +708,7 @@ impl ThreadManager { initial_history, auth_manager, parent_trace, + supports_openai_form_elicitation, )) .await } @@ -709,7 +720,9 @@ 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( self.state.environment_manager.as_ref(), &config.cwd, @@ -721,7 +734,7 @@ impl ThreadManager { config, initial_history, auth_manager, - self.agent_control(), + agent_control, session_source, /*parent_thread_id*/ None, /*forked_from_thread_id*/ None, @@ -733,6 +746,7 @@ impl ThreadManager { parent_trace, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, /*user_shell_override*/ None, )) .await @@ -742,7 +756,9 @@ 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( self.state.environment_manager.as_ref(), &config.cwd, @@ -751,7 +767,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, @@ -760,6 +776,7 @@ impl ThreadManager { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, /*user_shell_override*/ Some(user_shell_override), )) .await @@ -771,7 +788,9 @@ 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?; let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), @@ -784,7 +803,7 @@ impl ThreadManager { config, initial_history, auth_manager, - self.agent_control(), + agent_control, session_source, /*parent_thread_id*/ None, /*forked_from_thread_id*/ None, @@ -796,6 +815,7 @@ impl ThreadManager { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, /*user_shell_override*/ Some(user_shell_override), )) .await @@ -876,8 +896,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( @@ -906,6 +933,7 @@ impl ThreadManager { history: InitialHistory, thread_source: Option, parent_trace: Option, + supports_openai_form_elicitation: bool, ) -> CodexResult where S: Into, @@ -916,6 +944,7 @@ impl ThreadManager { history, thread_source, parent_trace, + supports_openai_form_elicitation, ) .await } @@ -927,6 +956,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. @@ -952,11 +982,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, @@ -965,13 +996,18 @@ impl ThreadManager { parent_trace, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + supports_openai_form_elicitation, /*user_shell_override*/ None, )) .await } 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)] @@ -1221,6 +1257,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, )) .await @@ -1258,6 +1295,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, )) .await @@ -1296,6 +1334,7 @@ impl ThreadManagerState { /*parent_trace*/ None, environments, /*thread_extension_init*/ ExtensionDataInit::default(), + /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, )) .await @@ -1317,6 +1356,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( @@ -1335,6 +1375,7 @@ impl ThreadManagerState { parent_trace, environments, thread_extension_init, + supports_openai_form_elicitation, user_shell_override, )) .await @@ -1358,6 +1399,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(_)); @@ -1425,9 +1467,11 @@ 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(), + 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 be952bbcbc..bfb7865bb2 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"); @@ -439,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(); @@ -463,6 +465,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 +480,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"); @@ -545,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"); @@ -567,6 +572,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 +599,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"); @@ -666,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 @@ -706,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 @@ -729,6 +738,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"); @@ -764,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 @@ -792,6 +803,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"); @@ -829,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 @@ -842,6 +855,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 +882,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"); @@ -920,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 @@ -947,6 +963,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 +980,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"); @@ -1022,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; @@ -1243,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 @@ -1254,6 +1274,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"); @@ -1350,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 @@ -1368,6 +1390,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"); @@ -1447,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 @@ -1458,6 +1482,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..eefea1eb8f 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"); @@ -4278,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 daf6f79b5c..42b9bf87ca 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -15,7 +15,9 @@ 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::TimeProvider; use codex_core::config::Config; use codex_core::resolve_installation_id; use codex_core::shell::Shell; @@ -39,6 +41,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 +261,8 @@ pub struct TestCodexBuilder { exec_server_url: Option, extensions: Arc>, user_instructions_provider: Option>, + supports_openai_form_elicitation: bool, + external_time_provider: Option>, } impl TestCodexBuilder { @@ -354,6 +359,16 @@ impl TestCodexBuilder { self } + pub fn with_openai_form_elicitation(mut self) -> Self { + self.supports_openai_form_elicitation = true; + 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"))) @@ -560,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(); @@ -574,6 +590,7 @@ impl TestCodexBuilder { path, auth_manager, user_shell_override, + self.supports_openai_form_elicitation, ), ) .await? @@ -585,6 +602,7 @@ impl TestCodexBuilder { path, auth_manager, /*parent_trace*/ None, + self.supports_openai_form_elicitation, )) .await? } @@ -594,11 +612,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 +1179,8 @@ pub fn test_codex() -> TestCodexBuilder { exec_server_url: None, 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/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/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/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/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/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/mod.rs b/codex-rs/core/tests/suite/mod.rs index 275458ffdc..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; @@ -100,6 +101,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/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/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(()) +} 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/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/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/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/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index ca11ff9072..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 { @@ -278,6 +279,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 }, 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. 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, })