use crate::driver_registry::DiscoveredModel;
use crate::error::{AgentLoopError, Result};
use crate::url_validation::is_blocked_ip;
use reqwest::StatusCode;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use serde::de::DeserializeOwned;
use std::sync::{Arc, OnceLock};
use std::time::Duration;
pub const AUDIO_CONTENT_PLACEHOLDER: &str = "[Audio content not supported]";
const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const HTTP_STREAM_READ_TIMEOUT: Duration = Duration::from_secs(300);
const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
const HTTP_POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(90);
const DNS_LOOKUP_TIMEOUT: Duration = Duration::from_secs(5);
struct SsrfGuardResolver;
type DnsBoxError = Box<dyn std::error::Error + Send + Sync>;
impl Resolve for SsrfGuardResolver {
fn resolve(&self, name: Name) -> Resolving {
let host = name.as_str().to_string();
Box::pin(async move {
let lookup = tokio::time::timeout(
DNS_LOOKUP_TIMEOUT,
tokio::net::lookup_host(format!("{host}:0")),
)
.await
.map_err(|_| -> DnsBoxError {
Box::new(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"DNS lookup timed out",
))
})?
.map_err(|e| -> DnsBoxError { Box::new(e) })?;
let addrs: Vec<std::net::SocketAddr> = lookup.collect();
for addr in &addrs {
if is_blocked_ip(addr.ip()) {
tracing::warn!(
host = %host,
resolved_ip = %addr.ip(),
"Provider HTTP client blocked: hostname resolves to private/internal address"
);
return Err(Box::new(std::io::Error::other(format!(
"host {host} resolves to blocked address {} (private/internal)",
addr.ip()
))) as DnsBoxError);
}
}
Ok(Box::new(addrs.into_iter()) as Addrs)
})
}
}
fn harden_builder(builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder {
builder
.redirect(reqwest::redirect::Policy::none())
.dns_resolver(Arc::new(SsrfGuardResolver))
}
pub fn shared_streaming_http_client() -> reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT
.get_or_init(|| {
harden_builder(
reqwest::Client::builder()
.connect_timeout(HTTP_CONNECT_TIMEOUT)
.read_timeout(HTTP_STREAM_READ_TIMEOUT)
.pool_idle_timeout(HTTP_POOL_IDLE_TIMEOUT),
)
.build()
.unwrap_or_else(|_| {
harden_builder(reqwest::Client::builder())
.build()
.unwrap_or_else(|_| reqwest::Client::new())
})
})
.clone()
}
pub fn shared_request_http_client() -> reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT
.get_or_init(|| {
harden_builder(
reqwest::Client::builder()
.connect_timeout(HTTP_CONNECT_TIMEOUT)
.timeout(HTTP_REQUEST_TIMEOUT)
.pool_idle_timeout(HTTP_POOL_IDLE_TIMEOUT),
)
.build()
.unwrap_or_else(|_| {
harden_builder(reqwest::Client::builder())
.build()
.unwrap_or_else(|_| reqwest::Client::new())
})
})
.clone()
}
#[derive(Debug, Clone)]
pub struct ParsedDataUrl {
pub media_type: String,
pub data: String,
}
pub fn parse_data_url(url: &str) -> Option<ParsedDataUrl> {
if !url.starts_with("data:") {
return None;
}
let parts: Vec<&str> = url.splitn(2, ',').collect();
if parts.len() != 2 {
return None;
}
let media_type = parts[0]
.trim_start_matches("data:")
.trim_end_matches(";base64")
.to_string();
let data = parts[1].to_string();
Some(ParsedDataUrl { media_type, data })
}
pub fn is_request_too_large(status: StatusCode, error_text: &str, extra_patterns: &[&str]) -> bool {
let error_lower = error_text.to_lowercase();
if status == StatusCode::PAYLOAD_TOO_LARGE {
return true;
}
if status.is_client_error() {
if error_lower.contains("input is too long") || error_lower.contains("maximum context") {
return true;
}
if error_lower.contains("exceeds the maximum")
&& (error_lower.contains("token") || error_lower.contains("context"))
{
return true;
}
for pattern in extra_patterns {
if error_lower.contains(pattern) {
return true;
}
}
}
false
}
pub const ANTHROPIC_TOO_LARGE_PATTERNS: &[&str] = &[
"prompt is too long",
"request size exceeded",
"context length",
"too many tokens",
];
pub const GEMINI_TOO_LARGE_PATTERNS: &[&str] = &[
"request payload size exceeds",
"content too large",
"token limit exceeded",
];
pub fn is_model_not_found(status: StatusCode, error_text: &str, patterns: &[&str]) -> bool {
if status != StatusCode::NOT_FOUND {
return false;
}
let error_lower = error_text.to_lowercase();
for pattern in patterns {
if error_lower.contains(pattern) {
return true;
}
}
false
}
pub const ANTHROPIC_NOT_FOUND_PATTERNS: &[&str] = &["not_found_error"];
pub const GEMINI_NOT_FOUND_PATTERNS: &[&str] = &["not_found", "model"];
pub async fn fetch_models<T, F>(
request: reqwest::RequestBuilder,
fetch_err_prefix: &str,
parse_err_prefix: &str,
none_on_statuses: &[StatusCode],
map: F,
) -> Result<Option<Vec<DiscoveredModel>>>
where
T: DeserializeOwned,
F: FnOnce(T) -> Vec<DiscoveredModel>,
{
let response = request
.send()
.await
.map_err(|e| AgentLoopError::llm(format!("{fetch_err_prefix}: {e}")))?;
let status = response.status();
if !status.is_success() {
let _ = response.bytes().await; if none_on_statuses.contains(&status) {
return Ok(None);
}
return Err(crate::openai_protocol::models_api_status_error(status));
}
let parsed: T = response
.json()
.await
.map_err(|e| AgentLoopError::llm(format!("{parse_err_prefix}: {e}")))?;
Ok(Some(map(parsed)))
}
pub mod thinking_budget {
pub const LOW: u32 = 1024;
pub const MEDIUM: u32 = 4096;
pub const HIGH: u32 = 16384;
pub const XHIGH: u32 = 32768;
pub fn from_effort(effort: &str) -> Option<u32> {
match effort.to_lowercase().as_str() {
"low" => Some(LOW),
"medium" => Some(MEDIUM),
"high" => Some(HIGH),
"xhigh" => Some(XHIGH),
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;
#[test]
fn test_parse_data_url_valid() {
let result = parse_data_url("data:image/png;base64,iVBOR").unwrap();
assert_eq!(result.media_type, "image/png");
assert_eq!(result.data, "iVBOR");
}
#[test]
fn test_parse_data_url_jpeg() {
let result = parse_data_url("data:image/jpeg;base64,/9j/4AAQ").unwrap();
assert_eq!(result.media_type, "image/jpeg");
assert_eq!(result.data, "/9j/4AAQ");
}
#[test]
fn test_parse_data_url_not_data() {
assert!(parse_data_url("https://example.com/image.png").is_none());
}
#[test]
fn test_parse_data_url_no_comma() {
assert!(parse_data_url("data:image/jpeg;base64").is_none());
}
#[test]
fn test_is_request_too_large_413() {
assert!(is_request_too_large(StatusCode::PAYLOAD_TOO_LARGE, "", &[]));
}
#[test]
fn test_is_request_too_large_generic() {
assert!(is_request_too_large(
StatusCode::BAD_REQUEST,
"input is too long",
&[]
));
}
#[test]
fn test_is_request_too_large_anthropic() {
assert!(is_request_too_large(
StatusCode::BAD_REQUEST,
"prompt is too long: 100000 tokens",
ANTHROPIC_TOO_LARGE_PATTERNS
));
}
#[test]
fn test_is_request_too_large_gemini() {
assert!(is_request_too_large(
StatusCode::BAD_REQUEST,
"request payload size exceeds limit",
GEMINI_TOO_LARGE_PATTERNS
));
}
#[test]
fn test_is_model_not_found_with_pattern() {
assert!(is_model_not_found(
StatusCode::NOT_FOUND,
r#"{"error":{"type":"not_found_error"}}"#,
ANTHROPIC_NOT_FOUND_PATTERNS
));
}
#[test]
fn test_is_model_not_found_no_match_without_pattern() {
assert!(!is_model_not_found(
StatusCode::NOT_FOUND,
"Endpoint not found",
ANTHROPIC_NOT_FOUND_PATTERNS
));
}
#[test]
fn test_is_model_not_found_not_404() {
assert!(!is_model_not_found(
StatusCode::BAD_REQUEST,
"model not found",
&[]
));
}
#[test]
fn test_is_model_not_found_gemini() {
assert!(is_model_not_found(
StatusCode::NOT_FOUND,
r#"{"error":{"status":"NOT_FOUND","message":"model foo"}}"#,
GEMINI_NOT_FOUND_PATTERNS
));
}
#[tokio::test]
async fn ssrf_resolver_blocks_loopback_literal() {
let resolver = SsrfGuardResolver;
let name = Name::from_str("127.0.0.1").unwrap();
let result = resolver.resolve(name).await;
assert!(result.is_err(), "loopback literal should be blocked");
}
#[tokio::test]
async fn ssrf_resolver_blocks_link_local_metadata_literal() {
let resolver = SsrfGuardResolver;
let name = Name::from_str("169.254.169.254").unwrap();
let result = resolver.resolve(name).await;
assert!(result.is_err(), "metadata IP should be blocked");
}
#[tokio::test]
async fn ssrf_resolver_blocks_private_rfc1918_literal() {
let resolver = SsrfGuardResolver;
for host in ["10.0.0.1", "192.168.1.1", "172.16.0.1"] {
let name = Name::from_str(host).unwrap();
assert!(
resolver.resolve(name).await.is_err(),
"{host} (RFC1918) should be blocked"
);
}
}
#[tokio::test]
async fn ssrf_resolver_allows_public_literal() {
let resolver = SsrfGuardResolver;
let name = Name::from_str("1.1.1.1").unwrap();
let result = resolver.resolve(name).await;
assert!(result.is_ok(), "public IP must be allowed");
}
#[test]
fn shared_clients_build_without_panicking() {
let _ = shared_streaming_http_client();
let _ = shared_request_http_client();
}
#[test]
fn test_thinking_budget_from_effort() {
assert_eq!(thinking_budget::from_effort("low"), Some(1024));
assert_eq!(thinking_budget::from_effort("medium"), Some(4096));
assert_eq!(thinking_budget::from_effort("HIGH"), Some(16384));
assert_eq!(thinking_budget::from_effort("xhigh"), Some(32768));
assert_eq!(thinking_budget::from_effort("unknown"), None);
}
}