Resolve model-selected search providers with policy-aware fallback (#59036)

## Why

Search discovery must account for the active model's supported tools while honoring host search restrictions. Automatic fallback must also preserve the provider that owns existing browsing references.

## What changed

- Resolve search from model metadata, client configuration, and managed requirements, exposing the selected provider's tools.
- Discover hosted search with the model's supported tools and carry the discovered invocation context into calls.
- Allow native fallback for eligible discovery failures only when search policy permits it and hosted search has not already been used. Do not replay failed hosted invocations through native search.
- Keep native fallback references on native search across later turns, discovery recovery, and resume. Release unused fallback offers and coordinate background calls to prevent conflicting provider use.

## Testing

Add app-server and core scenario coverage for provider selection, discovery failures, search restrictions, reference continuity, resume, and concurrent background search. Include a snapshot covering discovery and call history across a model switch.

GitOrigin-RevId: 0f22605722daaa424c4fa34660654cde5a6db049
This commit is contained in:
yitaoh-oai
2026-11-03 02:08:56 +00:00
committed by copyberry
parent 3c074bb44a
commit e839b828fb
15 changed files with 1749 additions and 11 deletions
@@ -27,6 +27,7 @@ use codex_app_server_protocol::WebSearchItem;
use codex_config::types::AuthCredentialsStoreMode;
use codex_features::Feature;
use codex_models_manager::model_info::model_info_from_slug;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS;
use core_test_support::responses;
use core_test_support::responses::strip_response_item_ids_from_json;
@@ -609,3 +610,404 @@ async fn search_request(server: &MockServer, search_path: &str) -> Result<wiremo
.find(|request| request.url.path() == search_path)
.context("expected standalone search request")
}
const XATU_V99: &str = "SonicBrowserToolV99";
const PAPER_MODEL_TOOLS: &[&str] = &["literature_search_query", "other_model_tool"];
#[derive(Clone, Copy, PartialEq, Eq)]
enum SearchPolicy {
Standard,
HostedOnly,
RequireNative,
RequireNativeClientMcp,
RequireNativeCached,
RequireNativeTelemetry,
RequireNativeDomainPolicy,
HostAllowedDomains(&'static [&'static str]),
Standalone(WebSearchMode),
Disabled,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum DiscoveryFailure {
Rpc(i32),
Http(u16),
}
/// Model Lab selection is resolved for the active turn on the same app-server used by both clients.
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::Standard, false; "xatu_literature_treatment")]
#[test_case::test_case(XATU_V99, (&[], false), None, SearchPolicy::Standard, false; "xatu_control")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), Some(DiscoveryFailure::Rpc(-32603)), SearchPolicy::Standard, false; "critical_discovery_error")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), Some(DiscoveryFailure::Rpc(-32602)), SearchPolicy::Standard, false; "invalid_discovery_request")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), Some(DiscoveryFailure::Http(503)), SearchPolicy::Standard, false; "service_unavailable_after_retries_allows_fallback")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), Some(DiscoveryFailure::Http(403)), SearchPolicy::Standard, false; "forbidden_discovery_disables_search")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::HostedOnly, false; "work_hosted_only")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), Some(DiscoveryFailure::Rpc(-32603)), SearchPolicy::HostedOnly, false; "work_prohibits_native_fallback")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::Standard, true; "failed_invocation_never_replays_on_native")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, false), None, SearchPolicy::Standard, false; "host_disables_literature_only")]
#[test_case::test_case(XATU_V99, (&["other_model_tool"], false), None, SearchPolicy::Standard, false; "model_with_other_supported_tools")]
#[test_case::test_case("SonicBrowserToolV66", (PAPER_MODEL_TOOLS, false), None, SearchPolicy::Standard, false; "inherit_connector_version")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::RequireNative, false; "managed_native_requirement")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::RequireNativeClientMcp, false; "native_requirement_overrides_client_mcp_without_model_preference")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::RequireNativeCached, false; "native_requirement_enforces_cached_turn")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::RequireNativeTelemetry, false; "native_requirement_with_telemetry_and_no_search_connector")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::RequireNativeDomainPolicy, false; "native_requirement_cannot_bypass_connector_domain_policy")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::Standalone(WebSearchMode::Cached), false; "standalone_model_honors_cached_turn")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::Standalone(WebSearchMode::Live), false; "desktop_installation_id_preserves_standalone_search")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), Some(DiscoveryFailure::Rpc(-32603)), SearchPolicy::HostAllowedDomains(&[]), false; "empty_host_and_native_domain_lists_deny_fallback")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), Some(DiscoveryFailure::Rpc(-32603)), SearchPolicy::HostAllowedDomains(&["example.com"]), false; "matching_query_filter_cannot_authorize_cached_navigation")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::Standalone(WebSearchMode::Disabled), false; "standalone_model_honors_disabled_turn")]
#[test_case::test_case(XATU_V99, (PAPER_MODEL_TOOLS, true), None, SearchPolicy::Disabled, false; "metadata_cannot_authorize_search")]
#[tokio::test]
async fn model_lab_search_provider(
pinned_type: &str,
(supported_tools, pinned_literature): (&[&str], bool),
error: Option<DiscoveryFailure>,
policy: SearchPolicy,
invocation_error: bool,
) -> Result<()> {
use core_test_support::apps_test_server::AppsTestServer;
use std::sync::Arc;
use std::sync::Mutex;
let server = responses::start_mock_server().await;
let tool = json!({
"name": "search_service.web_run", "description": "Search", "annotations": {"readOnlyHint": true},
"inputSchema": {"type":"object","properties":{"system2_search_query":{"type":"array"},"open":{"type":"array"}}},
"_meta": {"connector_id": codex_mcp::SEARCH_CONNECTOR, "connector_name":"search_service"}
});
let catalog_tools = if policy == SearchPolicy::RequireNativeTelemetry {
vec![]
} else {
vec![tool.clone()]
};
let apps =
AppsTestServer::mount_with_tools(&server, Arc::new(Mutex::new(catalog_tools))).await?;
let mut pinned = tool;
let search_query = if pinned_type == "SonicBrowserToolV66" {
let properties = pinned["inputSchema"]["properties"]
.as_object_mut()
.context("search schema")?;
let query = properties
.remove("system2_search_query")
.context("query schema")?;
properties.insert("system1_search_query".into(), query);
"system1_search_query"
} else {
"system2_search_query"
};
if pinned_literature {
pinned["inputSchema"]["properties"]["literature_search_query"] =
json!({"type":"array","items":{"type":"object"}});
}
let profile =
json!({"config":{"tool_type":pinned_type,"allow_literature_search":pinned_literature}});
pinned["_meta"]["_codex_apps"] = json!({"invocation_context":profile});
Mock::given(method("POST"))
.and(wiremock::matchers::body_partial_json(json!({
"method":"tools/list","params":{"_meta":{"x-codex-turn-metadata":{"mcp_request_meta":{
"openai/search_context":{"model_supported_tools":supported_tools}
}}}}
})))
.respond_with(move |r: &wiremock::Request| {
let req: Value = r.body_json().expect("JSON-RPC request");
let mut response = json!({"jsonrpc":"2.0","id":req["id"]});
match error {
Some(DiscoveryFailure::Http(status)) => return ResponseTemplate::new(status),
Some(DiscoveryFailure::Rpc(code)) => {
response["error"] = json!({"code":code,"message":"discovery failed"});
}
None => response["result"] = json!({"tools":[pinned]}),
}
ResponseTemplate::new(200).set_body_json(response)
})
.with_priority(1)
.mount(&server)
.await;
let namespace = "mcp__codex_apps__search_service";
let key = if pinned_literature {
"literature_search_query"
} else {
search_query
};
let commands = json!({(key): [{"q":"plasma biomarkers"}]});
let mut steps = Vec::new();
let native = matches!(
policy,
SearchPolicy::RequireNative
| SearchPolicy::RequireNativeClientMcp
| SearchPolicy::RequireNativeCached
| SearchPolicy::RequireNativeTelemetry
| SearchPolicy::Standalone(WebSearchMode::Cached | WebSearchMode::Live)
) || (matches!(
error,
Some(DiscoveryFailure::Rpc(-32603) | DiscoveryFailure::Http(503))
) && policy == SearchPolicy::Standard);
if let SearchPolicy::HostAllowedDomains(domains) = policy
&& !domains.is_empty()
{
// A stale model call must not reach native, even for a URL already in its cache.
steps.push(responses::sse(vec![
responses::ev_function_call_with_namespace(
"outside",
"web",
"run",
&json!({"open":[{"ref_id":"https://outside.example/paper"}]}).to_string(),
),
responses::ev_completed("outside"),
]));
}
if native {
mount_search_response(&server, "/v1/alpha/search", /*expected_calls*/ 1).await;
steps.push(responses::sse(vec![
responses::ev_function_call_with_namespace(
"native-search",
"web",
"run",
&commands.to_string(),
),
responses::ev_completed("native"),
]));
}
let allow_hosted = matches!(
policy,
SearchPolicy::Standard | SearchPolicy::HostedOnly | SearchPolicy::HostAllowedDomains(_)
);
let hosted = allow_hosted && error.is_none();
if hosted {
steps.push(responses::sse(vec![
responses::ev_tool_search_call(
"find-search",
&json!({"query": "search_service web_run"}),
),
responses::ev_completed("discovered"),
]));
steps.push(responses::sse(vec![
responses::ev_function_call_with_namespace(
"search-1",
namespace,
"_web_run",
&commands.to_string(),
),
responses::ev_completed("resp-1"),
]));
Mock::given(method("POST"))
.and(wiremock::matchers::body_partial_json(json!({
"method": "tools/call",
"params": {"name": "search_service.web_run"}
})))
.respond_with(move |r: &wiremock::Request| {
let body: Value = r.body_json().expect("JSON-RPC request");
ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": body["id"],
"result": {
"isError": invocation_error,
"content": [{"type": "text", "text": "search result"}]
}
}))
})
.with_priority(1)
.mount(&server)
.await;
}
steps.push(responses::sse(vec![
responses::ev_assistant_message("msg", "Done"),
responses::ev_completed("resp-final"),
]));
let responses = responses::mount_sse_sequence(&server, steps).await;
let home = TempDir::new()?;
let mut model = model_info_from_slug("mock-model");
model.web_search_provider = match policy {
SearchPolicy::RequireNativeClientMcp => None,
SearchPolicy::Standalone(_) => {
Some(codex_protocol::openai_models::WebSearchProvider::Standalone)
}
_ => Some(codex_protocol::openai_models::WebSearchProvider::SearchMcp),
};
model.supports_search_tool = true;
model.use_responses_lite = true;
model.experimental_supported_tools = supported_tools
.iter()
.map(|tool| (*tool).to_owned())
.collect();
let catalog = home.path().join("models.json");
std::fs::write(&catalog, serde_json::to_vec(&json!({"models":[model]}))?)?;
let mode = if matches!(policy, SearchPolicy::HostedOnly | SearchPolicy::Disabled) {
"disabled"
} else {
"live"
};
if matches!(
policy,
SearchPolicy::RequireNative
| SearchPolicy::RequireNativeClientMcp
| SearchPolicy::RequireNativeCached
| SearchPolicy::RequireNativeTelemetry
| SearchPolicy::RequireNativeDomainPolicy
) {
std::fs::write(
home.path().join("requirements.toml"),
"[feature_requirements]\nstandalone_web_search = true\n",
)?;
}
let client_provider = if policy == SearchPolicy::RequireNativeClientMcp {
"search_mcp"
} else {
"standalone"
};
let domain_config = match policy {
SearchPolicy::HostAllowedDomains(domains) => format!(
"\ntools.web_search.allowed_domains = {}",
serde_json::to_string(domains)?
),
_ => String::new(),
};
MockResponsesConfig::new(&server.uri())
.with_root_config(&format!(
"chatgpt_base_url = {}\nmodel_catalog_json = {}\nweb_search = \"{mode}\"\ntools.client_search_provider = \"{client_provider}\"{domain_config}",
serde_json::to_string(&apps.chatgpt_base_url)?,
serde_json::to_string(&catalog)?,
))
.with_model_provider("openai-custom")
.with_provider_name("OpenAI")
.with_provider_config("requires_openai_auth = true")
.with_features(&[
(Feature::Apps, !matches!(policy, SearchPolicy::Standalone(_))),
(Feature::StandaloneWebSearch, true),
(Feature::HostedSearchPolicy, policy == SearchPolicy::HostedOnly),
].into())
.with_provider_config("supports_websockets = false")
.with_provider_config("supports_standalone_web_search = true")
.write(home.path())?;
write_chatgpt_auth(
home.path(),
ChatGptAuthFixture::new("test-account-token").account_id("account-123"),
AuthCredentialsStoreMode::File,
)?;
let mut app = TestAppServer::builder()
.with_codex_home(home.path())
.with_env_overrides(&[("OPENAI_API_KEY", None)])
.build_initialized()
.await?;
let thread = app.start_thread(ThreadStartParams::default()).await?.thread;
app.start_turn_and_wait_for_completion(TurnStartParams {
thread_id: thread.id,
responsesapi_client_metadata: match policy {
SearchPolicy::HostedOnly | SearchPolicy::Disabled => Some(json!({"mode":"live"})),
SearchPolicy::RequireNativeCached => Some(json!({"mode":"cached"})),
SearchPolicy::RequireNativeTelemetry => {
Some(json!({"telemetry_attributes":{"source":"test"}}))
}
SearchPolicy::RequireNativeDomainPolicy => {
Some(json!({"allowed_domains":["example.org"]}))
}
SearchPolicy::HostAllowedDomains([]) => Some(json!({"allowed_domains":[]})),
SearchPolicy::HostAllowedDomains(domains) => {
Some(json!({"mode":"cached","allowed_domains":domains}))
}
SearchPolicy::Standalone(mode) => Some(json!({"mode":mode})),
_ => None,
}
.map(|context| {
let mut meta = json!({"openai/search_context":context});
if matches!(
policy,
SearchPolicy::Standalone(_) | SearchPolicy::RequireNativeTelemetry
) {
meta["openai/codex_installation_id"] = json!("desktop-installation");
}
[("mcp_request_meta".to_string(), meta.to_string())].into()
}),
input: vec![V2UserInput::Text {
text: "Search papers".into(),
text_elements: vec![],
}],
..Default::default()
})
.await?;
let requests = responses.requests();
let received = server
.received_requests()
.await
.context("recorded search requests")?;
let rpc: Vec<Value> = received
.iter()
.filter_map(|r| r.body_json::<Value>().ok())
.filter(|r| r.get("method").is_some())
.collect();
let scoped_lists = rpc
.iter()
.filter(|r| r["method"] == "tools/list")
.filter_map(|r| {
r.pointer("/params/_meta/x-codex-turn-metadata/mcp_request_meta/openai~1search_context")
})
.filter(|ctx| ctx["model_supported_tools"] == json!(supported_tools))
.count();
if error == Some(DiscoveryFailure::Http(503)) {
assert!(scoped_lists > 1, "fallback should follow a discovery retry");
} else {
assert_eq!(scoped_lists, usize::from(allow_hosted));
}
let request = requests[0].body_json();
let tools = request["input"]
.as_array()
.context("Lite input")?
.iter()
.find(|item| item["type"] == "additional_tools")
.context("Lite additional_tools")?;
let find_tool = |ns: &str, name: &str| responses::namespace_child_tool(tools, ns, name);
assert_eq!(find_tool("web", "run").is_some(), native);
assert!(!has_hosted_web_search(&request));
assert!(find_tool(namespace, "_web_run").is_none());
if hosted {
let found = requests[1].tool_search_output("find-search");
let tool = responses::namespace_child_tool(&found, namespace, "_web_run")
.context("profiled search tool in the lookup result")?;
assert_eq!(
tool.pointer("/parameters/properties/literature_search_query")
.is_some(),
pinned_literature
);
assert!(
requests[2]
.function_call_output("search-1")
.to_string()
.contains("search result")
);
}
assert_eq!(
received
.iter()
.filter(|r| r.url.path().ends_with("/alpha/search"))
.count(),
usize::from(native)
);
if matches!(
policy,
SearchPolicy::Standalone(WebSearchMode::Cached)
| SearchPolicy::RequireNativeCached
| SearchPolicy::RequireNativeTelemetry
) {
let body = search_request(&server, "/v1/alpha/search")
.await?
.body_json::<Value>()?;
assert_eq!(
body["settings"]["external_web_access"],
policy == SearchPolicy::RequireNativeTelemetry
);
}
let calls = core_test_support::apps_test_server::recorded_apps_tool_calls(&server).await;
assert_eq!(calls.len(), usize::from(hosted));
if let Some(call) = calls.first() {
assert_eq!(call["params"]["arguments"], commands);
assert_eq!(
call.pointer("/params/_meta/x-codex-turn-metadata/mcp_request_meta/openai~1search_context/model_supported_tools"),
Some(&json!(supported_tools))
);
assert_eq!(
call.pointer("/params/_meta/_codex_apps/invocation_context"),
Some(&profile)
);
}
Ok(())
}
+1
View File
@@ -111,6 +111,7 @@ pub(crate) mod mcp;
mod mcp_skill_dependencies;
mod mcp_tool_approval_templates;
mod mcp_tool_exposure;
mod model_search;
mod network_policy_decision;
pub use mcp::McpManager;
pub use mcp::McpServersDisabled;
+70
View File
@@ -488,6 +488,39 @@ async fn handle_approved_mcp_tool_call(
let mut tool_input = arguments_value
.clone()
.unwrap_or_else(|| JsonValue::Object(serde_json::Map::new()));
let is_search = prepared_call.is_host_owned_apps()
&& codex_mcp::is_model_search_tool(prepared_call.tool_info());
let search = if is_search {
step_context
.extension_data
.get::<crate::model_search::ModelSearch>()
} else {
None
};
// Code Mode can retain this step past a later turn's first native call.
// A call rejected here never reaches the peer or acquires MCP attribution.
if matches!(
search.as_deref(),
Some(crate::model_search::ModelSearch::Hosted { .. })
) && !crate::model_search::SearchOwner::admit(
&sess.services.thread_extension_data,
codex_protocol::openai_models::WebSearchProvider::SearchMcp,
) {
let result = notify_mcp_tool_call_skip(
sess,
turn_context,
call_id,
invocation,
item_metadata,
"This thread is using native search".to_string(),
/*already_started*/ true,
)
.await;
return HandledMcpToolCall {
result: CallToolResult::from_result(result),
tool_input,
};
}
let mut elicitation_type = None;
let mut request_error_code = None;
let result = async {
@@ -569,6 +602,32 @@ async fn handle_approved_mcp_tool_call(
serde_json::json!(turn_context.cyber_access_program),
);
}
let request_meta = if let Some(crate::model_search::ModelSearch::Hosted {
binding,
context,
}) = search.as_deref()
{
let mut meta = request_meta.unwrap_or_else(|| serde_json::json!({}));
let turn_meta = meta
.as_object_mut()
.ok_or_else(|| anyhow::anyhow!("invalid MCP metadata"))?
.entry(crate::X_CODEX_TURN_METADATA_HEADER)
.or_insert_with(|| serde_json::json!({}));
crate::model_search::with_search_context(turn_meta, context)?;
let pin = binding
.search_invocation_context(&prepared_call)
.ok_or_else(|| {
anyhow::anyhow!("Search MCP connection changed after discovery")
})?;
let apps = meta
.get_mut(MCP_TOOL_CODEX_APPS_META_KEY)
.and_then(serde_json::Value::as_object_mut)
.ok_or_else(|| anyhow::anyhow!("missing hosted Apps call metadata"))?;
apps.insert("invocation_context".to_string(), pin.clone());
Some(meta)
} else {
request_meta
};
let request_meta = with_mcp_tool_call_ids_meta(
request_meta,
&sess.thread_id.to_string(),
@@ -581,6 +640,17 @@ async fn handle_approved_mcp_tool_call(
request_meta,
)
.await?;
if is_search
&& let Some(usage) = step_context
.turn
.extension_data
.get::<crate::model_search::SearchTurnUsage>()
{
usage.record(
&step_context.turn,
codex_protocol::openai_models::WebSearchProvider::SearchMcp,
);
}
let mcp_call_trace = sess
.services
.rollout_thread_trace
+1 -1
View File
@@ -247,7 +247,7 @@ fn filter_non_codex_apps_mcp_tools_only(
})
}
fn filter_codex_apps_mcp_tools<'a>(
pub(crate) fn filter_codex_apps_mcp_tools<'a>(
mcp_tools: &'a [McpToolInfo],
config: &'a Config,
) -> impl Iterator<Item = &'a McpToolInfo> + 'a {
+456
View File
@@ -0,0 +1,456 @@
//! Resolve one search provider for the sampling step, preserving the host's search policy.
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::Weak;
use std::sync::atomic::AtomicU8;
use std::sync::atomic::Ordering;
use codex_extension_api::ExtensionData;
use codex_history::InitialHistory;
use codex_history::ResponseItemEnvelope;
use codex_history::RolloutItem;
use codex_mcp::McpBinding;
use codex_mcp::SearchDiscovery;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::mcp::McpAttributionStatus;
use codex_protocol::openai_models::WebSearchProvider;
use codex_tools::ToolSpec;
use serde_json::Map;
use serde_json::Value;
use serde_json::json;
use crate::session::step_settings::ResolvedStepSettings;
use crate::session::turn_context::TurnContext;
use crate::state::SessionServices;
use crate::turn_metadata::ExecutionMetadata;
#[derive(Clone)]
pub(crate) enum ModelSearch {
Hosted {
binding: Arc<McpBinding>,
context: Map<String, Value>,
},
ClientMcp,
Native {
mode: WebSearchMode,
fallback: bool,
},
Disabled,
}
#[derive(Default)]
pub(crate) struct SearchTurnUsage(AtomicU8);
impl SearchTurnUsage {
const MCP: u8 = 1;
const NATIVE: u8 = 2;
pub(crate) fn record(&self, turn: &TurnContext, provider: WebSearchProvider) {
let (flag, label) = match provider {
WebSearchProvider::SearchMcp => (Self::MCP, "search_mcp"),
WebSearchProvider::Standalone => (Self::NATIVE, "native"),
};
let previous = self.0.fetch_or(flag, Ordering::Relaxed);
if previous & flag == 0 {
turn.session_telemetry.counter(
"codex.model_search.turn",
/*inc*/ 1,
&[("provider", label)],
);
if previous != 0 {
turn.session_telemetry.counter(
"codex.model_search.turn",
/*inc*/ 1,
&[("provider", "mixed")],
);
}
}
}
}
struct CachedSearch {
base: Arc<McpBinding>,
request: Map<String, Value>,
discovery: SearchDiscovery,
}
/// Concurrent calls on the same provider share a reservation. An actual search
/// keeps it for the thread. An unused offer releases it when sampling finishes.
pub(crate) struct SearchOwner(WebSearchProvider);
impl SearchOwner {
fn reserve(store: &ExtensionData, provider: WebSearchProvider) -> Option<Arc<Self>> {
let slot = store.get_or_init(Mutex::<Weak<Self>>::default);
let mut slot = slot
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(owner) = slot.upgrade() {
return (owner.0 == provider).then_some(owner);
}
let owner = Arc::new(Self(provider));
*slot = Arc::downgrade(&owner);
Some(owner)
}
pub(crate) fn admit(store: &ExtensionData, provider: WebSearchProvider) -> bool {
let Some(owner) = Self::reserve(store, provider) else {
return false;
};
store.insert_arc(owner);
true
}
pub(crate) fn reserve_native_fallback(
step: &ExtensionData,
thread: &ExtensionData,
) -> Result<Option<Arc<Self>>, &'static str> {
if matches!(
step.get::<ModelSearch>().as_deref(),
Some(ModelSearch::Native { fallback: true, .. })
) {
return Self::reserve(thread, WebSearchProvider::Standalone)
.map(Some)
.ok_or("Search MCP is already in use in this thread");
}
Ok(None)
}
// Responses can execute its built-in search before the harness receives the
// first search event. Reserve it while that request can still issue a call.
pub(crate) fn reserve_server_fallback(
step: &ExtensionData,
thread: &ExtensionData,
tools: &mut Arc<[ToolSpec]>,
) -> Option<Arc<Self>> {
let is_search = |tool: &ToolSpec| matches!(tool, ToolSpec::WebSearch { .. });
if !tools.iter().any(is_search) {
return None;
}
match Self::reserve_native_fallback(step, thread) {
Ok(owner) => owner,
Err(_) => {
*tools = tools
.iter()
.filter(|tool| !is_search(tool))
.cloned()
.collect();
None
}
}
}
fn current(store: &ExtensionData) -> Option<WebSearchProvider> {
store.get::<Self>().map(|owner| owner.0)
}
pub(crate) fn restore(store: &ExtensionData, history: &InitialHistory) {
let used = |item: &ResponseItemEnvelope| {
item.metadata
.as_ref()
.is_some_and(|meta| meta.native_search_fallback_used)
};
if history.scan_rollout_items(|item| match item {
RolloutItem::ResponseItem(item) => used(item),
RolloutItem::Compacted(item) => item.replacement_history.iter().flatten().any(used),
_ => false,
}) {
Self::admit(store, WebSearchProvider::Standalone);
}
}
pub(crate) fn checkpoint(store: &ExtensionData, items: &mut [ResponseItemEnvelope]) {
if Self::current(store) == Some(WebSearchProvider::Standalone) {
for item in items {
item.metadata
.get_or_insert_default()
.native_search_fallback_used = true;
}
}
}
}
fn object(value: Option<&Value>) -> Option<Map<String, Value>> {
match value {
None => Some(Map::new()),
Some(Value::Object(map)) => Some(map.clone()),
Some(Value::String(encoded)) => serde_json::from_str(encoded).ok(),
_ => None,
}
}
impl crate::config::Config {
// Read the scalar override from the highest enabled configuration layer. Hosts can
// supply it through normal session flags without extending the shared Config struct.
pub(crate) fn client_search_provider(&self) -> Option<WebSearchProvider> {
self.config_layer_stack
.layers_high_to_low()
.find_map(|layer| layer.config.get("tools")?.get("client_search_provider"))?
.clone()
.try_into()
.ok()
}
pub(crate) fn requires_native_search(&self) -> bool {
let standalone = codex_features::Feature::StandaloneWebSearch;
if !self.features.enabled(standalone) {
return false;
}
let mut features = self.features.clone();
features.disable(standalone).is_err() || features.enabled(standalone)
}
}
fn requested_mode(context: &Map<String, Value>, default: WebSearchMode) -> Option<WebSearchMode> {
context.get("mode").map_or(Some(default), |value| {
serde_json::from_value(value.clone()).ok()
})
}
// These host metadata fields do not select connectors. Preserve the installation ID
// when narrowing discovery, and validate the search context separately for native.
fn unrestricted_connectors(request: &Map<String, Value>) -> bool {
request.keys().all(|key| {
matches!(
key.as_str(),
"openai/search_context" | "openai/codex_installation_id"
)
})
}
fn native_search_mode(turn: &TurnContext, request: &Map<String, Value>) -> Option<WebSearchMode> {
if !unrestricted_connectors(request) {
return None;
}
let context = object(request.get("openai/search_context"))?;
let configured = turn.config.web_search_mode.value();
let mut mode = configured.restrict_to(requested_mode(&context, configured)?);
// A host domain allowlist covers opens and clicks. Native filters only the query.
for (key, value) in &context {
match key.as_str() {
"mode" | "telemetry_attributes" | "model_supported_tools" => {}
"is_web_access_restricted" => {
if value.as_bool()? {
mode = mode.restrict_to(WebSearchMode::Cached);
}
}
"session_persistent" if value.as_bool()? || turn.config.ephemeral => {}
_ => return None,
}
}
turn.config.web_search_mode.can_set(&mode).ok()?;
(mode != WebSearchMode::Disabled).then_some(mode)
}
fn search_context(turn: &TurnContext, request: &Map<String, Value>) -> Option<Map<String, Value>> {
let native_mode = turn.config.web_search_mode.value();
if !turn.apps_enabled() {
return None;
}
let config = turn.config.web_search_config.as_ref();
if config.is_some_and(|config| {
config.user_location.is_some() || config.search_context_size.is_some()
}) {
return None;
}
let mut context = object(request.get("openai/search_context"))?;
let host_mode = requested_mode(&context, native_mode)?;
let mode = if native_mode == WebSearchMode::Disabled
&& turn
.config
.features
.enabled(codex_features::Feature::HostedSearchPolicy)
{
host_mode
} else {
native_mode.restrict_to(host_mode)
};
if mode == WebSearchMode::Disabled {
return None;
}
turn.config.web_search_mode.can_set(&mode).ok()?;
context.insert("mode".into(), json!(mode));
let restricted = context
.get("is_web_access_restricted")
.map_or(Some(false), Value::as_bool)?;
context.insert(
"is_web_access_restricted".into(),
json!(restricted || mode == WebSearchMode::Cached),
);
let persistent = context
.get("session_persistent")
.map_or(Some(true), Value::as_bool)?;
context.insert(
"session_persistent".into(),
json!(persistent && !turn.config.ephemeral),
);
if let Some(domains) = config
.and_then(|c| c.filters.as_ref()?.allowed_domains.as_ref())
.filter(|d| !d.is_empty())
{
// Distinct domain policies can contain wildcards. Don't guess at their intersection.
if context
.get("allowed_domains")
.is_some_and(|value| *value != json!(domains))
{
return None;
}
context.insert("allowed_domains".into(), json!(domains));
}
Some(context)
}
pub(crate) async fn resolve(
turn: &TurnContext,
settings: &ResolvedStepSettings,
base: &Arc<McpBinding>,
services: &SessionServices,
) -> Option<ModelSearch> {
let native_required = turn.config.requires_native_search();
let provider = if native_required {
WebSearchProvider::Standalone
} else if let Some(provider) = settings.model_info.web_search_provider {
provider
} else {
match turn.config.client_search_provider()? {
WebSearchProvider::SearchMcp => return Some(ModelSearch::ClientMcp),
WebSearchProvider::Standalone => WebSearchProvider::Standalone,
}
};
let observe = |decision, reason| {
turn.session_telemetry.counter(
"codex.model_search.discovery",
/*inc*/ 1,
&[("reason", reason)],
);
Some(decision)
};
let mut turn_meta = turn
.turn_metadata_state
.current_meta_value_for_mcp_request(ExecutionMetadata::from_settings(settings))
.unwrap_or_else(|| json!({}));
let Some(mut request) = object(turn_meta.get("mcp_request_meta")) else {
return observe(ModelSearch::Disabled, "invalid_request_metadata");
};
let native_mode = native_search_mode(turn, &request);
if native_required {
return match native_mode {
Some(mode) => observe(
ModelSearch::Native {
mode,
fallback: false,
},
"native_required",
),
None => observe(ModelSearch::Disabled, "fallback_restricted"),
};
}
if provider == WebSearchProvider::Standalone {
return Some(
native_mode.map_or(ModelSearch::Disabled, |mode| ModelSearch::Native {
mode,
fallback: false,
}),
);
}
if !crate::mcp_tool_exposure::filter_codex_apps_mcp_tools(base.tools(), &turn.config)
.any(codex_mcp::is_model_search_tool)
{
return observe(ModelSearch::Disabled, "connector_unavailable");
}
let usage = turn.extension_data.get_or_init(SearchTurnUsage::default);
let Some(mut context) = search_context(turn, &request) else {
return observe(ModelSearch::Disabled, "search_policy");
};
context.insert(
"model_supported_tools".into(),
json!(settings.model_info.experimental_supported_tools),
);
// Scoped discovery may narrow an unrestricted catalog. Keep existing allowlists,
// denylists, plugin selection, and all other host metadata intact.
if unrestricted_connectors(&request) {
request.insert(
"selected_connector_ids".into(),
json!([codex_mcp::SEARCH_CONNECTOR]),
);
}
request.insert("openai/search_context".into(), json!(&context));
turn_meta["mcp_request_meta"] = json!(&request);
// Only discovery is stable within the turn. A yielded search can finish on a later step.
let cached = turn
.extension_data
.get::<CachedSearch>()
.filter(|cached| Arc::ptr_eq(&cached.base, base) && cached.request == request);
let fresh = cached.is_none();
let cached = match cached {
Some(cached) => cached,
None => {
let cached = Arc::new(CachedSearch {
base: Arc::clone(base),
request,
discovery: base.discover_model_search(turn_meta).await,
});
turn.extension_data.insert_arc(Arc::clone(&cached));
cached
}
};
let finish = |decision, reason| {
if fresh {
observe(decision, reason)
} else {
Some(decision)
}
};
let owner = SearchOwner::current(&services.thread_extension_data);
let reason = match &cached.discovery {
SearchDiscovery::Ready(binding) => {
if owner != Some(WebSearchProvider::Standalone) {
return finish(
ModelSearch::Hosted {
binding: Arc::clone(binding),
context,
},
"hosted",
);
}
"native_browsing_state"
}
SearchDiscovery::Critical(reason) => *reason,
SearchDiscovery::Unavailable => return finish(ModelSearch::Disabled, "unavailable"),
};
let attribution = services.executed_tool_calls.mcp_attribution_snapshot();
if let Some(mode) = native_mode
&& owner != Some(WebSearchProvider::SearchMcp)
&& usage.0.load(Ordering::Relaxed) & SearchTurnUsage::MCP == 0
&& attribution.status != McpAttributionStatus::AttributionError
&& !attribution
.sources
.iter()
.any(|source| source.connector_id.as_deref() == Some(codex_mcp::SEARCH_CONNECTOR))
{
finish(
ModelSearch::Native {
mode,
fallback: true,
},
reason,
)
} else {
finish(ModelSearch::Disabled, "fallback_restricted")
}
}
/// Preserve all other per-call metadata while replacing the policy with the one used for discovery.
pub(crate) fn with_search_context(
request: &mut Value,
context: &Map<String, Value>,
) -> anyhow::Result<()> {
let envelope = request
.as_object_mut()
.ok_or_else(|| anyhow::anyhow!("invalid search turn metadata"))?;
let mut meta = object(envelope.get("mcp_request_meta"))
.ok_or_else(|| anyhow::anyhow!("invalid MCP request metadata"))?;
meta.insert("openai/search_context".into(), json!(context));
envelope.insert("mcp_request_meta".into(), Value::Object(meta));
Ok(())
}
+62 -3
View File
@@ -2870,6 +2870,24 @@ impl Session {
}
pub(crate) async fn emit_turn_item_started(&self, turn_context: &TurnContext, item: &TurnItem) {
if matches!(
item,
TurnItem::WebSearch(_)
| TurnItem::Extension(codex_extension_items::ExtensionItem::WebSearch(_))
) && let Some(usage) = turn_context
.extension_data
.get::<crate::model_search::SearchTurnUsage>()
{
// This turn selected Search MCP. Pin native only when a fallback call starts.
crate::model_search::SearchOwner::admit(
&self.services.thread_extension_data,
codex_protocol::openai_models::WebSearchProvider::Standalone,
);
usage.record(
turn_context,
codex_protocol::openai_models::WebSearchProvider::Standalone,
);
}
let started_at_ms = turn_context
.turn_timing_state
.record_item_started(item.id(), now_unix_timestamp_ms())
@@ -3901,6 +3919,10 @@ impl Session {
// Preserve any existing tool-specific override.
let policy: codex_utils_output_truncation::TruncationPolicy =
model_info.truncation_policy.into();
crate::model_search::SearchOwner::checkpoint(
&self.services.thread_extension_data,
&mut items,
);
for envelope in &mut items {
if matches!(
envelope.item,
@@ -4231,8 +4253,6 @@ impl Session {
let extension_data =
codex_extension_api::ExtensionData::new(turn_context.sub_id.clone());
extension_data.insert_arc(Arc::clone(&settings.model_info.model));
// Standalone and hosted search must use the same turn-scoped permission mode.
extension_data.insert(turn_context.config.web_search_mode.value());
if let Some(messages) = settings
.model_info
.model_messages
@@ -4249,16 +4269,51 @@ impl Session {
if let Some(messages) = model_messages.notes() {
extension_data.insert(messages.clone());
}
let mut required_servers = required_servers.to_vec();
if settings.model_info.web_search_provider
== Some(codex_protocol::openai_models::WebSearchProvider::SearchMcp)
&& turn_context.apps_enabled()
&& !turn_context.config.requires_native_search()
&& (turn_context.config.web_search_mode.value()
!= codex_protocol::config_types::WebSearchMode::Disabled
|| turn_context
.config
.features
.enabled(codex_features::Feature::HostedSearchPolicy))
{
required_servers.push(codex_mcp::CODEX_APPS_MCP_SERVER_NAME.to_owned());
}
let (mcp, prepared_recommendations) = tokio::join!(
// MCP refresh can be large; keep it off the sampling request's stack.
Box::pin(self.mcp_runtime_for_step(
turn_context.as_ref(),
&selected_capability_roots,
required_servers,
&required_servers,
required_plugins,
)),
turn::prepare_tool_recommendations(self.as_ref(), turn_context.as_ref()),
);
let search = crate::model_search::resolve(
turn_context.as_ref(),
&settings,
&mcp,
&self.services,
)
.await;
extension_data.insert(match &search {
Some(crate::model_search::ModelSearch::Native { mode, .. }) => *mode,
_ => turn_context.config.web_search_mode.value(),
});
let mcp = if let Some(search) = search {
let binding = match &search {
crate::model_search::ModelSearch::Hosted { binding, .. } => Arc::clone(binding),
_ => mcp,
};
extension_data.insert(search);
binding
} else {
mcp
};
extension_data.insert(selected_capability_roots);
// A step keeps the plugins from the environments it captured, even if shared MCP
// moves on to another environment. Its skill tools and the model use this same copy.
@@ -4527,6 +4582,10 @@ impl Session {
CompactedHistoryInstallation::Complete(metadata) => &metadata.compaction_model_hash,
CompactedHistoryInstallation::Prefix { model_info, .. } => &model_info.comp_hash,
};
crate::model_search::SearchOwner::checkpoint(
&self.services.thread_extension_data,
&mut items,
);
let mcp_revision = self
.services
.executed_tool_calls
+1
View File
@@ -1091,6 +1091,7 @@ impl Session {
thread_extension_init,
);
thread_extension_data.insert(crate::context::GuardianReviewEvidence::default());
crate::model_search::SearchOwner::restore(&thread_extension_data, &initial_history);
// Kick off independent async setup tasks in parallel to reduce startup latency.
//
// - initialize thread persistence with new or resumed session info
+5
View File
@@ -1872,6 +1872,11 @@ async fn run_sampling_request(
if first_attempt && turn_context.config.notify.is_some() {
notify_messages = legacy_notify_input_messages(&prompt.input);
}
let _native_search = crate::model_search::SearchOwner::reserve_server_fallback(
&step_context.extension_data,
&sess.services.thread_extension_data,
&mut prompt.tools,
);
let mut draining_for_mailbox = false;
let err = match try_run_sampling_request(
tool_runtime.clone(),
@@ -66,6 +66,17 @@ impl ToolExecutor<ToolInvocation> for ExtensionToolAdapter {
{
Box::pin(async move {
let step = &invocation.step_context;
let _native_search = if invocation.tool_name.namespace.as_deref() == Some("web")
&& invocation.tool_name.name == "run"
{
crate::model_search::SearchOwner::reserve_native_fallback(
&step.extension_data,
&invocation.session.services.thread_extension_data,
)
.map_err(|err| codex_tools::FunctionCallError::RespondToModel(err.to_string()))?
} else {
None
};
let selected: Vec<_> = match &invocation.environments {
ToolEnvironments::Arguments(environments) => {
environments.iter().map(|(_, env)| env.as_ref()).collect()
+37 -6
View File
@@ -159,6 +159,7 @@ pub(crate) fn build_tool_router(
session.services.mcp_manager.shared_definitions(),
&mut registry,
);
let search = step_store.get::<crate::model_search::ModelSearch>();
apply_mcp_tool_exposure_policy(
turn_context,
model_info,
@@ -172,12 +173,40 @@ pub(crate) fn build_tool_router(
extension_tool_executors(session, step_store),
&mut registry,
);
let native_mode = match search.as_deref() {
None => Some(turn_context.config.web_search_mode.value()),
Some(crate::model_search::ModelSearch::Native { mode, .. }) => Some(*mode),
_ => None,
};
if matches!(
search.as_deref(),
Some(
crate::model_search::ModelSearch::Native { .. }
| crate::model_search::ModelSearch::Disabled
)
) {
for tool in mcp
.tools()
.iter()
.filter(|tool| codex_mcp::is_model_search_tool(tool))
{
registry.remove(&tool.canonical_tool_name());
}
}
if native_mode.is_none() {
registry.remove(&ToolName::namespaced("web", "run"));
}
append_dynamic_tool_runtimes(dynamic_tools, &mut registry);
let hosted_specs = hosted_model_tool_specs(
turn_context,
model_info,
standalone_web_search_tool.as_slice(),
);
let hosted_specs = if let Some(mode) = native_mode {
hosted_model_tool_specs(
turn_context,
model_info,
mode,
standalone_web_search_tool.as_slice(),
)
} else {
Vec::new()
};
finalize_tool_router(
turn_context,
@@ -324,6 +353,7 @@ pub(crate) fn append_source_tools(
hosted_model_tool_specs(
turn_context,
model_info,
turn_context.config.web_search_mode.value(),
standalone_web_search_tool.as_slice(),
)
}
@@ -647,6 +677,7 @@ fn spec_for_model_request(
fn hosted_model_tool_specs(
turn_context: &TurnContext,
model_info: &ModelInfo,
mode: WebSearchMode,
registered_extension_tool_names: &[ToolName],
) -> Vec<ToolSpec> {
// Responses Lite accepts schemas for client-executed tools, not hosted Responses tools.
@@ -661,7 +692,7 @@ fn hosted_model_tool_specs(
// and the provider supports hosted search. `None` prevents emitting a hosted search tool.
let web_search_mode = (!standalone_web_search_available
&& turn_context.provider.capabilities().web_search)
.then_some(turn_context.config.web_search_mode.value());
.then_some(mode);
let web_search_config = web_search_mode
.as_ref()
.and(turn_context.config.web_search_config.as_ref());
+9 -1
View File
@@ -1142,7 +1142,12 @@ async fn environment_tools_follow_the_step_context() {
/*tool_suggest_candidates*/ None,
/*wait_for_environment_tool_config*/ None,
),
super::hosted_model_tool_specs(turn.as_ref(), turn.model_info(), &[]),
super::hosted_model_tool_specs(
turn.as_ref(),
turn.model_info(),
turn.config.web_search_mode.value(),
&[],
),
&Default::default(),
));
@@ -2273,6 +2278,7 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() {
super::hosted_model_tool_specs(
first_step_context.turn.as_ref(),
first_step_context.turn.model_info(),
first_step_context.turn.config.web_search_mode.value(),
&[],
),
&cache,
@@ -2307,6 +2313,7 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() {
super::hosted_model_tool_specs(
second_step_context.turn.as_ref(),
second_step_context.turn.model_info(),
second_step_context.turn.config.web_search_mode.value(),
&[],
),
&cache,
@@ -2372,6 +2379,7 @@ async fn tool_search_cache_rebuilds_when_deferred_world_state_changes() {
super::hosted_model_tool_specs(
step_context.turn.as_ref(),
step_context.turn.model_info(),
step_context.turn.config.web_search_mode.value(),
&[],
),
&cache,
@@ -1,8 +1,19 @@
//! Model-visible context for standalone web search.
use super::*;
use anyhow::Context;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::openai_models::WebSearchProvider;
use core_test_support::apps_test_server::AppsTestServer;
use core_test_support::apps_test_server::apps_enabled_builder;
use core_test_support::responses::ev_function_call;
use core_test_support::responses::ev_tool_search_call;
use core_test_support::test_codex::TestCodexBuilder;
use pretty_assertions::assert_eq;
use serde_json::Value;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use wiremock::Mock;
use wiremock::ResponseTemplate;
use wiremock::matchers::body_partial_json;
@@ -98,3 +109,545 @@ async fn literature_search_and_open() -> Result<()> {
));
Ok(())
}
fn literature_search_pin() -> Value {
json!({"config":{"tool_type":"SonicBrowserToolV99","allow_literature_search":true}})
}
async fn mount_model_search(
server: &wiremock::MockServer,
available: &Arc<AtomicBool>,
) -> Result<AppsTestServer> {
let mut tool = json!({
"name": "search_service.web_run", "description": "Find papers on the web.",
"annotations": {"readOnlyHint": true},
"inputSchema": {"type":"object","properties":{"system2_search_query":{"type":"array"},"open":{"type":"array"}}},
"_meta": {"connector_id": codex_mcp::SEARCH_CONNECTOR,"connector_name":"search_service"}
});
let apps =
AppsTestServer::mount_with_tools(server, Arc::new(Mutex::new(vec![tool.clone()]))).await?;
tool["inputSchema"]["properties"]["literature_search_query"] = json!({"type":"array"});
let pin = literature_search_pin();
tool["_meta"]["_codex_apps"] = json!({"invocation_context":pin});
// Host-owned Apps catalogs can exceed the generic MCP limit of 2,048 tools.
let mut catalog: Vec<_> = (0..2_048)
.map(|n| json!({"name":format!("other_{n}"),"inputSchema":{"type":"object"}}))
.collect();
catalog.push(tool);
let available = Arc::clone(available);
Mock::given(method("POST"))
.and(body_partial_json(json!({
"method":"tools/list","params":{"_meta":{"x-codex-turn-metadata":{"mcp_request_meta":{
"selected_connector_ids":[codex_mcp::SEARCH_CONNECTOR]
}}}}
})))
.respond_with(move |request: &wiremock::Request| {
let body: serde_json::Value = request.body_json().unwrap();
let mut response = json!({"jsonrpc":"2.0","id":body["id"]});
if available.load(Ordering::SeqCst) {
response["result"] = json!({"tools":catalog});
} else {
response["error"] = json!({"code":-32603,"message":"discovery unavailable"});
}
ResponseTemplate::new(200).set_body_json(response)
})
.with_priority(1)
.expect(2..)
.mount(server)
.await;
Ok(apps)
}
fn model_search_builder(apps: &AppsTestServer, tool_mode: ToolMode) -> TestCodexBuilder {
let auth = codex_core::test_support::auth_manager_from_auth(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
);
let mut extensions = ExtensionRegistryBuilder::new();
codex_web_search_extension::install(&mut extensions, Arc::clone(&auth));
apps_enabled_builder(apps.chatgpt_base_url.clone())
.with_auth_manager(auth)
.with_extensions(Arc::new(extensions.build()))
.with_config(move |config| {
configure_scenario_catalog(config);
if tool_mode == ToolMode::CodeMode {
config.features.enable(Feature::CodeMode).unwrap();
}
config
.features
.enable(Feature::StandaloneWebSearch)
.unwrap();
config.web_search_mode.set(WebSearchMode::Live).unwrap();
})
.with_model_info_override("native-paper", |model| {
model.web_search_provider = Some(WebSearchProvider::Standalone);
model.tool_mode = Some(ToolMode::Direct);
})
.with_model_info_override("hosted-paper", move |model| {
model.web_search_provider = Some(WebSearchProvider::SearchMcp);
model.tool_mode = Some(tool_mode);
if tool_mode == ToolMode::CodeMode {
model
.experimental_supported_tools
.push("test_sync_tool".into());
}
model.supports_search_tool = true;
model
.experimental_supported_tools
.push("literature_search_query".into());
})
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn model_selected_search_preserves_discovery_and_call_history() -> Result<()> {
use codex_protocol::protocol::ThreadSettingsOverrides;
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let discovery_available = Arc::new(AtomicBool::new(true));
let apps = mount_model_search(&server, &discovery_available).await?;
let pin = literature_search_pin();
let search = json!({"literature_search_query":[{"q":"FlashAttention-2"}]});
let reference_issued = Arc::new(AtomicBool::new(false));
let issued = Arc::clone(&reference_issued);
Mock::given(method("POST"))
.and(body_partial_json(json!({
"method":"tools/call","params":{"name":"search_service.web_run","arguments":search,
"_meta":{"_codex_apps":{"invocation_context":pin}}}
})))
.respond_with(move |request: &wiremock::Request| {
let body: serde_json::Value = request.body_json().unwrap();
issued.store(true, Ordering::SeqCst);
ResponseTemplate::new(200).set_body_json(json!({"jsonrpc":"2.0","id":body["id"],
"result":{"content":[{"type":"text","text":"[turn0academia0] FlashAttention-2 by Tri Dao."}]}}))
})
.with_priority(1)
.expect(1)
.mount(&server)
.await;
let open = json!({"open":[{"ref_id":"turn0academia0"}]});
Mock::given(method("POST"))
.and(body_partial_json(json!({
"method":"tools/call","params":{"name":"search_service.web_run","arguments":open,
"_meta":{"_codex_apps":{"invocation_context":pin}}}
})))
.respond_with(move |request: &wiremock::Request| {
let body: serde_json::Value = request.body_json().unwrap();
let text = if reference_issued.load(Ordering::SeqCst) {
"The paper reports a 2x speedup."
} else {
"Unknown reference."
};
ResponseTemplate::new(200).set_body_json(json!({"jsonrpc":"2.0","id":body["id"],
"result":{"content":[{"type":"text","text":text}]}}))
})
.with_priority(1)
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/v1/alpha/search"))
.respond_with(ResponseTemplate::new(400))
.expect(0)
.mount(&server)
.await;
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_tool_search_call("find-search", &json!({"query":"search_service web_run"})),
ev_completed("lookup"),
]),
sse(vec![
ev_function_call_with_namespace(
"paper-search",
"mcp__codex_apps__search_service",
"_web_run",
&search.to_string(),
),
ev_completed("search"),
]),
sse(vec![
ev_assistant_message("author", "The paper is by Tri Dao."),
ev_completed("first"),
]),
sse(vec![
ev_assistant_message("unavailable", "Search is temporarily unavailable."),
ev_completed("outage"),
]),
sse(vec![
ev_tool_search_call("find-open", &json!({"query":"search_service web_run"})),
ev_completed("lookup-open"),
]),
sse(vec![
ev_function_call_with_namespace(
"open",
"mcp__codex_apps__search_service",
"_web_run",
&open.to_string(),
),
ev_completed("opened"),
]),
sse(vec![
ev_assistant_message("speed", "It reports a 2x speedup."),
ev_completed("second"),
]),
sse(vec![
ev_assistant_message("recap", "Tri Dao reports a 2x speedup."),
ev_completed("third"),
]),
],
)
.await;
let test = model_search_builder(&apps, ToolMode::Direct)
.build_with_auto_env(&server)
.await?;
test.submit_turn("Find the FlashAttention-2 paper.").await?;
discovery_available.store(false, Ordering::SeqCst);
test.submit_turn("Open the paper and check its speedup.")
.await?;
discovery_available.store(true, Ordering::SeqCst);
test.submit_turn("Try opening the same paper again.")
.await?;
test.codex
.start_or_steer_turn(
TurnInputRequest::user_input(vec![text("Summarize the result.")]).with_thread_settings(
ThreadSettingsOverrides {
model: Some("native-paper".into()),
..Default::default()
},
),
)
.await?;
wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let requests = responses.requests();
assert_eq!(requests.len(), 8);
assert!(
requests[0]
.tool_by_name("mcp__codex_apps__search_service", "_web_run")
.is_none()
);
assert!(
requests[1]
.tool_search_output("find-search")
.to_string()
.contains("literature_search_query")
);
let output = requests[2].function_call_output("paper-search");
assert!(output.to_string().contains("Tri Dao"), "{output}");
assert!(requests[3].tool_by_name("web", "run").is_none());
let output = requests[6].function_call_output("open");
assert!(output.to_string().contains("2x"), "{output}");
insta::assert_snapshot!(context_snapshot::format_request_history_snapshot(
"MCP returns a paper reference. A discovery outage withholds native search, then MCP opens the same reference after recovery. Switching to a standalone model retains the history.",
&requests,
&ContextSnapshotOptions::default()
.rewrite_known_segments()
.include_request_settings(),
));
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn native_fallback_reference_survives_discovery_recovery_and_resume() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let available = Arc::new(AtomicBool::new(false));
let apps = mount_model_search(&server, &available).await?;
let mut test = model_search_builder(&apps, ToolMode::Direct)
.build_with_auto_env(&server)
.await?;
// Merely offering fallback during an outage must not pin a healthy thread to native.
let idle = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_assistant_message("hello", "Hello."),
ev_completed("hello"),
]),
sse(vec![
ev_tool_search_call("lookup", &json!({"query":"search_service web_run"})),
ev_completed("lookup"),
]),
sse(vec![
ev_assistant_message("ready", "Search is ready."),
ev_completed("ready"),
]),
],
)
.await;
test.submit_turn("Say hello.").await?;
available.store(true, Ordering::SeqCst);
test.submit_turn("Look up the search tool.").await?;
let requests = idle.requests();
assert!(requests[0].tool_by_name("web", "run").is_some());
assert!(
requests[2]
.tool_search_output("lookup")
.to_string()
.contains("literature_search_query")
);
let search = json!({"search_query":[{"q":"FlashAttention-2"}]});
let open = json!({"open":[{"ref_id":"turn0search0"}]});
let mut browser_session = None;
for (id, commands, output) in [
("search", search, "[turn0search0] FlashAttention-2"),
("open", open.clone(), "Tri Dao reports a 2x speedup."),
("reopen", open, "Tri Dao reports a 2x speedup."),
] {
let first_turn = id == "search";
available.store(!first_turn, Ordering::SeqCst);
if id == "reopen" {
test = model_search_builder(&apps, ToolMode::Direct)
.restart_with_auto_env(&server, &test)
.await?;
}
let native = Mock::given(method("POST"))
.and(path("/v1/alpha/search"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"output":output})))
.expect(1)
.mount_as_scoped(&server)
.await;
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_function_call_with_namespace(id, "web", "run", &commands.to_string()),
ev_completed("tool"),
]),
sse(vec![
ev_assistant_message("answer", output),
ev_completed("answer"),
]),
],
)
.await;
test.submit_turn(if first_turn {
"Find the paper."
} else {
"Open the paper you found."
})
.await?;
let requests = responses.requests();
assert!(requests[0].tool_by_name("web", "run").is_some());
let result = requests[1].function_call_output(id);
assert!(result.to_string().contains(output), "{result}");
let body: Value = native.received_requests().await[0].body_json()?;
assert_eq!(body["commands"], commands);
let session_id = body["id"].as_str().context("native search session ID")?;
assert_eq!(
browser_session
.get_or_insert_with(|| session_id.to_owned())
.as_str(),
session_id
);
}
assert!(
core_test_support::apps_test_server::recorded_apps_tool_calls(&server)
.await
.is_empty()
);
Ok(())
}
#[test_case::test_case(false; "unused_offer_releases_provider")]
#[test_case::test_case(true; "executed_call_keeps_provider")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_native_fallback(used: bool) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let available = Arc::new(AtomicBool::new(false));
let apps = mount_model_search(&server, &available).await?;
let test = model_search_builder(&apps, ToolMode::Direct)
.with_config(|config| {
config
.features
.disable(Feature::StandaloneWebSearch)
.unwrap();
})
.with_model_info_override("hosted-paper", |model| model.use_responses_lite = false)
.build_with_auto_env(&server)
.await?;
let mut events = Vec::new();
if used {
events.push(
core_test_support::responses::ev_web_search_call_added_partial("native", "in_progress"),
);
events.push(core_test_support::responses::ev_web_search_call_done(
"native",
"completed",
"FlashAttention-2",
));
}
events.extend([
ev_assistant_message("first", "Ready."),
ev_completed("first"),
]);
let responses = mount_sse_sequence(
&server,
vec![
sse(events),
sse(vec![
ev_assistant_message("second", "Ready."),
ev_completed("second"),
]),
],
)
.await;
test.submit_turn("Look for the paper.").await?;
available.store(true, Ordering::SeqCst);
test.submit_turn("Continue.").await?;
let native_offers: Vec<_> = responses
.requests()
.iter()
.map(|request| {
request.body_json()["tools"]
.as_array()
.unwrap()
.iter()
.any(|tool| tool["type"] == "web_search")
})
.collect();
assert_eq!(native_offers, vec![true, used]);
Ok(())
}
#[test_case::test_case(false; "mcp_finishes_before_native_is_used")]
#[test_case::test_case(true; "native_runs_before_retained_mcp_call")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn background_search_and_native_fallback(native_first: bool) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let available = Arc::new(AtomicBool::new(true));
let apps = mount_model_search(&server, &available).await?;
let test = model_search_builder(&apps, ToolMode::CodeMode)
.build_with_auto_env(&server)
.await?;
Mock::given(method("POST"))
.and(body_partial_json(json!({"method":"tools/call","params":{"name":"search_service.web_run","_meta":{"_codex_apps":{"invocation_context":literature_search_pin()}}}})))
.respond_with(|request: &wiremock::Request| {
let body: Value = request.body_json().unwrap();
ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc":"2.0", "id":body["id"], "result":{"content":[{"type":"text","text":"[turn0academia0] FlashAttention-2"}]}
}))
})
.with_priority(1).expect(u64::from(!native_first)).mount(&server).await;
// The barrier holds the original step's MCP call until the next turn has offered native.
let barrier = json!({"barrier":{"id":"background-search","participants":2,"timeout_ms":60000}});
let pending = format!(
r#"yield_control();
await tools.test_sync_tool({barrier});
text((await tools.mcp__codex_apps__search_service_web_run({{literature_search_query:[{{q:"FlashAttention-2"}}]}})).content);"#
);
let first = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_custom_tool_call("pending", "exec", &pending),
ev_completed("yielded"),
]),
sse(vec![
ev_assistant_message("message", "Searching."),
ev_completed("first"),
]),
],
)
.await;
test.submit_turn("Start a background paper search.").await?;
let yielded = first.requests()[1].custom_tool_call_output("pending");
let running = yielded["output"]
.as_str()
.or_else(|| yielded.pointer("/output/0/text")?.as_str())
.expect("yielded exec output");
let cell_id = running
.lines()
.next()
.unwrap()
.strip_prefix("Script running with cell ID ")
.expect("yielded exec cell");
available.store(false, Ordering::SeqCst);
let native_commands = json!({"search_query":[{"q":"FlashAttention-2"}]});
Mock::given(method("POST"))
.and(path("/v1/alpha/search"))
.and(body_partial_json(json!({"commands":native_commands})))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(json!({"output":"[turn0search0] native paper"})),
)
.expect(u64::from(native_first))
.mount(&server)
.await;
let mut steps = Vec::new();
if native_first {
steps.push(sse(vec![
ev_function_call_with_namespace("native", "web", "run", &native_commands.to_string()),
ev_completed("native"),
]));
}
steps.extend([
sse(vec![
ev_function_call("release", "test_sync_tool", &barrier.to_string()),
ev_completed("released"),
]),
sse(vec![
ev_function_call_with_namespace(
"wait-search",
"functions",
"wait",
&json!({"cell_id":cell_id,"yield_time_ms":10000}).to_string(),
),
ev_completed("waited"),
]),
sse(vec![
ev_assistant_message("answer", "The background result is available."),
ev_completed("done"),
]),
]);
let second = mount_sse_sequence(&server, steps).await;
test.submit_turn("Wait for the paper.").await?;
let requests = second.requests();
assert!(
requests[0].tool_by_name("web", "run").is_some(),
"outage initially allows native"
);
let last = requests.last().unwrap();
let result = last.function_call_output("wait-search").to_string();
if native_first {
assert!(
requests[1]
.function_call_output("native")
.to_string()
.contains("turn0search0")
);
assert!(result.contains("native search"), "{result}");
assert!(
last.tool_by_name("web", "run").is_some(),
"the skipped MCP call must not change ownership"
);
available.store(true, Ordering::SeqCst);
let recovered = mount_sse_sequence(
&server,
vec![sse(vec![
ev_assistant_message("answer", "Ready."),
ev_completed("recovered"),
])],
)
.await;
test.submit_turn("Continue using the same paper.").await?;
assert!(
recovered
.single_request()
.tool_by_name("web", "run")
.is_some()
);
} else {
assert!(result.contains("turn0academia0"), "{result}");
assert!(
last.tool_by_name("web", "run").is_none(),
"MCP owns the reference"
);
}
Ok(())
}
@@ -0,0 +1,137 @@
---
source: core/tests/suite/scenarios_web_search.rs
expression: "context_snapshot::format_request_history_snapshot(\"MCP returns a paper reference. A discovery outage withholds native search, then MCP opens the same reference after recovery. Switching to a standalone model retains the history.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)"
---
Scenario: MCP returns a paper reference. A discovery outage withholds native search, then MCP opens the same reference after recovery. Switching to a standalone model retains the history.
## Window 1
Settings:
include: ["reasoning.encrypted_content"]
model: "hosted-paper"
parallel_tool_calls: true
prompt_cache_key: "<PROMPT_CACHE_KEY 1>"
reasoning: {"effort":"medium"}
store: false
stream: true
text: {"verbosity":"low"}
tool_choice: "auto"
tools (9; hash=5429df93e52f82a6):
- function/exec_command: Runs a command in a PTY, returning output or a session ID for ongoing interaction.; args=[cmd, justification, login, max_output_tokens, prefix_rule, sandbox_permissions, shell, tty, workdir, yield_time_ms]; hash=c04108e51a8548c5
- function/write_stdin: Writes characters to an existing unified exec session and returns recent output.; args=[chars, max_output_tokens, session_id, yield_time_ms]; hash=7fb62d8a3e2df4e2
- function/list_mcp_resources: Lists resources provided by MCP servers. Resources allow servers to share data that provides cont...; args=[cursor, server]; hash=6792edaff13561bd
- function/list_mcp_resource_templates: Lists resource templates provided by MCP servers. Parameterized resource templates allow servers ...; args=[cursor, server]; hash=60f17c90ef7a4ebb
- function/read_mcp_resource: Read a specific resource from an MCP server given the server name and resource URI.; args=[server, uri]; hash=24246896b39008ae
- function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6
- custom/apply_patch: The `apply_patch` tool can be used to edit files. This is a FREEFORM tool, so do not wrap the pat...; hash=dea9f12b0111a4ba
- function/view_image: View a local image file from the filesystem when visual inspection is needed. Use this for images...; args=[detail, path]; hash=3b766d5bdec5d50e
- tool_search: # Tool discovery Searches over deferred tool metadata with BM25 and exposes matching tools for t...; args=[limit, query]; hash=b573ed69c0b0e6c7
-- request 1 (turn) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/developer[3]:
[01] <PERMISSIONS_INSTRUCTIONS>
[02] <APPS_INSTRUCTIONS>
[03] <tools>
Deferred tool namespaces:
- mcp__codex_apps__search_service: Tools for working with search_service.
- multi_agent_v1: Tools for spawning and managing sub-agents.
</tools>
02:message/user:
<ENVIRONMENT_CONTEXT>
03:message/user:
Find the FlashAttention-2 paper.
-- request 2 (turn) --
04:tool_search_call
05:tool_search_output
-- request 3 (turn) --
06:function_call/mcp__codex_apps__search_service._web_run:{"literature_search_query":[{"q":"FlashAttention-2"}]}
07:function_call_output:Wall time: <DURATION> seconds
Output: | [turn0academia0] FlashAttention-2 by Tri Dao.
-- request 4 (turn) --
08:message/assistant:
The paper is by Tri Dao.
09:message/developer:
<tools>
Removed deferred tool namespaces:
- mcp__codex_apps__search_service: Tools for working with search_service.
</tools>
10:message/user:
Open the paper and check its speedup.
-- request 5 (turn) --
11:message/assistant:
Search is temporarily unavailable.
12:message/developer:
<tools>
Added deferred tool namespaces:
- mcp__codex_apps__search_service: Tools for working with search_service.
</tools>
13:message/user:
Try opening the same paper again.
-- request 6 (turn) --
14:tool_search_call
15:tool_search_output
-- request 7 (turn) --
16:function_call/mcp__codex_apps__search_service._web_run:{"open":[{"ref_id":"turn0academia0"}]}
17:function_call_output:Wall time: <DURATION> seconds
Output: | The paper reports a 2x speedup.
## Window 2 (after request 7)
Settings: relative to window 1
model: "native-paper"
tools (10; hash=89fc167132311088):
+ namespace/web: Tools in the web namespace.; hash=0037b1ed8d50fab5
- function/run
-- request 8 (turn) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/developer[3]:
[01] <PERMISSIONS_INSTRUCTIONS>
[02] <APPS_INSTRUCTIONS>
[03] <tools>
Deferred tool namespaces:
- mcp__codex_apps__search_service: Tools for working with search_service.
- multi_agent_v1: Tools for spawning and managing sub-agents.
</tools>
02:message/user:
<ENVIRONMENT_CONTEXT>
03:message/user:
Find the FlashAttention-2 paper.
04:tool_search_call
05:tool_search_output
06:function_call/mcp__codex_apps__search_service._web_run:{"literature_search_query":[{"q":"FlashAttention-2"}]}
07:function_call_output:Wall time: <DURATION> seconds
Output: | [turn0academia0] FlashAttention-2 by Tri Dao.
08:message/assistant:
The paper is by Tri Dao.
09:message/developer:
<tools>
Removed deferred tool namespaces:
- mcp__codex_apps__search_service: Tools for working with search_service.
</tools>
10:message/user:
Open the paper and check its speedup.
11:message/assistant:
Search is temporarily unavailable.
12:message/developer:
<tools>
Added deferred tool namespaces:
- mcp__codex_apps__search_service: Tools for working with search_service.
</tools>
13:message/user:
Try opening the same paper again.
14:tool_search_call
15:tool_search_output
16:function_call/mcp__codex_apps__search_service._web_run:{"open":[{"ref_id":"turn0academia0"}]}
17:function_call_output:Wall time: <DURATION> seconds
Output: | The paper reports a 2x speedup.
18:message/assistant:
It reports a 2x speedup.
19:message/developer:
<MODEL_SWITCH>
20:message/developer:
<tools>
Removed deferred tool namespaces:
- mcp__codex_apps__search_service: Tools for working with search_service.
</tools>
21:message/user:
Summarize the result.
+4
View File
@@ -134,6 +134,10 @@ pub struct CodexHarnessMetadata {
)]
pub mcp_attribution: Option<McpAttribution>,
/// Native search ran during automatic fallback, so its browsing references stay on native.
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub native_search_fallback_used: bool,
/// Sender context captured by the host when this task message was accepted.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sender_user_messages: Option<Box<SenderUserMessages>>,