Skip to main content

systemprompt_api/services/gateway/image_fetch/
mod.rs

1//! Resolving caller-supplied image URLs to inline base64 for wires that
2//! cannot carry a URL.
3//!
4//! Gemini's `generateContent` has no URL image part at all — `inlineData` or a
5//! Files API handle are the only shapes it accepts — so the wire codec, which
6//! is synchronous and has no HTTP client, can only downgrade a URL image to
7//! text. This module does the fetch one layer up, in the dispatch pipeline,
8//! before the body is built, and rewrites the canonical request in place so
9//! the codec sees an image it can render.
10//!
11//! It is deliberately not a general-purpose fetcher. The URL comes from
12//! whoever sent the inference request, so every fetch is guarded by `guard`,
13//! bounded by a timeout, capped while the body streams, and accepted only if
14//! the server declares a MIME type Gemini takes.
15//!
16//! Copyright (c) systemprompt.io — Business Source License 1.1.
17//! See <https://systemprompt.io> for licensing details.
18
19mod guard;
20
21use base64::Engine as _;
22use base64::engine::general_purpose::STANDARD as BASE64;
23use systemprompt_models::net::{HTTP_CONNECT_TIMEOUT, trusted_http_hosts_from_env};
24
25use super::protocol::canonical::{CanonicalContent, CanonicalRequest, ImageSource};
26
27// Why: Gemini caps a whole `generateContent` request at 20 MB inline, and
28// base64 inflates by 4/3. A 5 MiB ceiling per image leaves a conversation room
29// for several images plus its text inside that budget, and is already well
30// above what any real photograph in a prompt weighs.
31pub const MAX_IMAGE_BYTES: usize = 5 * 1024 * 1024;
32
33// Why: the shapes Gemini documents for `inlineData`. A server declaring
34// anything else is either not serving an image or serving one the model cannot
35// decode; both are failures, not things to inline and hope.
36pub const ACCEPTED_MIME: [&str; 5] = [
37    "image/png",
38    "image/jpeg",
39    "image/webp",
40    "image/heic",
41    "image/heif",
42];
43
44const MAX_REDIRECTS: u8 = 3;
45const FETCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
46
47/// A caller-supplied image URL that could not be turned into inline data.
48///
49/// `caller_fault` separates "this URL was never going to work" — blocked host,
50/// wrong content type, too large — from a transport failure reaching an
51/// otherwise legitimate host, so the route layer can answer 400 or 502.
52#[derive(Debug, thiserror::Error)]
53#[error("image url {url} could not be inlined: {message}")]
54pub struct ImageFetchFailed {
55    pub url: String,
56    pub message: String,
57    pub caller_fault: bool,
58}
59
60/// Per-request bounds, so a test can point the fetcher at a loopback mock
61/// without the process-wide trust list that production reads from the
62/// environment.
63#[derive(Debug, Clone)]
64pub struct ImageFetchPolicy {
65    pub timeout: std::time::Duration,
66    pub max_bytes: usize,
67    pub max_redirects: u8,
68    pub trusted_hosts: Vec<String>,
69}
70
71impl Default for ImageFetchPolicy {
72    fn default() -> Self {
73        Self {
74            timeout: FETCH_TIMEOUT,
75            max_bytes: MAX_IMAGE_BYTES,
76            max_redirects: MAX_REDIRECTS,
77            trusted_hosts: trusted_http_hosts_from_env(),
78        }
79    }
80}
81
82/// Fetched bytes plus the MIME type the server declared for them.
83#[derive(Debug, Clone)]
84pub struct InlineImage {
85    pub media_type: String,
86    pub base64: String,
87}
88
89// Why: redirects are followed by hand so every hop is re-checked against the
90// guard. reqwest's own policy would resolve a redirect to 169.254.169.254
91// internally, and the only URL this code ever saw would be the innocent one.
92fn client() -> &'static reqwest::Client {
93    static CLIENT: std::sync::OnceLock<reqwest::Client> = std::sync::OnceLock::new();
94    CLIENT.get_or_init(|| {
95        reqwest::Client::builder()
96            .redirect(reqwest::redirect::Policy::none())
97            .connect_timeout(HTTP_CONNECT_TIMEOUT)
98            .build()
99            .unwrap_or_default()
100    })
101}
102
103// Why: the first failure aborts and returns rather than skipping the image. An
104// image the caller asked the model to look at is part of the prompt, and
105// answering about a prompt that quietly lost one of its inputs is the defect
106// this module exists to remove.
107pub async fn inline_url_images(
108    request: &mut CanonicalRequest,
109    policy: &ImageFetchPolicy,
110) -> Result<usize, ImageFetchFailed> {
111    let mut count = 0usize;
112    for message in &mut request.messages {
113        for content in &mut message.content {
114            let CanonicalContent::Image(ImageSource::Url { url, detail }) = content else {
115                continue;
116            };
117            let fetched = fetch(url, policy).await?;
118            *content = CanonicalContent::Image(ImageSource::Base64 {
119                media_type: fetched.media_type,
120                data: fetched.base64,
121                detail: *detail,
122            });
123            count += 1;
124        }
125    }
126    Ok(count)
127}
128
129// Why: the timeout wraps guard, connect, redirects and body read together, so
130// a host that stalls each step just under a per-step budget still cannot hold
131// the inference request open.
132pub async fn fetch(url: &str, policy: &ImageFetchPolicy) -> Result<InlineImage, ImageFetchFailed> {
133    let fail = |message: String, caller_fault: bool| ImageFetchFailed {
134        url: url.to_owned(),
135        message,
136        caller_fault,
137    };
138    tokio::time::timeout(policy.timeout, fetch_inner(url, policy))
139        .await
140        .map_or_else(
141            |_| Err(fail(format!("fetch exceeded {:?}", policy.timeout), false)),
142            |result| result.map_err(|(message, caller_fault)| fail(message, caller_fault)),
143        )
144}
145
146type FetchError = (String, bool);
147
148async fn fetch_inner(url: &str, policy: &ImageFetchPolicy) -> Result<InlineImage, FetchError> {
149    let mut next = guard::checked_url(url, &policy.trusted_hosts)
150        .await
151        .map_err(|e| (e, true))?;
152    for _ in 0..=policy.max_redirects {
153        let response = client()
154            .get(next.clone())
155            .send()
156            .await
157            .map_err(|e| (format!("request failed: {e}"), false))?;
158        if let Some(location) = redirect_target(&response) {
159            let joined = next
160                .join(&location)
161                .map_err(|e| (format!("invalid redirect target: {e}"), true))?;
162            next = guard::checked_url(joined.as_str(), &policy.trusted_hosts)
163                .await
164                .map_err(|e| (format!("redirect rejected: {e}"), true))?;
165            continue;
166        }
167        return read_image(response, policy).await;
168    }
169    Err((
170        format!("more than {} redirects", policy.max_redirects),
171        true,
172    ))
173}
174
175fn redirect_target(response: &reqwest::Response) -> Option<String> {
176    if !response.status().is_redirection() {
177        return None;
178    }
179    response
180        .headers()
181        .get(reqwest::header::LOCATION)
182        .and_then(|v| v.to_str().ok())
183        .map(ToOwned::to_owned)
184}
185
186async fn read_image(
187    mut response: reqwest::Response,
188    policy: &ImageFetchPolicy,
189) -> Result<InlineImage, FetchError> {
190    let status = response.status();
191    if !status.is_success() {
192        return Err((format!("host returned {status}"), true));
193    }
194    let media_type = declared_mime(&response)?;
195    // Why: the cap is enforced chunk by chunk rather than on the finished body,
196    // so a host advertising nothing and sending gigabytes is dropped after the
197    // first 5 MiB instead of being buffered whole and measured afterwards.
198    let mut body: Vec<u8> = Vec::new();
199    while let Some(chunk) = response
200        .chunk()
201        .await
202        .map_err(|e| (format!("read failed: {e}"), false))?
203    {
204        if body.len() + chunk.len() > policy.max_bytes {
205            return Err((format!("larger than {} bytes", policy.max_bytes), true));
206        }
207        body.extend_from_slice(&chunk);
208    }
209    if body.is_empty() {
210        return Err(("empty response body".to_owned(), true));
211    }
212    Ok(InlineImage {
213        media_type,
214        base64: BASE64.encode(&body),
215    })
216}
217
218fn declared_mime(response: &reqwest::Response) -> Result<String, FetchError> {
219    let raw = response
220        .headers()
221        .get(reqwest::header::CONTENT_TYPE)
222        .and_then(|v| v.to_str().ok())
223        .ok_or_else(|| ("no content-type".to_owned(), true))?;
224    let mime = raw
225        .split(';')
226        .next()
227        .unwrap_or_default()
228        .trim()
229        .to_ascii_lowercase();
230    if ACCEPTED_MIME.contains(&mime.as_str()) {
231        return Ok(mime);
232    }
233    Err((
234        format!("content-type {mime} is not an inlineable image"),
235        true,
236    ))
237}