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 HTTP_STREAM_POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(15);
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(|| {
crate::install_default_crypto_provider();
harden_builder(
reqwest::Client::builder()
.connect_timeout(HTTP_CONNECT_TIMEOUT)
.read_timeout(HTTP_STREAM_READ_TIMEOUT)
.pool_idle_timeout(HTTP_STREAM_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(|| {
crate::install_default_crypto_provider();
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 {
use crate::model::ReasoningEffort;
pub const MINIMAL: u32 = 1024;
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: ReasoningEffort) -> Option<u32> {
match effort {
ReasoningEffort::None => None,
ReasoningEffort::Minimal => Some(MINIMAL),
ReasoningEffort::Low => Some(LOW),
ReasoningEffort::Medium => Some(MEDIUM),
ReasoningEffort::High => Some(HIGH),
ReasoningEffort::Xhigh | ReasoningEffort::Max => Some(XHIGH),
}
}
}
const PROTECTED_REQUEST_HEADERS: &[&str] = &[
"host",
"content-length",
"transfer-encoding",
"connection",
"upgrade",
];
pub fn merge_request_headers(
base: Vec<(String, String)>,
extra: &[(String, String)],
) -> Vec<(String, String)> {
let mut merged = base;
for (name, value) in extra {
let name = name.trim();
if name.is_empty() {
continue;
}
if PROTECTED_REQUEST_HEADERS
.iter()
.any(|protected| name.eq_ignore_ascii_case(protected))
{
tracing::warn!(
header = %name,
"Ignoring connection-level header in extra_headers"
);
continue;
}
let mut replaced = false;
merged.retain_mut(|(existing, existing_value)| {
if !existing.eq_ignore_ascii_case(name) {
return true;
}
if replaced {
return false;
}
*existing_value = value.clone();
replaced = true;
true
});
if !replaced {
merged.push((name.to_string(), value.clone()));
}
}
merged
}
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;
#[test]
fn data_urls_preserve_media_and_complete_payload() {
for (url, media, data) in [
("data:image/png;base64,iVBOR", "image/png", "iVBOR"),
("data:image/jpeg;base64,/9j/4AAQ", "image/jpeg", "/9j/4AAQ"),
("data:text/plain,one,two", "text/plain", "one,two"),
("data:image/svg+xml;base64,PD4=", "image/svg+xml", "PD4="),
(
"data:text/plain;charset=utf-8;base64,",
"text/plain;charset=utf-8",
"",
),
] {
let parsed = parse_data_url(url).unwrap();
assert_eq!(
(parsed.media_type.as_str(), parsed.data.as_str()),
(media, data)
);
}
for invalid in [
"https://example.com/image.png",
"data:image/jpeg;base64",
"",
"image/png,data",
] {
assert!(parse_data_url(invalid).is_none(), "{invalid}");
}
}
#[test]
fn payload_classification_requires_status_and_token_context_or_provider_pattern() {
for (message, patterns, matches) in [
("INPUT IS TOO LONG", &[][..], true),
("maximum context reached", &[][..], true),
("request exceeds the maximum token count", &[][..], true),
("context exceeds the maximum", &[][..], true),
("rate exceeds the maximum", &[][..], false),
("authentication failed", &[][..], false),
(
"prompt is too long: 100000 tokens",
ANTHROPIC_TOO_LARGE_PATTERNS,
true,
),
("request size exceeded", ANTHROPIC_TOO_LARGE_PATTERNS, true),
(
"context length exceeded",
ANTHROPIC_TOO_LARGE_PATTERNS,
true,
),
("TOO MANY TOKENS", ANTHROPIC_TOO_LARGE_PATTERNS, true),
(
"request payload size exceeds limit",
GEMINI_TOO_LARGE_PATTERNS,
true,
),
("content too large", GEMINI_TOO_LARGE_PATTERNS, true),
("token limit exceeded", GEMINI_TOO_LARGE_PATTERNS, true),
("prompt is too long", GEMINI_TOO_LARGE_PATTERNS, false),
] {
for (status, expected) in [
(StatusCode::BAD_REQUEST, matches),
(StatusCode::PAYLOAD_TOO_LARGE, true),
(StatusCode::INTERNAL_SERVER_ERROR, false),
(StatusCode::OK, false),
] {
assert_eq!(
is_request_too_large(status, message, patterns),
expected,
"{status}: {message}"
);
}
}
assert!(is_request_too_large(StatusCode::PAYLOAD_TOO_LARGE, "", &[]));
}
#[test]
fn missing_model_classification_requires_404_and_provider_evidence() {
for (message, patterns, matches) in [
(
r#"{"error":{"type":"not_found_error"}}"#,
ANTHROPIC_NOT_FOUND_PATTERNS,
true,
),
("Endpoint not found", ANTHROPIC_NOT_FOUND_PATTERNS, false),
("NOT_FOUND", GEMINI_NOT_FOUND_PATTERNS, true),
("MODEL foo", GEMINI_NOT_FOUND_PATTERNS, true),
("missing endpoint", GEMINI_NOT_FOUND_PATTERNS, false),
("model not found", &[][..], false),
] {
for (status, expected) in [
(StatusCode::NOT_FOUND, matches),
(StatusCode::BAD_REQUEST, false),
(StatusCode::INTERNAL_SERVER_ERROR, false),
(StatusCode::OK, false),
] {
assert_eq!(
is_model_not_found(status, message, patterns),
expected,
"{status}: {message}"
);
}
}
}
#[tokio::test]
async fn resolver_rejects_private_addresses_and_returns_exact_public_address() {
for host in [
"127.0.0.1",
"169.254.169.254",
"10.0.0.1",
"192.168.1.1",
"172.16.0.1",
] {
let error = match SsrfGuardResolver
.resolve(Name::from_str(host).unwrap())
.await
{
Ok(_) => panic!("private address {host} accepted"),
Err(error) => error,
};
assert_eq!(
error.to_string(),
format!("host {host} resolves to blocked address {host} (private/internal)")
);
}
let addresses: Vec<_> = SsrfGuardResolver
.resolve(Name::from_str("1.1.1.1").unwrap())
.await
.unwrap()
.collect();
assert_eq!(addresses, ["1.1.1.1:0".parse().unwrap()]);
}
#[tokio::test]
async fn both_shared_http_clients_refuse_to_follow_redirects() {
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{method, path},
};
let destination = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(200))
.expect(0)
.mount(&destination)
.await;
let source = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/source"))
.respond_with(ResponseTemplate::new(302).insert_header("Location", destination.uri()))
.expect(2)
.mount(&source)
.await;
for client in [shared_streaming_http_client(), shared_request_http_client()] {
let response = client
.get(format!("{}/source", source.uri()))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::FOUND);
assert_eq!(response.headers()["location"], destination.uri());
}
destination.verify().await;
}
#[test]
fn thinking_efforts_have_literal_budgets_including_max_alias() {
use crate::model::ReasoningEffort::*;
for (effort, expected) in [
(None, Option::None),
(Minimal, Some(1024)),
(Low, Some(1024)),
(Medium, Some(4096)),
(High, Some(16384)),
(Xhigh, Some(32768)),
(Max, Some(32768)),
] {
assert_eq!(thinking_budget::from_effort(effort), expected, "{effort:?}");
}
}
#[test]
fn caller_overrides_replace_all_duplicates_preserving_order_and_other_headers() {
assert_eq!(
merge_request_headers(
vec![
("X-Token".into(), "old-a".into()),
("anthropic-version".into(), "2023-06-01".into()),
("x-token".into(), "old-b".into())
],
&[
(" x-TOKEN ".into(), "new".into()),
("Anthropic-Version".into(), "2024-01-01".into()),
("x-trace".into(), "first".into()),
("X-TRACE".into(), "last".into())
]
),
[
("X-Token".into(), "new".into()),
("anthropic-version".into(), "2024-01-01".into()),
("x-trace".into(), "last".into())
]
);
}
#[test]
fn caller_cannot_override_connection_headers_or_insert_blank_names() {
let base = vec![
("content-type".into(), "application/json".into()),
("host".into(), "configured.example".into()),
];
let extra = [
" Host ",
"CONTENT-LENGTH",
"transfer-Encoding",
"connection",
"UPGRADE",
"",
" \t",
]
.map(|name| (name.into(), "untrusted".into()));
assert_eq!(merge_request_headers(base.clone(), &extra), base);
}
}