diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c09b51b7cf..a98d142497 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -5526,6 +5526,8 @@ dependencies = [ "pretty_assertions", "schemars 0.8.22", "serde_json", + "tempfile", + "tokio", "url", ] diff --git a/codex-rs/app-server/tests/suite/v2/config_server_tests.rs b/codex-rs/app-server/tests/suite/v2/config_server_tests.rs index f6ca0981e8..f4854b8473 100644 --- a/codex-rs/app-server/tests/suite/v2/config_server_tests.rs +++ b/codex-rs/app-server/tests/suite/v2/config_server_tests.rs @@ -1,5 +1,6 @@ //! Exercises required HTTPS policy through real app-server startup and session RPCs. +use super::https_mock::tls_listener; use anyhow::Result; use app_test_support::TestAppServer; use codex_app_server_protocol::LoginAccountResponse; @@ -16,13 +17,9 @@ use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; use std::path::Path; use std::path::PathBuf; -use std::sync::Arc; use tempfile::TempDir; use tempfile::tempdir; -use tokio::net::TcpListener; use tokio::net::TcpStream; -use tokio_rustls::TlsAcceptor; -use tokio_rustls::rustls; use tokio_util::task::AbortOnDropHandle; use wiremock::Mock; use wiremock::MockServer; @@ -210,20 +207,7 @@ async fn https_mock_server( home: &Path, policy: &MockServer, ) -> Result<(String, PathBuf, AbortOnDropHandle<()>)> { - let cert = rcgen::generate_simple_self_signed(vec!["localhost".into(), "127.0.0.1".into()])?; - let ca_path = home.join("policy-ca.pem"); - std::fs::write(&ca_path, cert.cert.pem())?; - let tls = rustls::ServerConfig::builder_with_provider(Arc::new( - rustls::crypto::aws_lc_rs::default_provider(), - )) - .with_safe_default_protocol_versions()? - .with_no_client_auth() - .with_single_cert( - vec![cert.cert.der().clone()], - rustls::pki_types::PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der()).into(), - )?; - let acceptor = TlsAcceptor::from(Arc::new(tls)); - let listener = TcpListener::bind("127.0.0.1:0").await?; + let (listener, acceptor, ca_path) = tls_listener(home).await?; let origin_url = format!("https://{}", listener.local_addr()?); let origin = *policy.address(); let tls = AbortOnDropHandle::new(tokio::spawn(async move { diff --git a/codex-rs/app-server/tests/suite/v2/https_mock.rs b/codex-rs/app-server/tests/suite/v2/https_mock.rs new file mode 100644 index 0000000000..df26c82edc --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/https_mock.rs @@ -0,0 +1,27 @@ +//! Loopback TLS listener and its trusted test certificate. + +use anyhow::Result; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use tokio::net::TcpListener; +use tokio_rustls::TlsAcceptor; +use tokio_rustls::rustls; + +pub(super) async fn tls_listener(home: &Path) -> Result<(TcpListener, TlsAcceptor, PathBuf)> { + let cert = rcgen::generate_simple_self_signed(vec!["localhost".into(), "127.0.0.1".into()])?; + let ca_path = home.join("mock-ca.pem"); + std::fs::write(&ca_path, cert.cert.pem())?; + let tls = rustls::ServerConfig::builder_with_provider(Arc::new( + rustls::crypto::aws_lc_rs::default_provider(), + )) + .with_safe_default_protocol_versions()? + .with_no_client_auth() + .with_single_cert( + vec![cert.cert.der().clone()], + rustls::pki_types::PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der()).into(), + )?; + let acceptor = TlsAcceptor::from(Arc::new(tls)); + let listener = TcpListener::bind("127.0.0.1:0").await?; + Ok((listener, acceptor, ca_path)) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 032f5e5a32..2b14de3b4f 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -83,6 +83,7 @@ mod guardian_v2; mod history_notes_extension; mod hooks_list; mod host_skills; +mod https_mock; mod imagegen_extension; mod initialize; mod luna_reserve; diff --git a/codex-rs/app-server/tests/suite/v2/web_search.rs b/codex-rs/app-server/tests/suite/v2/web_search.rs index bff053b176..d6de23ce35 100644 --- a/codex-rs/app-server/tests/suite/v2/web_search.rs +++ b/codex-rs/app-server/tests/suite/v2/web_search.rs @@ -36,6 +36,9 @@ use wiremock::ResponseTemplate; use wiremock::matchers::method; use wiremock::matchers::path; +#[path = "web_search_routing.rs"] +mod routing; + // macOS and Windows Bazel CI can spend tens of seconds starting app-server // subprocesses or processing test RPCs under load. #[cfg(any(target_os = "macos", windows))] diff --git a/codex-rs/app-server/tests/suite/v2/web_search_routing.rs b/codex-rs/app-server/tests/suite/v2/web_search_routing.rs new file mode 100644 index 0000000000..c378e1f1dd --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/web_search_routing.rs @@ -0,0 +1,192 @@ +//! Standalone search uses the discovered workspace origin and credentials. + +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; +use app_test_support::TestAppServer; +use app_test_support::write_chatgpt_auth; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStatus; +use codex_app_server_protocol::UserInput; +use codex_config::types::AuthCredentialsStoreMode; +use codex_features::Feature; +use core_test_support::responses; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; +use tokio::io::copy_bidirectional; +use tokio::net::TcpStream; +use tokio::task::JoinSet; +use tokio::time::timeout; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::super::https_mock::tls_listener; +use super::search_request; + +const READ_TIMEOUT: Duration = Duration::from_secs(60); +const SEARCH_PATH: &str = "/backend-api/codex/alpha/search"; + +#[tokio::test] +async fn standalone_web_search_uses_workspace_routing() -> Result<()> { + let bootstrap = MockServer::start().await; + let routed = responses::start_mock_server().await; + let home = TempDir::new()?; + + let (listener, acceptor, cert_path) = tls_listener(home.path()).await?; + let routed_origin = format!("https://localhost:{}", listener.local_addr()?.port()); + let routed_address = *routed.address(); + let tls_bridge = async move { + let mut connections = JoinSet::new(); + loop { + tokio::select! { + connection = listener.accept() => { + let (stream, _) = connection.context("accept TLS bridge connection")?; + let acceptor = acceptor.clone(); + connections.spawn(async move { + let mut stream = acceptor.accept(stream).await.context("TLS bridge handshake")?; + let mut upstream = TcpStream::connect(routed_address).await.context("TLS bridge upstream connection")?; + copy_bidirectional(&mut stream, &mut upstream).await.context("TLS bridge transfer")?; + Ok::<(), anyhow::Error>(()) + }); + } + Some(result) = connections.join_next() => result.context("TLS bridge connection task")??, + } + } + }; + + Mock::given(method("GET")) + .and(path("/backend-api/wham/accounts/check")) + .and(header("chatgpt-account-id", "selected")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "accounts": [{ + "id": "selected", + "workspace_backend_origin": routed_origin, + "account_routing_override": "us_cr", + }], + }))) + .expect(1) + .mount(&bootstrap) + .await; + Mock::given(path("/backend-api/wham/config/bundle")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&bootstrap) + .await; + Mock::given(path(SEARCH_PATH)) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&bootstrap) + .await; + + Mock::given(method("POST")) + .and(path(SEARCH_PATH)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "output": "Routed search result", + }))) + .expect(1) + .mount(&routed) + .await; + + let call_id = "routed-search"; + let model_responses = vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call_with_namespace( + call_id, + "web", + "run", + &json!({"search_query": [{"q": "workspace routing"}]}).to_string(), + ), + responses::ev_completed("resp-1"), + ]), + responses::sse_completed("resp-2"), + ]; + let model_requests = responses::mount_sse_sequence(&routed, model_responses).await; + + MockResponsesConfig::new(&bootstrap.uri()) + .with_root_config(&format!( + "chatgpt_base_url = \"{}/backend-api\"", + bootstrap.uri() + )) + .with_model_provider("openai-custom") + .with_provider_name("OpenAI") + .with_provider_base_url(&format!("{}/backend-api/codex", bootstrap.uri())) + .with_provider_config("requires_openai_auth = true\nsupports_websockets = false") + .enable_feature(Feature::StandaloneWebSearch) + .write(home.path())?; + std::fs::write( + home.path().join("requirements.toml"), + "enforce_residency = \"us\"\n", + )?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("access-chatgpt").account_id("selected"), + AuthCredentialsStoreMode::File, + )?; + let mut env_overrides = vec![ + ("OPENAI_API_KEY", None), + ("CODEX_API_KEY", None), + ("CODEX_CA_CERTIFICATE", cert_path.to_str()), + ("SSL_CERT_FILE", cert_path.to_str()), + ]; + env_overrides.extend( + codex_network_proxy::PROXY_ENV_KEYS + .iter() + .map(|key| (*key, None)), + ); + let exercise = async { + let mut app = TestAppServer::builder() + .with_codex_home(home.path()) + .with_env_overrides(&env_overrides) + .build_initialized_with_timeout(READ_TIMEOUT) + .await?; + let thread = app.start_thread(ThreadStartParams::default()).await?.thread; + let completed = timeout( + READ_TIMEOUT, + app.start_turn_and_wait_for_completion(TurnStartParams { + thread_id: thread.id, + input: vec![UserInput::Text { + text: "Search the web".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }), + ) + .await??; + assert_eq!(completed.turn.status, TurnStatus::Completed); + let requests = model_requests.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[1].function_call_output(call_id)["output"], + json!([{"type": "input_text", "text": "Routed search result"}]) + ); + + let search = search_request(&routed, SEARCH_PATH).await?; + assert_eq!(search.headers["authorization"], "Bearer access-chatgpt"); + assert_eq!(search.headers["chatgpt-account-id"], "selected"); + assert_eq!(search.headers["x-openai-internal-codex-residency"], "us"); + assert_eq!(search.headers["x-openai-account-routing-override"], "us_cr"); + bootstrap.verify().await; + routed.verify().await; + Ok::<_, anyhow::Error>(app) + }; + let mut app = tokio::select! { + biased; + result = tls_bridge => return result, + result = exercise => result?, + }; + assert!( + timeout(READ_TIMEOUT, app.shutdown_gracefully()) + .await?? + .success() + ); + Ok(()) +} diff --git a/codex-rs/ext/web-search/Cargo.toml b/codex-rs/ext/web-search/Cargo.toml index a464b56bea..d72c779b61 100644 --- a/codex-rs/ext/web-search/Cargo.toml +++ b/codex-rs/ext/web-search/Cargo.toml @@ -31,3 +31,5 @@ url = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true } diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index 2394235bca..531cc364f9 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/codex-rs/ext/web-search/src/extension.rs @@ -19,6 +19,7 @@ use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolContributor; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; +use codex_model_provider::WorkspaceRoutingContext; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::config_types::WebSearchContextSize; @@ -130,7 +131,10 @@ impl ToolContributor for WebSearchExtension { ) -> Vec< Arc codex_extension_api::ToolExecutor>>, > { - let Some(config) = thread_store.get::() else { + let (Some(config), Some(workspace_routing)) = ( + thread_store.get::(), + thread_store.get::(), + ) else { return Vec::new(); }; if !config.available { @@ -144,6 +148,7 @@ impl ToolContributor for WebSearchExtension { config.provider.clone(), Some(self.auth_manager.clone()), ), + workspace_routing, settings: config.settings.clone(), supports_literature_search: step_store.get::().is_some_and(|model| { model @@ -165,6 +170,10 @@ pub fn install(registry: &mut ExtensionRegistryBuilder, auth_manager: Ar registry.tool_contributor(extension); } +#[cfg(test)] +#[path = "routing_tests.rs"] +mod routing_tests; + #[cfg(test)] mod tests { use codex_extension_api::ExtensionData; @@ -173,6 +182,7 @@ mod tests { use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; use codex_login::CodexAuth; + use codex_model_provider::WorkspaceRoutingContext; use codex_model_provider_info::ModelProviderInfo; use pretty_assertions::assert_eq; @@ -217,6 +227,9 @@ mod tests { let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("11111111-1111-4111-8111-111111111111"); let step_store = ExtensionData::new("step"); + thread_store.insert(WorkspaceRoutingContext::new( + "https://chatgpt.com/backend-api".to_string(), + )); thread_store.insert(WebSearchExtensionConfig { available: true, http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), diff --git a/codex-rs/ext/web-search/src/routing_tests.rs b/codex-rs/ext/web-search/src/routing_tests.rs new file mode 100644 index 0000000000..ec278d3b84 --- /dev/null +++ b/codex-rs/ext/web-search/src/routing_tests.rs @@ -0,0 +1,108 @@ +#![cfg(test)] + +//! Search must retain the host's routing context across thread configuration changes. + +use std::io; +use std::sync::Arc; + +use codex_core::config::ConfigBuilder; +use codex_extension_api::ConfigContributor; +use codex_extension_api::ConversationHistory; +use codex_extension_api::ExtensionData; +use codex_extension_api::FunctionCallError; +use codex_extension_api::NoopTurnItemEmitter; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStartInput; +use codex_extension_api::ToolCall; +use codex_extension_api::ToolCallSource; +use codex_extension_api::ToolContributor; +use codex_extension_api::ToolName; +use codex_extension_api::ToolPayload; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::WorkspaceRouting; +use codex_login::WorkspaceRoutingRequest; +use codex_login::WorkspaceRoutingResolver; +use codex_model_provider::ModelProviderFuture; +use codex_model_provider::WorkspaceRoutingContext; +use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::config_types::WebSearchMode; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TruncationPolicy; +use serde_json::json; + +use super::WebSearchExtension; + +#[tokio::test] +async fn config_changes_preserve_the_threads_established_routing() { + struct ReportBootstrap; + + impl WorkspaceRoutingResolver for ReportBootstrap { + fn resolve( + &self, + request: WorkspaceRoutingRequest, + ) -> ModelProviderFuture<'_, io::Result>> { + Box::pin(async move { Err(io::Error::other(request.chatgpt_base_url)) }) + } + } + + let home = tempfile::tempdir().unwrap(); + let mut config = ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .fallback_cwd(Some(home.path().to_path_buf())) + .build() + .await + .unwrap(); + config.web_search_mode.set(WebSearchMode::Cached).unwrap(); + config.model_provider = ModelProviderInfo::create_openai_provider(Some( + "http://127.0.0.1:0/backend-api/codex".into(), + )); + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let resolver: Arc = Arc::new(ReportBootstrap); + auth_manager.set_workspace_routing_resolver(Arc::downgrade(&resolver)); + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new("thread"); + // Use a thread bootstrap distinct from either config so reconstructing it is observable. + let thread_bootstrap = "https://thread.example/backend-api"; + thread_store.insert(WorkspaceRoutingContext::new(thread_bootstrap.into())); + let extension = WebSearchExtension { auth_manager }; + extension + .on_thread_start(ThreadStartInput { + config: &config, + session_source: &SessionSource::Exec, + persistent_thread_state_available: false, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &session_store, + thread_store: &thread_store, + }) + .await; + let mut changed = config.clone(); + changed.chatgpt_base_url = "https://changed.example/backend-api".into(); + extension.on_config_changed(&session_store, &thread_store, &config, &changed); + let tools = + extension.tools_for_step(&session_store, &thread_store, &ExtensionData::new("step")); + let call = ToolCall { + turn_id: "turn".into(), + call_id: "search".into(), + tool_name: ToolName::namespaced("web", "run"), + model: "test-model".into(), + codex_turn_metadata: None, + truncation_policy: TruncationPolicy::Bytes(1024), + source: ToolCallSource::Direct, + conversation_history: ConversationHistory::default(), + turn_item_emitter: Arc::new(NoopTurnItemEmitter), + attachment_store: codex_core::passthrough_image_store(), + ephemeral: false, + environments: Vec::new(), + payload: ToolPayload::Function { + arguments: json!({"search_query": [{"q": "test"}]}).to_string(), + }, + }; + assert!(matches!( + tools[0].handle(call).await, + Err(FunctionCallError::Fatal(message)) if message == thread_bootstrap + )); +} diff --git a/codex-rs/ext/web-search/src/tool.rs b/codex-rs/ext/web-search/src/tool.rs index 6ac78b060b..d7307bc2ae 100644 --- a/codex-rs/ext/web-search/src/tool.rs +++ b/codex-rs/ext/web-search/src/tool.rs @@ -1,3 +1,6 @@ +use std::sync::Arc; + +use codex_api::ReqwestTransport; use codex_api::SearchClient; use codex_api::SearchCommands; use codex_api::SearchQuery; @@ -19,9 +22,15 @@ use codex_extension_items::web_search::WebSearchAction; use codex_extension_items::web_search::WebSearchItem; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; +use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::add_originator_header; +use codex_login::default_client::create_client_for_route_async; use codex_login::default_client::create_transport_for_routes_async; +use codex_model_provider::ACCOUNT_ROUTING_HEADER; use codex_model_provider::SharedModelProvider; +use codex_model_provider::WorkspaceRoutingContext; +use codex_model_provider::prepare_provider_request; +use codex_model_provider::provider_uses_first_party_auth_path; use codex_protocol::models::WebSearchAction as CoreWebSearchAction; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::WebSearchBeginEvent; @@ -47,6 +56,7 @@ pub(crate) struct WebSearchTool { pub(crate) session_id: String, pub(crate) http_client_factory: HttpClientFactory, pub(crate) provider: SharedModelProvider, + pub(crate) workspace_routing: Arc, pub(crate) settings: SearchSettings, pub(crate) supports_literature_search: bool, pub(crate) originator: Option, @@ -102,18 +112,45 @@ impl WebSearchTool { ) -> Result, FunctionCallError> { let commands = parse_commands(call.function_arguments()?, self.supports_literature_search)?; let command_action = command_action(&commands); - let (provider, auth) = self - .provider - .api_provider_and_auth() - .await - .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; - let transport = create_transport_for_routes_async( - self.http_client_factory.clone(), - ClientRouteClass::Api, - ) + let provider = self.provider.as_ref(); + let first_party = provider_uses_first_party_auth_path(provider.info()); + let client = prepare_provider_request(provider, first_party, |_| async { + let routing_auth = if first_party { + provider.auth().await + } else { + None + }; + let (endpoint, api_auth) = provider.api_provider_and_auth().await?; + let resolved = provider + .route_responses_api_provider( + endpoint, + routing_auth.as_ref(), + &self.workspace_routing, + ) + .await?; + let endpoint = resolved.provider; + let transport = if resolved.redirect_policy == ClientRedirectPolicy::Reject + || endpoint.headers.contains_key(ACCOUNT_ROUTING_HEADER) + { + create_client_for_route_async( + self.http_client_factory.clone(), + endpoint.url_for_path("alpha/search"), + ClientRouteClass::Api, + ClientRedirectPolicy::Reject, + ) + .await + .map(ReqwestTransport::from_http_client) + } else { + create_transport_for_routes_async( + self.http_client_factory.clone(), + ClientRouteClass::Api, + ) + .await + }?; + Ok(SearchClient::new(transport, endpoint, api_auth)) + }) .await .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; - let client = SearchClient::new(transport, provider, auth); let request = SearchRequest { id: self.session_id.clone(), model: call.model.clone(),