mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
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:
BIN
Binary file not shown.
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
+137
@@ -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.
|
||||
@@ -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>>,
|
||||
|
||||
Reference in New Issue
Block a user