owlen/crates/owlen-core/tests/web_search_toggle.rs

use std::{any::Any, collections::HashMap, sync::Arc};

use async_trait::async_trait;
use futures::stream;
use owlen_core::{
    ChatStream, Provider, Result,
    config::Config,
    llm::ProviderConfig,
    session::SessionController,
    storage::StorageManager,
    types::{ChatRequest, ChatResponse, Message, ModelInfo},
    ui::NoOpUiController,
};
use serde_json::Value;
use tempfile::tempdir;

struct StubCloudProvider;

#[async_trait]
impl Provider for StubCloudProvider {
    fn name(&self) -> &str {
        "ollama_cloud"
    }

    async fn list_models(&self) -> Result<Vec<ModelInfo>> {
        Ok(vec![ModelInfo {
            id: "stub-cloud-model".to_string(),
            name: "Stub Cloud Model".to_string(),
            description: Some("Stub model for web toggle tests".to_string()),
            provider: self.name().to_string(),
            context_window: Some(8192),
            capabilities: vec!["chat".to_string(), "tools".to_string()],
            supports_tools: true,
        }])
    }

    async fn send_prompt(&self, _request: ChatRequest) -> Result<ChatResponse> {
        Ok(ChatResponse {
            message: Message::assistant(String::new()),
            usage: None,
            is_streaming: false,
            is_final: true,
        })
    }

    async fn stream_prompt(&self, _request: ChatRequest) -> Result<ChatStream> {
        Ok(Box::pin(stream::empty()))
    }

    async fn health_check(&self) -> Result<()> {
        Ok(())
    }

    fn as_any(&self) -> &(dyn Any + Send + Sync) {
        self
    }
}

#[tokio::test(flavor = "multi_thread")]
async fn toggling_web_search_updates_config_and_registry() {
    let temp_dir = tempdir().expect("temp dir");
    let storage = Arc::new(
        StorageManager::with_database_path(temp_dir.path().join("owlen-tests.db"))
            .await
            .expect("storage"),
    );

    let mut config = Config::default();
    config.privacy.encrypt_local_data = false;
    config.general.default_model = Some("stub-cloud-model".into());
    config.general.default_provider = "ollama_cloud".into();

    let mut provider_cfg = ProviderConfig {
        enabled: true,
        provider_type: "ollama_cloud".to_string(),
        base_url: Some("https://ollama.com".to_string()),
        api_key: Some("test-key".to_string()),
        api_key_env: None,
        extra: HashMap::new(),
    };
    provider_cfg.extra.insert(
        "web_search_endpoint".into(),
        Value::String("/api/web_search".into()),
    );
    config.providers.insert("ollama_cloud".into(), provider_cfg);

    let provider: Arc<dyn Provider> = Arc::new(StubCloudProvider);
    let ui = Arc::new(NoOpUiController);

    let mut session = SessionController::new(provider, config, storage, ui, false, None)
        .await
        .expect("session controller");

    assert!(
        !session
            .tool_registry()
            .tools()
            .iter()
            .any(|tool| tool == "web_search"),
        "web_search should be disabled by default"
    );

    session
        .set_tool_enabled("web_search", true)
        .await
        .expect("enable web_search");

    {
        let cfg = session.config_async().await;
        assert!(cfg.tools.web_search.enabled);
        assert!(cfg.privacy.enable_remote_search);
    }
    assert!(
        session
            .tool_registry()
            .tools()
            .iter()
            .any(|tool| tool == "web_search"),
        "web_search should be registered when enabled"
    );

    session
        .set_tool_enabled("web_search", false)
        .await
        .expect("disable web_search");

    {
        let cfg = session.config_async().await;
        assert!(!cfg.tools.web_search.enabled);
        assert!(!cfg.privacy.enable_remote_search);
    }
    assert!(
        !session
            .tool_registry()
            .tools()
            .iter()
            .any(|tool| tool == "web_search"),
        "web_search should be removed when disabled"
    );
}