From a98b94625e32d3fd62217364999eab33536bb61b Mon Sep 17 00:00:00 2001 From: soyeon-oai Date: Wed, 26 Aug 2026 22:45:07 +0000 Subject: [PATCH] Add trusted access context for MCP metadata (#40992) ## What changed - Add `TrustedAccessContext` to fetch account-bound cyber verified-access grants and attach them as host-owned `openai/entitlementContext` metadata. - Replace caller-supplied entitlement metadata and return an `unknown` status when authentication, identity consistency, response validation, size limits, or the lookup deadline cannot be verified. - Add `HttpResponseBodyStream::from_chunks` for lightweight in-memory streaming clients. ## Testing - Cover grant mapping, unsupported authentication, account and workspace changes, malformed or oversized responses, duplicate programs, and lookup timeouts. GitOrigin-RevId: 343b6a6608c98c45a34939ef507c2183516a11c1 --- codex-rs/codex-mcp/src/lib.rs | 2 + codex-rs/codex-mcp/src/trusted_access.rs | 231 ++++++ .../codex-mcp/src/trusted_access_tests.rs | 658 ++++++++++++++++++ .../src/client/http_response_body_stream.rs | 18 + 4 files changed, 909 insertions(+) create mode 100644 codex-rs/codex-mcp/src/trusted_access.rs create mode 100644 codex-rs/codex-mcp/src/trusted_access_tests.rs diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 9d24d2ddc0..254f992899 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -25,6 +25,7 @@ pub use runtime::SandboxState; pub use runtime::apply_http_headers_helper; pub use tool_catalog_cache::McpToolCatalogCache; pub use tools::ToolInfo; +pub use trusted_access::TrustedAccessContext; /// Backward-compatible name for the shared Codex Apps tools runtime. pub type CodexAppsToolsCache = codex_connectors::ConnectorRuntimeManager; @@ -112,3 +113,4 @@ pub(crate) mod runtime; pub(crate) mod server; mod tool_catalog_cache; pub(crate) mod tools; +mod trusted_access; diff --git a/codex-rs/codex-mcp/src/trusted_access.rs b/codex-rs/codex-mcp/src/trusted_access.rs new file mode 100644 index 0000000000..97622b646d --- /dev/null +++ b/codex-rs/codex-mcp/src/trusted_access.rs @@ -0,0 +1,231 @@ +use std::sync::Arc; +use std::time::Duration; + +use codex_exec_server::HttpClient; +use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRedirectPolicy; +use codex_exec_server::HttpRequestParams; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_protocol::auth::AuthMode; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Map; +use serde_json::Value; +use serde_json::json; + +const ENTITLEMENT_CONTEXT_KEY: &str = "openai/entitlementContext"; +const MAX_VERIFIED_ACCESS_RESPONSE_BYTES: usize = 1024 * 1024; +const TRUSTED_ACCESS_TIMEOUT: Duration = Duration::from_millis(2_500); + +#[derive(Deserialize)] +struct VerifiedAccessResponse { + programs: Vec, +} + +#[derive(Deserialize)] +struct VerifiedAccessProgram { + state: VerifiedAccessState, + grants: Vec, +} + +#[derive(Clone, Copy, Deserialize)] +#[serde(rename_all = "snake_case")] +enum VerifiedAccessState { + Active, + Inactive, + Unavailable, +} + +#[derive(Deserialize)] +struct VerifiedAccessGrant { + level: TrustedAccessLevel, + source: VerifiedAccessSource, +} + +#[derive(Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +enum TrustedAccessLevel { + Tac1, + Tac2, + Tac3, + Government, +} + +#[derive(Deserialize)] +#[serde(rename_all = "snake_case")] +enum VerifiedAccessSource { + Individual, + Organization, +} + +/// Fetches account-bound verified access for trusted, host-owned MCP metadata. +/// Callers must authorize the receiving plugin and tool before attaching it. +pub struct TrustedAccessContext { + auth: CodexAuth, + auth_manager: Arc, + chatgpt_base_url: String, + http_client: Arc, +} + +impl TrustedAccessContext { + pub fn new( + auth: CodexAuth, + auth_manager: Arc, + chatgpt_base_url: String, + http_client: Arc, + ) -> Self { + Self { + auth, + auth_manager, + chatgpt_base_url, + http_client, + } + } + + /// Replaces caller-supplied entitlement metadata with a fresh verified result. + pub async fn add_context(&self, meta: Option) -> Option { + let mut meta = match meta { + Some(Value::Object(meta)) => meta, + None => Map::new(), + other => return other, + }; + meta.remove(ENTITLEMENT_CONTEXT_KEY); + + let status = tokio::time::timeout(TRUSTED_ACCESS_TIMEOUT, self.fetch_status()) + .await + .ok() + .flatten() + .unwrap_or_else(|| { + json!({ + "schemaVersion": 1, + "status": "unknown", + "grants": [], + "stale": false + }) + }); + meta.insert( + ENTITLEMENT_CONTEXT_KEY.to_string(), + json!({ + "schemaVersion": 1, + "entitlements": { "cyber_trusted_access": status } + }), + ); + Some(Value::Object(meta)) + } + + async fn fetch_status(&self) -> Option { + let auth = self.auth_manager.auth().await?; + let expected_account_id = self + .auth + .get_account_id() + .filter(|account_id| !account_id.trim().is_empty())?; + let account_id = auth + .get_account_id() + .filter(|account_id| !account_id.trim().is_empty())?; + if !matches!( + self.auth.api_auth_mode(), + AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens + ) || !matches!( + auth.api_auth_mode(), + AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens + ) || account_id != expected_account_id + || auth.get_chatgpt_user_id() != self.auth.get_chatgpt_user_id() + || auth.is_workspace_account() != self.auth.is_workspace_account() + || auth.is_fedramp_account() != self.auth.is_fedramp_account() + { + return None; + } + + let headers = codex_model_provider::auth_provider_from_auth(&auth) + .to_auth_headers() + .iter() + .map(|(name, value)| { + Some(HttpHeader { + name: name.as_str().to_string(), + value: value.to_str().ok()?.to_string(), + value_env_var: None, + }) + }) + .collect::>>()?; + let (response, mut response_body) = self + .http_client + .http_request_stream(HttpRequestParams { + method: "GET".to_string(), + url: format!( + "{}/accounts/verified_access", + self.chatgpt_base_url.trim_end_matches('/') + ), + headers, + body: None, + timeout_ms: Some(TRUSTED_ACCESS_TIMEOUT.as_millis() as u64), + redirect_policy: HttpRedirectPolicy::Stop, + request_id: "trusted-access-status".to_string(), + stream_response: true, + }) + .await + .ok()?; + if response.status != 200 { + return None; + } + let mut response_bytes = Vec::new(); + while let Some(chunk) = response_body.recv().await.ok()? { + if chunk.len() > MAX_VERIFIED_ACCESS_RESPONSE_BYTES - response_bytes.len() { + return None; + } + response_bytes.extend_from_slice(&chunk); + } + let response: VerifiedAccessResponse = serde_json::from_slice(&response_bytes).ok()?; + let mut programs = response + .programs + .into_iter() + .filter(|program| program.get("program").and_then(Value::as_str) == Some("cyber")); + let program = programs.next()?; + if programs.next().is_some() { + return None; + } + let program: VerifiedAccessProgram = serde_json::from_value(program).ok()?; + + if self.auth_manager.auth_cached().is_none_or(|current| { + !matches!( + current.api_auth_mode(), + AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens + ) || current.get_account_id().as_deref() != Some(account_id.as_str()) + || current.get_chatgpt_user_id() != auth.get_chatgpt_user_id() + || current.is_workspace_account() != auth.is_workspace_account() + || current.is_fedramp_account() != auth.is_fedramp_account() + }) { + return None; + } + + let status = match program.state { + VerifiedAccessState::Active if !program.grants.is_empty() => "granted", + VerifiedAccessState::Inactive if program.grants.is_empty() => "not_granted", + VerifiedAccessState::Unavailable if program.grants.is_empty() => "unknown", + VerifiedAccessState::Active + | VerifiedAccessState::Inactive + | VerifiedAccessState::Unavailable => return None, + }; + let grants = program + .grants + .into_iter() + .map(|grant| { + let source = match grant.source { + VerifiedAccessSource::Individual => "user", + VerifiedAccessSource::Organization => "current_account", + }; + json!({ "level": grant.level, "source": source }) + }) + .collect::>(); + Some(json!({ + "schemaVersion": 1, + "status": status, + "grants": grants, + "stale": false + })) + } +} + +#[cfg(test)] +#[path = "trusted_access_tests.rs"] +mod tests; diff --git a/codex-rs/codex-mcp/src/trusted_access_tests.rs b/codex-rs/codex-mcp/src/trusted_access_tests.rs new file mode 100644 index 0000000000..63bc1d9d4a --- /dev/null +++ b/codex-rs/codex-mcp/src/trusted_access_tests.rs @@ -0,0 +1,658 @@ +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use codex_exec_server::ByteChunk; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRedirectPolicy; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthHeaders; +use codex_login::AuthKeyringBackendKind; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::ExternalAuth; +use codex_login::ExternalAuthFuture; +use codex_login::ExternalAuthRefreshContext; +use futures::FutureExt; +use futures::future::BoxFuture; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use tokio::sync::Notify; + +use super::MAX_VERIFIED_ACCESS_RESPONSE_BYTES; +use super::TrustedAccessContext; + +struct RecordingHttpClient { + requests: Mutex>, + status: u16, + response: Vec, + response_chunks: Option>>, + response_gate: Option<(Notify, Notify)>, +} + +impl RecordingHttpClient { + fn new(status: u16, response: Value) -> Self { + Self { + requests: Mutex::new(Vec::new()), + status, + response: serde_json::to_vec(&response).expect("serialize response"), + response_chunks: None, + response_gate: None, + } + } +} + +impl HttpClient for RecordingHttpClient { + fn http_request( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + self.requests.lock().expect("record request").push(params); + let response = HttpRequestResponse { + status: self.status, + headers: Vec::new(), + body: ByteChunk(self.response.clone()), + }; + async move { + if let Some((requested, release)) = &self.response_gate { + requested.notify_one(); + release.notified().await; + } + Ok(response) + } + .boxed() + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + self.requests.lock().expect("record request").push(params); + let response = HttpRequestResponse { + status: self.status, + headers: Vec::new(), + body: ByteChunk(Vec::new()), + }; + let chunks = self + .response_chunks + .clone() + .unwrap_or_else(|| vec![self.response.clone()]); + async move { + if let Some((requested, release)) = &self.response_gate { + requested.notify_one(); + release.notified().await; + } + Ok((response, HttpResponseBodyStream::from_chunks(chunks))) + } + .boxed() + } +} + +struct StaticExternalAuth(CodexAuth); + +impl ExternalAuth for StaticExternalAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.clone()) }) + } + + fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } +} + +struct SequencedExternalAuth(Mutex>); + +impl ExternalAuth for SequencedExternalAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + let mut sequence = self.0.lock().expect("auth sequence"); + Ok(if sequence.len() > 1 { + sequence.pop_front().expect("next auth") + } else { + sequence.front().expect("last auth").clone() + }) + }) + } + + fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } +} + +fn header_auth() -> CodexAuth { + CodexAuth::Headers(AuthHeaders::new( + [ + ( + "authorization".parse().unwrap(), + "Bearer synthetic-pat".parse().unwrap(), + ), + ( + "chatgpt-account-id".parse().unwrap(), + "account-a".parse().unwrap(), + ), + ] + .into_iter() + .collect(), + )) +} + +fn chatgpt_auth(account_id: &str) -> CodexAuth { + CodexAuth::from_external_chatgpt_tokens( + "header.e30.same", + account_id, + /*chatgpt_plan_type*/ None, + ) + .expect("test auth") +} + +fn fedramp_chatgpt_auth(account_id: &str) -> CodexAuth { + CodexAuth::from_external_chatgpt_tokens( + "header.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF9hY2NvdW50X2lzX2ZlZHJhbXAiOnRydWV9fQ.same", + account_id, + /*chatgpt_plan_type*/ None, + ) + .expect("FedRAMP test auth") +} + +fn context(auth: CodexAuth, client: Arc) -> TrustedAccessContext { + TrustedAccessContext::new( + auth.clone(), + AuthManager::from_auth_for_testing(auth), + "https://chatgpt.com/backend-api".to_string(), + client, + ) +} + +fn cyber_response(state: &str, grants: Value) -> Value { + json!({ "programs": [{ "program": "cyber", "state": state, "grants": grants }] }) +} + +fn expected_metadata(status: &str, grants: Value) -> Value { + json!({ "openai/entitlementContext": { + "schemaVersion": 1, + "entitlements": { "cyber_trusted_access": { + "schemaVersion": 1, + "status": status, + "grants": grants, + "stale": false + } } + } }) +} + +#[tokio::test] +async fn delivers_account_bound_verified_access_to_the_calling_plugin() { + let client = Arc::new(RecordingHttpClient::new( + /*status*/ 200, + cyber_response( + "active", + json!([ + { "level": "tac2", "source": "individual" }, + { "level": "government", "source": "organization" } + ]), + ), + )); + let mut context = context( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + client.clone(), + ); + context.chatgpt_base_url.push('/'); + let metadata = context + .add_context(Some(json!({ "threadId": "thread-1" }))) + .await; + let mut expected = expected_metadata( + "granted", + json!([ + { "level": "tac2", "source": "user" }, + { "level": "government", "source": "current_account" } + ]), + ); + expected["threadId"] = json!("thread-1"); + assert_eq!(metadata, Some(expected)); + + let mut requests = client.requests.lock().expect("inspect recorded requests"); + for request in requests.iter_mut() { + request + .headers + .sort_by(|left, right| left.name.cmp(&right.name)); + } + assert_eq!( + *requests, + vec![HttpRequestParams { + method: "GET".to_string(), + url: "https://chatgpt.com/backend-api/accounts/verified_access".to_string(), + headers: vec![ + HttpHeader { + name: "authorization".to_string(), + value: "Bearer Access Token".to_string(), + value_env_var: None, + }, + HttpHeader { + name: "chatgpt-account-id".to_string(), + value: "account_id".to_string(), + value_env_var: None, + }, + ], + body: None, + timeout_ms: Some(2_500), + redirect_policy: HttpRedirectPolicy::Stop, + request_id: "trusted-access-status".to_string(), + stream_response: true, + }] + ); +} + +#[tokio::test] +async fn rejects_auth_without_a_nonempty_account_id() -> anyhow::Result<()> { + for account_id in [None, Some(""), Some(" ")] { + let home = tempfile::tempdir()?; + std::fs::write( + home.path().join("auth.json"), + serde_json::to_vec(&json!({ + "auth_mode": "chatgpt", + "tokens": { + "id_token": "header.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lkIjoidXNlci1hIn19.signature", + "access_token": "synthetic-access-token", + "refresh_token": "synthetic-refresh-token", + "account_id": account_id + }, + "last_refresh": "2099-01-01T00:00:00Z" + }))?, + )?; + let auth = CodexAuth::from_auth_storage( + home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + &codex_login::test_support::transport_default_auth_route_config(), + ) + .await? + .expect("managed ChatGPT auth"); + let client = Arc::new(RecordingHttpClient::new( + /*status*/ 200, + cyber_response( + "active", + json!([{ "level": "tac1", "source": "individual" }]), + ), + )); + assert_eq!( + context(auth, client.clone()) + .add_context(/*meta*/ None) + .await, + Some(expected_metadata("unknown", json!([]))), + "account_id={account_id:?}" + ); + assert!(client.requests.lock().expect("inspect requests").is_empty()); + } + Ok(()) +} + +#[tokio::test] +async fn rejects_api_key_authentication_without_sending_credentials() { + let client = Arc::new(RecordingHttpClient::new( + /*status*/ 200, + json!({ "programs": [] }), + )); + let context = context(CodexAuth::from_api_key("test-api-key"), client.clone()); + + let metadata = context + .add_context(Some(json!({ + "threadId": "thread-1", + "openai/entitlementContext": { "forged": true } + }))) + .await; + + let mut expected = expected_metadata("unknown", json!([])); + expected["threadId"] = json!("thread-1"); + assert_eq!(metadata, Some(expected)); + assert!(client.requests.lock().expect("inspect requests").is_empty()); +} + +#[tokio::test] +async fn rejects_initial_unsupported_auth_after_switch_to_chatgpt() { + let client = Arc::new(RecordingHttpClient::new( + /*status*/ 200, + cyber_response( + "active", + json!([{ "level": "tac1", "source": "individual" }]), + ), + )); + let mut context = context(header_auth(), client.clone()); + context.auth_manager = AuthManager::from_auth_for_testing(chatgpt_auth("account-a")); + + assert_eq!( + context.add_context(/*meta*/ None).await, + Some(expected_metadata("unknown", json!([]))) + ); + assert!(client.requests.lock().expect("inspect requests").is_empty()); +} + +#[tokio::test] +async fn maps_verified_access_states_and_rejects_invalid_provider_responses() { + let cases = [ + ( + 200, + json!({ "programs": [ + { "program": "future", "state": "pending", "grants": [{ "level": "premium", "source": "subscription" }] }, + { "program": "cyber", "state": "active", "grants": [{ "level": "tac1", "source": "individual" }] } + ] }), + expected_metadata("granted", json!([{ "level": "tac1", "source": "user" }])), + ), + ( + 200, + cyber_response("inactive", json!([])), + expected_metadata("not_granted", json!([])), + ), + ( + 200, + cyber_response("unavailable", json!([])), + expected_metadata("unknown", json!([])), + ), + ( + 200, + json!({ "programs": [{ "program": "other", "state": "active", "grants": [] }] }), + expected_metadata("unknown", json!([])), + ), + ( + 200, + cyber_response( + "active", + json!([{ "level": "admin", "source": "individual" }]), + ), + expected_metadata("unknown", json!([])), + ), + ( + 200, + cyber_response("active", json!([])), + expected_metadata("unknown", json!([])), + ), + ( + 403, + json!({ "error": "forbidden" }), + expected_metadata("unknown", json!([])), + ), + ]; + + for (response_status, response, expected) in cases { + let response_description = response.to_string(); + let context = context( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + Arc::new(RecordingHttpClient::new(response_status, response)), + ); + let metadata = context.add_context(/*meta*/ None).await.expect("metadata"); + assert_eq!( + metadata, expected, + "HTTP {response_status}: {response_description}" + ); + } +} + +#[tokio::test] +async fn replaces_untrusted_entitlements_on_malformed_response() { + let mut client = RecordingHttpClient::new(/*status*/ 200, Value::Null); + client.response = b"{\"programs\":[".to_vec(); + let context = context(chatgpt_auth("account-a"), Arc::new(client)); + let mut expected = expected_metadata("unknown", json!([])); + expected["threadId"] = json!("thread-1"); + + assert_eq!( + context + .add_context(Some(json!({ + "threadId": "thread-1", + "openai/entitlementContext": { "forged": true } + }))) + .await, + Some(expected) + ); +} + +#[tokio::test] +async fn rejects_verified_access_response_larger_than_one_mebibyte() { + let valid_response = serde_json::to_vec(&cyber_response( + "active", + json!([{ "level": "tac1", "source": "individual" }]), + )) + .expect("serialize response"); + let mut client = RecordingHttpClient::new(/*status*/ 200, Value::Null); + client.response_chunks = Some(vec![ + vec![b' '; MAX_VERIFIED_ACCESS_RESPONSE_BYTES], + valid_response, + ]); + client.response = client + .response_chunks + .as_ref() + .expect("response chunks") + .concat(); + + assert_eq!( + context(chatgpt_auth("account-a"), Arc::new(client)) + .add_context(/*meta*/ None) + .await, + Some(expected_metadata("unknown", json!([]))) + ); +} + +#[tokio::test(start_paused = true)] +async fn returns_unknown_at_the_lookup_deadline() { + let mut client = RecordingHttpClient::new(/*status*/ 200, Value::Null); + client.response_gate = Some((Notify::new(), Notify::new())); + let client = Arc::new(client); + let context = context(chatgpt_auth("account-a"), client.clone()); + let lookup = context.add_context(/*meta*/ None); + tokio::pin!(lookup); + + assert!(futures::poll!(lookup.as_mut()).is_pending()); + assert_eq!(client.requests.lock().expect("recorded requests").len(), 1); + tokio::time::advance(Duration::from_millis(2_499)).await; + assert!(futures::poll!(lookup.as_mut()).is_pending()); + tokio::time::advance(Duration::from_millis(1)).await; + assert_eq!( + lookup + .now_or_never() + .expect("lookup must finish at the deadline"), + Some(expected_metadata("unknown", json!([]))) + ); +} + +#[tokio::test] +async fn rejects_duplicate_cyber_programs() { + let active = json!({ + "program": "cyber", "state": "active", + "grants": [{ "level": "tac1", "source": "individual" }] + }); + for other in [ + active.clone(), + json!({ "program": "cyber", "state": "inactive", "grants": [] }), + json!({ "program": "cyber" }), + ] { + for programs in [json!([active, other]), json!([other, active])] { + let context = context( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + Arc::new(RecordingHttpClient::new( + /*status*/ 200, + json!({ "programs": programs }), + )), + ); + assert_eq!( + context.add_context(/*meta*/ None).await, + Some(expected_metadata("unknown", json!([]))), + "duplicate programs: {programs}" + ); + } + } +} + +#[tokio::test] +async fn rejects_identity_changes_before_sending_credentials() { + for (description, initial_auth, selected_auth) in [ + ( + "account switch", + chatgpt_auth("account-a"), + chatgpt_auth("account-b"), + ), + ( + "standard to FedRAMP", + chatgpt_auth("account-a"), + fedramp_chatgpt_auth("account-a"), + ), + ( + "FedRAMP to standard", + fedramp_chatgpt_auth("account-a"), + chatgpt_auth("account-a"), + ), + ] { + let client = Arc::new(RecordingHttpClient::new( + /*status*/ 200, + cyber_response( + "active", + json!([{ "level": "tac1", "source": "individual" }]), + ), + )); + let mut context = context(initial_auth, client.clone()); + context.auth_manager = AuthManager::from_auth_for_testing(selected_auth); + assert_eq!( + context.add_context(/*meta*/ None).await, + Some(expected_metadata("unknown", json!([]))), + "identity change before the request: {description}" + ); + assert!( + client.requests.lock().expect("inspect requests").is_empty(), + "identity change before the request: {description}" + ); + } +} + +#[tokio::test] +async fn uses_the_checked_auth_snapshot_for_request_headers() -> anyhow::Result<()> { + let client = Arc::new(RecordingHttpClient::new( + /*status*/ 200, + cyber_response("inactive", json!([])), + )); + let context = context(chatgpt_auth("account-a"), client.clone()); + let refreshed = CodexAuth::from_external_chatgpt_tokens( + "header.e30.refreshed", + "account-a", + /*chatgpt_plan_type*/ None, + )?; + context + .auth_manager + .set_external_auth(Arc::new(SequencedExternalAuth(Mutex::new(VecDeque::from( + [context.auth.clone(), refreshed, header_auth()], + ))))) + .await?; + + assert_eq!( + context.add_context(/*meta*/ None).await, + Some(expected_metadata("not_granted", json!([]))) + ); + let requests = client.requests.lock().expect("recorded requests"); + assert_eq!(requests.len(), 1); + let authorization = requests[0] + .headers + .iter() + .find(|header| header.name == "authorization"); + assert_eq!( + authorization, + Some(&HttpHeader { + name: "authorization".to_string(), + value: "Bearer header.e30.refreshed".to_string(), + value_env_var: None, + }) + ); + Ok(()) +} + +#[tokio::test] +async fn rejects_identity_changes_while_request_is_in_flight() -> anyhow::Result<()> { + let personal = chatgpt_auth("account-a"); + let workspace = + CodexAuth::from_external_chatgpt_tokens("header.e30.same", "account-a", Some("team"))?; + for (description, initial_auth, selected_auth) in [ + ( + "account switch", + personal.clone(), + Some(chatgpt_auth("account-b")), + ), + ("logout", personal.clone(), None), + ( + "personal to workspace", + personal.clone(), + Some(workspace.clone()), + ), + ("workspace to personal", workspace, Some(personal)), + ( + "auth mode change", + chatgpt_auth("account-a"), + Some(header_auth()), + ), + ( + "standard to FedRAMP", + chatgpt_auth("account-a"), + Some(fedramp_chatgpt_auth("account-a")), + ), + ( + "FedRAMP to standard", + fedramp_chatgpt_auth("account-a"), + Some(chatgpt_auth("account-a")), + ), + ] { + let mut client = RecordingHttpClient::new( + /*status*/ 200, + cyber_response( + "active", + json!([{ "level": "tac1", "source": "individual" }]), + ), + ); + client.response_gate = Some((Notify::new(), Notify::new())); + let client = Arc::new(client); + let request_is_fedramp = initial_auth.is_fedramp_account(); + let context = context(initial_auth, client.clone()); + let auth_manager = &context.auth_manager; + auth_manager + .set_external_auth(Arc::new(StaticExternalAuth(context.auth.clone()))) + .await?; + + let (metadata, auth_change) = tokio::join!(context.add_context(/*meta*/ None), async { + let (requested, release) = client.response_gate.as_ref().expect("gated response"); + tokio::time::timeout(Duration::from_secs(5), requested.notified()).await?; + { + let requests = client.requests.lock().expect("inspect request"); + assert_eq!(requests.len(), 1); + assert!(requests[0].headers.iter().any(|header| { + header.name.eq_ignore_ascii_case("chatgpt-account-id") + && header.value == "account-a" + })); + assert_eq!( + requests[0] + .headers + .iter() + .any(|header| header.name.eq_ignore_ascii_case("x-openai-fedramp")), + request_is_fedramp + ); + } + if let Some(auth) = selected_auth { + auth_manager + .set_external_auth(Arc::new(StaticExternalAuth(auth))) + .await?; + } else { + auth_manager.clear_external_auth(); + } + release.notify_one(); + Ok::<(), anyhow::Error>(()) + }); + auth_change?; + + assert_eq!( + metadata, + Some(expected_metadata("unknown", json!([]))), + "auth change after the request: {description}" + ); + } + Ok(()) +} diff --git a/codex-rs/exec-server/src/client/http_response_body_stream.rs b/codex-rs/exec-server/src/client/http_response_body_stream.rs index 625c640812..be52e5cb6a 100644 --- a/codex-rs/exec-server/src/client/http_response_body_stream.rs +++ b/codex-rs/exec-server/src/client/http_response_body_stream.rs @@ -78,6 +78,24 @@ pub struct HttpResponseBodyStream { } impl HttpResponseBodyStream { + /// Creates an in-memory response stream from pre-buffered chunks. + /// + /// This is useful for [`crate::HttpClient`] implementations that already + /// own the response bytes, including lightweight test clients. + #[doc(hidden)] + pub fn from_chunks(chunks: Vec>) -> Self { + let body = futures::stream::iter( + chunks + .into_iter() + .map(|chunk| Ok::(chunk.into())), + ); + Self { + inner: HttpResponseBodyStreamInner::Local { + body: Box::pin(body), + }, + } + } + pub(super) fn local(response: HttpResponse) -> Self { Self { inner: HttpResponseBodyStreamInner::Local {