Attach verified access context to eligible plugin MCP calls (#41005)

## What changed

- Fetch ChatGPT account access and attach `cyber_trusted_access` under
  `openai/entitlementContext` when an installed or selected plugin explicitly
  requests it from a local, read-only stdio tool with no arguments.
- Reject caller-supplied entitlement context on direct MCP calls, and fall back
  to an `unknown` status when access cannot be verified or the account changes
  during the lookup.
- Count entitlement lookup time against the MCP tool-call timeout.

## Testing

- Cover eligible calls plus unowned, remote, HTTP, mutable, argument-bearing,
  API-key, undeclared, direct, and account-switch cases in core and app-server
  integration tests.

GitOrigin-RevId: 303b61585ef4530276c373d914a2d0ccef1f4ac2

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<String> = 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(())
+}

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,
     })

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;

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<ChildStdout>,
-    websocket_url: String,
+    pub(super) websocket_url: String,
 }
 
 impl ExecServerProcess {
-    async fn start() -> Result<Self> {
+    pub(super) async fn start() -> Result<Self> {
         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"])