diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 1cc32ea681..e0eff6147a 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2820,6 +2820,18 @@ dependencies = [ "zip 2.4.2", ] +[[package]] +name = "codex-credentialed-routes" +version = "0.0.0" +dependencies = [ + "anyhow", + "codex-backend-client", + "codex-login", + "codex-network-proxy", + "http 1.4.0", + "tracing", +] + [[package]] name = "codex-exec" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index e5f5025e91..42856e0f38 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -33,6 +33,7 @@ members = [ "connectors", "config", "context-fragments", + "credentialed-routes", "shell-command", "shell-escalation", "skills", @@ -165,6 +166,7 @@ codex-home = { path = "codex-home" } codex-config = { path = "config" } codex-connectors = { path = "connectors" } codex-context-fragments = { path = "context-fragments" } +codex-credentialed-routes = { path = "credentialed-routes" } codex-core = { path = "core" } codex-core-api = { path = "core-api" } codex-core-plugins = { path = "core-plugins" } diff --git a/codex-rs/credentialed-routes/BUILD.bazel b/codex-rs/credentialed-routes/BUILD.bazel new file mode 100644 index 0000000000..f040362f1d --- /dev/null +++ b/codex-rs/credentialed-routes/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "credentialed-routes", + crate_name = "codex_credentialed_routes", +) diff --git a/codex-rs/credentialed-routes/Cargo.toml b/codex-rs/credentialed-routes/Cargo.toml new file mode 100644 index 0000000000..dbabf86330 --- /dev/null +++ b/codex-rs/credentialed-routes/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "codex-credentialed-routes" +edition.workspace = true +version.workspace = true +license.workspace = true + +[lib] +name = "codex_credentialed_routes" +path = "src/lib.rs" +doctest = false + +[lints] +workspace = true + +[dependencies] +anyhow = { workspace = true } +codex-backend-client = { workspace = true } +codex-login = { workspace = true } +codex-network-proxy = { workspace = true } +http = { workspace = true } +tracing = { workspace = true } diff --git a/codex-rs/credentialed-routes/src/lib.rs b/codex-rs/credentialed-routes/src/lib.rs new file mode 100644 index 0000000000..c453398e36 --- /dev/null +++ b/codex-rs/credentialed-routes/src/lib.rs @@ -0,0 +1,101 @@ +use anyhow::Result; +use codex_backend_client::Client as BackendClient; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_network_proxy::ConfigReloader; +use codex_network_proxy::ConfigState; +use codex_network_proxy::CredentialedRoute; +use codex_network_proxy::CredentialedRouteProxyHeader; +use codex_network_proxy::CredentialedRoutesConfig; +use codex_network_proxy::CredentialedRoutesReloader; +use codex_network_proxy::CredentialedRoutesSource; +use http::HeaderMap; +use std::sync::Arc; +use tracing::debug; +use tracing::warn; + +/// Builds the initial proxy state and live reloader for credentialed routes. +pub async fn prepare_proxy_state( + base_state: ConfigState, + chatgpt_base_url: &str, + auth_manager: Arc, +) -> Result<(ConfigState, Arc, Vec)> { + let auth = auth_manager.auth().await; + let credentialed_routes = load_for_session(chatgpt_base_url, auth.as_ref()).await; + let reloader = Arc::new(CredentialedRoutesReloader::new( + base_state, + credentialed_routes.clone(), + source(chatgpt_base_url.to_string(), auth_manager), + )); + let state = ConfigReloader::reload_now(reloader.as_ref()).await?; + let route_prefixes = credentialed_routes.route_prefixes(); + Ok((state, reloader, route_prefixes)) +} + +/// Loads the initial credentialed routes for one Codex session. +pub async fn load_for_session( + chatgpt_base_url: &str, + auth: Option<&CodexAuth>, +) -> CredentialedRoutesConfig { + match fetch(chatgpt_base_url, auth).await { + Ok(credentialed_routes) => credentialed_routes, + Err(err) => { + warn!(error = %err, "failed to load credentialed routes for session"); + CredentialedRoutesConfig::default() + } + } +} + +/// Returns the live route source used by the managed proxy reloader. +pub fn source( + chatgpt_base_url: String, + auth_manager: Arc, +) -> Arc { + Arc::new(move || { + let chatgpt_base_url = chatgpt_base_url.clone(); + let auth_manager = Arc::clone(&auth_manager); + async move { + let auth = auth_manager.auth().await; + fetch(&chatgpt_base_url, auth.as_ref()).await + } + }) +} + +async fn fetch( + chatgpt_base_url: &str, + auth: Option<&CodexAuth>, +) -> Result { + let Some(auth) = auth.filter(|auth| auth.uses_codex_backend()) else { + return Ok(CredentialedRoutesConfig::default()); + }; + + let client = BackendClient::from_auth(chatgpt_base_url.to_string(), auth)?; + let response = client.list_credential_routes().await?; + debug!( + credentialed_routes = response.routes.len(), + "loaded credentialed routes for session" + ); + Ok(CredentialedRoutesConfig { + routes: response + .routes + .into_iter() + .map(|route| CredentialedRoute { + connector_id: route.connector_id, + link_id: route.link_id, + base_url: route.base_url, + }) + .collect(), + proxy_headers: credentialed_route_proxy_headers(client.credential_routes_proxy_headers()), + proxy_url: Some(client.credential_routes_proxy_url()), + }) +} + +fn credentialed_route_proxy_headers(headers: HeaderMap) -> Vec { + headers + .iter() + .map(|(name, value)| CredentialedRouteProxyHeader { + name: name.clone(), + value: value.clone(), + }) + .collect() +}