diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 6293bdd893..57d4742b64 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -610,8 +610,8 @@ process control and custom output caps do not apply to MXC. ### Gateway OAuth sign-in -Providers configured with `gateway_oauth` require a secondary OAuth credential in -addition to their primary authentication. Clients with a gateway sign-in UI set +A provider can obtain its gateway token through browser OAuth or a credential +export command. Clients with a gateway sign-in UI set `initialize.capabilities.explicitGatewayOauth: true`, complete initialization, and successfully call `account/gatewayOAuth/read` before sending authenticated requests, including startup `model/list` and inference requests. Repeat this probe on each @@ -630,6 +630,16 @@ With explicit login enabled, app-server refreshes existing credentials, but only `account/gatewayOAuth/login` RPC starts browser authorization. Requests needing sign-in fail promptly so the client can offer that flow. +Configure `[model_providers..gateway_oauth.credential_export]` to obtain the +gateway access token from a command on first use. Set `command`, optional `args`, +and optional `timeout_ms` (default 30000). The command prints one access token to +stdout. Codex caches it in memory and applies the configured gateway header or +cookie. Choose either `credential_export` or OAuth settings such as +`authorization_url`, `token_url`, and `client_id`. If the command fails, +the request returns the export error. Use direct +inference and catalog URLs because command credentials require same-origin catalog +access and reject HTTP redirects. + Clients that omit the capability or set it to `false`, including the TUI, retain automatic browser authorization after initialization. Startup credential reads cannot open a browser before initialization. Explicit opt-in is shared by gateway @@ -638,8 +648,9 @@ cannot be undone by a later connection that omits the capability. - `account/gatewayOAuth/read` returns the current effective `providerId`, `providerName`, `required`, `status`, and `error`. `required` indicates that this - provider uses gateway OAuth, including when already signed in. This operation - does not refresh tokens or open a browser; `status` is null for other providers. + provider uses gateway OAuth, including when already signed in. With a credential + exporter, `required` is false and `status` is null. This read does not refresh + tokens or open a browser. `notReady` means credentials are not ready. `succeeded` means saved credentials are locally usable, not that a gateway request has been verified. Reads observe usable replacement credentials saved by another process sharing the same home. diff --git a/codex-rs/app-server/src/gateway_oauth_notifications.rs b/codex-rs/app-server/src/gateway_oauth_notifications.rs index 91d8d5e43f..8002b6bfc0 100644 --- a/codex-rs/app-server/src/gateway_oauth_notifications.rs +++ b/codex-rs/app-server/src/gateway_oauth_notifications.rs @@ -46,7 +46,7 @@ pub(crate) fn spawn( .gateway_auth_manager() else { continue; }; - if manager.config() != &change.config { + if !matches!(manager.oauth_session(), Some(oauth) if oauth.config() == &change.config) { continue; } let (status, error) = status_fields(change.status); diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 4c46c69e25..8f75f86d07 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -117,7 +117,7 @@ pub(crate) struct AccountRequestProcessor { workspace_routing_fetches: Arc>, workspace_routing_shutdown: CancellationToken, gateway_login: Arc>>, - gateway_client: Arc>>>, + gateway_client: Arc>>>, _gateway_notifications: Arc>, pub(super) enterprise_login: Arc, } diff --git a/codex-rs/app-server/src/request_processors/account_processor/gateway_oauth.rs b/codex-rs/app-server/src/request_processors/account_processor/gateway_oauth.rs index c648eeea94..7899fcb0f5 100644 --- a/codex-rs/app-server/src/request_processors/account_processor/gateway_oauth.rs +++ b/codex-rs/app-server/src/request_processors/account_processor/gateway_oauth.rs @@ -7,7 +7,8 @@ use crate::transport::ConnectionId; use codex_app_server_protocol::GatewayOAuthCancelResponse; use codex_app_server_protocol::GatewayOAuthLoginResponse; use codex_app_server_protocol::GatewayOAuthReadResponse; -use codex_login::GatewayAuthManager; +use codex_login::OAuthSessionManager; +use codex_model_provider::GatewayAuthManager; use std::sync::PoisonError; pub(super) struct ActiveGatewayLogin { @@ -29,7 +30,7 @@ impl Drop for LoginGuard { impl AccountRequestProcessor { async fn gateway_client( &self, - ) -> Result<(Config, Option>), JSONRPCErrorError> { + ) -> Result<(Config, Option>), JSONRPCErrorError> { let config = self .config_manager .load_latest_config(/*fallback_cwd*/ None) @@ -43,11 +44,14 @@ impl AccountRequestProcessor { ) .gateway_auth_manager() .map_err(|error| internal_error(error.to_string()))?; + let session = manager + .as_deref() + .and_then(GatewayAuthManager::oauth_session); *self .gateway_client .lock() - .unwrap_or_else(PoisonError::into_inner) = manager.clone(); - Ok((config, manager)) + .unwrap_or_else(PoisonError::into_inner) = manager; + Ok((config, session)) } pub(crate) async fn gateway_oauth_read( diff --git a/codex-rs/app-server/src/request_processors/catalog_processor.rs b/codex-rs/app-server/src/request_processors/catalog_processor.rs index 80db727324..ca7d5c498c 100644 --- a/codex-rs/app-server/src/request_processors/catalog_processor.rs +++ b/codex-rs/app-server/src/request_processors/catalog_processor.rs @@ -216,7 +216,7 @@ impl CatalogRequestProcessor { .check_thread_model_provider(&config) .await .map_err(|err| config_load_error(&err))?; - if gateway.resolve_access_token().await.is_err() { + if let Err(error) = gateway.resolve_access_token().await { let current = self .config_manager .load_latest_config(/*fallback_cwd*/ None) @@ -232,9 +232,7 @@ impl CatalogRequestProcessor { "Model provider settings changed. Restart Codex to apply them, then retry fetching the model list", )); } - return Err(invalid_request( - "Gateway sign-in required or unavailable. Complete gateway sign-in and retry fetching the model list", - )); + return Err(invalid_request(error.to_string())); } } self.list_models(params, config, models_manager) diff --git a/codex-rs/app-server/tests/suite/v2/gateway_oauth_tests.rs b/codex-rs/app-server/tests/suite/v2/gateway_oauth_tests.rs index ac278cbf1d..9ad6995171 100644 --- a/codex-rs/app-server/tests/suite/v2/gateway_oauth_tests.rs +++ b/codex-rs/app-server/tests/suite/v2/gateway_oauth_tests.rs @@ -137,6 +137,71 @@ async fn model_list_requests_restart_after_gateway_provider_changes() -> Result< Ok(()) } +#[tokio::test] +async fn working_gateway_exporter_allows_model_list_without_an_oauth_session() -> Result<()> { + let oauth = MockServer::start().await; + let home = TempDir::new()?; + let token_file = home.path().join("exported-token"); + std::fs::write(&token_file, "")?; + let file_arg = serde_json::to_string(&token_file)?; + let command = if cfg!(windows) { + format!("command = 'cmd.exe', args = ['/D', '/C', 'type', {file_arg}]") + } else { + format!("command = 'cat', args = [{file_arg}]") + }; + std::fs::write( + home.path().join("config.toml"), + format!( + r#" +model_provider = "gateway" +[model_providers.gateway] +name = "Test gateway" +base_url = "{url}/v1" +[model_providers.gateway.gateway_oauth] +credential_export = {{ {command} }} +"#, + url = oauth.uri() + ), + )?; + let mut server = gateway_server(home.path()).await?; + let response = read_gateway(&mut server).await?; + assert!(!response.required); + assert_eq!(response.status, None); + // A local list still requires credentials, even when no remote catalog is fetched. + let id = server + .send_raw_request("model/list", Some(serde_json::json!({}))) + .await?; + let error = server + .read_stream_until_error_message(RequestId::Integer(id)) + .await?; + assert!( + error + .error + .message + .contains("Gateway credential export must print one access token") + ); + std::fs::write(&token_file, "exported")?; + for _ in 0..2 { + let _: codex_app_server_protocol::ModelListResponse = server + .request(|request_id| ClientRequest::ModelList { + request_id, + params: codex_app_server_protocol::ModelListParams::default(), + }) + .await?; + // The next preflight can use the token cached by the first. + std::fs::write(&token_file, "")?; + } + assert!( + oauth + .received_requests() + .await + .unwrap() + .iter() + .all(|request| !matches!(request.url.path(), "/authorize" | "/token")) + ); + Ok(()) +} + /// Account changes reload the catalog before validating a retained OAuth gateway. #[tokio::test] async fn model_list_leaves_oauth_gateway_after_account_change() -> Result<()> { diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index da77de3dd6..f605669d6b 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -492,35 +492,6 @@ ], "type": "object" }, - "AwsCredentialExportConfig": { - "additionalProperties": false, - "description": "Command used to export AWS signing credentials for a model provider.", - "properties": { - "args": { - "default": [], - "description": "Arguments passed to the credential export command.", - "items": { - "type": "string" - }, - "type": "array" - }, - "command": { - "description": "Executable to invoke directly, without a shell.", - "type": "string" - }, - "timeout_ms": { - "default": 30000, - "description": "Maximum time to wait for the credential export command to complete.", - "format": "uint64", - "minimum": 1.0, - "type": "integer" - } - }, - "required": [ - "command" - ], - "type": "object" - }, "AwsSigV4Config": { "additionalProperties": false, "description": "AWS SigV4 action selected by an exact-host, method, and path MITM hook.", @@ -1652,6 +1623,35 @@ ], "type": "string" }, + "CredentialExportConfig": { + "additionalProperties": false, + "description": "Command used to export credentials for a model provider.", + "properties": { + "args": { + "default": [], + "description": "Arguments passed to the credential export command.", + "items": { + "type": "string" + }, + "type": "array" + }, + "command": { + "description": "Executable to invoke directly, without a shell.", + "type": "string" + }, + "timeout_ms": { + "default": 30000, + "description": "Maximum time to wait for the credential export command to complete.", + "format": "uint64", + "minimum": 1.0, + "type": "integer" + } + }, + "required": [ + "command" + ], + "type": "object" + }, "CredentialProviderConfig": { "additionalProperties": false, "description": "Declarative description of an environment-backed credential family.", @@ -2016,6 +2016,88 @@ "type": "string" }, "GatewayOAuthConfig": { + "anyOf": [ + { + "$ref": "#/definitions/GatewayOAuthSettings" + }, + { + "additionalProperties": false, + "properties": { + "credential_export": { + "allOf": [ + { + "$ref": "#/definitions/CredentialExportConfig" + } + ], + "description": "Command that prints a gateway access token." + }, + "delivery": { + "allOf": [ + { + "$ref": "#/definitions/GatewayOAuthDelivery" + } + ], + "default": { + "kind": "header", + "name": "Authorization", + "scheme": "Bearer" + } + } + }, + "required": [ + "credential_export" + ], + "type": "object" + } + ] + }, + "GatewayOAuthDelivery": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "enum": [ + "header" + ], + "type": "string" + }, + "name": { + "type": "string" + }, + "scheme": { + "default": "Bearer", + "type": "string" + } + }, + "required": [ + "kind", + "name" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "enum": [ + "cookie" + ], + "type": "string" + }, + "name": { + "type": "string" + } + }, + "required": [ + "kind", + "name" + ], + "type": "object" + } + ] + }, + "GatewayOAuthSettings": { "additionalProperties": false, "properties": { "authorization_url": { @@ -2071,52 +2153,6 @@ ], "type": "object" }, - "GatewayOAuthDelivery": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "enum": [ - "header" - ], - "type": "string" - }, - "name": { - "type": "string" - }, - "scheme": { - "default": "Bearer", - "type": "string" - } - }, - "required": [ - "kind", - "name" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "enum": [ - "cookie" - ], - "type": "string" - }, - "name": { - "type": "string" - } - }, - "required": [ - "kind", - "name" - ], - "type": "object" - } - ] - }, "GhostSnapshotToml": { "additionalProperties": false, "properties": { @@ -3100,7 +3136,7 @@ "credential_export": { "allOf": [ { - "$ref": "#/definitions/AwsCredentialExportConfig" + "$ref": "#/definitions/CredentialExportConfig" } ], "description": "Optional command whose exported credentials replace the AWS SDK credential chain." @@ -3200,7 +3236,7 @@ "$ref": "#/definitions/GatewayOAuthConfig" } ], - "description": "Secondary OAuth credentials required by the provider's gateway." + "description": "Gateway access token from OAuth or a credential export command." }, "http_headers": { "additionalProperties": { diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 56aeec9903..704219bd6f 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -94,6 +94,7 @@ use codex_model_provider::ChatgptBaseUrl; use codex_model_provider::DEFAULT_CHATGPT_BASE_URL; use codex_model_provider::ProviderCapabilities; use codex_model_provider_info::GCP_PROVIDER_ID; +use codex_model_provider_info::GatewayOAuthConfig; use codex_model_provider_info::LEGACY_OLLAMA_CHAT_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OLLAMA_CHAT_PROVIDER_REMOVED_ERROR; @@ -1848,6 +1849,14 @@ impl Config { }; let mut factory = HttpClientFactory::new(outbound_proxy_policy) .with_network_policy(self.application_network_policy.clone()); + if self + .model_provider + .gateway_oauth + .as_ref() + .is_some_and(GatewayOAuthConfig::uses_credential_export) + { + factory = factory.with_sensitive_header_protection(); + } if !self.respect_system_proxy && self.features.enabled(Feature::SystemProxyFallback) { factory = factory.with_system_proxy_fallback(); } diff --git a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs index 1ebdd5486e..6facf10eb1 100644 --- a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs +++ b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs @@ -8,6 +8,7 @@ use codex_login::ExternalAuthFuture; use codex_login::ExternalAuthRefreshContext; use codex_model_provider_info::GatewayOAuthConfig; use codex_model_provider_info::GatewayOAuthDelivery; +use codex_model_provider_info::GatewayOAuthSettings; use codex_runtime_agent_instance::AgentHostAuthorization; use codex_runtime_agent_instance::RuntimeAuthManager; use codex_runtime_agent_instance::RuntimeRegistrationMetadata; @@ -97,7 +98,7 @@ async fn runtime_harness( } fn gateway_config(server: &MockServer, delivery: GatewayOAuthDelivery) -> GatewayOAuthConfig { - GatewayOAuthConfig { + GatewayOAuthConfig::OAuth(GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token", server.uri()), client_id: "client".into(), @@ -106,7 +107,7 @@ fn gateway_config(server: &MockServer, delivery: GatewayOAuthDelivery) -> Gatewa redirect_port: None, credentials_store: Default::default(), delivery, - } + }) } fn assertion_payload( diff --git a/codex-rs/core/tests/suite/gateway_auth.rs b/codex-rs/core/tests/suite/gateway_auth.rs index dcd56e256b..577cb63d57 100644 --- a/codex-rs/core/tests/suite/gateway_auth.rs +++ b/codex-rs/core/tests/suite/gateway_auth.rs @@ -29,6 +29,74 @@ use wiremock::ResponseTemplate; use wiremock::matchers::method; use wiremock::matchers::path; +fn gateway_export_settings(token_path: &std::path::Path) -> Result { + let path_arg = serde_json::to_string(token_path)?; + let command = if cfg!(windows) { + format!("command = 'cmd.exe'\nargs = ['/D', '/C', 'type', {path_arg}]") + } else { + format!("command = 'cat'\nargs = [{path_arg}]") + }; + Ok(format!( + "delivery = {{ kind = 'cookie', name = 'gateway' }}\n[gateway_oauth.credential_export]\n{command}" + )) +} + +async fn command_gateway(server: &MockServer, token_path: &std::path::Path) -> Result { + let provider: ModelProviderInfo = toml::from_str(&format!( + r#" +name = "gateway" +base_url = "{base}/v1" +requires_openai_auth = true +[gateway_oauth] +{gateway_settings} +"#, + base = server.uri(), + gateway_settings = gateway_export_settings(token_path)? + ))?; + test_codex() + .with_auth(CodexAuth::from_api_key("primary")) + .with_config(move |config| config.model_provider = provider) + .build_with_auto_env(server) + .await +} + +#[tokio::test] +async fn command_gateway_authenticates_without_an_oauth_grant_and_caches_the_token() -> Result<()> { + let server = MockServer::start().await; + let source = TempDir::new()?; + let token_path = source.path().join("access-token"); + std::fs::write(&token_path, "")?; + let test = command_gateway(&server, &token_path).await?; + let completed = + responses::sse_response(responses::sse(vec![responses::ev_completed("response")])); + let requests = + responses::mount_response_sequence(&server, vec![completed.clone(), completed]).await; + assert_eq!( + run_gateway_turn(&test).await?.as_deref(), + Some("Gateway credential export must print one access token"), + ); + std::fs::write(&token_path, "opaque,token\n")?; + assert_eq!(run_gateway_turn(&test).await?, None); + std::fs::remove_file(token_path)?; + assert_eq!(run_gateway_turn(&test).await?, None); + let headers: Vec<_> = requests + .requests() + .iter() + .map(|r| (r.header("authorization"), r.header("cookie"))) + .collect(); + assert_eq!( + headers, + vec![ + ( + Some("Bearer primary".to_string()), + Some("gateway=opaque,token".to_string()) + ); + 2 + ] + ); + Ok(()) +} + #[tokio::test] async fn inference_sends_chatgpt_and_gateway_credentials_to_the_configured_headers() -> Result<()> { sample_with_gateway( diff --git a/codex-rs/login/src/credential_export.rs b/codex-rs/login/src/credential_export.rs new file mode 100644 index 0000000000..10959d5610 --- /dev/null +++ b/codex-rs/login/src/credential_export.rs @@ -0,0 +1,74 @@ +//! Runs credential exporters directly with bounded output, timeout, and sanitized diagnostics. + +use std::io; +use std::ops::Deref; +use std::process::Stdio; + +use codex_model_provider_info::CredentialExportConfig; +use tokio::io::AsyncReadExt; +use tokio::process::Command; + +const MAX_CREDENTIAL_OUTPUT_BYTES: u64 = 64 * 1024; + +/// Captures stdout for the caller to parse in memory. Stderr is discarded. +pub async fn run_credential_export( + config: &CredentialExportConfig, + provider: &str, +) -> io::Result> { + config + .validate("credential_export") + .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))?; + let program = &config.command; + let mut command = Command::from(codex_utils_process::background_command(program)); + command + .args(config.args.iter().map(Deref::deref)) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .kill_on_drop(true); + + let mut child = command.spawn().map_err(|error| { + io::Error::new( + error.kind(), + format!( + "failed to run {provider} credential export command `{}`: {error}", + config.command + ), + ) + })?; + let stdout = child.stdout.take().ok_or_else(|| { + io::Error::other(format!( + "{provider} credential export stdout is unavailable" + )) + })?; + tokio::time::timeout(config.timeout(), async { + let mut output = Vec::new(); + stdout + .take(MAX_CREDENTIAL_OUTPUT_BYTES + 1) + .read_to_end(&mut output) + .await?; + if output.len() as u64 > MAX_CREDENTIAL_OUTPUT_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("{provider} credential export output exceeds 64 KiB"), + )); + } + let status = child.wait().await?; + if !status.success() { + return Err(io::Error::other(format!( + "{provider} credential export command exited with {status}" + ))); + } + Ok(output) + }) + .await + .map_err(|_| { + io::Error::new( + io::ErrorKind::TimedOut, + format!( + "{provider} credential export command timed out after {} ms", + config.timeout_ms + ), + ) + })? +} diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index a7cbdce3f6..e52f62e7ca 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -8,9 +8,11 @@ pub use auth::WorkspaceRoutingResolver; pub use auth::WorkspaceRoutingSession; mod callback_params; +mod credential_export; mod device_code_auth; mod oauth; mod oauth_session; +pub use credential_export::run_credential_export; mod outbound_proxy; mod pkce; mod server; @@ -83,7 +85,6 @@ pub use token_data::TokenData; pub use oauth_session::OAuthLoginControl as GatewayLoginControl; pub use oauth_session::OAuthSessionConfig as GatewayAuthConfig; pub use oauth_session::OAuthSessionError as GatewayAuthError; -pub use oauth_session::OAuthSessionManager as GatewayAuthManager; pub use oauth_session::OAuthSessionStatus as GatewayAuthStatus; pub use oauth_session::OAuthSessionStatusChange as GatewayAuthStatusChange; pub use oauth_session::subscribe_oauth_session_status as subscribe_gateway_auth_status; diff --git a/codex-rs/model-provider-info/src/gateway_oauth.rs b/codex-rs/model-provider-info/src/gateway_oauth.rs index c320599847..91bbbd18de 100644 --- a/codex-rs/model-provider-info/src/gateway_oauth.rs +++ b/codex-rs/model-provider-info/src/gateway_oauth.rs @@ -1,4 +1,4 @@ -//! Secondary OAuth credentials delivered alongside the provider's primary authentication. +//! Gateway credentials and their delivery alongside the provider's primary authentication. use codex_client::CHATGPT_ACCOUNT_ID_HEADER; use std::fmt; @@ -9,6 +9,7 @@ use schemars::JsonSchema; use serde::Deserialize; use serde::Serialize; +use crate::CredentialExportConfig; use crate::ModelProviderInfo; use crate::OAuthCredentialsStore; @@ -27,9 +28,25 @@ const RESERVED_GATEWAY_OAUTH_HEADERS: &[&str] = &[ const RESERVED_GATEWAY_OAUTH_HEADER_PREFIXES: &[&str] = &["x-codex-", "x-openai-", "sec-websocket-"]; +#[derive(Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] +#[serde( + untagged, + deny_unknown_fields, + expecting = "choose one gateway credential source: credential_export or OAuth settings (authorization_url, token_url, and client_id)" +)] +pub enum GatewayOAuthConfig { + OAuth(GatewayOAuthSettings), + Export { + /// Command that prints a gateway access token. + credential_export: CredentialExportConfig, + #[serde(default)] + delivery: GatewayOAuthDelivery, + }, +} + #[derive(Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] #[serde(deny_unknown_fields)] -pub struct GatewayOAuthConfig { +pub struct GatewayOAuthSettings { pub authorization_url: String, pub token_url: String, pub client_id: String, @@ -57,7 +74,7 @@ fn is_keyring(store: &OAuthCredentialsStore) -> bool { impl fmt::Debug for GatewayOAuthConfig { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("GatewayOAuthConfig") - .field("delivery", &self.delivery) + .field("delivery", self.delivery()) .finish_non_exhaustive() } } @@ -119,22 +136,49 @@ fn bearer_scheme() -> String { } impl GatewayOAuthConfig { + pub fn delivery(&self) -> &GatewayOAuthDelivery { + match self { + Self::OAuth(settings) => &settings.delivery, + Self::Export { delivery, .. } => delivery, + } + } + + pub fn uses_credential_export(&self) -> bool { + matches!(self, Self::Export { .. }) + } + /// Send the gateway bearer only to the configured endpoint. pub fn requires_gateway_origin(&self) -> bool { - self.delivery.is_authorization_bearer() + self.delivery().is_authorization_bearer() || self.uses_credential_export() } pub(crate) fn validate(&self, provider: &ModelProviderInfo) -> Result<(), String> { if provider.aws.is_some() || provider.is_amazon_bedrock() { return Err("provider gateway_oauth cannot be combined with AWS authentication".into()); } - validate_endpoint_url(&self.authorization_url, "gateway_oauth.authorization_url")?; - validate_endpoint_url(&self.token_url, "gateway_oauth.token_url")?; let base_url = provider .base_url .as_deref() .ok_or("gateway_oauth requires base_url")?; let gateway = validate_endpoint_url(base_url, "base_url")?; + match self { + Self::Export { + credential_export, .. + } => credential_export.validate("gateway_oauth.credential_export")?, + Self::OAuth(settings) => { + validate_endpoint_url( + &settings.authorization_url, + "gateway_oauth.authorization_url", + )?; + validate_endpoint_url(&settings.token_url, "gateway_oauth.token_url")?; + if settings.client_id.trim().is_empty() || settings.redirect_port == Some(0) { + return Err( + "gateway_oauth requires a nonempty client_id and a nonzero redirect_port" + .into(), + ); + } + } + } if self.requires_gateway_origin() && let Some(catalog) = provider.model_catalog_url.as_deref() && validate_endpoint_url(catalog, "model_catalog_url")?.origin() != gateway.origin() @@ -146,12 +190,7 @@ impl GatewayOAuthConfig { .into(), ); } - if self.client_id.trim().is_empty() || self.redirect_port == Some(0) { - return Err( - "gateway_oauth requires a nonempty client_id and a nonzero redirect_port".into(), - ); - } - let gateway_header = if self.delivery.is_authorization_bearer() { + let gateway_header = if self.delivery().is_authorization_bearer() { if provider.is_gcp() { return Err("Google Cloud uses the Authorization header for its own token. Remove gateway_oauth to use the built-in Google Cloud provider".into()); } @@ -169,7 +208,7 @@ impl GatewayOAuthConfig { } http::header::AUTHORIZATION } else { - self.delivery + self.delivery() .secondary_header() .map_err(|reason| format!("gateway_oauth.delivery {reason}"))? }; diff --git a/codex-rs/model-provider-info/src/gateway_oauth_tests.rs b/codex-rs/model-provider-info/src/gateway_oauth_tests.rs index f0165b251a..da519be61d 100644 --- a/codex-rs/model-provider-info/src/gateway_oauth_tests.rs +++ b/codex-rs/model-provider-info/src/gateway_oauth_tests.rs @@ -20,8 +20,8 @@ fn parses_header_and_cookie_delivery() { let header = provider("delivery = { kind = 'header', name = 'x-gateway-auth' }"); assert_eq!(header.validate(), Ok(())); assert_eq!( - header.gateway_oauth.as_ref().unwrap().delivery, - GatewayOAuthDelivery::Header { + header.gateway_oauth.as_ref().unwrap().delivery(), + &GatewayOAuthDelivery::Header { name: "x-gateway-auth".into(), scheme: "Bearer".into(), } @@ -47,7 +47,10 @@ fn rejects_unsafe_endpoints_and_delivery_collisions() { invalid.base_url = Some(url.into()); assert!(invalid.validate().is_err()); invalid = valid.clone(); - invalid.gateway_oauth.as_mut().unwrap().token_url = url.into(); + let Some(GatewayOAuthConfig::OAuth(settings)) = &mut invalid.gateway_oauth else { + unreachable!() + }; + settings.token_url = url.into(); let error = invalid.validate().unwrap_err(); assert!(!error.contains(url)); } @@ -77,6 +80,26 @@ fn rejects_unsafe_endpoints_and_delivery_collisions() { assert!(cookie.validate().is_err()); } +#[test] +fn credential_export_rejects_oauth_settings() { + for oauth in [ + "client_id = 'codex'", + "scopes = []", + "credentials_store = 'keyring'", + "authorization_url = 'https://example.test/authorize'\ntoken_url = 'https://example.test/token'\nclient_id = 'codex'", + ] { + let config = + format!("[gateway_oauth]\ncredential_export = {{ command = 'export-token' }}\n{oauth}"); + let error = toml::from_str::(&config).unwrap_err(); + assert!( + error + .to_string() + .contains("choose one gateway credential source"), + "{error}" + ); + } +} + #[test] fn gateway_bearer_requires_the_model_catalog_to_share_its_origin() { let other = "https://gateway.example.test:8443/catalog"; diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 7d0e1a8c81..efd48f17a9 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -41,6 +41,7 @@ pub use capabilities::ModelProviderCapabilities; pub use capabilities::RemoteCompactionSupport; pub use gateway_oauth::GatewayOAuthConfig; pub use gateway_oauth::GatewayOAuthDelivery; +pub use gateway_oauth::GatewayOAuthSettings; pub use oauth_session::OAuthCredentialsStore; pub use oauth_session::OAuthSessionConfig; @@ -72,7 +73,7 @@ pub fn read_managed_residency_requirement() -> Option { const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; -const DEFAULT_AWS_CREDENTIAL_EXPORT_TIMEOUT_MS: u64 = 30_000; +const DEFAULT_CREDENTIAL_EXPORT_TIMEOUT_MS: u64 = 30_000; const DEFAULT_AWS_AUTH_REFRESH_TIMEOUT_MS: u64 = 300_000; pub const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS: u64 = 15_000; /// Hard cap for user-configured `stream_max_retries`. @@ -176,7 +177,7 @@ pub struct ModelProviderInfo { pub experimental_bearer_token: Option, /// Command-backed bearer-token configuration for this provider. pub auth: Option, - /// Secondary OAuth credentials required by the provider's gateway. + /// Gateway access token from OAuth or a credential export command. pub gateway_oauth: Option, /// AWS SigV4 auth configuration for this provider. pub aws: Option, @@ -249,30 +250,48 @@ pub struct ModelProviderAwsAuthInfo { pub auth_refresh: Option, } -/// Command used to export AWS signing credentials for a model provider. +/// Command used to export credentials for a model provider. #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] #[schemars(deny_unknown_fields)] -pub struct AwsCredentialExportConfig { +pub struct CredentialExportConfig { /// Executable to invoke directly, without a shell. pub command: String, /// Arguments passed to the credential export command. #[serde(default)] pub args: Vec, /// Maximum time to wait for the credential export command to complete. - #[serde(default = "default_aws_credential_export_timeout_ms")] + #[serde(default = "default_credential_export_timeout_ms")] pub timeout_ms: NonZeroU64, } -impl AwsCredentialExportConfig { +pub type AwsCredentialExportConfig = CredentialExportConfig; + +impl CredentialExportConfig { + pub fn validate(&self, field: &str) -> Result<(), String> { + if self.command.trim().is_empty() { + return Err(format!("provider {field}.command must not be empty")); + } + let path = Path::new(&self.command); + let mut components = path.components(); + let bare = matches!((components.next(), components.next()), + (Some(Component::Normal(name)), None) if name == path.as_os_str()); + if !path.is_absolute() && !bare { + return Err(format!( + "provider {field}.command must be an absolute path or a bare executable name" + )); + } + Ok(()) + } + pub fn timeout(&self) -> Duration { Duration::from_millis(self.timeout_ms.get()) } } -fn default_aws_credential_export_timeout_ms() -> NonZeroU64 { - match NonZeroU64::new(DEFAULT_AWS_CREDENTIAL_EXPORT_TIMEOUT_MS) { +fn default_credential_export_timeout_ms() -> NonZeroU64 { + match NonZeroU64::new(DEFAULT_CREDENTIAL_EXPORT_TIMEOUT_MS) { Some(timeout_ms) => timeout_ms, - None => panic!("AWS credential export timeout must be non-zero"), + None => panic!("credential export timeout must be non-zero"), } } @@ -377,23 +396,7 @@ other non-default provider fields are not supported" .to_string(), ); } - if credential_export.command.trim().is_empty() { - return Err( - "provider aws.credential_export.command must not be empty".to_string() - ); - } - let command = Path::new(&credential_export.command); - let mut components = command.components(); - let is_bare_command = matches!( - (components.next(), components.next()), - (Some(Component::Normal(name)), None) if name == command.as_os_str() - ); - if !command.is_absolute() && !is_bare_command { - return Err( - "provider aws.credential_export.command must be an absolute path or a bare executable name" - .to_string(), - ); - } + credential_export.validate("aws.credential_export")?; } if let Some(auth_refresh) = aws.auth_refresh.as_ref() { diff --git a/codex-rs/model-provider/src/amazon_bedrock/credential_export.rs b/codex-rs/model-provider/src/amazon_bedrock/credential_export.rs index d6db7d93ed..2b1d567278 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/credential_export.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/credential_export.rs @@ -3,10 +3,6 @@ //! Exported secrets remain in memory; command output is bounded and never logged. use std::io; -use std::ops::Deref; -use std::path::Component; -use std::path::Path; -use std::process::Stdio; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; @@ -18,14 +14,11 @@ use codex_aws_auth::AwsAccessKeys; use codex_aws_auth::AwsCredentialsProvider; use codex_model_provider_info::AwsCredentialExportConfig; use serde::Deserialize; -use tokio::io::AsyncReadExt; -use tokio::process::Command; use tokio::sync::Mutex; use tokio::sync::MutexGuard; const DEFAULT_CREDENTIAL_LIFETIME: Duration = Duration::from_secs(60 * 60); const CREDENTIAL_REFRESH_WINDOW: Duration = Duration::from_secs(5 * 60); -const MAX_CREDENTIAL_OUTPUT_BYTES: u64 = 64 * 1024; #[derive(Debug)] pub(crate) struct AwsCredentialExport { @@ -158,70 +151,7 @@ impl AwsCredentialExport { } async fn export_credentials(&self) -> io::Result { - let program = Path::new(&self.config.command); - let mut components = program.components(); - let is_bare_command = matches!( - (components.next(), components.next()), - (Some(Component::Normal(name)), None) if name == program.as_os_str() - ); - if self.config.command.trim().is_empty() || (!program.is_absolute() && !is_bare_command) { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "AWS credential export command must be an absolute path or a bare executable name", - )); - } - - let mut command = Command::from(codex_utils_process::background_command(program)); - command - .args(self.config.args.iter().map(Deref::deref)) - .stdin(Stdio::null()) - .stdout(Stdio::piped()) - .stderr(Stdio::null()) - .kill_on_drop(true); - - let mut child = command.spawn().map_err(|error| { - io::Error::new( - error.kind(), - format!( - "failed to run AWS credential export command `{}`: {error}", - self.config.command - ), - ) - })?; - let stdout = child - .stdout - .take() - .ok_or_else(|| io::Error::other("AWS credential export stdout is unavailable"))?; - let output = tokio::time::timeout(self.config.timeout(), async { - let mut output = Vec::new(); - stdout - .take(MAX_CREDENTIAL_OUTPUT_BYTES + 1) - .read_to_end(&mut output) - .await?; - if output.len() as u64 > MAX_CREDENTIAL_OUTPUT_BYTES { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "AWS credential export output exceeds 64 KiB", - )); - } - let status = child.wait().await?; - if !status.success() { - return Err(io::Error::other(format!( - "AWS credential export command exited with {status}" - ))); - } - Ok(output) - }) - .await - .map_err(|_| { - io::Error::new( - io::ErrorKind::TimedOut, - format!( - "AWS credential export command timed out after {} ms", - self.config.timeout_ms - ), - ) - })??; + let output = codex_login::run_credential_export(&self.config, "AWS").await?; CachedCredentials::from_output(&output, SystemTime::now()) } } diff --git a/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs b/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs index 1f69d86283..c9697c761c 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs @@ -17,6 +17,7 @@ use super::super::AmazonBedrockModelProvider; use super::super::AwsAuthRecovery; use super::*; +const MAX_CREDENTIAL_OUTPUT_BYTES: usize = 64 * 1024; const SECRET: &str = "credential-export-secret"; const VALID_OUTPUT: &[u8] = br#"{"AccessKeyId":"exported","SecretAccessKey":"credential-export-secret","SessionToken":"credential-export-secret"}"#; @@ -183,10 +184,7 @@ fn credential_output_formats_expiration_and_sanitized_errors() { #[tokio::test] async fn credential_export_process_errors_are_bounded_and_sanitized() -> io::Result<()> { - let oversized = format!( - "{SECRET}{}", - "x".repeat(MAX_CREDENTIAL_OUTPUT_BYTES as usize) - ); + let oversized = format!("{SECRET}{}", "x".repeat(MAX_CREDENTIAL_OUTPUT_BYTES)); for (mode, output, timeout_ms, expected_kind) in [ ("missing", VALID_OUTPUT, 5_000, io::ErrorKind::NotFound), ("relative", VALID_OUTPUT, 5_000, io::ErrorKind::InvalidInput), @@ -239,7 +237,7 @@ async fn credential_export_process_errors_are_bounded_and_sanitized() -> io::Res #[tokio::test] async fn credential_export_reuses_cache_and_coalesces_refreshes() -> io::Result<()> { let mut output = VALID_OUTPUT.to_vec(); - output.resize(MAX_CREDENTIAL_OUTPUT_BYTES as usize, b' '); + output.resize(MAX_CREDENTIAL_OUTPUT_BYTES, b' '); let (directory, exporter) = exporter_fixture(&output, "success")?; let original = AwsAccessKeys { access_key_id: "exported".to_string(), diff --git a/codex-rs/model-provider/src/combined_auth.rs b/codex-rs/model-provider/src/combined_auth.rs index 6165dce5bf..536b057051 100644 --- a/codex-rs/model-provider/src/combined_auth.rs +++ b/codex-rs/model-provider/src/combined_auth.rs @@ -2,6 +2,7 @@ use std::sync::Arc; +use crate::GatewayAuthManager; use crate::auth::ResolvedProviderAuth; use codex_api::AuthError; use codex_api::AuthHeaderTelemetry; @@ -9,13 +10,9 @@ use codex_api::AuthHeadersFuture; use codex_api::AuthProvider; use codex_api::AuthProviderFuture; use codex_api::SharedAuthProvider; -use codex_login::GatewayAuthError; -use codex_login::GatewayAuthManager; -use codex_model_provider_info::GatewayOAuthConfig; use codex_model_provider_info::GatewayOAuthDelivery; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErr; -use codex_protocol::error::CodexErrorDetails; use codex_protocol::error::Result; use http::HeaderMap; use http::HeaderName; @@ -36,28 +33,8 @@ pub(crate) async fn compose_auth( let manager = manager .as_ref() .map_err(|error| CodexErr::InvalidRequest(error.clone()))?; - let token = manager.resolve_access_token().await.map_err(|error| { - if let Some(login_error) = error - .get_ref() - .and_then(|source| source.downcast_ref::()) - { - return match login_error { - GatewayAuthError::LoginRequired => { - CodexErrorDetails::GatewayLoginRequired(login_error.to_string()).into() - } - GatewayAuthError::LoginInProgress => { - CodexErr::InvalidRequest(login_error.to_string()) - } - }; - } - // Issuer errors may echo arbitrary credentials from configured URLs. Keep the - // diagnostic safe and bounded, preserving retries for transient failures. - CodexErr::Io(std::io::Error::new( - error.kind(), - "Gateway OAuth authentication failed; check the gateway configuration and credential store.", - )) - })?; - let (name, value) = gateway_header(config, &token)?; + let token = manager.resolve_access_token().await?; + let (name, value) = gateway_header(config.delivery(), &token)?; resolved.auth = Arc::new(CombinedAuth { primary: resolved.auth, name, @@ -66,29 +43,22 @@ pub(crate) async fn compose_auth( Ok(resolved) } -fn gateway_header(config: &GatewayOAuthConfig, token: &str) -> Result<(HeaderName, HeaderValue)> { +fn gateway_header( + delivery: &GatewayOAuthDelivery, + token: &str, +) -> Result<(HeaderName, HeaderValue)> { if token.is_empty() { return Err(CodexErr::InvalidRequest( - "gateway OAuth returned an empty token".into(), + "gateway access token is empty".into(), )); } - let (name, value) = match &config.delivery { + let (name, value) = match delivery { GatewayOAuthDelivery::Header { name, scheme } => ( HeaderName::from_bytes(name.as_bytes()) .map_err(|_| CodexErr::InvalidRequest("invalid gateway header".into()))?, format!("{scheme} {token}"), ), - GatewayOAuthDelivery::Cookie { name } => { - // RFC 6265 cookie-octet excludes separators, quotes, whitespace and non-ASCII. - if !token.bytes().all( - |byte| matches!(byte, 0x21 | 0x23..=0x2b | 0x2d..=0x3a | 0x3c..=0x5b | 0x5d..=0x7e), - ) { - return Err(CodexErr::InvalidRequest( - "gateway OAuth token is not a valid cookie value".into(), - )); - } - (http::header::COOKIE, format!("{name}={token}")) - } + GatewayOAuthDelivery::Cookie { name } => (http::header::COOKIE, format!("{name}={token}")), }; let mut value = HeaderValue::from_str(&value) .map_err(|_| CodexErr::InvalidRequest("invalid gateway OAuth token header".into()))?; diff --git a/codex-rs/model-provider/src/gateway_auth.rs b/codex-rs/model-provider/src/gateway_auth.rs new file mode 100644 index 0000000000..79c2661180 --- /dev/null +++ b/codex-rs/model-provider/src/gateway_auth.rs @@ -0,0 +1,87 @@ +//! Resolves credentials for a gateway, using its configured exporter or OAuth session. + +use std::sync::Arc; + +use crate::gateway_export::GatewayExport; +use codex_login::AuthManager; +use codex_login::GatewayAuthError; +use codex_login::OAuthSessionManager; +use codex_model_provider_info::GatewayOAuthConfig; +use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::error::CodexErr; +use codex_protocol::error::CodexErrorDetails; +use codex_protocol::error::Result; + +/// Resolves gateway request credentials and exposes the optional interactive OAuth session. +#[derive(Debug)] +pub struct GatewayAuthManager(GatewayCredentialSource); + +#[derive(Debug)] +enum GatewayCredentialSource { + Export(Arc), + OAuth(Arc), +} + +impl GatewayAuthManager { + pub(crate) fn new( + provider: &ModelProviderInfo, + config: &GatewayOAuthConfig, + owner: &Arc, + ) -> std::result::Result { + let state = crate::shared_state::process_shared_state(); + let source = match config { + GatewayOAuthConfig::Export { + credential_export, .. + } => GatewayCredentialSource::Export(state.gateway_export( + provider, + credential_export, + owner, + )), + GatewayOAuthConfig::OAuth(settings) => GatewayCredentialSource::OAuth( + state + .gateway_oauth_session(settings, &owner.runtime_config()) + .map_err(|_| "failed to create provider OAuth HTTP client".to_string())?, + ), + }; + Ok(Self(source)) + } + + /// The session used for explicit sign-in and OAuth status, when configured. + pub fn oauth_session(&self) -> Option> { + match &self.0 { + GatewayCredentialSource::OAuth(session) => Some(Arc::clone(session)), + GatewayCredentialSource::Export(_) => None, + } + } + + pub async fn resolve_access_token(&self) -> Result> { + match &self.0 { + GatewayCredentialSource::Export(export) => export.token().await, + GatewayCredentialSource::OAuth(session) => session + .resolve_access_token() + .await + .map(Arc::from) + .map_err(gateway_auth_error), + } + } +} + +fn gateway_auth_error(error: std::io::Error) -> CodexErr { + if let Some(login_error) = error + .get_ref() + .and_then(|source| source.downcast_ref::()) + { + return match login_error { + GatewayAuthError::LoginRequired => { + CodexErrorDetails::GatewayLoginRequired(login_error.to_string()).into() + } + GatewayAuthError::LoginInProgress => CodexErr::InvalidRequest(login_error.to_string()), + }; + } + // Issuer errors may echo arbitrary credentials from configured URLs. Keep the + // diagnostic safe and bounded, preserving retries for transient failures. + CodexErr::Io(std::io::Error::new( + error.kind(), + "Gateway OAuth authentication failed; check the gateway configuration and credential store.", + )) +} diff --git a/codex-rs/model-provider/src/gateway_export.rs b/codex-rs/model-provider/src/gateway_export.rs new file mode 100644 index 0000000000..48132abe64 --- /dev/null +++ b/codex-rs/model-provider/src/gateway_export.rs @@ -0,0 +1,64 @@ +//! Resolves and caches the configured gateway command's credential for this provider. + +use std::fmt; +use std::sync::Arc; + +use codex_http_client::NetworkPolicy; +use codex_model_provider_info::CredentialExportConfig; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result; +use tokio::sync::Mutex; + +pub(crate) struct GatewayExport { + command: CredentialExportConfig, + policy: NetworkPolicy, + cached: Mutex>>, +} + +impl fmt::Debug for GatewayExport { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("GatewayExport").finish_non_exhaustive() + } +} + +impl GatewayExport { + pub(crate) fn new(command: CredentialExportConfig, policy: NetworkPolicy) -> Self { + Self { + command, + policy, + cached: Mutex::default(), + } + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "concurrent cache misses must share a single credential export" + )] + pub(crate) async fn token(&self) -> Result> { + let mut cached = self.cached.lock().await; + if let Some(token) = cached.as_ref() { + return Ok(Arc::clone(token)); + } + let policy = self.policy.clone().for_current_account(); + let permit = policy + .acquire_for_unsupported_sdk() + .map_err(|error| CodexErr::Fatal(error.to_string()))?; + let output = permit + .run(codex_login::run_credential_export(&self.command, "Gateway")) + .await + .map_err(|error| CodexErr::Fatal(error.to_string()))? + .map_err(|error| CodexErr::InvalidRequest(error.to_string()))?; + let token = std::str::from_utf8(&output) + .ok() + .map(str::trim) + .filter(|value| !value.is_empty() && value.bytes().all(|b| b.is_ascii_graphic())) + .ok_or_else(|| { + CodexErr::InvalidRequest( + "Gateway credential export must print one access token".into(), + ) + })?; + let token = Arc::::from(token); + *cached = Some(Arc::clone(&token)); + Ok(token) + } +} diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 2282ca5b8a..9a84789afd 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -6,6 +6,8 @@ mod catalog_factory; mod chatgpt_base_url; mod client_auth; mod combined_auth; +mod gateway_auth; +mod gateway_export; mod gcp; mod models_endpoint; mod models_identity; @@ -45,6 +47,7 @@ pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; pub use codex_model_provider_info::GCP_PROVIDER_ID; pub use codex_model_provider_info::RemoteCompactionSupport; pub use codex_protocol::account::ProviderAccount; +pub use gateway_auth::GatewayAuthManager; pub use provider::ModelProvider; pub use provider::ModelProviderFuture; pub use provider::ProviderAccountError; diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 1e9ca0843e..072444db65 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -5,6 +5,7 @@ use std::pin::Pin; use std::sync::Arc; use std::time::Duration; +use crate::GatewayAuthManager; use codex_api::AgentIdentityTelemetry; use codex_api::ModelsClient; use codex_api::RequestTelemetry; @@ -20,11 +21,11 @@ use codex_http_client::HttpClientFactory; use codex_login::AuthEnvTelemetry; use codex_login::AuthManager; use codex_login::CodexAuth; -use codex_login::GatewayAuthManager; use codex_login::collect_auth_env_telemetry; use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::create_client_for_route_async; use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; +use codex_model_provider_info::GatewayOAuthConfig; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::ModelCatalogOriginator; use codex_models_manager::manager::ModelsEndpointClient; @@ -65,7 +66,12 @@ impl OpenAiModelsEndpoint { auth_manager: Option>, gateway_auth_manager: Option, String>>, ) -> Self { - let redirect_policy = if provider_info.model_catalog_url.is_some() { + let redirect_policy = if provider_info.model_catalog_url.is_some() + || provider_info + .gateway_oauth + .as_ref() + .is_some_and(GatewayOAuthConfig::uses_credential_export) + { ClientRedirectPolicy::Reject } else { ClientRedirectPolicy::Default diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index a0d1e53773..c562d6e40d 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -4,6 +4,7 @@ use std::path::PathBuf; use std::pin::Pin; use std::sync::Arc; +use crate::GatewayAuthManager; use codex_api::ApiError; use codex_api::Provider; use codex_api::SharedAuthProvider; @@ -11,7 +12,6 @@ use codex_api::TransportError; use codex_features::RuntimeAuthMode; use codex_login::AuthManager; use codex_login::CodexAuth; -use codex_login::GatewayAuthManager; use codex_login::WorkspaceRoutingRequest; use codex_login::default_client::ClientRedirectPolicy; use codex_model_provider_info::ModelProviderInfo; @@ -484,9 +484,7 @@ impl ConfiguredModelProvider { let manager = auth_manager .as_ref() .ok_or_else(|| "gateway_oauth requires auth runtime configuration".to_string())?; - crate::shared_state::process_shared_state() - .gateway_auth(config, &manager.runtime_config()) - .map_err(|_| "failed to create provider OAuth HTTP client".to_string()) + GatewayAuthManager::new(&info, config, manager).map(Arc::new) }); let auth_manager = auth_manager_for_provider(auth_manager, &info); let runtime_for_destination = |destination: Option<&str>| { diff --git a/codex-rs/model-provider/src/request_setup.rs b/codex-rs/model-provider/src/request_setup.rs index 26a0bbd009..407b89f97d 100644 --- a/codex-rs/model-provider/src/request_setup.rs +++ b/codex-rs/model-provider/src/request_setup.rs @@ -17,6 +17,7 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::create_client_for_route; +use codex_model_provider_info::GatewayOAuthConfig; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::Result; use codex_runtime_agent_instance::RuntimeAuth; @@ -158,7 +159,16 @@ pub async fn resolve_request( .map(|state| state.generation) .or(state.map(|state| state.generation)), api_provider, - redirect_policy, + redirect_policy: if provider + .info() + .gateway_oauth + .as_ref() + .is_some_and(GatewayOAuthConfig::uses_credential_export) + { + ClientRedirectPolicy::Reject + } else { + redirect_policy + }, api_auth: resolved_auth.auth, agent_identity_telemetry: resolved_auth.agent_identity_telemetry, runtime_id, diff --git a/codex-rs/model-provider/src/shared_state.rs b/codex-rs/model-provider/src/shared_state.rs index 83849fa08b..587df43729 100644 --- a/codex-rs/model-provider/src/shared_state.rs +++ b/codex-rs/model-provider/src/shared_state.rs @@ -3,12 +3,16 @@ use std::sync::Mutex; use std::sync::OnceLock; use std::sync::Weak; +use crate::gateway_export::GatewayExport; use codex_keyring_store::DefaultKeyringStore; +use codex_login::AuthManager; use codex_login::AuthRuntimeConfig; use codex_login::GatewayAuthConfig; -use codex_login::GatewayAuthManager; -use codex_model_provider_info::GatewayOAuthConfig; +use codex_login::OAuthSessionManager; +use codex_model_provider_info::CredentialExportConfig; +use codex_model_provider_info::GatewayOAuthSettings; use codex_model_provider_info::ModelProviderAwsAuthInfo; +use codex_model_provider_info::ModelProviderInfo; use crate::amazon_bedrock::AwsAuthRecovery; use crate::amazon_bedrock::AwsCredentialExport; @@ -18,26 +22,34 @@ use crate::amazon_bedrock::AwsCredentialExport; pub(crate) struct ModelProviderSharedState { aws_credential_exports: Mutex)>>, aws_auth_recoveries: Mutex)>>, - gateway_managers: Mutex< + gateway_exports: Mutex>, + oauth_sessions: Mutex< Vec<( GatewayAuthConfig, AuthRuntimeConfig, - Weak, + Weak, )>, >, } +#[derive(Debug)] +struct GatewayExportEntry { + owner: Weak, + provider: ModelProviderInfo, + export: Weak, +} + pub(crate) fn process_shared_state() -> &'static ModelProviderSharedState { static STATE: OnceLock = OnceLock::new(); STATE.get_or_init(ModelProviderSharedState::default) } impl ModelProviderSharedState { - pub(crate) fn gateway_auth( + pub(crate) fn gateway_oauth_session( &self, - config: &GatewayOAuthConfig, + config: &GatewayOAuthSettings, runtime: &AuthRuntimeConfig, - ) -> std::io::Result> { + ) -> std::io::Result> { let oauth = GatewayAuthConfig { authorization_url: config.authorization_url.clone(), device_authorization_url: None, @@ -49,7 +61,7 @@ impl ModelProviderSharedState { credentials_store: config.credentials_store, }; let mut managers = self - .gateway_managers + .oauth_sessions .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); managers.retain(|(_, _, manager)| manager.strong_count() != 0); @@ -62,7 +74,7 @@ impl ModelProviderSharedState { { return Ok(manager); } - let manager = Arc::new(GatewayAuthManager::new( + let manager = Arc::new(OAuthSessionManager::new( oauth.clone(), runtime.codex_home.clone(), runtime.auth_route_config.http_client_factory(), @@ -72,6 +84,38 @@ impl ModelProviderSharedState { Ok(manager) } + // The catalog preflight and endpoint construct separate providers for the same owner. + pub(crate) fn gateway_export( + &self, + provider: &ModelProviderInfo, + command: &CredentialExportConfig, + owner: &Arc, + ) -> Arc { + let owner_key = Arc::downgrade(owner); + let mut entries = self + .gateway_exports + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + entries.retain(|entry| entry.export.strong_count() != 0 && entry.owner.strong_count() != 0); + if let Some(export) = entries + .iter() + .find(|entry| entry.owner.ptr_eq(&owner_key) && &entry.provider == provider) + .and_then(|entry| entry.export.upgrade()) + { + return export; + } + let export = Arc::new(GatewayExport::new( + command.clone(), + owner.application_network_policy(), + )); + entries.push(GatewayExportEntry { + owner: owner_key, + provider: provider.clone(), + export: Arc::downgrade(&export), + }); + export + } + pub(crate) fn aws_credential_export( &self, aws: &ModelProviderAwsAuthInfo, diff --git a/codex-rs/model-provider/src/shared_state_test_support.rs b/codex-rs/model-provider/src/shared_state_test_support.rs index 5a6d1fdafa..23a926c2cf 100644 --- a/codex-rs/model-provider/src/shared_state_test_support.rs +++ b/codex-rs/model-provider/src/shared_state_test_support.rs @@ -9,7 +9,7 @@ use codex_keyring_store::CredentialStoreError; use codex_keyring_store::KeyringStore; use codex_login::AuthManager; use codex_login::GatewayAuthConfig; -use codex_login::GatewayAuthManager; +use codex_login::OAuthSessionManager; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OAuthCredentialsStore; use codex_secrets::LocalSecretsNamespace; @@ -64,11 +64,12 @@ pub fn seed_gateway_auth( info: &ModelProviderInfo, primary: &AuthManager, token: serde_json::Value, -) -> Arc { - let config = info - .gateway_oauth - .clone() - .expect("gateway test configuration"); +) -> Arc { + let Some(codex_model_provider_info::GatewayOAuthConfig::OAuth(config)) = + info.gateway_oauth.clone() + else { + panic!("expected gateway OAuth test configuration") + }; let runtime = primary.runtime_config(); let oauth = GatewayAuthConfig { authorization_url: config.authorization_url, @@ -85,7 +86,7 @@ pub fn seed_gateway_auth( "gateway-test-passphrase".to_string(), )))); let manager = Arc::new( - GatewayAuthManager::new( + OAuthSessionManager::new( oauth.clone(), runtime.codex_home.clone(), runtime.auth_route_config.http_client_factory(), @@ -130,7 +131,7 @@ pub fn seed_gateway_auth( } } process_shared_state() - .gateway_managers + .oauth_sessions .lock() .expect("gateway test manager registry lock") .push((oauth, runtime, Arc::downgrade(&manager))); diff --git a/codex-rs/model-provider/src/shared_state_tests.rs b/codex-rs/model-provider/src/shared_state_tests.rs index 655466ce8f..4180d4bde3 100644 --- a/codex-rs/model-provider/src/shared_state_tests.rs +++ b/codex-rs/model-provider/src/shared_state_tests.rs @@ -18,7 +18,9 @@ use codex_login::CodexAuth; use codex_login::auth::AgentIdentityAuthPolicy; use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::create_client_for_route; +use codex_model_provider_info::GatewayOAuthConfig; use codex_model_provider_info::GatewayOAuthDelivery; +use codex_model_provider_info::GatewayOAuthSettings; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OAuthCredentialsStore; use codex_models_manager::manager::RefreshStrategy; @@ -68,7 +70,7 @@ async fn gateway_credentials_accompany_primary_auth_in_models_and_responses() { name: "gateway_session".into(), }, ] { - let config = GatewayOAuthConfig { + let config = GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token", server.uri()), client_id: "test".into(), @@ -80,7 +82,7 @@ async fn gateway_credentials_accompany_primary_auth_in_models_and_responses() { delivery: delivery.clone(), }; let info = ModelProviderInfo { - gateway_oauth: Some(config.clone()), + gateway_oauth: Some(GatewayOAuthConfig::OAuth(config)), ..ModelProviderInfo::create_openai_provider(Some(server.uri())) }; let mut primary_headers = HeaderMap::new(); @@ -100,7 +102,12 @@ async fn gateway_credentials_accompany_primary_auth_in_models_and_responses() { ); let provider = create_model_provider(info.clone(), Some(primary)); assert!(Arc::ptr_eq( - &provider.gateway_auth_manager().unwrap().unwrap(), + &provider + .gateway_auth_manager() + .unwrap() + .unwrap() + .oauth_session() + .unwrap(), &manager, )); let models = provider.models_manager_without_cache(/*config_model_catalog*/ None); @@ -210,7 +217,7 @@ async fn gateway_preserves_runtime_auth_and_rejects_header_collisions() { home.path().to_path_buf(), ); let info = ModelProviderInfo { - gateway_oauth: Some(GatewayOAuthConfig { + gateway_oauth: Some(GatewayOAuthConfig::OAuth(GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token", server.uri()), client_id: "client".into(), @@ -222,7 +229,7 @@ async fn gateway_preserves_runtime_auth_and_rejects_header_collisions() { name: "x-gateway-auth".into(), scheme: "Bearer".into(), }, - }), + })), ..ModelProviderInfo::create_openai_provider(Some(server.uri())) }; let _gateway = seed_gateway_auth( @@ -312,7 +319,7 @@ async fn gateway_refresh_preserves_primary_auth_and_hides_issuer_errors() { home.path().to_path_buf(), ); let info = ModelProviderInfo { - gateway_oauth: Some(GatewayOAuthConfig { + gateway_oauth: Some(GatewayOAuthConfig::OAuth(GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token?credential=query-secret", server.uri()), client_id: "client".into(), @@ -324,7 +331,7 @@ async fn gateway_refresh_preserves_primary_auth_and_hides_issuer_errors() { name: "x-gateway-auth".into(), scheme: "Bearer".into(), }, - }), + })), ..ModelProviderInfo::create_openai_provider(Some(server.uri())) }; let body = if succeeds { @@ -407,7 +414,7 @@ async fn models_cache_is_reused_only_for_matching_gateway_configuration() { codex_login::default_client::RESIDENCY_HEADER_NAME.to_string(), "us".into(), )])), - gateway_oauth: Some(GatewayOAuthConfig { + gateway_oauth: Some(GatewayOAuthConfig::OAuth(GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token", server.uri()), client_id: "client".into(), @@ -420,7 +427,7 @@ async fn models_cache_is_reused_only_for_matching_gateway_configuration() { name: "x-gateway-auth".into(), scheme: "Bearer".into(), }, - }), + })), ..ModelProviderInfo::create_openai_provider(Some(server.uri())) }; let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); @@ -429,8 +436,10 @@ async fn models_cache_is_reused_only_for_matching_gateway_configuration() { .models .remove(0); for resource in ["first", "second"] { - info.gateway_oauth.as_mut().unwrap().resource = - Some(format!("{}/{resource}", server.uri())); + let Some(GatewayOAuthConfig::OAuth(settings)) = &mut info.gateway_oauth else { + unreachable!() + }; + settings.resource = Some(format!("{}/{resource}", server.uri())); model.slug = format!("catalog-{resource}"); let _gateway = seed_gateway_auth( &info, @@ -489,7 +498,7 @@ async fn models_observe_gateway_rotation(instance: ProviderInstance) { ); let info = ModelProviderInfo { model_catalog_url: Some(format!("{}/models", server.uri()).into()), - gateway_oauth: Some(GatewayOAuthConfig { + gateway_oauth: Some(GatewayOAuthConfig::OAuth(GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token", server.uri()), client_id: "client".into(), @@ -501,7 +510,7 @@ async fn models_observe_gateway_rotation(instance: ProviderInstance) { name: "x-gateway-auth".into(), scheme: "Bearer".into(), }, - }), + })), ..ModelProviderInfo::create_openai_provider(Some(server.uri())) }; let gateway = seed_gateway_auth( @@ -604,7 +613,7 @@ async fn gateway_setup_errors_are_reported_when_authentication_is_requested() { let server = MockServer::start().await; let mut info = ModelProviderInfo { - gateway_oauth: Some(GatewayOAuthConfig { + gateway_oauth: Some(GatewayOAuthConfig::OAuth(GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token", server.uri()), client_id: "client".into(), @@ -616,7 +625,7 @@ async fn gateway_setup_errors_are_reported_when_authentication_is_requested() { name: "x-gateway-auth".into(), scheme: "Bearer".into(), }, - }), + })), ..ModelProviderInfo::create_openai_provider(Some(server.uri())) }; let provider = create_model_provider(info.clone(), /*auth_manager*/ None); @@ -648,7 +657,10 @@ async fn gateway_setup_errors_are_reported_when_authentication_is_requested() { "failed to create provider OAuth HTTP client" ); - info.gateway_oauth.as_mut().unwrap().client_id.clear(); + let Some(GatewayOAuthConfig::OAuth(settings)) = &mut info.gateway_oauth else { + unreachable!() + }; + settings.client_id.clear(); let primary = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("primary")); let error = create_model_provider(info, Some(primary)) .api_auth() @@ -672,7 +684,7 @@ async fn gateway_login_required_invalidates_the_last_successful_catalog() { ); let info = ModelProviderInfo { model_catalog_url: Some(format!("{}/models", server.uri()).into()), - gateway_oauth: Some(GatewayOAuthConfig { + gateway_oauth: Some(GatewayOAuthConfig::OAuth(GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token", server.uri()), client_id: "client".into(), @@ -684,7 +696,7 @@ async fn gateway_login_required_invalidates_the_last_successful_catalog() { name: "x-gateway-auth".into(), scheme: "Bearer".into(), }, - }), + })), ..ModelProviderInfo::create_openai_provider(Some(server.uri())) }; let login_control = codex_login::GatewayLoginControl::for_runtime(&primary.runtime_config()); @@ -757,7 +769,7 @@ async fn gateway_login_states_stop_request_retries() { home.path().to_path_buf(), ); let info = ModelProviderInfo { - gateway_oauth: Some(GatewayOAuthConfig { + gateway_oauth: Some(GatewayOAuthConfig::OAuth(GatewayOAuthSettings { authorization_url: format!("{}/authorize", server.uri()), token_url: format!("{}/token", server.uri()), client_id: "client".into(), @@ -769,7 +781,7 @@ async fn gateway_login_states_stop_request_retries() { name: "x-gateway-auth".into(), scheme: "Bearer".into(), }, - }), + })), ..ModelProviderInfo::create_openai_provider(Some(server.uri())) }; let login_control = codex_login::GatewayLoginControl::for_runtime(&primary.runtime_config()); @@ -795,7 +807,8 @@ async fn gateway_login_states_stop_request_retries() { ); let cancel = tokio::sync::Notify::new(); let (url_tx, url_rx) = tokio::sync::oneshot::channel(); - let login = gateway.login_with_browser(cancel.notified(), |_| { + let oauth = gateway.oauth_session().unwrap(); + let login = oauth.login_with_browser(cancel.notified(), |_| { let _ = url_tx.send(()); }); let check = async {