Skip to main content

gproxy_protocol/adapt/generate/resources/
mod.rs

1//! Materialize declared foreign URLs/native IDs through scoped ResourceAccess.
2//! Bytes and factual MIME stay in concrete vendor fields; no content IR is used.
3
4mod chat;
5mod claude;
6mod gemini;
7mod responses;
8use crate::{
9    Dialect,
10    capability::{ResourceAccess, ResourceReference},
11    codec::{self, CodecLimits},
12    transform::{TransformError, TransformErrorKind},
13};
14use base64::{Engine as _, engine::general_purpose::STANDARD};
15use std::time::SystemTime;
16
17pub struct GenerationResources<'a, R: ResourceAccess> {
18    pub access: &'a R,
19    pub scope: &'a R::Scope,
20    pub limits: CodecLimits,
21    pub max_references: usize,
22    pub now: SystemTime,
23    /// The dialect the request is converted into. A reference that dialect
24    /// reads as it stands is left in place rather than fetched.
25    pub target: Dialect,
26}
27
28/// Whether a media reference converted from `source` reaches `target` as a
29/// reference the upstream reads itself, so no bytes need fetching. Only
30/// URLs to dialects with a URL form for that media, and file IDs between the
31/// OpenAI dialects, survive conversion; every other reference is fetched.
32/// Tests pin this against what the direct transforms actually keep.
33pub fn reference_passes_through(
34    source: Dialect,
35    target: Dialect,
36    image: bool,
37    reference: &ResourceReference,
38) -> bool {
39    let target = target.pair_dialect();
40    match reference {
41        ResourceReference::Url(_) if image => target != Dialect::Gemini,
42        ResourceReference::Url(_) => matches!(target, Dialect::OpenAi | Dialect::Claude),
43        ResourceReference::Id(_) => {
44            !image
45                && source.family() == crate::WireFamily::OpenAi
46                && target.family() == crate::WireFamily::OpenAi
47        }
48        #[allow(unreachable_patterns)]
49        _ => false,
50    }
51}
52
53pub(super) struct Budget<'a, 'b, R: ResourceAccess> {
54    ctx: &'a GenerationResources<'b, R>,
55    source: Dialect,
56    remaining: u64,
57    references: usize,
58}
59
60struct Media {
61    bytes: bytes::Bytes,
62    mime: String,
63    filename: Option<String>,
64}
65
66impl Media {
67    fn data_uri(&self) -> String {
68        format!("data:{};base64,{}", self.mime, STANDARD.encode(&self.bytes))
69    }
70}
71
72impl<R: ResourceAccess> GenerationResources<'_, R> {
73    fn budget(&self, source: Dialect) -> Budget<'_, '_, R> {
74        Budget {
75            ctx: self,
76            source,
77            remaining: self
78                .limits
79                .max_body_bytes
80                .min(self.access.limits().read_bytes),
81            references: 0,
82        }
83    }
84}
85
86impl<R: ResourceAccess> Budget<'_, '_, R> {
87    /// Fetch a reference the target cannot read, or `None` to leave it.
88    async fn read(
89        &mut self,
90        reference: ResourceReference,
91        image: bool,
92    ) -> Result<Option<Media>, TransformError> {
93        if reference_passes_through(self.source, self.ctx.target, image, &reference) {
94            return Ok(None);
95        }
96        if self.references >= self.ctx.max_references || self.remaining == 0 {
97            return Err(limit());
98        }
99        self.references += 1;
100        let result = self.ctx.access.read(self.ctx.scope, &reference).await?;
101        if result
102            .metadata
103            .length
104            .is_some_and(|length| length > self.remaining)
105        {
106            return Err(limit());
107        }
108        if result
109            .metadata
110            .expires_at
111            .is_some_and(|expiry| expiry <= self.ctx.now)
112        {
113            return Err(TransformError::invalid_result(
114                "resource.expiry",
115                "resource is expired",
116            ));
117        }
118        let mime = result
119            .metadata
120            .mime
121            .ok_or_else(|| TransformError::missing_metadata("resource.mime"))?;
122        if mime.trim().is_empty() || mime.contains([';', ',', '\r', '\n']) {
123            return Err(TransformError::invalid_result(
124                "resource.mime",
125                "exact MIME type required",
126            ));
127        }
128        let mut limits = self.ctx.limits;
129        limits.max_body_bytes = limits.max_body_bytes.min(self.remaining);
130        limits.max_buffer_bytes = limits.max_buffer_bytes.min(self.remaining);
131        let bytes = codec::read_http_body(result.body, limits)
132            .await
133            .map_err(|e| {
134                TransformError::new(
135                    if e.kind() == codec::CodecErrorKind::Limit {
136                        TransformErrorKind::Limit
137                    } else {
138                        TransformErrorKind::Host
139                    },
140                    "resource.body",
141                    e.to_string(),
142                )
143            })?;
144        if result
145            .metadata
146            .length
147            .is_some_and(|n| n != bytes.len() as u64)
148        {
149            return Err(TransformError::invalid_result(
150                "resource.length",
151                "declared length differs from actual bytes",
152            ));
153        }
154        self.remaining = self
155            .remaining
156            .checked_sub(bytes.len() as u64)
157            .ok_or_else(limit)?;
158        if (image || matches!(mime.as_str(), "image/png" | "image/jpeg" | "image/webp"))
159            && crate::transform::images::inspect_image(&bytes)?.mime() != mime
160        {
161            return Err(TransformError::invalid_result(
162                "resource.mime",
163                "image bytes differ from MIME",
164            ));
165        }
166        Ok(Some(Media {
167            bytes,
168            mime,
169            filename: result.metadata.filename,
170        }))
171    }
172}
173
174fn limit() -> TransformError {
175    TransformError::new(
176        TransformErrorKind::Limit,
177        "generation.resources",
178        "resource count or aggregate byte budget exceeded",
179    )
180}