magi-code 0.96.1

use super::{ToolResult, ToolResultDisplay, ToolRuntime, url_fetch};
use crate::{
    cancellation::AgentCancellation,
    output::{is_credential_like_key, redact_sensitive_text},
};
use chrono::{Duration as ChronoDuration, Utc};
use reqwest::{Url, blocking::Client, redirect::Policy};
use serde::Deserialize;
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use std::{collections::VecDeque, io::Read, time::Duration};

pub(super) const OUTPUT_MAX_BYTES: usize = 48 * 1024;
const RESPONSE_MAX_BYTES: usize = 2 * 1024 * 1024;
const PAGE_MAX_BYTES: usize = 40_000;
const CACHE_MAX_PAGES: usize = 64;

pub(super) const MAX_DOMAINS: usize = 20;
pub(super) const MAX_DOMAIN_CHARS: usize = 253;

#[derive(Debug, Clone, Copy, Deserialize)]
#[serde(rename_all = "lowercase")]
pub(super) enum RecencyFilter {
    Day,
    Week,
    Month,
    Year,
}

impl RecencyFilter {
    fn start_date(self) -> String {
        let days = match self {
            Self::Day => 1,
            Self::Week => 7,
            Self::Month => 30,
            Self::Year => 365,
        };
        (Utc::now().date_naive() - ChronoDuration::days(days)).to_string()
    }
}

#[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub(super) enum Operation {
    Search,
    Open,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(super) struct WebArgs {
    pub operation: Operation,
    pub query: Option<String>,
    pub url: Option<String>,
    pub reference: Option<String>,
    pub limit: Option<usize>,
    #[serde(rename = "domainFilter")]
    pub domain_filter: Option<Vec<String>>,
    #[serde(rename = "recencyFilter")]
    pub recency_filter: Option<RecencyFilter>,
}

impl WebArgs {
    pub(super) fn validate(mut self) -> anyhow::Result<Self> {
        for value in [&mut self.query, &mut self.url, &mut self.reference]
            .into_iter()
            .flatten()
        {
            *value = value.trim().to_owned();
            anyhow::ensure!(
                !value.is_empty(),
                "query, url and reference must not be blank"
            );
        }
        anyhow::ensure!(
            self.query.as_ref().is_none_or(|q| q.chars().count() <= 512),
            "query exceeds 512 characters"
        );
        anyhow::ensure!(
            self.url.as_ref().is_none_or(|u| u.len() <= 2048),
            "url exceeds 2048 bytes"
        );
        anyhow::ensure!(
            self.reference.as_ref().is_none_or(|r| r.len() <= 80),
            "reference exceeds 80 bytes"
        );
        anyhow::ensure!(
            self.limit.is_none_or(|n| (1..=10).contains(&n)),
            "limit must be between 1 and 10"
        );
        if let Some(domains) = &mut self.domain_filter {
            anyhow::ensure!(
                domains.len() <= MAX_DOMAINS,
                "domainFilter must contain at most {MAX_DOMAINS} domains"
            );
            for domain in domains {
                *domain = domain.trim().to_owned();
                validate_domain_filter(domain)?;
            }
        }
        anyhow::ensure!(
            self.operation == Operation::Search
                || (self.domain_filter.is_none() && self.recency_filter.is_none()),
            "domainFilter and recencyFilter are search-only; omit them for open"
        );
        match self.operation {
            Operation::Search => anyhow::ensure!(
                self.query.is_some() && self.url.is_none() && self.reference.is_none(),
                "search requires query; url and reference must be omitted"
            ),
            Operation::Open => anyhow::ensure!(
                self.query.is_none()
                    && (self.url.is_some() != self.reference.is_some())
                    && self.limit.is_none(),
                "open requires exactly one of url or reference; query and limit must be omitted"
            ),
        }
        Ok(self)
    }
}

fn validate_domain_filter(domain: &str) -> anyhow::Result<()> {
    let domain = domain.strip_prefix('-').unwrap_or(domain);
    anyhow::ensure!(
        !domain.is_empty() && domain.len() <= MAX_DOMAIN_CHARS,
        "domainFilter domain portion must contain 1 to {MAX_DOMAIN_CHARS} characters"
    );
    for label in domain.split('.') {
        anyhow::ensure!(
            !label.is_empty()
                && label.len() <= 63
                && !label.starts_with('-')
                && !label.ends_with('-')
                && label
                    .bytes()
                    .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-'),
            "domainFilter entries must be bare ASCII domains with labels of 1 to 63 characters; no schemes, ports, paths, whitespace, or leading/trailing label hyphens"
        );
    }
    Ok(())
}

#[derive(Debug, Clone)]
struct Page {
    reference: String,
    url: String,
    title: String,
    text: String,
    truncated: bool,
}

#[derive(Debug, Default)]
pub(super) struct WebCache {
    pages: VecDeque<Page>,
}

impl WebCache {
    fn insert(&mut self, url: String, title: String, text: String) -> Page {
        self.insert_with_truncation(url, title, text, false)
    }

    fn insert_with_truncation(
        &mut self,
        url: String,
        title: String,
        text: String,
        truncated: bool,
    ) -> Page {
        let text = redact_sensitive_text(&text);
        let truncated = truncated || text.len() > PAGE_MAX_BYTES;
        let text = clip(&text, PAGE_MAX_BYTES).to_owned();
        let title = clip(&redact_sensitive_text(&title), 512).to_owned();
        let mut digest = Sha256::new();
        for value in [&url, &title, &text] {
            digest.update((value.len() as u64).to_le_bytes());
            digest.update(value.as_bytes());
        }
        digest.update([u8::from(truncated)]);
        let digest = digest.finalize();
        let reference = format!(
            "web:{}",
            digest
                .iter()
                .map(|byte| format!("{byte:02x}"))
                .collect::<String>()
        );
        if let Some(index) = self.pages.iter().position(|p| p.reference == reference) {
            // A result returned again is fresh, even when its content is unchanged.
            self.pages.remove(index);
        }
        let page = Page {
            reference,
            url,
            title,
            text,
            truncated,
        };
        if self.pages.len() == CACHE_MAX_PAGES {
            self.pages.pop_front();
        }
        self.pages.push_back(page.clone());
        page
    }

    fn get(&self, reference: &str) -> anyhow::Result<Page> {
        self.pages.iter().find(|p| p.reference == reference).cloned()
            .ok_or_else(|| anyhow::anyhow!("web cache miss: reference expired or belongs to another runtime; open the public URL again"))
    }
}

#[derive(Deserialize)]
struct Response {
    results: Vec<SearchResult>,
}
#[derive(Deserialize)]
struct SearchResult {
    url: String,
    title: Option<String>,
    text: Option<String>,
}

impl ToolRuntime {
    pub(super) fn web_for_context(
        &self,
        args: WebArgs,
        context: &crate::output::ToolDispatchContext,
    ) -> anyhow::Result<ToolResult> {
        let args = args.validate()?;
        context.cancellation.check()?;
        if args.operation != Operation::Search
            || context.hook_context.provider_id.as_deref()
                != Some(crate::providers::OPENAI_CODEX_PROVIDER)
        {
            return self.web(args, &context.cancellation);
        }
        anyhow::ensure!(
            args.domain_filter.is_none() && args.recency_filter.is_none(),
            "Codex web search does not support domainFilter or recencyFilter; omit them. No Exa fallback is performed."
        );
        let model = context
            .hook_context
            .model_id
            .as_deref()
            .filter(|model| !model.trim().is_empty())
            .ok_or_else(|| anyhow::anyhow!("Codex web search requires the active model"))?;
        let paths = self.view_image_paths.as_ref().ok_or_else(|| {
            anyhow::anyhow!("Codex web search requires configured credential paths")
        })?;
        let crate::auth::ProviderCredential::OAuth { access, account_id } =
            crate::auth::codex_credential_from_store(paths)?
        else {
            anyhow::bail!("Codex web search requires OAuth credentials")
        };
        let provider = crate::providers::OpenAiCodexProvider::new(
            model,
            access,
            account_id,
            crate::providers::ReqwestHttpTransport,
        )
        .with_store_auth_refresh(paths.clone());
        let output = provider.web_search(
            args.query.as_deref().unwrap_or_default(),
            args.limit.unwrap_or(5),
            &context.cancellation,
        )?;
        let content = serde_json::to_string_pretty(&output)?;
        Ok(result(
            true,
            content,
            json!({"provider":"openai-codex", "kind":"search_synthesis", "cached":false}),
        ))
    }

    pub(super) fn web(
        &self,
        args: WebArgs,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<ToolResult> {
        let args = args.validate()?;
        cancellation.check()?;
        if let Some(reference) = &args.reference {
            let page = self
                .web_cache
                .lock()
                .map_err(|_| anyhow::anyhow!("web cache lock poisoned"))?
                .get(reference)?;
            return Ok(page_result(&page));
        }
        if let Some(url) = &args.url {
            validate_url(url, cancellation)?;
            let fetched = url_fetch::fetch_url(
                url_fetch::UrlFetchInput {
                    url: url.clone(),
                    max_tokens: None,
                }
                .validate()?,
                cancellation,
            )?;
            cancellation.check()?;
            return self.cache_fetched_page(url.clone(), fetched);
        }
        let key = std::env::var("EXA_API_KEY")
            .ok()
            .filter(|s| !s.trim().is_empty())
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "EXA_API_KEY is required for web search; set it in the process environment"
                )
            })?;
        let client = Client::builder()
            .connect_timeout(Duration::from_secs(5))
            .timeout(Duration::from_secs(20))
            .redirect(Policy::none())
            .user_agent(concat!("magi-code/", env!("CARGO_PKG_VERSION")))
            .build()?;
        self.web_with_client(
            args,
            &client,
            "https://api.exa.ai",
            key.trim(),
            cancellation,
        )
    }

    fn web_with_client(
        &self,
        args: WebArgs,
        client: &Client,
        base: &str,
        key: &str,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<ToolResult> {
        anyhow::ensure!(
            args.operation == Operation::Search,
            "Exa only supports search"
        );
        let mut body = json!({"query": args.query, "type": "auto", "numResults": args.limit.unwrap_or(5), "contents": {"text": {"maxCharacters": 10000}}});
        let mut include_domains = Vec::new();
        let mut exclude_domains = Vec::new();
        for domain in args.domain_filter.as_deref().unwrap_or_default() {
            if let Some(excluded) = domain.strip_prefix('-') {
                exclude_domains.push(excluded);
            } else {
                include_domains.push(domain.as_str());
            }
        }
        if !include_domains.is_empty() {
            body["includeDomains"] = json!(include_domains);
        }
        if !exclude_domains.is_empty() {
            body["excludeDomains"] = json!(exclude_domains);
        }
        if let Some(recency) = args.recency_filter {
            body["startPublishedDate"] = json!(recency.start_date());
        }
        cancellation.check()?;
        let response = client
            .post(format!("{base}/search"))
            .header("x-api-key", key)
            .json(&body)
            .send();
        cancellation.check()?;
        let response = match response {
            Ok(response) => response,
            Err(error) => {
                return Ok(result(
                    false,
                    "EXA request failed".into(),
                    json!({"timeout": error.is_timeout()}),
                ));
            }
        };
        let status = response.status();
        if !status.is_success() {
            // Never echo provider error bodies or credentials.
            return Ok(result(
                false,
                format!("EXA API error {status}"),
                json!({"rate_limited": status.as_u16() == 429}),
            ));
        }
        anyhow::ensure!(
            response
                .content_length()
                .is_none_or(|n| n <= RESPONSE_MAX_BYTES as u64),
            "EXA response exceeds byte limit"
        );
        let mut bytes = Vec::new();
        let read_result = response
            .take((RESPONSE_MAX_BYTES + 1) as u64)
            .read_to_end(&mut bytes);
        cancellation.check()?;
        read_result.map_err(|_| anyhow::anyhow!("EXA response body read failed"))?;
        anyhow::ensure!(
            bytes.len() <= RESPONSE_MAX_BYTES,
            "EXA response exceeds byte limit"
        );
        let response: Response = serde_json::from_slice(&bytes)
            .map_err(|_| anyhow::anyhow!("EXA returned malformed JSON"))?;
        let mut pages = Vec::new();
        let max_results = args.limit.unwrap_or(5);
        let mut omitted = response.results.len().saturating_sub(max_results);
        for source in response.results.into_iter().take(max_results) {
            cancellation.check()?;
            if validate_url(&source.url, cancellation).is_err() {
                cancellation.check()?;
                omitted += 1;
                continue;
            }
            let Some(text) = source.text.filter(|t| !t.trim().is_empty()) else {
                omitted += 1;
                continue;
            };
            let page = self
                .web_cache
                .lock()
                .map_err(|_| anyhow::anyhow!("web cache lock poisoned"))?
                .insert(source.url, source.title.unwrap_or_default(), text);
            pages.push(page);
        }
        cancellation.check()?;
        let mut content = String::new();
        for page in &pages {
            content.push_str(&format!(
                "[{}] {}\nURL: {}\n{}\n\n",
                page.reference,
                page.title,
                page.url,
                clip(&page.text, 1200)
            ));
        }
        if pages.is_empty() {
            content.push_str("No usable public results returned.");
        }
        Ok(result(
            true,
            content,
            json!({"results": pages.len(), "omitted": omitted, "truncated": pages.iter().any(|page| page.truncated || page.text.len() > 1200), "scope": "bounded Exa page text; open a reference to inspect cached text"}),
        ))
    }
    fn cache_fetched_page(
        &self,
        url: String,
        fetched: url_fetch::UrlFetchOutput,
    ) -> anyhow::Result<ToolResult> {
        if !fetched.success {
            let mut metadata = fetched.metadata;
            metadata["provider"] = json!("direct");
            return Ok(result(false, fetched.content, metadata));
        }
        let page = self
            .web_cache
            .lock()
            .map_err(|_| anyhow::anyhow!("web cache lock poisoned"))?
            .insert_with_truncation(
                url,
                fetched.metadata["title"]
                    .as_str()
                    .unwrap_or_default()
                    .to_owned(),
                fetched.content,
                fetched.metadata["truncated"].as_bool().unwrap_or(false),
            );
        let mut output = page_result(&page);
        output.metadata["provider"] = json!("direct");
        output.metadata["cache"] = json!(false);
        Ok(output)
    }
}

fn validate_url(url: &str, cancellation: &AgentCancellation) -> anyhow::Result<()> {
    anyhow::ensure!(url.len() <= 2048, "public URL exceeds 2048 bytes");
    let parsed = Url::parse(url).map_err(|_| anyhow::anyhow!("invalid public URL"))?;
    anyhow::ensure!(
        parsed.username().is_empty() && parsed.password().is_none(),
        "URL credentials are not allowed"
    );
    anyhow::ensure!(
        !parsed
            .query_pairs()
            .any(|(name, _)| is_credential_like_key(&name)),
        "sensitive URL parameters are not allowed"
    );
    if let Some(fragment) = parsed.fragment() {
        // Fragments are submitted to Exa too; decode parameter names just like queries.
        let mut fragment_url = parsed.clone();
        fragment_url.set_query(Some(fragment));
        anyhow::ensure!(
            !fragment_url
                .query_pairs()
                .any(|(name, _)| is_credential_like_key(&name)),
            "sensitive URL parameters are not allowed"
        );
    }
    anyhow::ensure!(
        redact_sensitive_text(url) == url,
        "sensitive URL parameters are not allowed"
    );
    url_fetch::validate_public_url(&parsed, cancellation)
}

fn page_result(page: &Page) -> ToolResult {
    result(
        true,
        format!(
            "[{}] {}\nURL: {}\n\n{}",
            page.reference, page.title, page.url, page.text
        ),
        json!({"provider": "web-cache", "reference": page.reference, "url": page.url, "cache": true, "truncated": page.truncated, "scope": "bounded page text, not necessarily the complete page"}),
    )
}

fn clip(text: &str, max: usize) -> &str {
    let mut end = text.len().min(max);
    while !text.is_char_boundary(end) {
        end -= 1;
    }
    &text[..end]
}

fn result(success: bool, content: String, mut metadata: Value) -> ToolResult {
    let content = redact_sensitive_text(&content);
    if metadata.get("provider").is_none() {
        metadata["provider"] = json!("exa");
    }
    if content.len() > OUTPUT_MAX_BYTES {
        metadata["truncated"] = json!(true);
    }
    ToolResult {
        tool_name: "web".into(),
        success,
        content: clip(&content, OUTPUT_MAX_BYTES).to_owned(),
        metadata,
        display: ToolResultDisplay::default(),
    }
}