openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bind reusable backend clients to runtime authentication (#59255)

Browse Source 
## Why

Reusable clients need fresh authentication for each request while preserving the credential source and transport policy selected when the client was created.

## What changed

- Add runtime-aware backend client construction and per-request authentication. Allow token refresh within the original credential source, reject credential replacement, and validate destination and account/organization selectors.
- Let runtime authentication finish persisting rotated tokens when the requesting task is cancelled.
- Apply asynchronous authentication to file creation and finalization requests, generating fresh assertions while keeping credentials off blob uploads.
- Preserve direct sandbox routing without enabling legacy custom-CA fallback for protected clients that reject redirects.
- Exclude login account metadata from model cache identities when provider bearer credentials override the login, and record authentication selection outcomes and duration.

## Testing

Add regression coverage for credential refresh and replacement, organization selection, destination restrictions, transport policy preservation, cancellation during token rotation, credential stripping on redirects, and file upload authentication in both runtime authentication modes.

GitOrigin-RevId: d045595e869dfcdfd9f705a6228cf0eb4c2d2c12
This commit is contained in:
cooper-oai
2026-11-03 23:35:00 +00:00
committed by copyberry
parent 1176913a6f
commit 8f32b72151
38 changed files with 1512 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+7
View File
@@ -2436,18 +2436,23 @@ name = "codex-backend-client"
version = "0.0.0"
dependencies = [
"anyhow",
"base64 0.22.1",
"codex-api",
"codex-backend-openapi-models",
"codex-features",
"codex-http-client",
"codex-login",
"codex-model-provider",
"codex-protocol",
"codex-runtime-agent-instance",
"http 1.4.0",
"pretty_assertions",
"serde",
"serde_json",
"tempfile",
"tokio",
"url",
"uuid",
"wiremock",
]
@@ -4773,6 +4778,7 @@ name = "codex-runtime-agent-instance"
version = "0.0.0"
dependencies = [
"base64 0.22.1",
"bytes",
"codex-agent-identity",
"codex-api",
"codex-features",
@@ -4780,6 +4786,7 @@ dependencies = [
"codex-login",
"codex-utils-redacted-string",
"ed25519-dalek",
"futures",
"http 1.4.0",
"pretty_assertions",
"rand 0.9.3",
codex-rs/app-server/src/turn_cost_worker.rs
+2 -1
View File
@@ -454,7 +454,8 @@ impl WorkerRuntime {
.map_err(|error| RequestError::Other(error.into()))?;
let endpoint = provider.url_for_path("analytics/codex/turn-costs");
let client = BackendClient::new(provider.base_url.clone(), http_client_factory)
.with_auth_provider(auth);
.with_auth_provider(auth)
.map_err(RequestError::Other)?;
client
.query_api_key_turn_costs_at(&endpoint, turn_ids, &provider.headers)
.await
codex-rs/backend-client/Cargo.toml
+5
View File
@@ -20,13 +20,18 @@ http = { workspace = true }
url = { workspace = true }
codex-backend-openapi-models = { path = "../codex-backend-openapi-models" }
codex-api = { workspace = true }
codex-features = { workspace = true }
codex-http-client = { workspace = true }
codex-login = { workspace = true }
codex-model-provider = { workspace = true }
codex-runtime-agent-instance = { workspace = true }
codex-protocol = { workspace = true }
tokio = { workspace = true, features = ["time"] }
[dev-dependencies]
base64 = { workspace = true }
uuid = { workspace = true }
pretty_assertions = "1"
tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt"] }
wiremock = { workspace = true }
codex-rs/backend-client/src/analytics_session.rs
+2 -1
View File
@@ -55,7 +55,8 @@ impl AnalyticsSession {
.with_auth_provider(codex_model_provider::auth_provider_from_auth_manager(
Arc::clone(&auth_manager),
&auth,
));
))
.map_err(|_| "Couldn't initialize local Analytics authentication.")?;
let mut session = Self {
client,
auth_manager,
codex-rs/backend-client/src/analytics_session_tests.rs
+2 -1
View File
@@ -87,7 +87,8 @@ async fn assert_unauthorized_after_recovery(outcome: RefreshOutcome, expected_re
.with_auth_provider(codex_model_provider::auth_provider_from_auth_manager(
Arc::clone(&auth_manager),
&auth,
));
))
.unwrap();
let session = AnalyticsSession {
account: AnalyticsAccount {
id: "account-a".into(),
codex-rs/backend-client/src/client.rs
+97 -27
View File
@@ -9,7 +9,9 @@ use crate::types::RateLimitStatusPayload;
use crate::types::TokenUsageProfile;
use crate::types::TurnAttemptsSiblingTurnsResponse;
use anyhow::Result;
use codex_api::AuthError;
use codex_api::SharedAuthProvider;
use codex_features::RuntimeAuthMode;
use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER;
use codex_http_client::ClientRouteClass;
use codex_http_client::HttpClientFactory;
@@ -19,12 +21,14 @@ use codex_http_client::RouteAwareClientPool;
use codex_http_client::RouteAwareRequestError;
use codex_login::CodexAuth;
use codex_login::default_client::get_codex_user_agent;
use codex_model_provider::ClientAuth;
use codex_protocol::account::PlanType as AccountPlanType;
use codex_protocol::protocol::CreditsSnapshot;
use codex_protocol::protocol::RateLimitReachedType;
use codex_protocol::protocol::RateLimitSnapshot;
use codex_protocol::protocol::RateLimitWindow;
use codex_protocol::protocol::SpendControlLimitSnapshot;
use codex_runtime_agent_instance::RuntimeAuthManager;
use http::Method;
use http::StatusCode;
use http::header::CACHE_CONTROL;
@@ -36,6 +40,7 @@ use http::header::USER_AGENT;
use serde::Serialize;
use serde::de::DeserializeOwned;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;
pub(crate) mod analytics;
@@ -113,6 +118,15 @@ impl From<anyhow::Error> for RequestError {
}
}
impl From<AuthError> for RequestError {
fn from(error: AuthError) -> Self {
match error {
AuthError::Policy(denied) => Self::Policy(denied),
other => Self::Other(other.into()),
}
}
}
impl From<codex_http_client::HttpError> for RequestError {
fn from(error: codex_http_client::HttpError) -> Self {
match error {
@@ -161,7 +175,7 @@ impl PathStyle {
pub struct Client {
base_url: String,
http: RouteAwareClientPool,
auth_provider: SharedAuthProvider,
auth: ClientAuth,
user_agent: Option<HeaderValue>,
chatgpt_account_id: Option<String>,
chatgpt_account_is_fedramp: bool,
@@ -173,7 +187,7 @@ impl fmt::Debug for Client {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Client")
.field("base_url", &self.base_url)
.field("auth_provider", &"<provider>")
.field("auth", &"<provider>")
.field("user_agent", &self.user_agent)
.field("chatgpt_account_id", &self.chatgpt_account_id)
.field(
@@ -239,7 +253,7 @@ impl Client {
Self {
base_url,
http,
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
auth: ClientAuth::new(codex_model_provider::unauthenticated_auth_provider()),
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
@@ -253,17 +267,47 @@ impl Client {
auth: &CodexAuth,
http_client_factory: HttpClientFactory,
) -> Self {
Self::new(base_url, http_client_factory)
.with_user_agent(get_codex_user_agent())
.with_auth_provider(codex_model_provider::auth_provider_from_auth(auth))
let mut client =
Self::new(base_url, http_client_factory).with_user_agent(get_codex_user_agent());
client.auth = ClientAuth::new(codex_model_provider::auth_provider_from_auth(auth));
client
}
pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Self {
self.auth_provider = auth;
self
/// Build a client using the captured credential and the request's organization/project headers.
pub async fn from_auth_with_runtime(
base_url: impl Into<String>,
auth: &CodexAuth,
http_client_factory: HttpClientFactory,
runtime: Option<&(Arc<RuntimeAuthManager>, RuntimeAuthMode)>,
headers: &HeaderMap,
) -> Result<Self> {
let base_url = base_url.into();
let runtime = codex_model_provider::runtime_auth_for_url(runtime, &base_url)?;
Self::from_auth_with_selected_runtime(base_url, auth, http_client_factory, runtime, headers)
.await
}
pub fn with_bearer_token(self, token: String) -> Self {
async fn from_auth_with_selected_runtime(
base_url: String,
auth: &CodexAuth,
http_client_factory: HttpClientFactory,
runtime: Option<&(Arc<RuntimeAuthManager>, RuntimeAuthMode)>,
headers: &HeaderMap,
) -> Result<Self> {
let (auth, http_client_factory) =
ClientAuth::for_auth(auth, runtime, http_client_factory, headers).await?;
let mut client = Self::new(base_url, http_client_factory);
client.auth = auth;
Ok(client.with_user_agent(get_codex_user_agent()))
}
/// Configure provider auth, returning an error for a runtime-bound client.
pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Result<Self> {
self.auth.set_provider(auth)?;
Ok(self)
}
pub fn with_bearer_token(self, token: String) -> Result<Self> {
self.with_auth_provider(std::sync::Arc::new(
codex_model_provider::BearerAuthProvider::new(token),
))
@@ -291,14 +335,19 @@ impl Client {
self
}
fn headers(&self) -> Result<HeaderMap> {
let mut h = HeaderMap::new();
async fn headers(&self) -> Result<HeaderMap, RequestError> {
self.headers_with(HeaderMap::new()).await
}
async fn headers_with(&self, mut h: HeaderMap) -> Result<HeaderMap, RequestError> {
if let Some(ua) = &self.user_agent {
h.insert(USER_AGENT, ua.clone());
} else {
h.insert(USER_AGENT, HeaderValue::from_static("codex-cli"));
}
self.auth_provider.add_auth_headers(&mut h)?;
let selected = self.auth.resolve(&h).await?;
h = selected.auth.apply_auth_headers(h).await?;
if let Some(acc) = &self.chatgpt_account_id
&& let Ok(hv) = HeaderValue::from_str(acc)
{
@@ -309,6 +358,9 @@ impl Client {
{
h.insert(name, HeaderValue::from_static("true"));
}
selected
.validate_headers(&h)
.map_err(|error| RequestError::Other(error.into()))?;
Ok(h)
}
@@ -322,6 +374,7 @@ impl Client {
method: &str,
url: &str,
) -> Result<(String, String)> {
self.auth.validate_destination(&self.base_url, url)?;
let res = req.send().await?;
let status = res.status();
let ct = res
@@ -343,6 +396,9 @@ impl Client {
method: &str,
url: &str,
) -> std::result::Result<(String, String), RequestError> {
self.auth
.validate_destination(&self.base_url, url)
.map_err(|error| RequestError::Other(error.into()))?;
let res = req.send().await?;
let status = res.status();
let content_type = res
@@ -368,7 +424,9 @@ impl Client {
&self,
url: &str,
) -> std::result::Result<(String, String), RequestError> {
let request = self.request(Method::GET, url).headers(self.headers()?);
let request = self
.request(Method::GET, url)
.headers(self.headers().await?);
if !self.http.allows_system_proxy_fallback() {
return self.exec_request_detailed(request, "GET", url).await;
}
@@ -395,7 +453,7 @@ impl Client {
.http
.clone()
.with_outbound_proxy_policy(OutboundProxyPolicy::RespectSystemProxy);
let request = http.get(url).headers(self.headers()?);
let request = http.get(url).headers(self.headers().await?);
self.exec_request_detailed(request, "GET", url).await
}
@@ -435,7 +493,9 @@ impl Client {
pub async fn get_token_usage_profile(&self) -> Result<TokenUsageProfile> {
let url = self.token_usage_profile_url();
let req = self.request(Method::GET, &url).headers(self.headers()?);
let req = self
.request(Method::GET, &url)
.headers(self.headers().await?);
let (body, ct) = self.exec_request(req, "GET", &url).await?;
self.decode_json(&url, &ct, &body)
}
@@ -454,7 +514,7 @@ impl Client {
let url = self.send_add_credits_nudge_email_url();
let req = self
.request(Method::POST, &url)
.headers(self.headers()?)
.headers(self.headers().await?)
.header(CONTENT_TYPE, HeaderValue::from_static("application/json"))
.json(&SendAddCreditsNudgeEmailRequest { credit_type });
self.exec_request_detailed(req, "POST", &url).await?;
@@ -469,7 +529,9 @@ impl Client {
cursor: Option<&str>,
) -> Result<PaginatedListTaskListItem> {
let url = self.list_tasks_url(limit, task_filter, environment_id, cursor)?;
let req = self.request(Method::GET, &url).headers(self.headers()?);
let req = self
.request(Method::GET, &url)
.headers(self.headers().await?);
let (body, ct) = self.exec_request(req, "GET", &url).await?;
self.decode_json::<PaginatedListTaskListItem>(&url, &ct, &body)
}
@@ -521,7 +583,9 @@ impl Client {
PathStyle::CodexApi => format!("{}/api/codex/tasks/{}", self.base_url, task_id),
PathStyle::ChatGptApi => format!("{}/wham/tasks/{}", self.base_url, task_id),
};
let req = self.request(Method::GET, &url).headers(self.headers()?);
let req = self
.request(Method::GET, &url)
.headers(self.headers().await?);
let (body, ct) = self.exec_request(req, "GET", &url).await?;
let parsed: CodeTaskDetailsResponse = self.decode_json(&url, &ct, &body)?;
Ok((parsed, body, ct))
@@ -542,7 +606,9 @@ impl Client {
self.base_url, task_id, turn_id
),
};
let req = self.request(Method::GET, &url).headers(self.headers()?);
let req = self
.request(Method::GET, &url)
.headers(self.headers().await?);
let (body, ct) = self.exec_request(req, "GET", &url).await?;
self.decode_json::<TurnAttemptsSiblingTurnsResponse>(&url, &ct, &body)
}
@@ -581,7 +647,7 @@ impl Client {
let url = self.user_settings_url();
let req = self
.request(Method::GET, &url)
.headers(self.headers()?)
.headers(self.headers().await?)
.header(
CACHE_CONTROL,
HeaderValue::from_static("no-cache, no-store"),
@@ -597,7 +663,7 @@ impl Client {
let url = self.workspace_messages_url();
let req = self
.request(Method::GET, &url)
.headers(self.headers()?)
.headers(self.headers().await?)
.header(CACHE_CONTROL, HeaderValue::from_static("no-store"));
let (body, ct) = self.exec_request_detailed(req, "GET", &url).await?;
self.decode_json::<CodexWorkspaceMessagesResponse>(&url, &ct, &body)
@@ -613,7 +679,7 @@ impl Client {
};
let req = self
.request(Method::POST, &url)
.headers(self.headers()?)
.headers(self.headers().await?)
.header(CONTENT_TYPE, HeaderValue::from_static("application/json"))
.json(&request_body);
let (body, ct) = self.exec_request(req, "POST", &url).await?;
@@ -1269,8 +1335,8 @@ mod tests {
);
}
#[test]
fn authenticated_user_settings_client_uses_active_workspace_headers() {
#[tokio::test]
async fn authenticated_user_settings_client_uses_active_workspace_headers() {
let auth = CodexAuth::from_external_chatgpt_tokens(
"e30.e30.c2ln",
"workspace-123",
@@ -1282,7 +1348,7 @@ mod tests {
&auth,
HttpClientFactory::new(codex_http_client::OutboundProxyPolicy::ReqwestDefault),
);
let headers = client.headers().unwrap();
let headers = client.headers().await.unwrap();
assert_eq!(
[
@@ -1304,7 +1370,7 @@ mod tests {
HttpClientFactory::new(codex_http_client::OutboundProxyPolicy::ReqwestDefault),
ClientRouteClass::Api,
),
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
auth: ClientAuth::new(codex_model_provider::unauthenticated_auth_provider()),
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
@@ -1313,3 +1379,7 @@ mod tests {
}
}
}
#[cfg(test)]
#[path = "client/runtime_auth_tests.rs"]
mod runtime_auth_tests;
codex-rs/backend-client/src/client/analytics.rs
+1 -1
View File
@@ -96,7 +96,7 @@ impl Client {
}
let request = self
.request(Method::GET, url.as_str())
.headers(self.headers()?);
.headers(self.headers().await?);
let (body, _) = self
.exec_request_detailed(request, "GET", url.as_str())
.await?;
codex-rs/backend-client/src/client/chatgpt_turn_cost.rs
+1 -1
View File
@@ -57,7 +57,7 @@ impl Client {
};
let request = self
.request(Method::POST, &url)
.headers(self.headers()?)
.headers(self.headers().await?)
.json(&TurnCostsRequest {
threads: threads
.iter()
codex-rs/backend-client/src/client/plan_history.rs
+3 -1
View File
@@ -66,7 +66,9 @@ impl Client {
PathStyle::ChatGptApi => "wham",
};
let url = format!("{}/{prefix}/usage/plan_limit_history?days=7", self.base_url);
let request = self.request(Method::GET, &url).headers(self.headers()?);
let request = self
.request(Method::GET, &url)
.headers(self.headers().await?);
let (body, _) = match self.exec_request_detailed(request, "GET", &url).await {
Ok(response) => response,
Err(error) if error.status().is_some_and(|status| status.as_u16() == 404) => {
codex-rs/backend-client/src/client/profile.rs
+3 -1
View File
@@ -61,7 +61,9 @@ impl Client {
/// Read the full profile using the same route as account token usage.
pub async fn get_account_profile(&self) -> Result<AccountProfile, RequestError> {
let url = self.token_usage_profile_url();
let request = self.request(Method::GET, &url).headers(self.headers()?);
let request = self
.request(Method::GET, &url)
.headers(self.headers().await?);
let (body, _) = self.exec_request_detailed(request, "GET", &url).await?;
serde_json::from_str(&body)
.map_err(|_| RequestError::Other(anyhow::anyhow!("Invalid account profile response.")))
codex-rs/backend-client/src/client/rate_limit_resets.rs
+7 -3
View File
@@ -71,7 +71,9 @@ impl Client {
supports_luna_reserve: bool,
) -> Result<RateLimitStatusWithResetCredits> {
let url = self.rate_limit_status_url();
let mut req = self.request(Method::GET, &url).headers(self.headers()?);
let mut req = self
.request(Method::GET, &url)
.headers(self.headers().await?);
if supports_luna_reserve {
req = req.header("x-openai-codex-luna-reserve", HeaderValue::from_static("1"));
}
@@ -81,7 +83,9 @@ impl Client {
pub async fn list_rate_limit_reset_credits(&self) -> Result<RateLimitResetCreditsDetails> {
let url = self.rate_limit_reset_credits_url();
let req = self.request(Method::GET, &url).headers(self.headers()?);
let req = self
.request(Method::GET, &url)
.headers(self.headers().await?);
let (body, ct) = self.exec_request(req, "GET", &url).await?;
self.decode_json(&url, &ct, &body)
}
@@ -111,7 +115,7 @@ impl Client {
let url = self.consume_rate_limit_reset_credit_url();
let req = self
.request(Method::POST, &url)
.headers(self.headers()?)
.headers(self.headers().await?)
.header(CONTENT_TYPE, HeaderValue::from_static("application/json"))
.json(&ConsumeRateLimitResetCreditRequest {
redeem_request_id,
codex-rs/backend-client/src/client/rate_limit_resets_tests.rs
+2 -1
View File
@@ -5,6 +5,7 @@ use crate::types::ConsumeRateLimitResetCreditCode;
use crate::types::RateLimitResetCreditDetails;
use crate::types::RateLimitResetCreditsDetails;
use crate::types::RateLimitResetCreditsSummary;
use codex_model_provider::ClientAuth;
use pretty_assertions::assert_eq;
use wiremock::Mock;
use wiremock::MockServer;
@@ -178,7 +179,7 @@ fn test_client(base_url: &str, path_style: PathStyle) -> Client {
),
codex_http_client::ClientRouteClass::Api,
),
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
auth: ClientAuth::new(codex_model_provider::unauthenticated_auth_provider()),
user_agent: None,
chatgpt_account_id: None,
chatgpt_account_is_fedramp: false,
codex-rs/backend-client/src/client/runtime_auth_tests.rs
+568
View File
@@ -0,0 +1,568 @@
//! Verify the backend client resolves a fresh runtime proof for each HTTP request.
use super::*;
use codex_features::RuntimeAuthMode;
use codex_http_client::DestinationPolicy;
use codex_http_client::NetworkPolicyController;
use codex_http_client::NetworkPolicyDenied;
use codex_http_client::OutboundProxyPolicy;
use codex_login::AuthCredentialsStoreMode;
use codex_login::AuthKeyringBackendKind;
use codex_login::AuthManager;
use codex_login::AuthRouteConfig;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_runtime_agent_instance::AgentHostAuthorization;
use codex_runtime_agent_instance::RuntimeAuthManager;
use codex_runtime_agent_instance::RuntimeRegistrationMetadata;
use codex_runtime_agent_instance::RuntimeRole;
use pretty_assertions::assert_eq;
use std::sync::Arc;
use uuid::Uuid;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
fn runtime_provider(
server: &MockServer,
credentials: Arc<AuthManager>,
mode: RuntimeAuthMode,
) -> (Arc<RuntimeAuthManager>, RuntimeAuthMode) {
(
Arc::new(RuntimeAuthManager::new(
credentials,
AgentHostAuthorization::InstallationId {
installation_id: Uuid::nil(),
},
RuntimeRegistrationMetadata {
harness: Some("codex".into()),
harness_version: None,
role: RuntimeRole::Orchestrator,
},
format!("{}/register", server.uri()).parse().unwrap(),
)),
mode,
)
}
#[tokio::test]
async fn account_runtime_auth_signs_each_request_and_strips_credentials_on_redirect() {
for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] {
let server = MockServer::start().await;
let destination = MockServer::start().await;
Mock::given(path("/register"))
.respond_with(|request: &wiremock::Request| {
let body: serde_json::Value = request.body_json().unwrap();
let organization = body["actor_auth"]["organization_id"].as_str().unwrap_or("org-a");
ResponseTemplate::new(200).set_body_json(serde_json::json!({
"runtime_agent_instance_id": format!("runtime-{organization}"), "agent_host_id": "host-a",
"actor": {"kind": "platform", "user_id": "user-a", "organization_id": organization, "project_id": null},
}))
})
.expect(2).mount(&server).await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({"accounts": [], "account_ordering": []})),
)
.expect(2)
.mount(&destination)
.await;
Mock::given(path("/turn-costs"))
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []})),
)
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(307).insert_header("location", destination.uri()))
.expect(2)
.mount(&server)
.await;
let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault);
let auth = CodexAuth::from_api_key("key-a");
let runtime = runtime_provider(
&server,
AuthManager::from_auth_for_testing(auth.clone()),
mode,
);
let client = Client::from_auth_with_selected_runtime(
server.uri(),
&auth,
factory,
Some(&runtime),
&HeaderMap::new(),
)
.await
.unwrap();
client.get_accounts_check().await.unwrap();
client.get_accounts_check().await.unwrap();
assert!(
client
.clone()
.with_chatgpt_account_id("different-workspace")
.get_accounts_check()
.await
.is_err()
);
let stale_headers = HeaderMap::from_iter([
(
HeaderName::from_static("openai-organization"),
HeaderValue::from_static("org-b"),
),
(
http::header::AUTHORIZATION,
HeaderValue::from_static("Bearer stale"),
),
(
HeaderName::from_static("x-openai-agent-assertion"),
HeaderValue::from_static("stale"),
),
]);
client
.query_api_key_turn_costs_at(
&format!("{}/turn-costs", server.uri()),
&[],
&stale_headers,
)
.await
.unwrap();
assert!(
client
.query_api_key_turn_costs_at(&destination.uri(), &[], &stale_headers)
.await
.is_err()
);
let redirected = destination.received_requests().await.unwrap();
assert_eq!(redirected.len(), 2);
for request in redirected {
assert!(!request.headers.contains_key("authorization"));
assert!(!request.headers.contains_key("x-openai-agent-assertion"));
}
let requests = server.received_requests().await.unwrap();
let cost_request = requests
.iter()
.find(|request| request.url.path() == "/turn-costs")
.unwrap();
assert_eq!(cost_request.headers["openai-organization"], "org-b");
let proofs: Vec<_> = requests
.iter()
.filter(|r| r.url.path() != "/register")
.map(|request| {
let authorization = request.headers["authorization"].to_str().unwrap();
match mode {
RuntimeAuthMode::DualHeader => {
assert_eq!(authorization, "Bearer key-a");
request.headers["x-openai-agent-assertion"]
.to_str()
.unwrap()
}
RuntimeAuthMode::AssertionOnly => {
assert!(!request.headers.contains_key("x-openai-agent-assertion"));
authorization.strip_prefix("AgentAssertion ").unwrap()
}
}
})
.collect();
assert_eq!(proofs.len(), 3);
assert_ne!(proofs[0], proofs[1]);
assert_ne!(proofs[1], proofs[2]);
assert_ne!(proofs[2], "stale");
}
}
#[tokio::test]
async fn client_registers_in_requested_organization_when_default_is_denied() {
for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/register"))
.respond_with(|request: &wiremock::Request| {
let body: serde_json::Value = request.body_json().unwrap();
if body["actor_auth"]["organization_id"] != "org-b" {
return ResponseTemplate::new(403);
}
ResponseTemplate::new(200).set_body_json(serde_json::json!({
"runtime_agent_instance_id": "runtime-org-b",
"agent_host_id": "host-a",
"actor": {
"kind": "platform",
"user_id": "user-a",
"organization_id": "org-b",
"project_id": null,
},
}))
})
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/turn-costs"))
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []})),
)
.expect(1)
.mount(&server)
.await;
let auth = CodexAuth::from_api_key("key-a");
let runtime = runtime_provider(
&server,
AuthManager::from_auth_for_testing(auth.clone()),
mode,
);
let headers = HeaderMap::from_iter([(
HeaderName::from_static("openai-organization"),
HeaderValue::from_static("org-b"),
)]);
let client = Client::from_auth_with_selected_runtime(
server.uri(),
&auth,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
Some(&runtime),
&headers,
)
.await
.unwrap();
client
.query_api_key_turn_costs_at(&format!("{}/turn-costs", server.uri()), &[], &headers)
.await
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 2);
assert_eq!(requests[1].url.path(), "/turn-costs");
assert_eq!(requests[1].headers["openai-organization"], "org-b");
let authorization = requests[1].headers["authorization"].to_str().unwrap();
match mode {
RuntimeAuthMode::DualHeader => {
assert_eq!(authorization, "Bearer key-a");
assert!(requests[1].headers.contains_key("x-openai-agent-assertion"));
}
RuntimeAuthMode::AssertionOnly => {
assert!(authorization.starts_with("AgentAssertion "));
assert!(!requests[1].headers.contains_key("x-openai-agent-assertion"));
}
}
server.verify().await;
}
}
#[tokio::test]
async fn unsupported_openai_backend_rejects_runtime_auth_before_registration() {
for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] {
let server = MockServer::start().await;
let auth = CodexAuth::from_api_key("key-a");
let runtime = runtime_provider(
&server,
AuthManager::from_auth_for_testing(auth.clone()),
mode,
);
let result = Client::from_auth_with_runtime(
"https://future.openai.com/backend-api",
&auth,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
Some(&runtime),
&HeaderMap::new(),
)
.await;
let Err(error) = result else {
panic!("unsupported OpenAI backend accepted runtime auth");
};
assert!(matches!(
error.downcast_ref::<CodexErr>().map(CodexErr::details),
Some(CodexErrorDetails::UnsupportedOperation(_))
));
assert!(server.received_requests().await.unwrap().is_empty());
}
}
#[tokio::test]
async fn custom_backend_keeps_its_credential_without_registering_a_runtime() {
for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] {
let server = MockServer::start().await;
Mock::given(path("/turn-costs"))
.respond_with(
ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []})),
)
.expect(1)
.mount(&server)
.await;
let auth = CodexAuth::from_api_key("custom-backend-key");
let runtime = runtime_provider(
&server,
AuthManager::from_auth_for_testing(auth.clone()),
mode,
);
let client = Client::from_auth_with_runtime(
server.uri(),
&auth,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
Some(&runtime),
&HeaderMap::new(),
)
.await
.unwrap();
client
.query_api_key_turn_costs_at(
&format!("{}/turn-costs", server.uri()),
&[],
&HeaderMap::new(),
)
.await
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].headers["authorization"],
"Bearer custom-backend-key"
);
assert!(!requests[0].headers.contains_key("x-openai-agent-assertion"));
server.verify().await;
}
}
#[tokio::test]
async fn reusable_runtime_client_keeps_the_selected_accounts_transport_policy() {
let server = MockServer::start().await;
Mock::given(path("/register"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"runtime_agent_instance_id": "runtime-policy", "agent_host_id": "host-a",
"actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": null},
})))
.expect(1).mount(&server).await;
Mock::given(path("/turn-costs"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []})))
.expect(2)
.mount(&server)
.await;
let controller = NetworkPolicyController::default();
let policy = controller.policy();
let home = tempfile::tempdir().unwrap();
codex_login::login_with_api_key(
home.path(),
"key-a",
AuthCredentialsStoreMode::Ephemeral,
AuthKeyringBackendKind::Direct,
)
.unwrap();
let credentials = Arc::new(
AuthManager::new(
home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::Ephemeral,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
AuthRouteConfig::from_http_client_factory(
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
.with_network_policy(policy.clone()),
),
)
.await,
);
assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted));
let auth = credentials.auth().await.unwrap();
let runtime = runtime_provider(&server, credentials, RuntimeAuthMode::AssertionOnly);
let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault);
let client = Client::from_auth_with_selected_runtime(
server.uri(),
&auth,
factory.clone(),
Some(&runtime),
&HeaderMap::new(),
)
.await
.unwrap();
let url = format!("{}/turn-costs", server.uri());
client
.query_api_key_turn_costs_at(&url, &[], &HeaderMap::new())
.await
.unwrap();
policy.invalidate();
assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted));
let error = client
.query_api_key_turn_costs_at(&url, &[], &HeaderMap::new())
.await
.unwrap_err();
assert!(matches!(
error,
RequestError::Policy(NetworkPolicyDenied::Revoked)
));
let fresh = Client::from_auth_with_selected_runtime(
server.uri(),
&auth,
factory,
Some(&runtime),
&HeaderMap::new(),
)
.await
.unwrap();
fresh
.query_api_key_turn_costs_at(&url, &[], &HeaderMap::new())
.await
.unwrap();
policy.invalidate();
assert!(controller.publish(
policy.revision(),
DestinationPolicy::Restricted {
allowed_hosts: Default::default()
}
));
let selectors = HeaderMap::from_iter([(
HeaderName::from_static("openai-organization"),
HeaderValue::from_static("org-b"),
)]);
let error = fresh
.query_api_key_turn_costs_at(&url, &[], &selectors)
.await
.unwrap_err();
assert!(
matches!(
error,
RequestError::Policy(NetworkPolicyDenied::Destination)
),
"{error}"
);
server.verify().await;
}
struct RotatingAuth(std::sync::RwLock<CodexAuth>);
impl codex_login::ExternalAuth for RotatingAuth {
fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> {
Box::pin(async { Ok(self.0.read().unwrap().clone()) })
}
fn refresh(
&self,
_: codex_login::ExternalAuthRefreshContext,
) -> codex_login::ExternalAuthFuture<'_, CodexAuth> {
self.resolve()
}
}
#[tokio::test]
async fn reusable_runtime_client_refreshes_credentials_without_accepting_a_new_session() {
use base64::Engine;
let oauth = |session, token| {
let claims = serde_json::json!({
"iss": "https://auth.openai.com/", "session_id": session, "jti": token,
"https://api.openai.com/auth": {
"user_id": "user-a", "chatgpt_account_id": "workspace-a"
}
});
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string());
CodexAuth::from_external_chatgpt_tokens(
&format!("header.{payload}.signature"),
"workspace-a",
/*chatgpt_plan_type*/ None,
)
.unwrap()
};
for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] {
let server = MockServer::start().await;
Mock::given(path("/register"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"runtime_agent_instance_id": "runtime-refresh", "agent_host_id": "host-a",
"actor": {"kind": "chatgpt", "user_id": "user-a", "account_id": "workspace-a"}
})))
.mount(&server)
.await;
Mock::given(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"plan_type": "plus",
"rate_limit": {
"allowed": true, "limit_reached": false,
"primary_window": {
"used_percent": 0, "limit_window_seconds": 300,
"reset_after_seconds": 60, "reset_at": 2000000000
}
}
})))
.expect(2)
.mount(&server)
.await;
let initial = oauth("login-a", "initial");
let refreshed = oauth("login-a", "refreshed");
let credentials = AuthManager::from_auth_for_testing(initial.clone());
let source = Arc::new(RotatingAuth(std::sync::RwLock::new(initial.clone())));
credentials.set_external_auth(source.clone()).await.unwrap();
let runtime = runtime_provider(&server, credentials, mode);
let client = Client::from_auth_with_selected_runtime(
server.uri(),
&initial,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
Some(&runtime),
&HeaderMap::new(),
)
.await
.unwrap();
assert!(
client
.clone()
.with_bearer_token("replacement-bearer".into())
.is_err()
);
assert!(
client
.clone()
.with_auth_provider(Arc::new(codex_model_provider::BearerAuthProvider::new(
"replacement-provider".to_string(),
)))
.is_err()
);
client.get_rate_limits().await.unwrap();
*source.0.write().unwrap() = refreshed.clone();
client.get_rate_limits().await.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(
requests
.iter()
.filter(|r| r.url.path() == "/register")
.count(),
1
);
let requests: Vec<_> = requests
.iter()
.filter(|r| r.url.path() == "/api/codex/usage")
.collect();
assert_eq!(requests.len(), 2);
let mut proofs = Vec::new();
for (request, expected) in requests.iter().zip([initial, refreshed]) {
let authorization = request.headers["authorization"].to_str().unwrap();
let proof = match mode {
RuntimeAuthMode::DualHeader => {
assert_eq!(
authorization,
format!("Bearer {}", expected.get_token().unwrap())
);
request.headers["x-openai-agent-assertion"]
.to_str()
.unwrap()
}
RuntimeAuthMode::AssertionOnly => {
authorization.strip_prefix("AgentAssertion ").unwrap()
}
};
let claims: serde_json::Value = serde_json::from_slice(
&base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(proof.split('.').nth(2).unwrap())
.unwrap(),
)
.unwrap();
assert_eq!(claims["runtime_agent_instance"], "runtime-refresh");
proofs.push(proof);
}
assert_ne!(proofs[0], proofs[1]);
*source.0.write().unwrap() = oauth("another-login", "replacement");
assert!(client.get_rate_limits().await.is_err());
}
}
codex-rs/backend-client/src/client/task_usage.rs
+1 -1
View File
@@ -103,7 +103,7 @@ impl Client {
}
let request = self
.request(Method::POST, &url)
.headers(self.headers()?)
.headers(self.headers().await?)
.json(&Query { threads });
let (body, _) = self.exec_request_detailed(request, "POST", &url).await?;
let response: TaskUsageResponse = serde_json::from_str(&body)
codex-rs/backend-client/src/client/thread_usage.rs
+1 -1
View File
@@ -82,7 +82,7 @@ impl Client {
let url = self.thread_usage_url();
let request = self
.request(Method::POST, &url)
.headers(self.headers()?)
.headers(self.headers().await?)
.header(CONTENT_TYPE, HeaderValue::from_static("application/json"))
.json(&ThreadUsageQueryRequest { thread_ids });
let (body, _) = self.exec_request_detailed(request, "POST", &url).await?;
codex-rs/backend-client/src/client/turn_usage.rs
+1 -2
View File
@@ -85,8 +85,7 @@ impl Client {
turn_ids: &[String],
provider_headers: &HeaderMap,
) -> Result<Vec<ApiKeyTurnCost>, RequestError> {
let mut headers = provider_headers.clone();
headers.extend(self.headers()?);
let headers = self.headers_with(provider_headers.clone()).await?;
let request = self
.request(Method::POST, url)
.headers(headers)
codex-rs/backend-client/src/client_request_tests.rs
+3 -1
View File
@@ -119,7 +119,8 @@ async fn migrated_requests_preserve_query_auth_and_json_body() {
)
.with_auth_provider(Arc::new(codex_model_provider::BearerAuthProvider::new(
"request-token".to_string(),
)));
)))
.unwrap();
let tasks = client
.list_tasks(
@@ -195,6 +196,7 @@ fn fallback_client(base_url: &str) -> Client {
.with_auth_provider(Arc::new(codex_model_provider::BearerAuthProvider::new(
"bootstrap-token".to_string(),
)))
.unwrap()
.with_chatgpt_account_id("workspace-123")
.with_user_agent("bootstrap-test")
}
codex-rs/cloud-config/src/config_server.rs
+9 -7
View File
@@ -191,13 +191,15 @@ impl BundleSource for ConfigServerSource {
});
}
*identity = credentials.authorization_id;
fetch_bundle(
&self
.client
.clone()
.with_bearer_token(credentials.access_token),
)
.await
let client = self
.client
.clone()
.with_bearer_token(credentials.access_token)
.map_err(|_| BundleRequestError::Fatal {
error: auth_failure(),
metric_reason: Some("auth"),
})?;
fetch_bundle(&client).await
}
async fn recover_auth(
codex-rs/cloud-tasks-client/src/http.rs
+7 -3
View File
@@ -42,9 +42,13 @@ impl HttpClient {
self
}
pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Self {
self.backend = self.backend.clone().with_auth_provider(auth);
self
pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Result<Self> {
self.backend = self
.backend
.clone()
.with_auth_provider(auth)
.map_err(|error| CloudTaskError::Msg(error.to_string()))?;
Ok(self)
}
pub fn with_chatgpt_account_id(mut self, account_id: impl Into<String>) -> Self {
codex-rs/cloud-tasks/src/lib.rs
+1 -1
View File
@@ -111,7 +111,7 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result<BackendContext>
}
let auth_provider = codex_model_provider::auth_provider_from_auth(&auth);
http = http.with_auth_provider(auth_provider);
http = http.with_auth_provider(auth_provider)?;
if let Some(acc) = auth.get_account_id() {
append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}"));
}
codex-rs/codex-api/src/files.rs
+6 -5
View File
@@ -8,6 +8,7 @@ use codex_http_client::RequestBuilder;
use codex_http_client::RouteAwareClientPool;
use codex_http_client::RouteAwareRequestError;
use futures::Stream;
use http::HeaderMap;
use http::Method;
use http::StatusCode;
use http::header::CONTENT_LENGTH;
@@ -154,7 +155,7 @@ where
"file_size": file_size_bytes,
"use_case": OPENAI_FILE_USE_CASE,
});
let request = authorized_request(client_pool, auth, Method::POST, &create_url)?;
let request = authorized_request(client_pool, auth, Method::POST, &create_url).await?;
let create_request = match hosted_upload {
Some(context) => serde_json::json!({
"file_name": file_name.as_str(),
@@ -365,7 +366,8 @@ where
};
let finalize_started_at = Instant::now();
loop {
let finalize_response = authorized_request(client_pool, auth, Method::POST, &finalize_url)?
let finalize_response = authorized_request(client_pool, auth, Method::POST, &finalize_url)
.await?
.json(&finalize_request)
.send()
.await
@@ -432,14 +434,13 @@ where
}
}
fn authorized_request(
async fn authorized_request(
client_pool: &RouteAwareClientPool,
auth: &dyn AuthProvider,
method: Method,
url: &str,
) -> Result<RequestBuilder, OpenAiFileError> {
let mut headers = http::HeaderMap::new();
auth.add_auth_headers(&mut headers)?;
let headers = auth.apply_auth_headers(HeaderMap::new()).await?;
Ok(client_pool
.request(method, url)
codex-rs/core/tests/suite/client_websockets_runtime_auth.rs
+61
View File
@@ -499,3 +499,64 @@ async fn gateway_refresh_failure_blocks_runtime_inference(mode: RuntimeAuthMode)
.all(|request| matches!(request.url.path(), "/register" | "/token"))
);
}
#[test_case(RuntimeAuthMode::DualHeader; "dual_header")]
#[test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn model_runtime_auth_strips_credentials_on_cross_origin_redirects(mode: RuntimeAuthMode) {
skip_if_no_network!();
let server = MockServer::start().await;
let destination = MockServer::start().await;
Mock::given(path("/register"))
.respond_with(registered)
.expect(1)
.mount(&server)
.await;
Mock::given(path("/v1/responses"))
.respond_with(ResponseTemplate::new(307).insert_header("location", destination.uri()))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(responses::sse(vec![
ev_response_created("resp-redirected"),
ev_completed("resp-redirected"),
])),
)
.expect(1)
.mount(&destination)
.await;
let mut provider =
ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri())));
provider.supports_websockets = false;
let (harness, _) = runtime_harness(&server, provider, mode).await;
complete(&mut harness.client.new_session(), &harness)
.await
.unwrap();
let requests = server.received_requests().await.unwrap();
let request = requests
.iter()
.find(|r| r.url.path() == "/v1/responses")
.unwrap();
assertion_payload(
request.headers["authorization"].to_str().unwrap(),
request
.headers
.get("x-openai-agent-assertion")
.map(|v| v.to_str().unwrap().to_string()),
mode,
"key-a",
);
let redirected = destination.received_requests().await.unwrap();
assert_eq!(redirected.len(), 1);
assert_eq!(redirected[0].body, request.body);
assert!(!redirected[0].headers.contains_key("authorization"));
assert!(
!redirected[0]
.headers
.contains_key("x-openai-agent-assertion")
);
}
codex-rs/http-client/src/route_aware_client_pool.rs
+16 -7
View File
@@ -56,6 +56,7 @@ pub struct RouteAwareClientPool {
pub(crate) client_builder: HttpClientBuilder,
default_headers: HeaderMap,
custom_ca_fallback: CustomCaFallback,
direct_proxy: bool,
clients: Arc<Mutex<HashMap<OutboundProxyRoute, TransportClient>>>,
client_build: Arc<tokio::sync::Mutex<()>>,
rustls_clients: Option<RustlsClientCache>,
@@ -289,6 +290,7 @@ impl RouteAwareClientPool {
client_builder,
default_headers,
custom_ca_fallback: CustomCaFallback::Disabled,
direct_proxy: false,
clients: Arc::new(Mutex::new(HashMap::new())),
client_build: Arc::default(),
rustls_clients: None,
@@ -357,6 +359,12 @@ impl RouteAwareClientPool {
self
}
/// Uses direct routing without changing destination policy or TLS configuration.
pub fn with_direct_proxy(mut self) -> Self {
self.direct_proxy = true;
self
}
/// Preserves the legacy sandbox client's direct routing and custom-CA fallback.
pub fn with_legacy_direct_proxy_and_custom_ca_fallback(mut self) -> Self {
self.custom_ca_fallback = CustomCaFallback::LegacyDirect;
@@ -478,13 +486,14 @@ impl RouteAwareClientPool {
F: FnOnce(String) -> Fut,
Fut: Future<Output = io::Result<OutboundProxyRoute>>,
{
let route = if self.custom_ca_fallback == CustomCaFallback::LegacyDirect {
OutboundProxyRoute::Direct
} else {
resolve_route(request_url.to_string())
.await
.map_err(RouteAwareClientPoolError::Resolve)?
};
let route =
if self.direct_proxy || self.custom_ca_fallback == CustomCaFallback::LegacyDirect {
OutboundProxyRoute::Direct
} else {
resolve_route(request_url.to_string())
.await
.map_err(RouteAwareClientPoolError::Resolve)?
};
if let Some(rustls_clients) = self.rustls_clients.as_ref()
&& let Ok(url) = reqwest::Url::parse(request_url)
&& rustls_clients.requires_rustls(&url, &route)
codex-rs/http-client/src/route_aware_client_pool_tests.rs
+26 -9
View File
@@ -377,7 +377,12 @@ async fn legacy_custom_ca_fallback_is_limited_to_reqwest_default() {
.expect("invalid CA fixture should be written");
for ca_env in ["CODEX_CA_CERTIFICATE", "SSL_CERT_FILE"] {
for policy_name in ["reqwest-default", "respect-system-proxy"] {
for policy_name in [
"reqwest-default",
"respect-system-proxy",
"strict-reqwest-default",
"strict-direct",
] {
let output = std::process::Command::new(
std::env::current_exe().expect("test executable should be available"),
)
@@ -403,18 +408,24 @@ async fn legacy_custom_ca_fallback_is_limited_to_reqwest_default() {
};
let outbound_proxy_policy = match policy_name.as_str() {
"reqwest-default" => OutboundProxyPolicy::ReqwestDefault,
"reqwest-default" | "strict-reqwest-default" | "strict-direct" => {
OutboundProxyPolicy::ReqwestDefault
}
"respect-system-proxy" => OutboundProxyPolicy::RespectSystemProxy,
_ => panic!("unexpected test proxy policy: {policy_name}"),
};
let pool = RouteAwareClientPool::with_chatgpt_cloudflare_cookies(
HttpClientFactory::new(outbound_proxy_policy),
ClientRouteClass::Other,
)
.with_legacy_custom_ca_fallback();
);
let pool = match policy_name.as_str() {
"strict-reqwest-default" => pool,
"strict-direct" => pool.with_direct_proxy(),
_ => pool.with_legacy_custom_ca_fallback(),
};
match outbound_proxy_policy {
OutboundProxyPolicy::ReqwestDefault => {
match policy_name.as_str() {
"reqwest-default" => {
let (address, server) = spawn_response_server(vec![
"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_string(),
]);
@@ -428,13 +439,19 @@ async fn legacy_custom_ca_fallback_is_limited_to_reqwest_default() {
let requests = server.join().expect("response server should finish");
assert_eq!(requests.len(), 1);
}
OutboundProxyPolicy::RespectSystemProxy => {
_ => {
let error = pool
.client_for_url_with_resolver("http://127.0.0.1/update", |_| async {
Ok(OutboundProxyRoute::Direct)
if policy_name == "strict-direct" {
Err(io::Error::other(
"direct routing must not need proxy discovery",
))
} else {
Ok(OutboundProxyRoute::Direct)
}
})
.await
.expect_err("system-proxy routes should reject invalid custom CAs");
.expect_err("strict routes should reject invalid custom CAs");
assert!(matches!(error, RouteAwareClientPoolError::Build(_)));
}
codex-rs/login/src/auth/default_client.rs
+35 -8
View File
@@ -263,11 +263,19 @@ pub fn create_client_for_route(
route_class: ClientRouteClass,
redirect_policy: ClientRedirectPolicy,
) -> Result<HttpClient, BuildRouteAwareHttpClientError> {
let legacy_fallback = redirect_policy == ClientRedirectPolicy::Default
|| !http_client_factory.protects_sensitive_headers();
let builder = match redirect_policy {
ClientRedirectPolicy::Default => default_http_client_builder(),
ClientRedirectPolicy::Reject => default_http_client_builder().without_redirects(),
};
create_client_for_route_with_builder(http_client_factory, request_url, route_class, builder)
create_client_for_route_with_builder(
http_client_factory,
request_url,
route_class,
builder,
legacy_fallback,
)
}
fn create_client_for_route_with_builder(
@@ -275,9 +283,11 @@ fn create_client_for_route_with_builder(
request_url: &str,
route_class: ClientRouteClass,
builder: HttpClientBuilder,
legacy_fallback: bool,
) -> Result<HttpClient, BuildRouteAwareHttpClientError> {
if http_client_factory.network_policy().is_managed()
|| http_client_factory.protects_sensitive_headers()
|| !legacy_fallback
{
let mut pool = codex_http_client::RouteAwareClientPool::with_builder(
http_client_factory.clone(),
@@ -285,11 +295,17 @@ fn create_client_for_route_with_builder(
builder,
);
if is_sandboxed() {
pool = pool.with_legacy_direct_proxy_and_custom_ca_fallback();
} else if matches!(
http_client_factory.outbound_proxy_policy(),
OutboundProxyPolicy::ReqwestDefault
) {
pool = if legacy_fallback {
pool.with_legacy_direct_proxy_and_custom_ca_fallback()
} else {
pool.with_direct_proxy()
};
} else if legacy_fallback
&& matches!(
http_client_factory.outbound_proxy_policy(),
OutboundProxyPolicy::ReqwestDefault
)
{
pool = pool.with_legacy_custom_ca_fallback();
}
return Ok(pool.into_client());
@@ -324,6 +340,8 @@ pub async fn create_client_for_route_async(
route_class: ClientRouteClass,
redirect_policy: ClientRedirectPolicy,
) -> std::io::Result<HttpClient> {
let legacy_fallback = redirect_policy == ClientRedirectPolicy::Default
|| !http_client_factory.protects_sensitive_headers();
create_client_for_route_async_with_builder(
http_client_factory,
request_url,
@@ -332,6 +350,7 @@ pub async fn create_client_for_route_async(
ClientRedirectPolicy::Default => default_http_client_builder(),
ClientRedirectPolicy::Reject => default_http_client_builder().without_redirects(),
},
legacy_fallback,
)
.await
}
@@ -347,6 +366,7 @@ pub async fn create_client_for_route_without_request_logging_async(
request_url,
route_class,
|| default_http_client_builder().without_request_logging(),
/*legacy_fallback*/ true,
)
.await
}
@@ -356,6 +376,7 @@ async fn create_client_for_route_async_with_builder(
request_url: String,
route_class: ClientRouteClass,
builder: impl FnOnce() -> HttpClientBuilder + Send + 'static,
legacy_fallback: bool,
) -> std::io::Result<HttpClient> {
let permit = ROUTE_AWARE_CLIENT_BUILD_PERMIT
.acquire()
@@ -368,6 +389,7 @@ async fn create_client_for_route_async_with_builder(
&request_url,
route_class,
builder(),
legacy_fallback,
)
.map_err(std::io::Error::from)
})
@@ -413,8 +435,13 @@ pub async fn create_transport_for_routes_async(
},
);
if sandboxed {
pool = pool.with_legacy_direct_proxy_and_custom_ca_fallback();
} else if transport_default_proxy {
pool = match redirect_policy {
ClientRedirectPolicy::Default => {
pool.with_legacy_direct_proxy_and_custom_ca_fallback()
}
ClientRedirectPolicy::Reject => pool.with_direct_proxy(),
};
} else if transport_default_proxy && redirect_policy == ClientRedirectPolicy::Default {
pool = pool.with_legacy_custom_ca_fallback();
}
ReqwestTransport::from_route_aware_client_pool(pool)
codex-rs/login/src/auth/default_client_tests.rs
+58
View File
@@ -355,3 +355,61 @@ fn test_macos() {
.unwrap();
assert!(re.is_match(&user_agent));
}
#[tokio::test]
async fn signed_reject_routes_preserve_invalid_ca_errors() {
const CHILD: &str = "CODEX_TEST_SIGNED_ROUTE_INVALID_CA";
if std::env::var_os(CHILD).is_none() {
let dir = tempfile::tempdir().unwrap();
let ca = dir.path().join("invalid.pem");
std::fs::write(&ca, "invalid certificate").unwrap();
let output = std::process::Command::new(std::env::current_exe().unwrap())
.arg("signed_reject_routes_preserve_invalid_ca_errors")
.arg("--nocapture")
.env(CHILD, "1")
.env("CODEX_CA_CERTIFICATE", ca)
.env_remove("SSL_CERT_FILE")
.output()
.unwrap();
assert!(
output.status.success(),
"{}",
String::from_utf8_lossy(&output.stderr)
);
assert!(String::from_utf8_lossy(&output.stdout).contains("1 passed"));
return;
}
let url = "https://127.0.0.1:9/";
let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
.with_sensitive_header_protection();
let sync_client = create_client_for_route(
&factory,
url,
ClientRouteClass::Api,
ClientRedirectPolicy::Reject,
)
.unwrap();
let async_client = create_client_for_route_async(
factory,
url.into(),
ClientRouteClass::Api,
ClientRedirectPolicy::Reject,
)
.await
.unwrap();
for client in [sync_client, async_client] {
let error = client.get(url).send().await.unwrap_err();
assert!(
matches!(
error,
codex_http_client::HttpError::Route(
codex_http_client::RouteAwareClientPoolError::Build(
codex_http_client::BuildRouteAwareHttpClientError::CustomCa(_)
)
)
),
"{error}"
);
}
}
codex-rs/model-provider/src/auth.rs
+55 -1
View File
@@ -440,7 +440,61 @@ pub(crate) fn uses_fedramp_api_key(
&& auth.is_some_and(CodexAuth::is_fedramp_account))
}
/// Builds request-header auth for a first-party Codex auth snapshot.
/// Select the shared runtime without changing the identity of captured request data.
pub async fn auth_provider_for_request(
auth: &CodexAuth,
runtime: Option<&(Arc<RuntimeAuthManager>, RuntimeAuthMode)>,
headers: &HeaderMap,
) -> Result<ResolvedProviderAuth, AuthError> {
observe_auth_selection(runtime.map(|(_, mode)| *mode), "request", async {
match runtime {
Some((manager, mode)) => {
let manager = Arc::clone(manager);
let auth = auth.clone();
let selection = crate::runtime_auth::resource_selection(headers)?;
let selected = crate::runtime_auth::finish_runtime_selection(async move {
manager.auth_for(&auth, &selection).await
})
.await?;
Ok(ResolvedProviderAuth::for_runtime(selected, *mode))
}
None => {
let mut resolved = ResolvedProviderAuth::new(auth_provider_from_auth(auth));
resolved.credential = Some(auth.clone());
Ok(resolved)
}
}
})
.await
}
/// Record bounded auth-mode and outcome labels without changing the selection result.
pub(crate) async fn observe_auth_selection<T, E>(
mode: Option<RuntimeAuthMode>,
operation: &'static str,
selection: impl Future<Output = Result<T, E>>,
) -> Result<T, E> {
let started = std::time::Instant::now();
let result = selection.await;
let mode = match mode {
None => "ordinary",
Some(RuntimeAuthMode::DualHeader) => "dual_header",
Some(RuntimeAuthMode::AssertionOnly) => "assertion_only",
};
let _ = codex_otel::record_global_operation(
"codex.auth.selection",
"codex.auth.selection.duration",
started.elapsed(),
&[
("mode", mode),
("operation", operation),
("outcome", if result.is_ok() { "success" } else { "error" }),
],
);
result
}
/// Builds request-header auth for a first-party Codex credential.
pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider {
match auth {
CodexAuth::AgentIdentity(auth) => {
codex-rs/model-provider/src/client_auth.rs
+120
View File
@@ -0,0 +1,120 @@
//! Authentication for clients that reuse a connection across requests.
use crate::ResolvedProviderAuth;
use codex_api::AuthError;
use codex_api::SharedAuthProvider;
use codex_features::RuntimeAuthMode;
use codex_http_client::HttpClientFactory;
use codex_login::CodexAuth;
use codex_protocol::error::Result as CodexResult;
use codex_runtime_agent_instance::RuntimeAuthManager;
use http::HeaderMap;
use std::sync::Arc;
/// Keeps a reusable client's credential source fixed while allowing token refresh.
/// Each request resolves fresh auth using the same source as the client transport.
#[derive(Clone)]
pub struct ClientAuth {
source: ClientAuthSource,
}
#[derive(Clone)]
enum ClientAuthSource {
Provider(SharedAuthProvider),
Runtime {
manager: Arc<RuntimeAuthManager>,
mode: RuntimeAuthMode,
source_generation: u64,
},
}
impl ClientAuth {
pub fn new(provider: SharedAuthProvider) -> Self {
Self {
source: ClientAuthSource::Provider(provider),
}
}
/// Select the auth source and matching transport for a reusable client's request headers.
/// The caller selects runtime auth only for a supported OpenAI destination.
pub async fn for_auth(
auth: &CodexAuth,
runtime: Option<&(Arc<RuntimeAuthManager>, RuntimeAuthMode)>,
factory: HttpClientFactory,
headers: &HeaderMap,
) -> Result<(Self, HttpClientFactory), AuthError> {
let resolved = crate::auth_provider_for_request(auth, runtime, headers).await?;
let factory = resolved.bind_http_client_factory(factory)?;
let source = match (runtime, resolved.runtime) {
(Some((manager, mode)), Some(selected)) => ClientAuthSource::Runtime {
manager: Arc::clone(manager),
mode: *mode,
source_generation: selected.change_state().source_generation,
},
_ => ClientAuthSource::Provider(resolved.auth),
};
Ok((Self { source }, factory))
}
/// Replace ordinary authentication; reject changes to a runtime-bound client.
pub fn set_provider(&mut self, provider: SharedAuthProvider) -> Result<(), AuthError> {
match &mut self.source {
ClientAuthSource::Provider(current) => {
*current = provider;
Ok(())
}
ClientAuthSource::Runtime { .. } => Err(AuthError::Build(
"cannot replace credentials on a runtime-bound client".to_string(),
)),
}
}
/// Resolve one request, preserving the source that authenticated the client.
pub async fn resolve(&self, headers: &HeaderMap) -> Result<ResolvedProviderAuth, AuthError> {
let mode = match &self.source {
ClientAuthSource::Provider(_) => None,
ClientAuthSource::Runtime { mode, .. } => Some(*mode),
};
crate::auth::observe_auth_selection(mode, "client_request", async {
match &self.source {
ClientAuthSource::Provider(provider) => {
Ok(ResolvedProviderAuth::new(Arc::clone(provider)))
}
ClientAuthSource::Runtime {
manager,
mode,
source_generation,
} => {
let selected =
crate::select_runtime_auth_for_source(manager, *source_generation, headers)
.await?;
Ok(ResolvedProviderAuth::for_runtime(selected, *mode))
}
}
})
.await
}
/// Allow OpenAI requests and the original origin, never an unrelated provider.
pub fn validate_destination(&self, base_url: &str, url: &str) -> CodexResult<()> {
if matches!(&self.source, ClientAuthSource::Provider(_)) {
return Ok(());
}
crate::validate_runtime_destination(Some(url))?;
if codex_runtime_agent_instance::is_openai_api_url(url) {
return Ok(());
}
let origin = |value| {
url::Url::parse(value)
.map(|url| url.origin())
.map_err(std::io::Error::other)
};
if origin(url)? != origin(base_url)? {
return Err(std::io::Error::other(
"runtime-authenticated clients cannot send requests to another provider",
)
.into());
}
Ok(())
}
}
codex-rs/model-provider/src/lib.rs
+5
View File
@@ -4,6 +4,7 @@ mod bearer_auth_provider;
mod capabilities;
mod catalog_factory;
mod chatgpt_base_url;
mod client_auth;
mod combined_auth;
mod gcp;
mod models_endpoint;
@@ -27,6 +28,7 @@ pub use auth::AgentIdentitySessionFallback;
pub use auth::ProviderAuthContext;
pub use auth::ProviderAuthScope;
pub use auth::ResolvedProviderAuth;
pub use auth::auth_provider_for_request;
pub use auth::auth_provider_from_auth;
pub use auth::auth_provider_from_auth_manager;
pub use auth::unauthenticated_auth_provider;
@@ -36,6 +38,7 @@ pub use capabilities::ProviderCapabilities;
pub use catalog_factory::ModelCatalogFactory;
pub use chatgpt_base_url::ChatgptBaseUrl;
pub use chatgpt_base_url::DEFAULT_CHATGPT_BASE_URL;
pub use client_auth::ClientAuth;
pub use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID;
pub use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_PROVIDER_ID;
pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL;
@@ -57,6 +60,8 @@ pub use provider::resolve_api_provider;
pub use request_preparation::prepare_provider_request;
pub use request_setup::ResolvedProviderRequest;
pub use request_setup::resolve_request;
pub use runtime_auth::runtime_auth_for_url;
pub use runtime_auth::select_runtime_auth_for_source;
pub use runtime_auth::validate_runtime_destination;
pub use runtime_auth::validate_runtime_provider;
pub use runtime_auth::validate_runtime_selectors;
codex-rs/model-provider/src/models_endpoint_refresh_tests.rs
+77 -19
View File
@@ -20,22 +20,43 @@ use wiremock::matchers::body_partial_json;
use wiremock::matchers::method;
use wiremock::matchers::path;
#[derive(Clone, Copy)]
enum Caller {
Catalog,
Startup,
CapturedCredential,
SourceGeneration,
}
#[tokio::test]
async fn catalog_timeout_preserves_refreshed_tokens_with_and_without_runtime_auth() {
check_catalog_timeout_preserves_refreshed_tokens(/*refresh_after_auth_change*/ false).await;
check_timeout_preserves_refreshed_tokens(Caller::Catalog).await;
}
#[tokio::test]
async fn catalog_startup_timeout_preserves_refreshed_tokens() {
check_catalog_timeout_preserves_refreshed_tokens(/*refresh_after_auth_change*/ true).await;
check_timeout_preserves_refreshed_tokens(Caller::Startup).await;
}
async fn check_catalog_timeout_preserves_refreshed_tokens(refresh_after_auth_change: bool) {
#[tokio::test]
async fn cancelled_runtime_client_selection_persists_the_rotated_token() {
check_timeout_preserves_refreshed_tokens(Caller::CapturedCredential).await;
check_timeout_preserves_refreshed_tokens(Caller::SourceGeneration).await;
}
async fn check_timeout_preserves_refreshed_tokens(caller: Caller) {
for mode in [
None,
Some(RuntimeAuthMode::DualHeader),
Some(RuntimeAuthMode::AssertionOnly),
] {
if matches!(
caller,
Caller::CapturedCredential | Caller::SourceGeneration
) && mode != Some(RuntimeAuthMode::DualHeader)
{
continue;
}
let server = MockServer::start().await;
let home = tempfile::tempdir().unwrap();
let token = |expires| {
@@ -155,23 +176,60 @@ async fn check_catalog_timeout_preserves_refreshed_tokens(refresh_after_auth_cha
Some(Arc::clone(&credentials)),
/*gateway_auth_manager*/ None,
)
.with_runtime_auth(runtime_auth),
.with_runtime_auth(runtime_auth.clone()),
);
if refresh_after_auth_change {
let manager = OpenAiModelsManager::new_without_cache(
endpoint.clone(),
Some(Arc::clone(&credentials)),
);
manager.refresh_after_auth_change(factory.clone()).await;
} else {
let error = endpoint
.list_models("test", factory.clone())
.await
.unwrap_err();
assert!(matches!(
error.details(),
codex_protocol::error::CodexErrorDetails::RequestTimeout
));
match caller {
Caller::CapturedCredential => {
let auth = credentials.auth_cached().unwrap();
assert!(
timeout(
MODELS_REFRESH_TIMEOUT,
crate::auth_provider_for_request(
&auth,
runtime_auth.as_ref(),
&http::HeaderMap::new()
)
)
.await
.is_err()
);
}
Caller::SourceGeneration => {
let generation = credentials
.auth_change_state_receiver()
.borrow()
.source_generation;
let runtime = &runtime_auth.as_ref().unwrap().0;
assert!(
timeout(
MODELS_REFRESH_TIMEOUT,
crate::select_runtime_auth_for_source(
runtime,
generation,
&http::HeaderMap::new()
)
)
.await
.is_err()
);
}
Caller::Startup => {
let manager = OpenAiModelsManager::new_without_cache(
endpoint.clone(),
Some(Arc::clone(&credentials)),
);
manager.refresh_after_auth_change(factory.clone()).await;
}
Caller::Catalog => {
let error = endpoint
.list_models("test", factory.clone())
.await
.unwrap_err();
assert!(matches!(
error.details(),
codex_protocol::error::CodexErrorDetails::RequestTimeout
));
}
}
assert_eq!(
server
codex-rs/model-provider/src/models_identity.rs
+7
View File
@@ -16,6 +16,13 @@ pub(crate) fn identity(
provider_info: &ModelProviderInfo,
auth: Option<&CodexAuth>,
) -> CoreResult<String> {
// Provider bearer overrides replace the login credential and its account metadata.
let auth =
if provider_info.env_key.is_some() || provider_info.experimental_bearer_token.is_some() {
None
} else {
auth
};
let mut api_provider = resolve_api_provider(provider_info, auth)?;
let mut digest = Sha256::new();
// Length-prefix every field to avoid ambiguity between adjacent values.
codex-rs/model-provider/src/models_identity_tests.rs
+27
View File
@@ -113,3 +113,30 @@ fn cache_identity_tracks_catalog_url() {
assert_ne!(bundled, first);
assert_ne!(first, second);
}
#[test]
fn provider_key_catalog_identity_does_not_depend_on_unused_login() {
let mut provider =
ModelProviderInfo::create_openai_provider(Some("https://api.openai.com/v1".into()));
// Keep this input stable while the residency test changes the process default.
provider.http_headers = Some(std::collections::HashMap::from([(
codex_login::default_client::RESIDENCY_HEADER_NAME.to_string(),
"us".into(),
)]));
provider.requires_openai_auth = false;
provider.experimental_bearer_token = Some("provider-key".into());
let identity_without_login = identity(&provider, /*auth*/ None).unwrap();
let selected_key = CodexAuth::from_api_key("provider-key");
let unused_login = chatgpt_auth("user@example.com", "user", "workspace", "team", "token");
for auth in [&selected_key, &unused_login] {
assert_eq!(
identity(&provider, Some(auth)).unwrap(),
identity_without_login
);
}
provider.experimental_bearer_token = Some("different-provider-key".into());
assert_ne!(
identity(&provider, /*auth*/ None).unwrap(),
identity_without_login
);
}
codex-rs/model-provider/src/request_setup.rs
+37 -25
View File
@@ -33,31 +33,42 @@ pub(crate) async fn resolve_request_auth(
destination: Option<String>,
) -> codex_protocol::error::Result<ResolvedProviderAuth> {
tokio::spawn(async move {
if let Some((runtime, mode)) = runtime {
crate::validate_runtime_destination(destination.as_deref())?;
crate::validate_runtime_provider(&info, manager.as_ref(), &runtime)?;
let headers = info.build_header_map()?;
let selected = crate::runtime_auth::select_runtime_auth(&runtime, &headers)
.await
.map_err(|error| codex_api::map_api_error(error.into()))?;
let resolved = ResolvedProviderAuth::for_runtime(selected, mode);
resolved.validate_headers(&headers)?;
tracing::debug!(runtime_auth_mode = ?mode, "selected runtime authentication");
return Ok(resolved);
}
crate::auth::observe_auth_selection(
runtime.as_ref().map(|(_, mode)| *mode),
"model",
async move {
if let Some((runtime, mode)) = runtime {
crate::validate_runtime_destination(destination.as_deref())?;
crate::validate_runtime_provider(&info, manager.as_ref(), &runtime)?;
let headers = info.build_header_map()?;
let selected = crate::runtime_auth::select_runtime_auth(&runtime, &headers)
.await
.map_err(|error| codex_api::map_api_error(error.into()))?;
let resolved = ResolvedProviderAuth::for_runtime(selected, mode);
resolved.validate_headers(&headers)?;
tracing::debug!(runtime_auth_mode = ?mode, "selected runtime authentication");
return Ok(resolved);
}
let credential = match &manager {
Some(manager) => manager.auth().await,
None => None,
};
let mut resolved = match scope {
Some(scope) if crate::provider_uses_first_party_auth_path(&info) => {
resolve_provider_auth_for_scope(manager, credential.as_ref(), &info, scope).await?
}
_ => ResolvedProviderAuth::new(resolve_provider_auth(credential.as_ref(), &info)?),
};
resolved.credential = credential;
Ok(resolved)
let credential = match &manager {
Some(manager) => manager.auth().await,
None => None,
};
let mut resolved = match scope {
Some(scope) if crate::provider_uses_first_party_auth_path(&info) => {
resolve_provider_auth_for_scope(manager, credential.as_ref(), &info, scope)
.await?
}
_ => ResolvedProviderAuth::new(resolve_provider_auth(
credential.as_ref(),
&info,
)?),
};
resolved.credential = credential;
Ok(resolved)
},
)
.await
})
.await?
}
@@ -85,9 +96,10 @@ impl ResolvedProviderRequest {
} else {
self.redirect_policy
};
let request_url = self.api_provider.url_for_path(endpoint);
let client = create_client_for_route(
&self.http_client_factory,
&self.api_provider.url_for_path(endpoint),
&request_url,
ClientRouteClass::Api,
redirect_policy,
)
codex-rs/model-provider/src/runtime_auth.rs
+48 -6
View File
@@ -1,6 +1,7 @@
//! Keep model requests and catalog refreshes bound to the same registered actor.
use codex_api::AuthError;
use codex_features::RuntimeAuthMode;
use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER;
use codex_http_client::OPENAI_ORGANIZATION_HEADER;
use codex_http_client::OPENAI_PROJECT_HEADER;
@@ -13,6 +14,7 @@ use codex_protocol::error::Result as CodexResult;
use codex_runtime_agent_instance::ActorInfo;
use codex_runtime_agent_instance::PlatformResourceSelection;
use codex_runtime_agent_instance::RuntimeAuth;
use codex_runtime_agent_instance::RuntimeAuthError;
use codex_runtime_agent_instance::RuntimeAuthManager;
use http::HeaderMap;
use std::sync::Arc;
@@ -22,6 +24,39 @@ pub(crate) async fn select_runtime_auth(
manager: &RuntimeAuthManager,
headers: &HeaderMap,
) -> Result<RuntimeAuth, AuthError> {
manager
.auth_with_selection(&resource_selection(headers)?)
.await
.map_err(AuthError::from)
}
/// Select fresh request auth without allowing a reusable client's source to change.
pub async fn select_runtime_auth_for_source(
manager: &Arc<RuntimeAuthManager>,
source_generation: u64,
headers: &HeaderMap,
) -> Result<RuntimeAuth, AuthError> {
let manager = Arc::clone(manager);
let selection = resource_selection(headers)?;
finish_runtime_selection(
async move { manager.auth_for_source(source_generation, &selection).await },
)
.await
}
// Once token rotation begins, let it finish persisting even if the request is cancelled.
pub(crate) async fn finish_runtime_selection(
selection: impl Future<Output = Result<RuntimeAuth, RuntimeAuthError>> + Send + 'static,
) -> Result<RuntimeAuth, AuthError> {
tokio::spawn(selection)
.await
.map_err(|error| AuthError::Build(format!("runtime auth task failed: {error}")))?
.map_err(AuthError::from)
}
pub(crate) fn resource_selection(
headers: &HeaderMap,
) -> Result<PlatformResourceSelection, AuthError> {
let selector = |name| {
headers
.get(name)
@@ -33,14 +68,10 @@ pub(crate) async fn select_runtime_auth(
})
.transpose()
};
let selection = PlatformResourceSelection {
Ok(PlatformResourceSelection {
organization_id: selector(OPENAI_ORGANIZATION_HEADER)?,
project_id: selector(OPENAI_PROJECT_HEADER)?,
};
manager
.auth_with_selection(&selection)
.await
.map_err(AuthError::from)
})
}
/// A runtime may not replace an independently selected provider credential.
@@ -74,6 +105,17 @@ pub fn validate_runtime_destination(destination: Option<&str>) -> CodexResult<()
Ok(())
}
/// Select runtime auth for supported OpenAI URLs; reject unsupported OpenAI URLs.
pub fn runtime_auth_for_url<'a>(
runtime: Option<&'a (Arc<RuntimeAuthManager>, RuntimeAuthMode)>,
url: &str,
) -> CodexResult<Option<&'a (Arc<RuntimeAuthManager>, RuntimeAuthMode)>> {
if runtime.is_some() {
validate_runtime_destination(Some(url))?;
}
Ok(runtime.filter(|_| classify_openai_api_url(url) == OpenAiDestination::Supported))
}
/// Account and organization selectors must agree with the registered actor.
/// API project access is authorized by the backend for each request.
pub fn validate_runtime_selectors(headers: &HeaderMap, actor: &ActorInfo) -> CodexResult<()> {
codex-rs/runtime-agent-instance/Cargo.toml
+2
View File
@@ -30,6 +30,8 @@ uuid = { workspace = true, features = ["serde", "v4"] }
zeroize = { workspace = true }
[dev-dependencies]
bytes = { workspace = true }
futures = { workspace = true }
pretty_assertions = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
codex-rs/runtime-agent-instance/src/lib.rs
+3
View File
@@ -28,3 +28,6 @@ pub use registration::RegistrationErrorCode;
pub use registration::RuntimeRegistrationClient;
pub use runtime::RuntimeAgentInstance;
pub use runtime::RuntimeRegistration;
/// Harness-owned lazy runtime, shared before any request requires registration.
pub type RuntimeAuthSlot = std::sync::Arc<std::sync::OnceLock<std::sync::Arc<RuntimeAuthManager>>>;
codex-rs/runtime-agent-instance/src/manager.rs
+69 -6
View File
@@ -135,16 +135,25 @@ impl RuntimeAuthManager {
}
/// Select auth and a runtime for the requested organization and project.
#[expect(
clippy::await_holding_invalid_type,
reason = "runtime initialization serializes registration and refresh for one credential source"
)]
pub async fn auth_with_selection(
&self,
selection: &PlatformResourceSelection,
) -> Result<RuntimeAuth, RuntimeAuthError> {
self.select_auth(selection, /*expected_source*/ None).await
}
#[expect(
clippy::await_holding_invalid_type,
reason = "runtime initialization serializes registration and refresh for one credential source"
)]
async fn select_auth(
&self,
selection: &PlatformResourceSelection,
expected_source: Option<u64>,
) -> Result<RuntimeAuth, RuntimeAuthError> {
// This initializer's lock can span IO. No AuthManager lock is held across registration.
let mut state = self.state.lock().await;
self.check_source(expected_source)?;
let (auth, http_client_factory) = self
.credentials
.auth_with_http_client_factory()
@@ -152,7 +161,13 @@ impl RuntimeAuthManager {
.ok_or(RuntimeAuthError::Unavailable)?;
let mut recovery = self.credentials.unauthorized_recovery();
let mut result = self
.auth_once(&mut state, auth, http_client_factory, selection)
.auth_once(
&mut state,
auth,
http_client_factory,
selection,
expected_source,
)
.await;
while matches!(
&result,
@@ -164,6 +179,7 @@ impl RuntimeAuthManager {
))
) && recovery.has_next()
{
self.check_source(expected_source)?;
recovery.next().await.map_err(|error| match error {
RefreshTokenError::Policy(denied) => RuntimeAuthError::Policy(denied),
RefreshTokenError::Permanent(_) | RefreshTokenError::Transient(_) => {
@@ -176,7 +192,13 @@ impl RuntimeAuthManager {
.await
.ok_or(RuntimeAuthError::Unavailable)?;
result = self
.auth_once(&mut state, auth, http_client_factory, selection)
.auth_once(
&mut state,
auth,
http_client_factory,
selection,
expected_source,
)
.await;
}
result
@@ -188,10 +210,12 @@ impl RuntimeAuthManager {
auth: CodexAuth,
http_client_factory: HttpClientFactory,
selection: &PlatformResourceSelection,
expected_source: Option<u64>,
) -> Result<RuntimeAuth, RuntimeAuthError> {
let actor_auth = actor_authorization(&auth, selection)?;
let http_client_factory = http_client_factory.with_sensitive_header_protection();
let change_state = self.auth_state(&auth)?;
self.check_source(expected_source)?;
// Registered runtimes can serve different projects. Pending and denied
// registrations retain their original scope and idempotency key.
if state.as_ref().is_some_and(|cached| {
@@ -269,6 +293,45 @@ impl RuntimeAuthManager {
}
}
/// Resolve auth for work whose routing, cache key, or payload uses a captured credential.
/// A same-source refresh is allowed; a different source must not take over that work.
/// This binds credential identity, not when the caller captured it. Restoring the exact
/// same credential can select a new runtime after logout; a different source cannot.
pub async fn auth_for(
&self,
expected: &CodexAuth,
selection: &PlatformResourceSelection,
) -> Result<RuntimeAuth, RuntimeAuthError> {
let source = self
.credentials
.auth_change_state_for(expected)
.ok_or(RuntimeAuthError::Changed)?;
self.auth_for_source(source.source_generation, selection)
.await
}
/// Select another request within a reusable client's original credential source.
pub async fn auth_for_source(
&self,
source_generation: u64,
selection: &PlatformResourceSelection,
) -> Result<RuntimeAuth, RuntimeAuthError> {
self.select_auth(selection, Some(source_generation)).await
}
fn check_source(&self, expected: Option<u64>) -> Result<(), RuntimeAuthError> {
if expected.is_some_and(|expected| {
self.credentials
.auth_change_state_receiver()
.borrow()
.source_generation
!= expected
}) {
return Err(RuntimeAuthError::Changed);
}
Ok(())
}
fn auth_state(&self, auth: &CodexAuth) -> Result<AuthChangeState, RuntimeAuthError> {
let state = self
.credentials
codex-rs/runtime-agent-instance/src/manager_tests.rs
+137
View File
@@ -571,3 +571,140 @@ async fn registration_failure_reports_retryability() {
);
}
}
#[tokio::test]
async fn captured_request_auth_allows_refresh_but_rejects_source_replacement() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(registered)
.mount(&server)
.await;
let initial = oauth(Some("login-a"), "old");
let refreshed = oauth(Some("login-a"), "new");
let (manager, source, _) = setup(&server, initial.clone(), refreshed.clone()).await;
let runtime = manager.auth().await.unwrap().runtime().id().to_string();
*source.current.write().unwrap() = refreshed.clone();
let selected = manager
.auth_for(&initial, &PlatformResourceSelection::default())
.await
.unwrap();
let source_generation = selected.change_state().source_generation;
let request_auth = selected.into_auth_provider(RuntimeAuthMode::DualHeader);
let headers = request_auth
.apply_auth_headers(HeaderMap::new())
.await
.unwrap();
assert_eq!(
headers[AUTHORIZATION],
format!("Bearer {}", refreshed.get_token().unwrap())
);
let proof = headers["x-openai-agent-assertion"].to_str().unwrap();
let payload: Value = serde_json::from_slice(
&base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(proof.split('.').nth(2).unwrap())
.unwrap(),
)
.unwrap();
assert_eq!(payload["runtime_agent_instance"], runtime);
*source.current.write().unwrap() = oauth(Some("another-login"), "replacement");
assert!(matches!(
manager
.auth_for_source(source_generation, &PlatformResourceSelection::default())
.await,
Err(RuntimeAuthError::Changed)
));
assert!(
manager
.auth_for(&initial, &PlatformResourceSelection::default())
.await
.is_err()
);
assert_eq!(server.received_requests().await.unwrap().len(), 1);
}
#[tokio::test]
async fn file_upload_runtime_auth_signs_control_requests_but_not_blob_uploads() {
use wiremock::matchers::path;
for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] {
let server = MockServer::start().await;
let blob = MockServer::start().await;
Mock::given(path("/agent/register"))
.respond_with(registered)
.expect(1)
.mount(&server)
.await;
Mock::given(path("/files"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"file_id": "file-a", "upload_url": blob.uri(),
})))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/files/file-a/uploaded"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"status": "retry"})))
.up_to_n_times(1)
.with_priority(1)
.mount(&server)
.await;
Mock::given(path("/files/file-a/uploaded"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"status": "success", "download_url": blob.uri(), "file_name": "hello.txt", "file_size_bytes": 5,
}))).expect(1).mount(&server).await;
Mock::given(method("PUT"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&blob)
.await;
let initial = CodexAuth::from_api_key("key-a");
let (manager, _, _) = setup(&server, initial.clone(), initial).await;
let provider = manager.auth().await.unwrap().into_auth_provider(mode);
let pool =
codex_http_client::RouteAwareClientPool::new_without_redirects_or_request_logging(
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
codex_http_client::ClientRouteClass::Api,
);
let uploaded = codex_api::upload_openai_file(
&server.uri(),
provider.as_ref(),
&pool,
"hello.txt".into(),
/*file_size_bytes*/ 5,
|| async {
Ok(futures::stream::iter([Ok::<_, std::io::Error>(
bytes::Bytes::from_static(b"hello"),
)]))
},
/*hosted_upload*/ None,
)
.await
.unwrap();
assert_eq!(uploaded.uri, "sediment://file-a");
let requests = server.received_requests().await.unwrap();
let mut proofs = std::collections::HashSet::new();
for request in requests
.iter()
.filter(|r| r.url.path().starts_with("/files"))
{
let authorization = request.headers[AUTHORIZATION].to_str().unwrap();
let proof = match mode {
RuntimeAuthMode::DualHeader => {
assert_eq!(authorization, "Bearer key-a");
request.headers["x-openai-agent-assertion"]
.to_str()
.unwrap()
}
RuntimeAuthMode::AssertionOnly => {
assert!(!request.headers.contains_key("x-openai-agent-assertion"));
authorization.strip_prefix("AgentAssertion ").unwrap()
}
};
assert!(proofs.insert(proof));
}
assert_eq!(proofs.len(), 3);
let blobs = blob.received_requests().await.unwrap();
assert_eq!(blobs.len(), 1);
assert!(!blobs[0].headers.contains_key(AUTHORIZATION));
assert!(!blobs[0].headers.contains_key("x-openai-agent-assertion"));
assert_eq!(blobs[0].body, b"hello");
}
}