diff --git a/agents-docs/CONFIG.md b/agents-docs/CONFIG.md index adb8fe39..456603b9 100644 --- a/agents-docs/CONFIG.md +++ b/agents-docs/CONFIG.md @@ -43,18 +43,43 @@ The Rust CLI looks for `config.toml` in: ```toml max_chars = 8000 -profile = "balanced" -skip_providers = ["exa"] +min_chars = 200 +exa_results = 5 +tavily_results = 3 +output_limit = 10 +log_level = "info" + +[routing] +min_free_quality_to_skip_paid = 0.70 + +# Pre-warms the cache for the top N tracked domains at CLI startup; disabled with `enabled = false`. +[routing.prewarm] +enabled = true # warm cache for top domains on CLI startup +top_n_domains = 20 # number of top tracked domains to pre-warm +max_concurrency = 4 # parallel pre-warm requests (respects provider rate limits) [cache.ttl] firecrawl = 21600 # 6 hours exa = 14400 # 4 hours +tavily = 14400 # 4 hours +serper = 7200 # 2 hours jina = 7200 # 2 hours +mistral = 28800 # 8 hours duckduckgo = 3600 # 1 hour llms_txt = 28800 # 8 hours synthesis = 43200 # 12 hours default = 3600 # fallback for any unlisted provider +[semantic_cache] +enabled = true +path = ".do-wdr_cache" +threshold = 0.85 +max_entries = 10000 + +[providers.exa_mcp.rate_limit] +requests_per_second = 10 +burst = 20 + [api_keys] tavily = "your-key-here" serper = "your-key-here" diff --git a/cli/src/startup.rs b/cli/src/startup.rs index d193e467..b2322b50 100644 --- a/cli/src/startup.rs +++ b/cli/src/startup.rs @@ -16,7 +16,7 @@ pub async fn prewarm_cache(resolver: Arc, config: &Config) -> Result<( } let top_n = config.routing.prewarm.top_n_domains; - let max_concurrency = config.routing.prewarm.max_concurrency.max(1); + let max_concurrency = config.routing.prewarm.max_concurrency; let domains = { let routing_memory = resolver.routing_memory(); @@ -35,52 +35,90 @@ pub async fn prewarm_cache(resolver: Arc, config: &Config) -> Result<( max_concurrency ); - let semaphore = Arc::new(Semaphore::new(max_concurrency)); - let mut join_set = JoinSet::new(); - - for domain in domains { - let url = format!("https://{}", domain); + let resolve = move |url: String| { let resolver = resolver.clone(); + async move { + resolver + .resolve_url(&url) + .await + .map(|_| ()) + .map_err(anyhow::Error::from) + } + }; - let permit = Arc::clone(&semaphore).acquire_owned().await?; + prewarm_domains(domains, max_concurrency, resolve).await +} - join_set.spawn(async move { - tracing::debug!("Pre-warming domain: {}", domain); +/// Resolve a batch of domains with bounded concurrency, awaiting completion. +/// +/// The `resolve` closure is injected so tests can exercise the semaphore and +/// timeout loop offline without a network-backed `Resolver`. +pub async fn prewarm_domains( + domains: Vec, + max_concurrency: usize, + resolve: F, +) -> anyhow::Result<()> +where + F: Fn(String) -> Fut + Clone + Send + Sync + 'static, + Fut: Future> + Send + 'static, +{ + // A zero limit would deadlock the spawn loop on the first acquire; treat + // it as an explicit serialization request like the CLI-side default. + let max_concurrency = max_concurrency.max(1); - let result = tokio::time::timeout(PER_TASK_TIMEOUT, resolver.resolve_url(&url)).await; - match result { - Ok(Ok(_)) => { - tracing::debug!("Pre-warm succeeded for domain: {}", domain); - } - Ok(Err(e)) => { - tracing::warn!("Pre-warm failed for domain {}: {}", domain, e); - } - Err(_) => { - tracing::warn!("Pre-warm timed out for domain: {}", domain); - } - } - drop(permit); - }); - } + // The overall budget covers spawning AND joining. Without this, tasks that + // each burn PER_TASK_TIMEOUT would let the spawn phase run for many minutes + // before the join timeout ever started counting. + let completed = tokio::time::timeout(OVERALL_TIMEOUT, async { + let semaphore = Arc::new(Semaphore::new(max_concurrency)); + let mut join_set = JoinSet::new(); - let deadline = tokio::time::sleep(OVERALL_TIMEOUT); - tokio::pin!(deadline); + for domain in domains { + let url = format!("https://{}", domain); + let resolve = resolve.clone(); - loop { - tokio::select! { - Some(result) = join_set.join_next() => { - if let Err(e) = result { - tracing::warn!("Pre-warm task panicked: {}", e); + let permit = Arc::clone(&semaphore).acquire_owned().await?; + + join_set.spawn(async move { + tracing::debug!("Pre-warming domain: {}", domain); + + let result = tokio::time::timeout(PER_TASK_TIMEOUT, resolve(url)).await; + match result { + Ok(Ok(_)) => { + tracing::debug!("Pre-warm succeeded for domain: {}", domain); + } + Ok(Err(e)) => { + tracing::warn!("Pre-warm failed for domain {}: {}", domain, e); + } + Err(_) => { + tracing::warn!("Pre-warm timed out for domain: {}", domain); + } } + drop(permit); + }); + } + + while let Some(result) = join_set.join_next().await { + if let Err(e) = result { + tracing::warn!("Pre-warm task panicked: {}", e); } - _ = &mut deadline => { - tracing::warn!("Pre-warming timed out after {}s, abandoning remaining tasks", OVERALL_TIMEOUT.as_secs()); - join_set.abort_all(); - break; - } - else => break, } - } + + Ok::<(), anyhow::Error>(()) + }) + .await; + + match completed { + Ok(result) => result, + Err(_) => { + // Dropping the timed-out future aborts all in-flight tasks. + tracing::warn!( + "Pre-warming timed out after {}s, abandoning remaining tasks", + OVERALL_TIMEOUT.as_secs() + ); + Ok(()) + } + }?; tracing::info!("Cache pre-warming completed"); Ok(()) diff --git a/cli/tests/config_prewarm.rs b/cli/tests/config_prewarm.rs new file mode 100644 index 00000000..cc99aef3 --- /dev/null +++ b/cli/tests/config_prewarm.rs @@ -0,0 +1,47 @@ +//! Round-trip of `[routing.prewarm]` from TOML into `PrewarmConfig`. +//! +//! The committed `config.toml` carries an explicit `[routing.prewarm]` block. +//! These tests prove the section deserializes into the typed config and that +//! the committed values match the same defaults used when the section is +//! left unset (per `cli/src/config/defaults.rs`). + +use do_wdr_lib::config::Config; + +#[test] +fn prewarm_config_defaults_match_committed_block() { + let config = Config::default(); + assert!(config.routing.prewarm.enabled); + assert_eq!(config.routing.prewarm.top_n_domains, 20); + assert_eq!(config.routing.prewarm.max_concurrency, 4); +} + +#[test] +fn prewarm_config_roundtrip_from_toml() { + // Explicit [routing.prewarm] values deserialize into PrewarmConfig. + let content = r#" +[routing.prewarm] +enabled = false +top_n_domains = 7 +max_concurrency = 2 +"#; + let config: Config = toml::from_str(content).unwrap(); + assert!(!config.routing.prewarm.enabled); + assert_eq!(config.routing.prewarm.top_n_domains, 7); + assert_eq!(config.routing.prewarm.max_concurrency, 2); +} + +#[test] +fn committed_config_toml_parses_with_prewarm() { + // Guard the actual repo file, not just an inline sample: a typo (or wrong + // type) in the committed [routing.prewarm] block must fail here, since no + // test otherwise reads the repo-root config.toml. + let path = concat!(env!("CARGO_MANIFEST_DIR"), "/../config.toml"); + let config = Config::from_file(path).expect("committed config.toml must parse"); + assert!(config.routing.prewarm.enabled); + assert_eq!(config.routing.prewarm.top_n_domains, 20); + assert_eq!(config.routing.prewarm.max_concurrency, 4); + // Canary for key placement: log_level is a TOP-LEVEL Config field; nested + // under [routing] it deserializes as "" (serde default) instead of "info", + // so this assertion catches a re-introduced misnesting. + assert_eq!(config.log_level, "info"); +} diff --git a/cli/tests/startup_prewarm.rs b/cli/tests/startup_prewarm.rs new file mode 100644 index 00000000..7c861e04 --- /dev/null +++ b/cli/tests/startup_prewarm.rs @@ -0,0 +1,62 @@ +//! Guards of the cache pre-warming core — offline, no network. +//! +//! `prewarm_domains` is the injected-resolve core behind `prewarm_cache` +//! (`cli/src/startup.rs`). These tests pin the "empty domains → graceful +//! skip" guard and the "every provided domain is resolved once" contract. + +use do_wdr_lib::config::Config; +use do_wdr_lib::resolver::Resolver; +use do_wdr_lib::startup::{prewarm_cache, prewarm_domains}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +#[tokio::test] +async fn prewarm_empty_domains_skips_resolve() { + let counter = Arc::new(AtomicUsize::new(0)); + let counter_c = counter.clone(); + let resolve = move |_url: String| { + counter_c.fetch_add(1, Ordering::SeqCst); + async move { Ok(()) } + }; + + // No tracked domains (mirrors prewarm_cache's `domains.is_empty()` guard): + // the resolve fn must never be called and the call must succeed gracefully. + let result = prewarm_domains(Vec::new(), 4, resolve).await; + + assert!(result.is_ok()); + assert_eq!(counter.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn prewarm_cache_disabled_returns_without_work() { + // Covers the one guard the core tests cannot reach: prewarm_cache with + // enabled=false must short-circuit before touching the resolver. The + // explicit budget turns a regression (e.g. the guard dropping) into a + // failure instead of a hang. + let mut config = Config::default(); + config.routing.prewarm.enabled = false; + config.semantic_cache.enabled = false; + let resolver = Arc::new(Resolver::with_config(config.clone()).await); + + tokio::time::timeout(Duration::from_secs(10), prewarm_cache(resolver, &config)) + .await + .expect("enabled=false prewarm must return promptly") + .expect("enabled=false prewarm must succeed"); +} + +#[tokio::test] +async fn prewarm_invokes_resolve_for_each_domain() { + let counter = Arc::new(AtomicUsize::new(0)); + let counter_c = counter.clone(); + let resolve = move |_url: String| { + counter_c.fetch_add(1, Ordering::SeqCst); + async move { Ok(()) } + }; + + let domains = vec!["a.example".to_string(), "b.example".to_string()]; + let result = prewarm_domains(domains, 2, resolve).await; + + assert!(result.is_ok()); + assert_eq!(counter.load(Ordering::SeqCst), 2); +} diff --git a/cli/tests/startup_semaphore.rs b/cli/tests/startup_semaphore.rs new file mode 100644 index 00000000..bcb5fd3c --- /dev/null +++ b/cli/tests/startup_semaphore.rs @@ -0,0 +1,63 @@ +//! Concurrency limits of the cache pre-warming core — offline, no network. +//! +//! Verifies the semaphore bounds peak in-flight work to `max_concurrency` and +//! that the acquire/join loop completes without deadlock or the overall timeout. + +use do_wdr_lib::startup::prewarm_domains; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +#[tokio::test] +async fn max_concurrency_respected() { + let peak = Arc::new(AtomicUsize::new(0)); + let in_flight = Arc::new(AtomicUsize::new(0)); + let peak_c = peak.clone(); + let in_flight_c = in_flight.clone(); + let resolve = move |_url: String| { + let peak = peak_c.clone(); + let in_flight = in_flight_c.clone(); + async move { + let current = in_flight.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(current, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(10)).await; + in_flight.fetch_sub(1, Ordering::SeqCst); + Ok(()) + } + }; + + let domains: Vec = (0..6).map(|i| format!("{}.example", i)).collect(); + let result = prewarm_domains(domains, 2, resolve).await; + + assert!(result.is_ok()); + let observed_peak = peak.load(Ordering::SeqCst); + assert!( + observed_peak <= 2, + "peak concurrency {observed_peak} exceeded limit 2" + ); + assert_eq!( + in_flight.load(Ordering::SeqCst), + 0, + "all tasks must complete and release their permit" + ); +} + +#[tokio::test] +async fn prewarm_without_deadlock_or_timeout() { + let resolve = move |_url: String| async move { + tokio::time::sleep(Duration::from_millis(1)).await; + Ok(()) + }; + + // 20 domains, 4-permit semaphore: the acquire loop must drain the queue and + // join every task. A deadlocked acquire would trip this explicit budget + // instead of silently hanging the test run. + let domains: Vec = (0..20).map(|i| format!("{}.example", i)).collect(); + tokio::time::timeout( + Duration::from_secs(10), + prewarm_domains(domains, 4, resolve), + ) + .await + .expect("prewarm must complete within the 10s budget (no deadlock)") + .expect("prewarm must succeed"); +} diff --git a/config.toml b/config.toml index e895ae6f..77281c81 100644 --- a/config.toml +++ b/config.toml @@ -1,6 +1,3 @@ -[routing] -min_free_quality_to_skip_paid = 0.70 - max_chars = 8000 min_chars = 200 exa_results = 5 @@ -8,6 +5,14 @@ tavily_results = 3 output_limit = 10 log_level = "info" +[routing] +min_free_quality_to_skip_paid = 0.70 + +[routing.prewarm] +enabled = true # warm cache for top domains on CLI startup +top_n_domains = 20 # number of top tracked domains to pre-warm +max_concurrency = 4 # parallel pre-warm requests (respects provider rate limits) + [cache.ttl] firecrawl = 21600 # 6 hours exa = 14400 # 4 hours