use std::sync::Arc;use codex_api::AllowedCaller;use codex_api::ApproximateLocation;use codex_api::LocationType;use codex_api::SearchContextSize;use codex_api::SearchFilters;use codex_api::SearchSettings;use codex_core::config::Config;use codex_extension_api::ConfigContributor;use codex_extension_api::ExtensionData;use codex_extension_api::ExtensionFuture;use codex_extension_api::ExtensionRegistryBuilder;use codex_extension_api::ThreadLifecycleContributor;use codex_extension_api::ThreadStartInput;use codex_extension_api::ToolContributor;use codex_login::AuthManager;use codex_model_provider::create_model_provider;use codex_model_provider_info::ModelProviderInfo;use codex_protocol::config_types::WebSearchContextSize;use codex_protocol::config_types::WebSearchMode;use crate::tool::WebSearchTool;#[derive(Clone)]struct WebSearchExtension {    auth_manager: Arc<AuthManager>,}#[derive(Clone)]struct WebSearchExtensionConfig {    available: bool,    provider: ModelProviderInfo,    settings: SearchSettings,}impl From<&Config> for WebSearchExtensionConfig {    fn from(config: &Config) -> Self {        let web_search_mode = config.web_search_mode.value();        Self {            // Core selects this executor per turn using the feature flag or model metadata.            available: config.model_provider.is_openai()                && web_search_mode != WebSearchMode::Disabled,            provider: config.model_provider.clone(),            settings: search_settings(config, web_search_mode),        }    }}fn search_settings(config: &Config, web_search_mode: WebSearchMode) -> SearchSettings {    let web_search_config = config.web_search_config.as_ref();    SearchSettings {        user_location: web_search_config            .and_then(|config| config.user_location.as_ref())            .map(|location| ApproximateLocation {                r#type: LocationType::Approximate,                country: location.country.clone(),                region: location.region.clone(),                city: location.city.clone(),                timezone: location.timezone.clone(),            }),        search_context_size: web_search_config            .and_then(|config| config.search_context_size)            .map(|size| match size {                WebSearchContextSize::Low => SearchContextSize::Low,                WebSearchContextSize::Medium => SearchContextSize::Medium,                WebSearchContextSize::High => SearchContextSize::High,            }),        filters: web_search_config            .and_then(|config| config.filters.as_ref())            .map(|filters| SearchFilters {                allowed_domains: filters.allowed_domains.clone(),                blocked_domains: None,            }),        allowed_callers: Some(vec![AllowedCaller::Direct]),        external_web_access: Some(match web_search_mode {            WebSearchMode::Live => true,            WebSearchMode::Cached | WebSearchMode::Disabled => false,        }),        ..Default::default()    }}impl ThreadLifecycleContributor<Config> for WebSearchExtension {    fn on_thread_start<'a>(        &'a self,        input: ThreadStartInput<'a, Config>,    ) -> ExtensionFuture<'a, ()> {        Box::pin(async move {            input                .thread_store                .insert(WebSearchExtensionConfig::from(input.config));        })    }}impl ConfigContributor<Config> for WebSearchExtension {    fn on_config_changed(        &self,        _session_store: &ExtensionData,        thread_store: &ExtensionData,        _previous_config: &Config,        new_config: &Config,    ) {        thread_store.insert(WebSearchExtensionConfig::from(new_config));    }}impl ToolContributor for WebSearchExtension {    fn tools(        &self,        session_store: &ExtensionData,        thread_store: &ExtensionData,    ) -> Vec<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {        let Some(config) = thread_store.get::<WebSearchExtensionConfig>() else {            return Vec::new();        };        if !config.available {            return Vec::new();        }        vec![Arc::new(WebSearchTool {            session_id: session_store.level_id().to_string(),            provider: create_model_provider(                config.provider.clone(),                Some(self.auth_manager.clone()),            ),            settings: config.settings.clone(),        })]    }}pub fn install(registry: &mut ExtensionRegistryBuilder<Config>, auth_manager: Arc<AuthManager>) {    let extension = Arc::new(WebSearchExtension { auth_manager });    registry.thread_lifecycle_contributor(extension.clone());    registry.config_contributor(extension.clone());    registry.tool_contributor(extension);}#[cfg(test)]mod tests {    use codex_extension_api::ExtensionData;    use codex_extension_api::ExtensionRegistryBuilder;    use codex_extension_api::ToolName;    use codex_login::CodexAuth;    use codex_model_provider_info::ModelProviderInfo;    use pretty_assertions::assert_eq;    use super::AuthManager;    use super::Config;    use super::WebSearchExtensionConfig;    use super::install;    use crate::tool::RUN_TOOL_NAME;    use crate::tool::WEB_NAMESPACE;    #[test]    fn installed_extension_contributes_web_run_when_enabled() {        let mut builder = ExtensionRegistryBuilder::<Config>::new();        install(            &mut builder,            AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")),        );        let registry = builder.build();        let session_store = ExtensionData::new("session");        let thread_store = ExtensionData::new("11111111-1111-4111-8111-111111111111");        thread_store.insert(WebSearchExtensionConfig {            available: true,            provider: ModelProviderInfo::create_openai_provider(/*base_url*/ None),            settings: Default::default(),        });        let tool_names = registry            .tool_contributors()            .iter()            .flat_map(|contributor| contributor.tools(&session_store, &thread_store))            .map(|tool| (tool.tool_name(), tool.supports_parallel_tool_calls()))            .collect::<Vec<_>>();        assert_eq!(            tool_names,            vec![(ToolName::namespaced(WEB_NAMESPACE, RUN_TOOL_NAME), true)]        );    }}