diff --git a/src/api.rs b/src/api.rs index c0ee303..449034a 100644 --- a/src/api.rs +++ b/src/api.rs @@ -13,6 +13,7 @@ use serde_json::{Map, Value}; use tokio::time::sleep; use crate::error::KagiError; +use crate::http; use crate::parser::parse_assistant_thread_list; #[cfg(test)] use crate::types::ApiMeta; @@ -31,11 +32,6 @@ use crate::types::{ TranslateWarning, TranslationSuggestionsResponse, WordInsightsResponse, }; -const USER_AGENT: &str = concat!( - "kagi-cli/", - env!("CARGO_PKG_VERSION"), - " (+https://github.com/Microck/kagi-cli)" -); const KAGI_SUMMARIZE_URL: &str = "https://kagi.com/api/v0/summarize"; const KAGI_SUBSCRIBER_SUMMARIZE_URL: &str = "https://kagi.com/mother/summary_labs"; const KAGI_NEWS_LATEST_URL: &str = "https://news.kagi.com/api/batches/latest"; @@ -2179,11 +2175,7 @@ where } fn build_client() -> Result { - Client::builder() - .user_agent(USER_AGENT) - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|error| KagiError::Network(format!("failed to build HTTP client: {error}"))) + http::client_30s() } fn map_transport_error(error: reqwest::Error) -> KagiError { diff --git a/src/http.rs b/src/http.rs new file mode 100644 index 0000000..b8dfba6 --- /dev/null +++ b/src/http.rs @@ -0,0 +1,41 @@ +use std::sync::OnceLock; +use std::time::Duration; + +use reqwest::Client; + +use crate::error::KagiError; + +const USER_AGENT: &str = concat!( + "kagi-cli/", + env!("CARGO_PKG_VERSION"), + " (+https://github.com/Microck/kagi-cli)" +); + +static CLIENT_20S: OnceLock> = OnceLock::new(); +static CLIENT_30S: OnceLock> = OnceLock::new(); + +pub fn client_20s() -> Result { + cached_client(&CLIENT_20S, Duration::from_secs(20)) +} + +pub fn client_30s() -> Result { + cached_client(&CLIENT_30S, Duration::from_secs(30)) +} + +fn cached_client( + slot: &OnceLock>, + timeout: Duration, +) -> Result { + let result = slot.get_or_init(|| { + Client::builder() + .user_agent(USER_AGENT) + .timeout(timeout) + .build() + .map_err(|error| format!("failed to build HTTP client: {error}")) + }); + + result + .as_ref() + .cloned() + .map_err(|error| KagiError::Network(error.clone())) +} diff --git a/src/main.rs b/src/main.rs index 13483b3..fd50072 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ mod auth; mod auth_wizard; mod cli; mod error; +mod http; mod parser; mod quick; mod search; @@ -53,7 +54,7 @@ struct SearchRequestOptions { no_personalized: bool, } -#[tokio::main] +#[tokio::main(flavor = "current_thread")] async fn main() { if let Err(error) = run().await { eprintln!("{error}"); @@ -746,10 +747,9 @@ async fn run_batch_search( let semaphore_clone = Arc::clone(&semaphore); let credentials_clone = credentials.clone(); let options_clone = options.clone(); - let format_clone = format.clone(); let query_clone = query.clone(); - let handle: tokio::task::JoinHandle> = + let handle: tokio::task::JoinHandle> = tokio::spawn(async move { let _permit = semaphore_clone.acquire().await; rate_limiter_clone.acquire().await?; @@ -758,19 +758,7 @@ async fn run_batch_search( let response = execute_search_request(&request, credentials_clone).await?; - let output = match format_clone.as_str() { - "pretty" => format_pretty_response(&response, use_color), - "compact" => serde_json::to_string(&response).map_err(|error| { - KagiError::Parse(format!("failed to serialize search response: {error}")) - })?, - "markdown" => format_markdown_response(&response), - "csv" => format_csv_response(&response), - _ => serde_json::to_string_pretty(&response).map_err(|error| { - KagiError::Parse(format!("failed to serialize search response: {error}")) - })?, - }; - - Ok((query, output)) + Ok((query, response)) }); handles.push(handle); @@ -804,24 +792,18 @@ async fn run_batch_search( if format == "json" || format == "compact" { // For machine-readable formats, create a proper JSON envelope let queries: Vec = results.iter().map(|(query, _)| query.clone()).collect(); - let mut results_json = serde_json::json!({ + let results_payload = results + .into_iter() + .map(|(_, response)| serde_json::to_value(response)) + .collect::, _>>() + .map_err(|error| { + KagiError::Parse(format!("failed to serialize batch search response: {error}")) + })?; + let results_json = serde_json::json!({ "queries": queries, - "results": [] + "results": results_payload }); - let results_array = results_json["results"].as_array_mut().unwrap(); - - for (query, output) in results { - // Parse the individual JSON output and add to array - let parsed: serde_json::Value = serde_json::from_str(&output).map_err(|e| { - KagiError::Parse(format!( - "failed to parse batch result for '{}': {}", - query, e - )) - })?; - results_array.push(parsed); - } - if format == "compact" { println!("{}", serde_json::to_string(&results_json)?); } else { @@ -829,7 +811,15 @@ async fn run_batch_search( } } else { // For human-readable formats, output with headers - for (query, output) in results { + for (query, response) in results { + let output = match format.as_str() { + "pretty" => format_pretty_response(&response, use_color), + "markdown" => format_markdown_response(&response), + "csv" => format_csv_response(&response), + _ => serde_json::to_string_pretty(&response).map_err(|error| { + KagiError::Parse(format!("failed to serialize search response: {error}")) + })?, + }; println!("=== Results for: {} ===", query); println!("{}", output); println!(); diff --git a/src/quick.rs b/src/quick.rs index 251c0d5..24e2d35 100644 --- a/src/quick.rs +++ b/src/quick.rs @@ -3,12 +3,12 @@ use scraper::Html; use serde::Deserialize; use crate::error::KagiError; +use crate::http; use crate::search::{SearchRequest, validate_lens_value}; use crate::types::{ QuickMessage, QuickMeta, QuickReferenceCollection, QuickReferenceItem, QuickResponse, }; -const USER_AGENT: &str = "kagi-cli/0.1.0 (+https://github.com/)"; const KAGI_QUICK_ANSWER_URL: &str = "https://kagi.com/mother/context"; pub async fn execute_quick( @@ -452,11 +452,7 @@ fn format_client_error_suffix(body: &str) -> String { } fn build_client() -> Result { - Client::builder() - .user_agent(USER_AGENT) - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|error| KagiError::Network(format!("failed to build HTTP client: {error}"))) + http::client_30s() } fn map_transport_error(error: reqwest::Error) -> KagiError { diff --git a/src/search.rs b/src/search.rs index 77ddb96..a00c899 100644 --- a/src/search.rs +++ b/src/search.rs @@ -2,12 +2,12 @@ use reqwest::{Client, StatusCode, header}; use serde::Deserialize; use crate::error::KagiError; +use crate::http; use crate::parser::parse_search_results; use crate::types::{SearchResponse, SearchResult}; const KAGI_SEARCH_URL: &str = "https://kagi.com/html/search"; const KAGI_API_SEARCH_URL: &str = "https://kagi.com/api/v0/search"; -const USER_AGENT: &str = "kagi-cli/0.1.0 (+https://github.com/)"; const UNAUTHENTICATED_MARKERS: [&str; 3] = [ "Kagi Search - A Premium Search Engine", "Welcome to Kagi", @@ -406,11 +406,7 @@ fn looks_unauthenticated(body: &str) -> bool { } fn build_client() -> Result { - Client::builder() - .user_agent(USER_AGENT) - .timeout(std::time::Duration::from_secs(20)) - .build() - .map_err(|error| KagiError::Network(format!("failed to build HTTP client: {error}"))) + http::client_20s() } fn map_transport_error(error: reqwest::Error) -> KagiError {