diff --git a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs index 37c650448b..8b2ccbcd4d 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs @@ -97,7 +97,7 @@ impl AppsRequestProcessor { codex_apps_tools_cache_key: cache_key.clone(), client_mcp_extensions: ClientMcpExtensions::default(), auth: auth.clone(), - codex_apps_auth_manager, + auth_manager: codex_apps_auth_manager, elicitation_reviewer: None, elicitation_lifecycle: None, }) diff --git a/codex-rs/app-server/tests/suite/v2/daybreak_access.rs b/codex-rs/app-server/tests/suite/v2/daybreak_access.rs new file mode 100644 index 0000000000..923212b4c5 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/daybreak_access.rs @@ -0,0 +1,229 @@ +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; +use app_test_support::TestAppServer; +use app_test_support::write_chatgpt_auth; +use codex_app_server_protocol::CapabilityRootLocation; +use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::ItemCompletedNotification; +use codex_app_server_protocol::McpServerToolCallParams; +use codex_app_server_protocol::McpServerToolCallResponse; +use codex_app_server_protocol::SelectedCapabilityRoot; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::UserInput; +use codex_config::types::AuthCredentialsStoreMode; +use codex_features::Feature; +use codex_utils_path_uri::PathUri; +use core_test_support::responses; +use core_test_support::stdio_server_bin; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; +use test_case::test_case; +use wiremock::Mock; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[derive(Clone, Copy)] +enum AccessCall { + Installed, + Selected, + Unowned, + ApiKey, + Unlisted, +} + +#[test_case(AccessCall::Installed; "installed direct call")] +#[test_case(AccessCall::Selected; "selected model and direct calls")] +#[test_case(AccessCall::Unowned; "unowned direct call")] +#[test_case(AccessCall::ApiKey; "api key direct call")] +#[test_case(AccessCall::Unlisted; "unlisted direct call")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn daybreak_access_respects_plugin_provenance(case: AccessCall) -> Result<()> { + let server = responses::start_mock_server().await; + let home = TempDir::new()?; + let selected_package = TempDir::new()?; + let selected = matches!(case, AccessCall::Selected); + let plugin_root = if selected { + selected_package.path().to_path_buf() + } else { + home.path().join("plugins/cache/test/sample/local") + }; + std::fs::create_dir_all(plugin_root.join(".codex-plugin"))?; + std::fs::write( + plugin_root.join(".codex-plugin/plugin.json"), + r#"{"name":"sample"}"#, + )?; + let mut mcp_config = json!({ + "command":stdio_server_bin()?, "cwd":plugin_root, "environment_id":"local", + "default_tools_approval_mode":"approve", "env":{"MCP_TEST_DAYBREAK_READ_ONLY":"true"} + }); + if matches!(case, AccessCall::Unlisted) { + mcp_config["env"] = json!({}); + } + std::fs::write( + plugin_root.join(".mcp.json"), + serde_json::to_vec(&json!({"mcpServers":{"sample":mcp_config}}))?, + )?; + let extra_config = match case { + AccessCall::Installed | AccessCall::ApiKey | AccessCall::Unlisted => { + "[plugins.\"sample@test\"]\nenabled = true\n".to_string() + } + AccessCall::Selected => String::new(), + AccessCall::Unowned => toml::to_string(&json!({"mcp_servers":{"sample":mcp_config}}))?, + }; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Plugins) + .disable_feature(Feature::Apps) + .with_root_config(&format!( + "chatgpt_base_url = \"{}\"\ncli_auth_credentials_store = \"file\"", + server.uri() + )) + .with_extra_config(&extra_config) + .write(home.path())?; + let api_key = matches!(case, AccessCall::ApiKey).then_some("daybreak-test-api-key"); + if api_key.is_none() { + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("daybreak-test-token").account_id("daybreak-account"), + AuthCredentialsStoreMode::File, + )?; + } + let mut app = TestAppServer::builder() + .with_codex_home(home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", api_key)]) + // Daybreak metadata is local-host-only, even when the test runner provides a remote executor. + .without_auto_env() + .build_initialized_with_timeout(std::time::Duration::from_secs(/*secs*/ 30)) + .await?; + let selected_capability_roots = if selected { + Some(vec![SelectedCapabilityRoot { + id: "selected-daybreak".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "local".to_string(), + path: PathUri::from_host_native_path(&plugin_root)?, + }, + }]) + } else { + None + }; + let response: ThreadStartResponse = app + .request(|request_id| ClientRequest::ThreadStart { + request_id, + params: ThreadStartParams { + selected_capability_roots, + ..Default::default() + }, + }) + .await?; + let thread = response.thread; + Mock::given(method("GET")) + .and(path("/accounts/verified_access")) + .and(header("authorization", "Bearer daybreak-test-token")) + .and(header("chatgpt-account-id", "daybreak-account")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "programs": [{"program":"cyber", "state":"active", "grants":[ + {"level":"tac1", "source":"individual"} + ]}] + }))) + .expect(match case { + AccessCall::Selected => 1, + AccessCall::Installed + | AccessCall::Unowned + | AccessCall::ApiKey + | AccessCall::Unlisted => 0, + }) + .mount(&server) + .await; + let expected = json!({ + "schemaVersion":1, "entitlements":{"cyber_trusted_access":{ + "schemaVersion":1,"status":"granted","grants":[{"level":"tac1","source":"user"}],"stale":false + }} + }); + let mut results = Vec::new(); + // Selected capabilities are materialized by the first turn, as in the executor MCP tests. + if selected { + let _responses = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_tool_search_call( + "search", + &json!({"query":"get_codex_security_daybreak_access"}), + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "daybreak-call", + "mcp__sample", + "get_codex_security_daybreak_access", + "{}", + ), + responses::ev_completed("resp-2"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("done", "done"), + responses::ev_completed("resp-3"), + ]), + ], + ) + .await; + app.start_turn_and_wait_for_completion(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![UserInput::Text { + text: "Check Daybreak access".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + loop { + let event: ItemCompletedNotification = app.read_notification("item/completed").await?; + if let ThreadItem::McpToolCall { + result, plugin_id, .. + } = event.item + { + assert_eq!(plugin_id.as_deref(), Some("selected-daybreak")); + results.push(( + result + .expect("successful model call") + .structured_content + .expect("echoed metadata"), + Some(expected), + )); + break; + } + } + } + let response: McpServerToolCallResponse = app + .request(|request_id| ClientRequest::McpServerToolCall { + request_id, + params: McpServerToolCallParams { + thread_id: thread.id.clone(), + server: "sample".to_string(), + tool: "get_daybreak_access".to_string(), + arguments: Some(json!({})), + meta: Some(json!({ + "source":"direct-test", + "openai/entitlementContext":{"untrusted":"caller-supplied"} + })), + }, + }) + .await?; + let direct_metadata = response.structured_content.expect("echoed metadata"); + assert_eq!(direct_metadata["source"], "direct-test"); + results.push((direct_metadata, None)); + for (metadata, expected) in results { + assert_eq!(metadata.get("openai/entitlementContext"), expected.as_ref()); + assert_eq!(metadata["threadId"], thread.id); + assert!(!metadata.to_string().contains("daybreak-test-token")); + } + server.verify().await; + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index d654961905..7ea4813a71 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -22,6 +22,7 @@ mod connection_handling_websocket_unix; mod curated_mcp_sync; mod current_time; mod cyber_access_program; +mod daybreak_access; mod dynamic_tools; mod environment_add; mod environment_info; diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index ba3f8a3aa2..c6d79fea0d 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -170,7 +170,7 @@ impl fmt::Debug for McpBinding { /// one [`McpBinding`]. #[derive(Clone)] pub struct PreparedMcpCall { - _connections: Arc, + connections: Arc, client: Arc, config: Arc, catalog_revision: u64, @@ -201,7 +201,7 @@ impl PreparedMcpCall { let server_name = tool_info.server_name.clone(); config.permission_profile_for_server(&server_name)?; Some(Self { - _connections: connections, + connections, client, config, catalog_revision, @@ -321,10 +321,40 @@ impl PreparedMcpCall { )); } let (arguments, meta) = prepare().await?; + let timeout_deadline = + effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout); + let add_trusted_access_context = self.connections.add_trusted_access_context( + &self.tool_info, + &self.server_metadata, + arguments.as_ref(), + meta, + ); + let meta = match effective_timeout.zip(timeout_deadline) { + Some((timeout, deadline)) => { + tokio::time::timeout_at(deadline, add_trusted_access_context) + .await + .map_err(|_| { + anyhow::anyhow!("timed out awaiting tools/call after {timeout:.0?}") + })? + } + None => add_trusted_access_context.await, + }; + let remaining_timeout = match effective_timeout.zip(timeout_deadline) { + Some((timeout, deadline)) => { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err(anyhow::anyhow!( + "timed out awaiting tools/call after {timeout:.0?}" + )); + } + Some(remaining) + } + None => None, + }; let result = self .client .client - .call_tool(tool_name.clone(), arguments, meta, effective_timeout) + .call_tool(tool_name.clone(), arguments, meta, remaining_timeout) .await .with_context(|| format!("tool call failed for `{}/{tool_name}`", self.server_name))?; drop(current_revision); diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index d0c50bfa56..7485565422 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -52,6 +52,8 @@ use crate::tool_catalog_cache::McpToolCatalogCacheContext; use crate::tools::ToolFilter; use crate::tools::ToolInfo; use crate::tools::filter_tools; +use crate::trusted_access::ENTITLEMENT_CONTEXT_KEY; +use crate::trusted_access::TrustedAccessContext; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; @@ -192,6 +194,7 @@ pub(crate) struct McpConnectionSet { prefix_mcp_tool_names: bool, non_prefixed_mcp_tool_servers: Vec, elicitation_requests: ElicitationRequestManager, + pub(crate) trusted_access: Option, } impl McpConnectionSet { @@ -203,6 +206,7 @@ impl McpConnectionSet { input: McpRuntimeInput, elicitation_router: ElicitationRequestRouter, ) -> Self { + let trusted_access = TrustedAccessContext::from_runtime(&input); let McpRuntimeInput { startup_policy, config, @@ -218,7 +222,7 @@ impl McpConnectionSet { codex_apps_tools_cache_key, client_mcp_extensions, auth, - codex_apps_auth_manager, + auth_manager, elicitation_reviewer, elicitation_lifecycle, } = input; @@ -271,7 +275,7 @@ impl McpConnectionSet { let static_chatgpt_auth_provider = auth .filter(|auth| auth.uses_codex_backend()) .map(codex_model_provider::auth_provider_from_auth); - let codex_apps_auth_provider = codex_apps_auth_manager.and_then(|auth_manager| { + let codex_apps_auth_provider = auth_manager.and_then(|auth_manager| { auth.filter(|auth| auth.uses_codex_backend()).map(|auth| { codex_model_provider::auth_provider_from_auth_manager(auth_manager, auth) }) @@ -682,6 +686,7 @@ impl McpConnectionSet { prefix_mcp_tool_names, non_prefixed_mcp_tool_servers, elicitation_requests: elicitation_requests.clone(), + trusted_access, }; let summary_publication_gate = publication_gate; tokio::spawn(async move { @@ -742,6 +747,7 @@ impl McpConnectionSet { prefix_mcp_tool_names, non_prefixed_mcp_tool_servers: Vec::new(), elicitation_requests: ElicitationRequestManager::default(), + trusted_access: None, } } @@ -867,7 +873,7 @@ impl McpConnectionSet { tool: &str, environment_id: Option<&str>, arguments: Option, - meta: Option, + mut meta: Option, requested_timeout: Option, wait_for_server: bool, ) -> Result { @@ -911,6 +917,10 @@ impl McpConnectionSet { } (server_timeout, requested_timeout) => server_timeout.or(requested_timeout), }; + // Direct callers cannot supply host-owned entitlement metadata, even for unlisted tools. + if let Some(serde_json::Value::Object(meta)) = meta.as_mut() { + meta.remove(ENTITLEMENT_CONTEXT_KEY); + } let result: rmcp::model::CallToolResult = client .client .call_tool(tool.to_string(), arguments, meta, effective_timeout) diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index e168840c15..eb024b173e 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -36,8 +36,14 @@ use codex_connectors::ConnectorRuntimeContext; use codex_connectors::ConnectorRuntimeContextKey; use codex_connectors::ConnectorRuntimeFetchSource; use codex_connectors::ConnectorRuntimeManager; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; use codex_exec_server_test_support::environment_manager_without_environments; use codex_login::AuthHeaders; +use codex_login::AuthManager; use codex_login::CodexAuth; use codex_protocol::ToolName; use codex_protocol::approvals::ElicitationRequest; @@ -115,6 +121,7 @@ impl McpConnectionSet { /*lifecycle*/ None, ElicitationRequestRouter::default(), ), + trusted_access: None, } } @@ -238,6 +245,24 @@ fn create_test_server_info(title: &str) -> McpServerInfo { struct TestInProcessTransportFactory; +struct PendingHttpClient; + +impl HttpClient for PendingHttpClient { + fn http_request( + &self, + _params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + futures::future::pending().boxed() + } + + fn http_request_stream( + &self, + _params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + futures::future::pending().boxed() + } +} + impl InProcessTransportFactory for TestInProcessTransportFactory { fn open(&self) -> BoxFuture<'static, io::Result> { async { @@ -438,6 +463,81 @@ async fn create_test_managed_client(tools: Vec) -> ManagedClient { } } +#[tokio::test(start_paused = true)] +async fn prepared_call_timeout_includes_trusted_access_lookup() { + let mut tool = create_test_tool("docs", "access"); + tool.tool.annotations = Some(rmcp::model::ToolAnnotations::new().read_only(true)); + let mut tool_meta = rmcp::model::MetaObject::new(); + tool_meta.insert( + "openai/requestedEntitlements".to_string(), + serde_json::json!(["cyber_trusted_access"]), + ); + tool.tool.meta = Some(tool_meta); + + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionSet::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + let mut config = crate::mcp::tests::test_mcp_config(std::env::temp_dir()); + let mut catalog = crate::ResolvedMcpCatalog::builder(); + catalog.register(crate::McpServerRegistration::from_plugin( + "docs".to_string(), + crate::McpPluginAttribution::new("docs@test".to_string(), "Docs".to_string()), + /*plugin_order*/ 0, + serde_json::from_value(serde_json::json!({ "command": "docs" })) + .expect("plugin MCP config"), + )); + config.mcp_server_catalog = catalog.build(); + config + .server_permission_profiles + .insert("docs".to_string(), PermissionProfile::default()); + manager.tool_plugin_provenance = Arc::new(crate::tool_plugin_provenance(&config)); + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + manager.trusted_access = Some(TrustedAccessContext::new( + auth.clone(), + AuthManager::from_auth_for_testing(auth), + "https://chatgpt.com/backend-api".to_string(), + Arc::new(PendingHttpClient), + )); + let manager = Arc::new(manager); + let server_metadata = McpServerMetadata { + environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + pollutes_memory: true, + origin: Some(McpServerOrigin::Stdio), + supports_parallel_tool_calls: false, + default_tools_approval_mode: None, + tool_approval_modes: HashMap::new(), + }; + let prepared = crate::PreparedMcpCall::new( + manager, + Arc::new(create_test_managed_client(vec![tool.clone()]).await), + Arc::new(config), + /*catalog_revision*/ 0, + Arc::new(RwLock::new(0)), + tool, + server_metadata, + Some("docs@test".to_string()), + /*selected_plugin_server*/ false, + ) + .expect("docs should retain its permission profile"); + + let started = tokio::time::Instant::now(); + let error = prepared + .call( + Some(serde_json::json!({})), + /*meta*/ None, + Some(Duration::from_secs(1)), + ) + .await + .expect_err("trusted access lookup should consume the call timeout"); + + assert_eq!(started.elapsed(), Duration::from_secs(1)); + assert!(format!("{error:#}").contains("timed out awaiting tools/call after 1s")); +} + async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { AsyncManagedClient { client: futures::future::ready::>(Ok( @@ -1887,7 +1987,7 @@ async fn codex_apps_extension_does_not_share_host_owned_tools_cache() -> anyhow: codex_apps_tools_cache_key: cache_key, client_mcp_extensions: ClientMcpExtensions::default(), auth: None, - codex_apps_auth_manager: None, + auth_manager: None, elicitation_reviewer: None, elicitation_lifecycle: None, }, @@ -3817,7 +3917,7 @@ async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() - ), client_mcp_extensions: ClientMcpExtensions::default(), auth: Some(hosted_auth.clone()), - codex_apps_auth_manager: None, + auth_manager: None, elicitation_reviewer: None, elicitation_lifecycle: None, }, @@ -3936,7 +4036,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ), client_mcp_extensions: ClientMcpExtensions::default(), auth: None, - codex_apps_auth_manager: None, + auth_manager: None, elicitation_reviewer: None, elicitation_lifecycle: None, }, @@ -4333,7 +4433,7 @@ async fn reconcile_reusable_server( ), client_mcp_extensions: ClientMcpExtensions::default(), auth: None, - codex_apps_auth_manager: None, + auth_manager: None, elicitation_reviewer: None, elicitation_lifecycle: None, }, @@ -4787,7 +4887,7 @@ async fn reconciliation_replaces_connection_when_protocol_mode_changes() { ), client_mcp_extensions: ClientMcpExtensions::default(), auth: None, - codex_apps_auth_manager: None, + auth_manager: None, elicitation_reviewer: None, elicitation_lifecycle: None, }, @@ -4845,7 +4945,7 @@ async fn reconciliation_reuses_legacy_stdio_server_when_modern_protocol_is_enabl ), client_mcp_extensions: ClientMcpExtensions::default(), auth: None, - codex_apps_auth_manager: None, + auth_manager: None, elicitation_reviewer: None, elicitation_lifecycle: None, }, diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 05d8a8de68..b1261863b7 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -431,7 +431,7 @@ pub async fn read_mcp_resource( codex_apps_tools_cache_key: connector_runtime_context_key(auth), client_mcp_extensions: ClientMcpExtensions::default(), auth: auth.cloned(), - codex_apps_auth_manager: None, + auth_manager: None, elicitation_reviewer: None, elicitation_lifecycle: None, }, @@ -506,7 +506,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( codex_apps_tools_cache_key: connector_runtime_context_key(auth), client_mcp_extensions: ClientMcpExtensions::default(), auth: auth.cloned(), - codex_apps_auth_manager: None, + auth_manager: None, elicitation_reviewer: None, elicitation_lifecycle: None, }, diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 57d5299177..bdd07128e1 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -80,7 +80,7 @@ pub struct McpRuntimeInput { pub codex_apps_tools_cache_key: ConnectorRuntimeContextKey, pub client_mcp_extensions: ClientMcpExtensions, pub auth: Option, - pub codex_apps_auth_manager: Option>, + pub auth_manager: Option>, pub elicitation_reviewer: Option, pub elicitation_lifecycle: Option, } @@ -672,7 +672,7 @@ impl McpRuntimeContext { self.local_process_cwd.clone() } - fn local_http_client(&self) -> Arc { + pub(crate) fn local_http_client(&self) -> Arc { Arc::clone(&self.local_http_client) } diff --git a/codex-rs/codex-mcp/src/trusted_access.rs b/codex-rs/codex-mcp/src/trusted_access.rs index 97622b646d..baed4690cb 100644 --- a/codex-rs/codex-mcp/src/trusted_access.rs +++ b/codex-rs/codex-mcp/src/trusted_access.rs @@ -1,6 +1,11 @@ use std::sync::Arc; use std::time::Duration; +use crate::connection_manager::McpConnectionSet; +use crate::runtime::McpRuntimeInput; +use crate::server::McpServerMetadata; +use crate::server::McpServerOrigin; +use crate::tools::ToolInfo; use codex_exec_server::HttpClient; use codex_exec_server::HttpHeader; use codex_exec_server::HttpRedirectPolicy; @@ -14,10 +19,55 @@ use serde_json::Map; use serde_json::Value; use serde_json::json; -const ENTITLEMENT_CONTEXT_KEY: &str = "openai/entitlementContext"; +pub(crate) const ENTITLEMENT_CONTEXT_KEY: &str = "openai/entitlementContext"; const MAX_VERIFIED_ACCESS_RESPONSE_BYTES: usize = 1024 * 1024; +const REQUESTED_ENTITLEMENTS_KEY: &str = "openai/requestedEntitlements"; +const CYBER_TRUSTED_ACCESS_ENTITLEMENT: &str = "cyber_trusted_access"; const TRUSTED_ACCESS_TIMEOUT: Duration = Duration::from_millis(2_500); +impl McpConnectionSet { + /// Installed and task-selected plugins may request supported advisory entitlement metadata. + /// Model calls use the local, read-only, zero-argument boundary. + pub(crate) async fn add_trusted_access_context( + &self, + tool: &ToolInfo, + server: &McpServerMetadata, + arguments: Option<&Value>, + meta: Option, + ) -> Option { + if tool + .tool + .meta + .as_deref() + .and_then(|meta| meta.get(REQUESTED_ENTITLEMENTS_KEY)) + .and_then(Value::as_array) + .is_some_and(|entitlements| { + entitlements.iter().all(Value::is_string) + && entitlements.iter().any(|entitlement| { + entitlement.as_str() == Some(CYBER_TRUSTED_ACCESS_ENTITLEMENT) + }) + }) + && self + .plugin_id_for_mcp_server_name(&tool.server_name) + .is_some() + && arguments.is_none_or(|arguments| arguments.as_object().is_some_and(Map::is_empty)) + && server.environment_id == codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID + && matches!(server.origin, Some(McpServerOrigin::Stdio)) + && tool + .tool + .annotations + .as_ref() + .and_then(|annotations| annotations.read_only_hint) + == Some(true) + && let Some(context) = self.trusted_access.as_ref() + { + context.add_context(meta).await + } else { + meta + } + } +} + #[derive(Deserialize)] struct VerifiedAccessResponse { programs: Vec, @@ -69,6 +119,22 @@ pub struct TrustedAccessContext { } impl TrustedAccessContext { + pub(crate) fn from_runtime(input: &McpRuntimeInput) -> Option { + let auth = input.auth.as_ref()?; + if !matches!( + auth.api_auth_mode(), + AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens + ) { + return None; + } + Some(Self::new( + auth.clone(), + input.auth_manager.clone()?, + input.config.chatgpt_base_url.clone(), + input.runtime_context.local_http_client(), + )) + } + pub fn new( auth: CodexAuth, auth_manager: Arc, diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index ef3b7839fe..ba6ca1a544 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -268,7 +268,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( codex_apps_tools_cache_key: connector_runtime_context_key(auth.as_ref()), client_mcp_extensions: ClientMcpExtensions::default(), auth: auth.clone(), - codex_apps_auth_manager, + auth_manager: codex_apps_auth_manager, elicitation_reviewer: None, elicitation_lifecycle: None, }) diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index b241d35165..421dcf6af2 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -358,10 +358,6 @@ impl Session { }) .collect(), ); - let codex_apps_auth_manager = - codex_mcp::host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()) - .then(|| Arc::clone(&self.services.auth_manager)); - McpRuntimeInput { startup_policy: if matches!(desired.session_source, SessionSource::SubAgent(_)) { McpStartupPolicy::LazyWhenCached @@ -381,7 +377,7 @@ impl Session { codex_apps_tools_cache_key: connector_runtime_context_key(auth.as_ref()), client_mcp_extensions: self.services.client_mcp_extensions.for_mcp_servers(), auth, - codex_apps_auth_manager, + auth_manager: Some(Arc::clone(&self.services.auth_manager)), elicitation_reviewer, elicitation_lifecycle: Some(self.mcp_elicitation_lifecycle()), } diff --git a/codex-rs/core/tests/suite/daybreak_access.rs b/codex-rs/core/tests/suite/daybreak_access.rs new file mode 100644 index 0000000000..34cf689a97 --- /dev/null +++ b/codex-rs/core/tests/suite/daybreak_access.rs @@ -0,0 +1,448 @@ +use std::sync::Arc; + +use anyhow::Result; +use codex_core::TurnInputRequest; +use codex_features::Feature; +use codex_login::CodexAuth; +use codex_login::ExternalAuth; +use codex_login::ExternalAuthFuture; +use codex_login::ExternalAuthRefreshContext; +use codex_protocol::protocol::EventMsg; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::stdio_server_bin; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use core_test_support::wait_for_mcp_server; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use tempfile::TempDir; +use test_case::test_case; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::net::TcpListener; +use tokio::sync::oneshot; +use wiremock::Mock; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[derive(Clone, Copy)] +enum AccessCase { + InstalledStdio, + UndeclaredDaybreakName, + ExternalTokens, + #[cfg(unix)] + RemoteStdio, + UnownedStdio, + InstalledHttp, + MutableStdio, + NonemptyArguments, + ApiKey, +} + +#[test_case(AccessCase::InstalledStdio; "installed local plugin")] +#[test_case(AccessCase::UndeclaredDaybreakName; "undeclared Daybreak tool name")] +#[test_case(AccessCase::ExternalTokens; "external ChatGPT tokens")] +#[cfg_attr(unix, test_case(AccessCase::RemoteStdio; "installed remote stdio plugin"))] +#[test_case(AccessCase::UnownedStdio; "unowned local server")] +#[test_case(AccessCase::InstalledHttp; "HTTP plugin in default environment")] +#[test_case(AccessCase::MutableStdio; "mutable plugin tool")] +#[test_case(AccessCase::NonemptyArguments; "nonempty arguments")] +#[test_case(AccessCase::ApiKey; "API key session")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn daybreak_metadata_follows_the_actual_plugin_call(case: AccessCase) -> Result<()> { + skip_if_no_network!(Ok(())); + #[cfg(unix)] + let remote = if matches!(case, AccessCase::RemoteStdio) { + Some(super::multi_exec_server_sandbox::ExecServerProcess::start().await?) + } else { + None + }; + #[cfg(unix)] + let remote_url = remote.as_ref().map(|remote| remote.websocket_url.clone()); + #[cfg(not(unix))] + let remote_url: Option = None; + let tool_name = if matches!( + case, + AccessCase::UndeclaredDaybreakName | AccessCase::InstalledHttp + ) { + "get_daybreak_access" + } else { + "get_codex_security_daybreak_access" + }; + let server = responses::start_mock_server().await; + Mock::given(method("GET")) + .and(path("/accounts/verified_access")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "programs": [{"program":"cyber", "state":"active", "grants":[ + {"level":"tac1", "source":"individual"} + ]}] + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/remote-mcp")) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("MCP JSON-RPC request"); + let result = match body["method"].as_str() { + Some("initialize") => json!({ + "protocolVersion": body["params"]["protocolVersion"], + "capabilities":{"tools":{}}, + "serverInfo":{"name":"daybreak-test", "version":"1"} + }), + Some("tools/list") => json!({"tools":[{ + "name":"get_daybreak_access", "description":"Check Daybreak access", + "inputSchema":{"type":"object", "properties":{}}, + "annotations":{"readOnlyHint":true} + }]}), + Some("tools/call") => json!({ + "content":[], "structuredContent":body["params"]["_meta"] + }), + _ => return ResponseTemplate::new(202), + }; + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc":"2.0", "id":body["id"], "result":result + })) + }) + .mount(&server) + .await; + + let arguments = if matches!(case, AccessCase::NonemptyArguments) { + r#"{"refresh":true}"# + } else { + "{}" + }; + let home = Arc::new(TempDir::new()?); + let plugin_root = super::plugins::write_sample_plugin_manifest_and_config(&home); + let server_config = if matches!(case, AccessCase::InstalledHttp) { + json!({"url":format!("{}/remote-mcp", server.uri())}) + } else { + json!({ + "command":stdio_server_bin()?, + "cwd":plugin_root, + "environment_id": if remote_url.is_some() { "remote" } else { "local" }, + "default_tools_approval_mode":"approve", + "env":{"MCP_TEST_DAYBREAK_READ_ONLY": + if matches!(case, AccessCase::MutableStdio) { "false" } else { "true" }} + }) + }; + if matches!(case, AccessCase::UnownedStdio) { + let config_path = home.path().join("config.toml"); + let existing = std::fs::read_to_string(&config_path)?; + let extra = toml::to_string(&json!({"mcp_servers":{"sample":server_config}}))?; + std::fs::write(config_path, format!("{existing}\n{extra}"))?; + } else { + std::fs::write( + plugin_root.join(".mcp.json"), + serde_json::to_vec(&json!({ + "mcpServers":{"sample":server_config} + }))?, + )?; + } + let auth = if matches!(case, AccessCase::ApiKey) { + CodexAuth::from_api_key("test-api-key") + } else if matches!(case, AccessCase::ExternalTokens) { + CodexAuth::from_external_chatgpt_tokens( + "header.e30.external", + "external-account", + /*chatgpt_plan_type*/ None, + )? + } else { + CodexAuth::create_dummy_chatgpt_auth_for_testing() + }; + let base_url = server.uri(); + let mut builder = test_codex() + .with_home(home) + .with_auth(auth) + .with_config(move |config| config.chatgpt_base_url = base_url); + if let Some(remote_url) = &remote_url { + builder = builder.with_exec_server_url(remote_url.clone()); + } + let test = builder.build_with_remote_and_local_env(&server).await?; + if remote_url.is_some() { + assert!( + test.thread_manager + .environment_manager() + .get_environment("remote") + .expect("remote executor") + .is_remote() + ); + let (config, _) = test.codex.current_mcp_config_and_runtime_context().await; + assert_eq!( + config.mcp_server_catalog.configured_servers()["sample"].environment_id, + "remote" + ); + } + wait_for_mcp_server(&test.codex, "sample").await?; + let mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_tool_search_call("search", &json!({"query":tool_name})), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "daybreak-call", + "mcp__sample", + tool_name, + arguments, + ), + responses::ev_completed("resp-2"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("done", "done"), + responses::ev_completed("resp-3"), + ]), + ], + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Call the sample Daybreak access tool".to_string(), + text_elements: Vec::new(), + }])) + .await?; + let end = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::McpToolCallEnd(_)) + }) + .await; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.codex.shutdown_and_wait().await?; + let EventMsg::McpToolCallEnd(end) = end else { + unreachable!() + }; + let metadata = end + .result + .expect("successful MCP call") + .structured_content + .expect("echoed metadata"); + let expected_context = match case { + AccessCase::InstalledStdio | AccessCase::ExternalTokens => { + Some(json!({"schemaVersion":1, "entitlements":{ + "cyber_trusted_access":{"schemaVersion":1,"status":"granted","grants":[ + {"level":"tac1","source":"user"} + ],"stale":false} + }})) + } + #[cfg(unix)] + AccessCase::RemoteStdio => None, + AccessCase::UndeclaredDaybreakName + | AccessCase::UnownedStdio + | AccessCase::InstalledHttp + | AccessCase::MutableStdio + | AccessCase::NonemptyArguments + | AccessCase::ApiKey => None, + }; + assert_eq!( + metadata.get("openai/entitlementContext"), + expected_context.as_ref() + ); + assert_eq!( + end.plugin_id.as_deref(), + if matches!(case, AccessCase::UnownedStdio) { + None + } else { + Some("sample@test") + } + ); + assert!(metadata["threadId"].is_string()); + assert!(!metadata.to_string().contains("Access Token")); + assert!(!metadata.to_string().contains("header.e30.external")); + let requests = server.received_requests().await.expect("recorded requests"); + let access_requests: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == "/accounts/verified_access") + .collect(); + assert_eq!( + access_requests.len(), + usize::from(matches!( + case, + AccessCase::InstalledStdio | AccessCase::ExternalTokens + )) + ); + if let Some(request) = access_requests.first() { + let (authorization, account_id) = if matches!(case, AccessCase::ExternalTokens) { + ("Bearer header.e30.external", "external-account") + } else { + ("Bearer Access Token", "account_id") + }; + assert_eq!(request.headers["authorization"], authorization); + assert_eq!(request.headers["chatgpt-account-id"], account_id); + } + assert!( + mock.requests()[2] + .function_call_output("daybreak-call") + .is_object() + ); + Ok(()) +} + +struct SelectedAuth(CodexAuth); + +impl ExternalAuth for SelectedAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.clone()) }) + } + + fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.clone()) }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn daybreak_discards_in_flight_account_changes_without_apps() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let listener = TcpListener::bind("127.0.0.1:0").await?; + let base_url = format!("http://{}", listener.local_addr()?); + let (request_tx, request_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let response_task = tokio::spawn(async move { + let (mut stream, request) = loop { + let (stream, _) = listener.accept().await?; + let mut stream = BufReader::new(stream); + let mut request = String::new(); + loop { + let mut line = String::new(); + let read = stream.read_line(&mut line).await?; + if read == 0 || line == "\r\n" { + break; + } + request.push_str(&line); + } + if request.starts_with("GET /accounts/verified_access ") { + break (stream, request); + } + stream + .get_mut() + .write_all( + b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await?; + }; + request_tx.send(request).expect("request receiver"); + release_rx.await?; + let body = json!({"programs":[{"program":"cyber","state":"active","grants":[ + {"level":"tac1","source":"individual"} + ]}]}) + .to_string(); + stream.get_mut().write_all(format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ).as_bytes()).await?; + Ok::<_, anyhow::Error>(()) + }); + let home = Arc::new(TempDir::new()?); + let plugin_root = super::plugins::write_sample_plugin_manifest_and_config(&home); + std::fs::write( + plugin_root.join(".mcp.json"), + serde_json::to_vec(&json!({ + "mcpServers":{"sample":{"command":stdio_server_bin()?,"cwd":plugin_root, + "environment_id":"local","default_tools_approval_mode":"approve", + "env":{"MCP_TEST_DAYBREAK_READ_ONLY":"true"}}} + }))?, + )?; + let test = test_codex() + .with_home(home) + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config.chatgpt_base_url = base_url; + config + .features + .disable(Feature::Apps) + .expect("disable Apps"); + }) + .build_with_remote_and_local_env(&server) + .await?; + wait_for_mcp_server(&test.codex, "sample").await?; + responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_tool_search_call( + "search", + &json!({"query":"get_codex_security_daybreak_access"}), + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "daybreak-call", + "mcp__sample", + "get_codex_security_daybreak_access", + "{}", + ), + responses::ev_completed("resp-2"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("done", "done"), + responses::ev_completed("resp-3"), + ]), + ], + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Call the sample Daybreak access tool".to_string(), + text_elements: Vec::new(), + }])) + .await?; + let call = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::McpToolCallEnd(_)) + }); + let switch = async { + let request = + tokio::time::timeout(std::time::Duration::from_secs(10), request_rx).await??; + assert!(request.starts_with("GET /accounts/verified_access ")); + assert!( + request + .to_ascii_lowercase() + .contains("chatgpt-account-id: account_id") + ); + test.thread_manager + .auth_manager() + .set_external_auth(Arc::new(SelectedAuth( + CodexAuth::from_external_chatgpt_tokens( + "header.e30.changed", + "other-account", + /*chatgpt_plan_type*/ None, + )?, + ))) + .await?; + release_tx.send(()).expect("release response"); + Ok::<_, anyhow::Error>(()) + }; + let (result, switched) = tokio::join!(call, switch); + switched?; + response_task.await??; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.codex.shutdown_and_wait().await?; + let EventMsg::McpToolCallEnd(result) = result else { + unreachable!() + }; + assert_eq!( + result + .result + .expect("successful call") + .structured_content + .expect("metadata")["openai/entitlementContext"], + json!({ + "schemaVersion":1,"entitlements":{"cyber_trusted_access":{ + "schemaVersion":1,"status":"unknown","grants":[],"stale":false + }} + }) + ); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mcp_auth_refresh.rs b/codex-rs/core/tests/suite/mcp_auth_refresh.rs index 1552312644..6a9f3ec62c 100644 --- a/codex-rs/core/tests/suite/mcp_auth_refresh.rs +++ b/codex-rs/core/tests/suite/mcp_auth_refresh.rs @@ -108,7 +108,7 @@ async fn hosted_plugin_runtime_ps_mcp_tool_calls_use_current_auth_manager_token( codex_apps_tools_cache_key: codex_mcp::codex_apps_tools_cache_key(Some(&expected_auth)), client_mcp_extensions: ClientMcpExtensions::default(), auth: Some(expected_auth.clone()), - codex_apps_auth_manager: Some(Arc::clone(&auth_manager)), + auth_manager: Some(Arc::clone(&auth_manager)), elicitation_reviewer: None, elicitation_lifecycle: None, }) diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 8437a41d5d..a34c08a902 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -62,6 +62,7 @@ mod context_annotations; mod current_time_reminder; mod cyber_access_program; mod cyber_exec_policy; +mod daybreak_access; mod deprecation_notice; mod exec; mod exec_policy; diff --git a/codex-rs/core/tests/suite/multi_exec_server_sandbox.rs b/codex-rs/core/tests/suite/multi_exec_server_sandbox.rs index 840bda2797..8928df9f6f 100644 --- a/codex-rs/core/tests/suite/multi_exec_server_sandbox.rs +++ b/codex-rs/core/tests/suite/multi_exec_server_sandbox.rs @@ -43,15 +43,15 @@ const SECOND_CALL_ID: &str = "write-from-second"; const EXEC_SERVER_START_TIMEOUT: Duration = Duration::from_secs(30); const TURN_COMPLETE_TIMEOUT: Duration = Duration::from_secs(30); -struct ExecServerProcess { +pub(super) struct ExecServerProcess { _codex_home: TempDir, child: Child, _stdout: BufReader, - websocket_url: String, + pub(super) websocket_url: String, } impl ExecServerProcess { - async fn start() -> Result { + pub(super) async fn start() -> Result { let codex_home = TempDir::new()?; let mut child = Command::new(codex_utils_cargo_bin::cargo_bin("codex")?) .args(["exec-server", "--listen", "ws://127.0.0.1:0"]) 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 b049531711..79a2e1c066 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -79,6 +79,28 @@ impl TestToolServer { Arc::new(sandbox_meta_schema), ); sandbox_meta_tool.annotations = Some(ToolAnnotations::new().read_only(true)); + let entitlement_tools = std::env::var("MCP_TEST_DAYBREAK_READ_ONLY") + .ok() + .into_iter() + .flat_map(|read_only| { + ["get_codex_security_daybreak_access", "get_daybreak_access"].map(|name| { + let mut tool = sandbox_meta_tool.clone(); + tool.name = Cow::Borrowed(name); + tool.description = + Some(Cow::Borrowed("Return requested account access metadata.")); + tool.annotations = Some(ToolAnnotations::new().read_only(read_only == "true")); + if name == "get_codex_security_daybreak_access" { + let mut meta = MetaObject::new(); + meta.insert( + "openai/requestedEntitlements".to_string(), + json!(["cyber_trusted_access"]), + ); + tool.meta = Some(meta); + } + tool + }) + }) + .collect::>(); #[expect(clippy::expect_used)] let thread_hint_schema: JsonObject = serde_json::from_value(json!({ @@ -124,6 +146,7 @@ impl TestToolServer { Self::image_scenario_tool(), sandbox_meta_tool, ]; + tools.extend(entitlement_tools); if std::env::var_os("MCP_TEST_ENABLE_NODE_REPL_JS").is_some() { #[expect(clippy::expect_used)] let schema: JsonObject = serde_json::from_value(json!({ @@ -662,9 +685,9 @@ impl ServerHandler for TestToolServer { .supports_openai_form_elicitation .load(Ordering::Relaxed), }))), - "sandbox_meta" => Ok(Self::structured_result(serde_json::Value::Object( - context.meta.0.0, - ))), + "sandbox_meta" | "get_codex_security_daybreak_access" | "get_daybreak_access" => Ok( + Self::structured_result(serde_json::Value::Object(context.meta.0.0)), + ), "cwd" => { let cwd = std::env::current_dir() .map(|path| path.to_string_lossy().into_owned())