use super::request::ImageRequestContext;
use crate::{
transform::{Report, TransformError, TransformErrorKind},
wire::{
DeclaredFields, gemini as g,
openai::{images as o, responses as r},
},
};
use base64::Engine;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ImageMetadata {
pub format: o::ImageOutputFormat,
pub width: u32,
pub height: u32,
}
impl ImageMetadata {
pub fn mime(&self) -> &'static str {
match self.format {
o::ImageOutputFormat::Png => "image/png",
o::ImageOutputFormat::Jpeg => "image/jpeg",
o::ImageOutputFormat::Webp => "image/webp",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GeneratedImagePart {
pub bytes: bytes::Bytes,
pub metadata: ImageMetadata,
pub revised_prompt: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ImageUsageFacts {
Responses(r::response::ResponseUsage),
Gemini(g::UsageMetadata),
}
#[derive(Debug, Clone, PartialEq)]
pub struct ImageCallResult {
pub image: GeneratedImagePart,
pub usage: Option<ImageUsageFacts>,
pub response_id: Option<String>,
pub model: Option<String>,
pub report: Report,
}
fn invalid(detail: &str) -> TransformError {
TransformError::invalid_result("image.output", detail)
}
fn u32be(b: &[u8]) -> u32 {
u32::from_be_bytes(b.try_into().unwrap())
}
fn u32le(b: &[u8]) -> u32 {
u32::from_le_bytes(b.try_into().unwrap())
}
fn crc32(b: &[u8]) -> u32 {
let mut c = !0u32;
for x in b {
c ^= u32::from(*x);
for _ in 0..8 {
c = (c >> 1) ^ ((0u32.wrapping_sub(c & 1)) & 0xedb88320);
}
}
!c
}
pub fn inspect_image(b: &[u8]) -> Result<ImageMetadata, TransformError> {
let (format, width, height) = if b.starts_with(b"\x89PNG\r\n\x1a\n") {
let mut pos = 8usize;
let mut dimensions = None;
let mut data = false;
let mut ended = false;
while pos < b.len() {
if b.len() - pos < 12 {
return Err(invalid("truncated PNG chunk"));
}
let len = u32be(&b[pos..pos + 4]) as usize;
let end = pos
.checked_add(12)
.and_then(|p| p.checked_add(len))
.filter(|p| *p <= b.len())
.ok_or_else(|| invalid("truncated PNG payload"))?;
let kind = &b[pos + 4..pos + 8];
if crc32(&b[pos + 4..end - 4]) != u32be(&b[end - 4..end]) {
return Err(invalid("invalid PNG CRC"));
}
if pos == 8 {
if kind != b"IHDR" || len != 13 {
return Err(invalid("PNG requires IHDR"));
}
let width = u32be(&b[pos + 8..pos + 12]);
let height = u32be(&b[pos + 12..pos + 16]);
let depth = b[pos + 16];
let color = b[pos + 17];
if width > 0x7fffffff
|| height > 0x7fffffff
|| !match color {
0 => matches!(depth, 1 | 2 | 4 | 8 | 16),
2 | 4 | 6 => matches!(depth, 8 | 16),
3 => matches!(depth, 1 | 2 | 4 | 8),
_ => false,
}
|| b[pos + 18] != 0
|| b[pos + 19] != 0
|| b[pos + 20] > 1
{
return Err(invalid("invalid PNG header"));
}
dimensions = Some((width, height));
} else if kind == b"IHDR" {
return Err(invalid("duplicate PNG header"));
}
if kind == b"IDAT" && len > 0 {
data = true;
}
if kind == b"IEND" {
if len != 0 || end != b.len() {
return Err(invalid("invalid PNG end"));
}
ended = true;
}
pos = end;
}
if !data || !ended {
return Err(invalid("PNG image data or end missing"));
}
let (w, h) = dimensions.ok_or_else(|| invalid("PNG dimensions missing"))?;
(o::ImageOutputFormat::Png, w, h)
} else if b.starts_with(&[0xff, 0xd8]) {
let mut pos = 2;
let mut dims = None;
let mut scan = false;
let mut ended = false;
while pos < b.len() {
if b[pos] != 0xff {
return Err(invalid("invalid JPEG marker"));
}
while pos < b.len() && b[pos] == 0xff {
pos += 1;
}
if pos >= b.len() {
return Err(invalid("truncated JPEG marker"));
}
let marker = b[pos];
pos += 1;
if marker == 0xd9 {
ended = pos == b.len();
break;
}
if marker == 0x00 || marker == 0xd8 {
return Err(invalid("unexpected JPEG marker"));
}
if marker == 0x01 || (0xd0..=0xd7).contains(&marker) {
continue;
}
if b.len() - pos < 2 {
return Err(invalid("truncated JPEG segment"));
}
let len = u16::from_be_bytes([b[pos], b[pos + 1]]) as usize;
let end = pos
.checked_add(len)
.filter(|p| len >= 2 && *p <= b.len())
.ok_or_else(|| invalid("invalid JPEG segment length"))?;
if matches!(marker,0xc0..=0xc3|0xc5..=0xc7|0xc9..=0xcb|0xcd..=0xcf) {
if len < 8 {
return Err(invalid("truncated JPEG frame"));
}
let h = u16::from_be_bytes([b[pos + 3], b[pos + 4]]) as u32;
let w = u16::from_be_bytes([b[pos + 5], b[pos + 6]]) as u32;
if dims.replace((w, h)).is_some() {
return Err(invalid("multiple JPEG frames unsupported"));
}
}
pos = end;
if marker == 0xda {
scan = true;
while pos < b.len() {
if b[pos] == 0xff {
if pos + 1 >= b.len() {
return Err(invalid("truncated JPEG scan"));
}
if b[pos + 1] == 0x00 || (0xd0..=0xd7).contains(&b[pos + 1]) {
pos += 2;
continue;
}
break;
}
pos += 1;
}
}
}
if !scan || !ended {
return Err(invalid("JPEG scan or end missing"));
}
let (w, h) = dims.ok_or_else(|| invalid("JPEG dimensions missing"))?;
(o::ImageOutputFormat::Jpeg, w, h)
} else if b.len() >= 12 && &b[..4] == b"RIFF" && &b[8..12] == b"WEBP" {
if u32le(&b[4..8]) as u64 + 8 != b.len() as u64 {
return Err(invalid("WebP length mismatch"));
}
let mut pos = 12;
let mut dims = None;
while pos < b.len() {
if b.len() - pos < 8 {
return Err(invalid("truncated WebP chunk"));
}
let size = u32le(&b[pos + 4..pos + 8]) as usize;
let end = pos
.checked_add(8)
.and_then(|p| p.checked_add(size))
.filter(|p| *p <= b.len())
.ok_or_else(|| invalid("truncated WebP payload"))?;
let data = &b[pos + 8..end];
let kind = &b[pos..pos + 4];
let next = if kind == b"VP8 " {
if data.len() < 10 || data[0] & 1 != 0 || &data[3..6] != b"\x9d\x01\x2a" {
return Err(invalid("invalid WebP VP8 frame"));
}
Some((
u32::from(u16::from_le_bytes([data[6], data[7]]) & 0x3fff),
u32::from(u16::from_le_bytes([data[8], data[9]]) & 0x3fff),
))
} else if kind == b"VP8L" {
if data.len() < 5 || data[0] != 0x2f {
return Err(invalid("invalid WebP lossless frame"));
}
let bits = u32le(&data[1..5]);
Some(((bits & 0x3fff) + 1, ((bits >> 14) & 0x3fff) + 1))
} else if kind == b"ANIM" || kind == b"ANMF" {
return Err(invalid("animated WebP is not a generated still image"));
} else {
None
};
if let Some(d) = next
&& dims.replace(d).is_some()
{
return Err(invalid("multiple WebP images"));
}
pos = end
.checked_add(size % 2)
.filter(|p| *p <= b.len())
.ok_or_else(|| invalid("missing WebP padding"))?;
}
let (w, h) = dims.ok_or_else(|| invalid("WebP image data missing"))?;
(o::ImageOutputFormat::Webp, w, h)
} else {
return Err(invalid("unsupported or malformed PNG/JPEG/WebP image"));
};
if width == 0 || height == 0 {
return Err(invalid("zero image dimensions"));
}
Ok(ImageMetadata {
format,
width,
height,
})
}
pub fn decode_image(
encoded: &str,
mime: Option<&str>,
max_bytes: u64,
) -> Result<GeneratedImagePart, TransformError> {
if !encoded.len().is_multiple_of(4) {
return Err(invalid("invalid padded image base64 length"));
}
let padding = if encoded.ends_with("==") {
2
} else if encoded.ends_with('=') {
1
} else {
0
};
let decoded_bound = (encoded.len() as u64 / 4)
.saturating_mul(3)
.saturating_sub(padding);
if decoded_bound > max_bytes {
return Err(TransformError::new(
TransformErrorKind::Limit,
"image.bytes",
"decoded image exceeds cap",
));
}
let bytes = base64::engine::general_purpose::STANDARD
.decode(encoded)
.map_err(|_| invalid("invalid image base64"))?;
if bytes.len() as u64 > max_bytes {
return Err(TransformError::new(
TransformErrorKind::Limit,
"image.bytes",
"decoded image exceeds cap",
));
}
let metadata = inspect_image(&bytes)?;
if mime.is_some_and(|m| m != metadata.mime()) {
return Err(invalid("declared MIME contradicts image bytes"));
}
Ok(GeneratedImagePart {
bytes: bytes.into(),
metadata,
revised_prompt: None,
})
}
pub fn image_response_from_responses(
body: &r::response::GenerateContentResponseBody,
_context: &ImageRequestContext,
max_bytes: u64,
) -> Result<ImageCallResult, TransformError> {
if body.id.trim().is_empty()
|| body.model.trim().is_empty()
|| body.created_at < 0
|| body.status != Some(r::response::ResponseStatus::Completed)
|| body.error.is_some()
|| body.incomplete_details.is_some()
{
return Err(invalid(
"Responses generation did not complete successfully",
));
}
let mut out = None;
let mut report = Report::default();
for item in &body.output {
match item {
r::response::ResponseOutputItem::ImageGenerationCall(call) => {
if call.status != r::input::ImageGenerationStatus::Completed
|| call.id.trim().is_empty()
{
return Err(invalid("image tool did not complete"));
}
let mut image = decode_image(
call.result
.as_deref()
.ok_or_else(|| invalid("missing image result"))?,
None,
max_bytes,
)?;
image.revised_prompt = call.revised_prompt.clone().flatten();
if out.replace(image).is_some() {
return Err(invalid("expected exactly one image per call"));
}
}
r::response::ResponseOutputItem::Reasoning(_) => {
report.omitted("output.reasoning", "Images has no reasoning representation");
}
r::response::ResponseOutputItem::Message(message) => {
if message.status != r::input::OutputMessageStatus::Completed
|| message
.content
.iter()
.any(|p| matches!(p, r::input::OutputContent::Refusal(_)))
{
return Err(invalid("incomplete or refused image response message"));
}
report.omitted(
"output.message",
"ancillary text has no Images representation and is not a revised prompt",
);
}
_ => return Err(invalid("unexpected non-image tool output")),
}
}
let usage = body
.usage
.as_ref()
.and_then(Option::as_ref)
.map(|u| {
if [
u.input_tokens,
u.output_tokens,
u.total_tokens,
u.input_tokens_details.cache_write_tokens,
u.input_tokens_details.cached_tokens,
u.output_tokens_details.reasoning_tokens,
]
.iter()
.any(|n| *n < 0)
|| u.input_tokens.checked_add(u.output_tokens) != Some(u.total_tokens)
|| u.input_tokens_details
.cache_write_tokens
.checked_add(u.input_tokens_details.cached_tokens)
.is_none_or(|v| v > u.input_tokens)
|| u.output_tokens_details.reasoning_tokens > u.output_tokens
{
return Err(invalid("invalid Responses usage"));
}
Ok(ImageUsageFacts::Responses(u.clone().into_declared()))
})
.map(crate::transform::optional)
.transpose()?
.flatten();
Ok(ImageCallResult {
image: out.ok_or_else(|| invalid("missing image result"))?,
usage,
response_id: Some(body.id.clone()),
model: Some(body.model.clone()),
report,
})
}
pub fn image_response_from_gemini(
body: &g::GenerateContentResponseBody,
_context: &ImageRequestContext,
max_bytes: u64,
) -> Result<ImageCallResult, TransformError> {
if body
.prompt_feedback
.as_ref()
.and_then(|p| p.block_reason)
.is_some_and(|v| v != g::BlockReason::Unspecified)
{
return Err(invalid("Gemini blocked image generation"));
}
let candidates = body.candidates.as_deref().unwrap_or_default();
if candidates.len() != 1 {
return Err(invalid("expected exactly one Gemini candidate"));
}
let candidate = &candidates[0];
if candidate.finish_reason != Some(g::FinishReason::Stop)
|| candidate.index.is_some_and(|i| i != 0)
{
return Err(invalid("Gemini candidate did not finish normally"));
}
let content = candidate
.content
.as_ref()
.ok_or_else(|| invalid("missing Gemini image content"))?;
if content.role.as_deref().is_some_and(|r| r != "model") {
return Err(invalid("invalid Gemini candidate role"));
}
let mut out = None;
let mut report = Report::default();
for part in content.parts.as_deref().unwrap_or_default() {
if part.function_call.is_some()
|| part.function_response.is_some()
|| part.file_data.is_some()
|| part.executable_code.is_some()
|| part.code_execution_result.is_some()
|| part.tool_call.is_some()
|| part.tool_response.is_some()
|| part.video_metadata.is_some()
|| part.media_processing.is_some()
|| part.audio_transcription.is_some()
|| part.speech_metadata.is_some()
{
return Err(invalid("unexpected non-image Gemini tool or media output"));
}
if part.text.is_some() && part.inline_data.is_some() {
return Err(invalid("Gemini part contains multiple payloads"));
}
if part.thought == Some(true) || part.text.is_some() {
report.omitted(
"candidate.text/thought",
"ancillary text has no Images representation and is not a revised prompt",
);
continue;
}
let data = part
.inline_data
.as_ref()
.ok_or_else(|| invalid("unexpected empty Gemini output"))?;
let image = decode_image(&data.data, Some(&data.mime_type), max_bytes)?;
if out.replace(image).is_some() {
return Err(invalid("expected exactly one Gemini image"));
}
}
let usage = body
.usage_metadata
.as_ref()
.map(|u| {
if [
u.prompt_token_count,
u.candidates_token_count,
u.cached_content_token_count,
u.tool_use_prompt_token_count,
u.thoughts_token_count,
u.total_token_count,
]
.into_iter()
.flatten()
.any(|n| n < 0)
{
return Err(invalid("negative Gemini usage"));
}
for details in [
&u.prompt_tokens_details,
&u.cache_tokens_details,
&u.candidates_tokens_details,
&u.tool_use_prompt_tokens_details,
]
.into_iter()
.flatten()
{
if details.iter().any(|d| d.token_count.is_some_and(|n| n < 0)) {
return Err(invalid("negative Gemini modality token count"));
}
}
Ok(ImageUsageFacts::Gemini(u.clone().into_declared()))
})
.map(crate::transform::optional)
.transpose()?
.flatten();
Ok(ImageCallResult {
image: out.ok_or_else(|| invalid("missing Gemini image"))?,
usage,
response_id: body.response_id.clone(),
model: body.model_version.clone(),
report,
})
}