Route standalone web search through the thread's workspace context (#58140)

## Why

Standalone web search could send requests to the bootstrap backend instead of the discovered workspace backend, bypassing workspace routing.

## What changed

Use the thread's shared `WorkspaceRoutingContext` when preparing search requests so they use the resolved backend, authentication, residency, and account routing headers. Reject redirects when required by the resolved routing policy or account routing header, and preserve the established routing context across configuration changes.

## Testing

Add an app-server integration test that verifies search reaches the discovered HTTPS backend with the expected credentials and routing headers, without contacting the bootstrap search endpoint. Add a regression test that verifies configuration changes preserve the thread's routing context.

GitOrigin-RevId: 67160f1b3e54231dc100ce3ed4cfcea965b56d26
This commit is contained in:
sgarg-oai
2026-10-29 18:09:55 +00:00
committed by copyberry
parent 7d764a236e
commit f32b6cfd2b
10 changed files with 398 additions and 29 deletions
+2
View File
@@ -5526,6 +5526,8 @@ dependencies = [
"pretty_assertions",
"schemars 0.8.22",
"serde_json",
"tempfile",
"tokio",
"url",
]
@@ -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 {
@@ -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))
}
@@ -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;
@@ -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))]
@@ -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(())
}
+2
View File
@@ -31,3 +31,5 @@ url = { workspace = true }
[dev-dependencies]
pretty_assertions = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true }
+14 -1
View File
@@ -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<dyn for<'call> codex_extension_api::ToolExecutor<codex_extension_api::ToolCall<'call>>>,
> {
let Some(config) = thread_store.get::<WebSearchExtensionConfig>() else {
let (Some(config), Some(workspace_routing)) = (
thread_store.get::<WebSearchExtensionConfig>(),
thread_store.get::<WorkspaceRoutingContext>(),
) 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::<ModelInfo>().is_some_and(|model| {
model
@@ -165,6 +170,10 @@ pub fn install(registry: &mut ExtensionRegistryBuilder<Config>, 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),
@@ -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<Option<WorkspaceRouting>>> {
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<dyn WorkspaceRoutingResolver> = 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
));
}
+47 -10
View File
@@ -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<WorkspaceRoutingContext>,
pub(crate) settings: SearchSettings,
pub(crate) supports_literature_search: bool,
pub(crate) originator: Option<String>,
@@ -102,18 +112,45 @@ impl WebSearchTool {
) -> Result<Box<dyn ToolOutput>, 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(),