yoagent 0.26.5

//! Google Generative AI (Gemini) provider.
//!
//! Uses the `streamGenerateContent` endpoint with SSE streaming.
//! API key is passed as a query parameter.
//!
//! # Prompt caching: implicit only, deliberately
//!
//! This provider sends no cache directives and ignores [`CacheStrategy`].
//! Gemini caches implicitly on its own, and the `cachedContentTokenCount` this
//! module reads back into [`Usage::cache_read`] reports the result of that
//! automatic behaviour — it is telemetry, not an acknowledgement of anything
//! the client asked for.
//!
//! Gemini's *explicit* caching is a separate resource with a different
//! lifecycle: you create a `CachedContent` object, receive a handle, reference
//! it by name on later requests, and manage its TTL and its own billing line.
//! That does not fit behind [`CacheStrategy`], which describes where to place
//! markers inside a single request. Wiring it here would mean either
//! misrepresenting a stateful resource as a per-request flag, or silently
//! creating server-side objects with a lifetime the caller cannot see.
//!
//! It is worth doing as its own API — with explicit create/reference/expire
//! surface — and it is not worth pretending the current enum can express it.
//! See yologdev/yoagent#123.

use super::traits::*;
use crate::types::*;
use futures::StreamExt;
use serde::Deserialize;
use tokio::sync::mpsc;
use tracing::{debug, warn};

pub struct GoogleProvider;

#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
impl StreamProvider for GoogleProvider {
    fn protocol(&self) -> Option<crate::provider::ApiProtocol> {
        Some(crate::provider::ApiProtocol::GoogleGenerativeAi)
    }

    async fn stream(
        &self,
        config: StreamConfig,
        tx: mpsc::UnboundedSender<StreamEvent>,
        cancel: tokio_util::sync::CancellationToken,
    ) -> Result<Message, ProviderError> {
        let model_config = config
            .model_config
            .as_ref()
            .ok_or_else(|| ProviderError::Other("ModelConfig required".into()))?;

        let base_url = &model_config.base_url;
        // The key goes in a header, never the URL: reqwest puts the URL in
        // its error text, which reaches logs, retry events and transcripts.
        let url = format!(
            "{}/v1beta/models/{}:streamGenerateContent?alt=sse",
            base_url, config.model
        );
        // Trimmed: the URL parser used to strip a trailing CR/LF (a key read
        // from a CRLF `.env`), and a header value rejects it. Skipped when the
        // caller sets `x-goog-api-key` in `ModelConfig.headers` (a header is
        // appended, not replaced, so it would be sent twice) or has no key.
        // An `Authorization` header is left alone: it may be a proxy's own.
        let key = config.api_key.trim();
        let user_key = model_config
            .headers
            .keys()
            .any(|k| k.eq_ignore_ascii_case("x-goog-api-key"));
        let key = if key.is_empty() || user_key {
            None
        } else {
            let mut value = reqwest::header::HeaderValue::from_str(key).map_err(|_| {
                ProviderError::Auth(
                    "the Gemini API key contains a character not allowed in an HTTP header \
                     (check GEMINI_API_KEY / GOOGLE_API_KEY, or the configured key)"
                        .into(),
                )
            })?;
            value.set_sensitive(true);
            Some(value)
        };

        let body = build_request_body(&config);
        debug!("Google GenAI request: model={}", config.model);

        let client = reqwest::Client::new();
        let mut request = client.post(&url).header("content-type", "application/json");
        if let Some(key) = key {
            request = request.header("x-goog-api-key", key);
        }

        for (k, v) in &model_config.headers {
            request = request.header(k, v);
        }

        // Google streams JSON chunks separated by newlines, not SSE.
        // With alt=sse, it does use SSE format.
        let response = request
            .json(&body)
            .send()
            .await
            .map_err(|e| ProviderError::Network(e.to_string()))?;

        if !response.status().is_success() {
            let status = response.status();
            let retry_after = parse_retry_after(response.headers());
            let body = response.text().await.unwrap_or_default();
            return Err(ProviderError::classify_with_retry_after(
                status.as_u16(),
                &format!("Google API error {}: {}", status, body),
                retry_after,
            ));
        }

        parse_stream(response, &config.model, &model_config.provider, tx, cancel).await
    }
}

/// Parse a Gemini `streamGenerateContent?alt=sse` response into the turn's
/// message. Shared with Vertex AI, which streams the same format.
pub(crate) async fn parse_stream(
    response: reqwest::Response,
    model: &str,
    provider: &str,
    tx: mpsc::UnboundedSender<StreamEvent>,
    cancel: tokio_util::sync::CancellationToken,
) -> Result<Message, ProviderError> {
    let mut content: Vec<Content> = Vec::new();
    let mut usage = Usage::default();
    let mut stop_reason = StopReason::Stop;
    let mut error_message: Option<String> = None;
    // Every complete Gemini response ends with a `finishReason` (or a
    // `promptFeedback.blockReason` when the prompt itself was blocked). A
    // stream that ends without one was cut off: an error, never a finished
    // (empty or truncated) answer.
    let mut finished = false;

    let _ = tx.send(StreamEvent::Start);

    // Parse SSE stream
    let mut stream = response.bytes_stream();
    let mut buffer = SseBuffer::default();

    loop {
        tokio::select! {
            _ = cancel.cancelled() => {
                return Err(ProviderError::Cancelled);
            }
            chunk = stream.next() => {
                match chunk {
                    None if finished => break,
                    None => {
                        let err = ProviderError::Network(
                            "Gemini stream ended without a finishReason".into(),
                        );
                        warn!("Google stream error: {}", err);
                        return Err(err);
                    }
                    // The response was already complete (a finishReason
                    // arrived): keep it rather than retry and bill it twice.
                    Some(Err(e)) if finished => {
                        warn!("Google stream error after finishReason; keeping the response: {}", e);
                        break;
                    }
                    Some(Err(e)) => {
                        // Match the other providers: a transport failure is an
                        // error (and retryable), not a silently truncated turn.
                        let provider_err = ProviderError::Network(e.to_string());
                        warn!("Google stream error: {}", provider_err);
                        return Err(provider_err);
                    }
                    Some(Ok(bytes)) => {
                        // A transport chunk may end inside a UTF-8 code point.
                        // Decode only once the complete SSE event is buffered.
                        buffer.push(&bytes);

                        // Process complete SSE events (handle both \n\n and \r\n\r\n)
                        while let Some(data) = buffer.next_data() {
                            if data.is_empty() {
                                continue;
                            }

                            // Google reports mid-stream failures as
                            // {"error": {...}} payloads, which would otherwise
                            // deserialize into an empty chunk and vanish.
                            if is_error_payload(&data) {
                                let provider_err = classify_sse_error_event(&data);
                                warn!("Google in-stream error: {}", provider_err);
                                return Err(provider_err);
                            }

                            let chunk: GoogleChunk = match serde_json::from_str(&data) {
                                Ok(c) => c,
                                Err(e) => {
                                    warn!("Failed to parse Google chunk: {}", e);
                                    continue;
                                }
                            };

                            // The prompt itself was blocked: no candidates follow.
                            if let Some(reason) = chunk
                                .prompt_feedback
                                .as_ref()
                                .and_then(|f| f.block_reason.as_deref())
                            {
                                finished = true;
                                warn!("Gemini blocked the prompt (blockReason={})", reason);
                                error_message = Some(format!(
                                    "Prompt blocked by Gemini (blockReason: {})",
                                    reason
                                ));
                                stop_reason = StopReason::Refusal;
                            }

                            // Process candidates
                            for candidate in &chunk.candidates.unwrap_or_default() {
                                if let Some(c) = &candidate.content {
                                    for part in &c.parts {
                                        if let Some(text) = part_text(part) {
                                            if part.thought.unwrap_or(false) {
                                                // Thought summary part → Thinking content.
                                                let think_idx = content.iter().position(|c| matches!(c, Content::Thinking { .. }));
                                                let idx = match think_idx {
                                                    Some(i) => i,
                                                    None => {
                                                        content.push(Content::thinking(String::new()));
                                                        content.len() - 1
                                                    }
                                                };
                                                if let Some(Content::Thinking { thinking, .. }) = content.get_mut(idx) {
                                                    thinking.push_str(text);
                                                }
                                                let _ = tx.send(StreamEvent::ThinkingDelta {
                                                    content_index: idx,
                                                    delta: text.to_string(),
                                                });
                                                continue;
                                            }
                                            let text_idx = content.iter().position(|c| matches!(c, Content::Text { .. }));
                                            let idx = match text_idx {
                                                Some(i) => i,
                                                None => {
                                                    content.push(Content::Text { text: String::new() });
                                                    content.len() - 1
                                                }
                                            };
                                            if let Some(Content::Text { text: t }) = content.get_mut(idx) {
                                                t.push_str(text);
                                            }
                                            let _ = tx.send(StreamEvent::TextDelta {
                                                content_index: idx,
                                                delta: text.to_string(),
                                            });
                                        }
                                        if let Some(fc) = &part.function_call {
                                            let id = fc.id.clone().unwrap_or_else(|| format!("google-fc-{}", content.len()));
                                            let args = fc.args.clone().unwrap_or(serde_json::Value::Object(Default::default()));
                                            let metadata = part.thought_signature.as_ref().map(|sig| {
                                                serde_json::json!({"thought_signature": sig})
                                            });
                                            let idx = content.len();
                                            content.push(Content::ToolCall {
                                                id: id.clone(),
                                                name: fc.name.clone(),
                                                arguments: args,
                                                provider_metadata: metadata,
                                            });
                                            let _ = tx.send(StreamEvent::ToolCallStart {
                                                content_index: idx,
                                                id,
                                                name: fc.name.clone(),
                                            });
                                            let _ = tx.send(StreamEvent::ToolCallEnd { content_index: idx });
                                            stop_reason = StopReason::ToolUse;
                                        }
                                    }
                                }
                                if let Some(reason) = &candidate.finish_reason {
                                    finished = true;
                                    // Don't override ToolUse -- Gemini returns "STOP"
                                    // even when it emits function calls
                                    if stop_reason != StopReason::ToolUse {
                                        stop_reason = match reason.as_str() {
                                            "STOP" => StopReason::Stop,
                                            "MAX_TOKENS" | "RECITATION" => StopReason::Length,
                                            // A tool call Gemini could not produce or
                                            // was not allowed to make: a failed turn,
                                            // with Gemini's own explanation when given.
                                            "MALFORMED_FUNCTION_CALL" | "UNEXPECTED_TOOL_CALL"
                                            | "TOO_MANY_TOOL_CALLS" => {
                                                let detail = candidate
                                                    .finish_message
                                                    .as_deref()
                                                    .map(|m| format!(": {m}"))
                                                    .unwrap_or_default();
                                                warn!("Gemini tool call failed (finishReason={}){}", reason, detail);
                                                error_message = Some(format!(
                                                    "Gemini failed to make a valid tool call (finishReason: {reason}){detail}. Retrying the request usually works."
                                                ));
                                                StopReason::Error
                                            }
                                            "SAFETY" | "PROHIBITED_CONTENT" | "BLOCKLIST"
                                            | "SPII" | "IMAGE_SAFETY" => {
                                                warn!(
                                                    "Gemini blocked the response (finishReason={})",
                                                    reason
                                                );
                                                error_message = Some(format!(
                                                    "Response blocked by Gemini safety filters (finishReason: {})",
                                                    reason
                                                ));
                                                StopReason::Refusal
                                            }
                                            other => {
                                                warn!("Gemini finished with an unhandled finishReason={}", other);
                                                StopReason::Stop
                                            }
                                        };
                                    }
                                }
                            }

                            // Process usage
                            if let Some(u) = &chunk.usage_metadata {
                                // promptTokenCount includes cached tokens;
                                // keep `input` as the uncached remainder so
                                // downstream sums don't double-count.
                                usage.input = u
                                    .prompt_token_count
                                    .unwrap_or(0)
                                    .saturating_sub(u.cached_content_token_count.unwrap_or(0));
                                // Thinking tokens bill as output ("Output price
                                // (including thinking tokens)") but are counted
                                // apart, in `thoughtsTokenCount`.
                                usage.output = u
                                    .candidates_token_count
                                    .unwrap_or(0)
                                    .saturating_add(u.thoughts_token_count.unwrap_or(0));
                                usage.total_tokens = u.total_token_count.unwrap_or(0);
                                usage.cache_read = u.cached_content_token_count.unwrap_or(0);
                            }
                        }
                    }
                }
            }
        }
    }

    let message = Message::Assistant {
        content,
        stop_reason,
        model: model.to_string(),
        provider: provider.to_string(),
        usage,
        timestamp: now_ms(),
        error_message,
    };

    let _ = tx.send(StreamEvent::Done {
        message: message.clone(),
    });
    Ok(message)
}

/// Bytes of an SSE stream not yet split into events. Bytes, not text: a
/// transport chunk may end inside a UTF-8 code point, so each event is
/// decoded only once it is complete.
#[derive(Default)]
struct SseBuffer {
    bytes: Vec<u8>,
    /// How far a separator has been searched for: a search resumes there
    /// (less the three bytes a split `\r\n\r\n` can leave), so a large
    /// event arriving in small chunks is scanned once, not once per chunk.
    searched: usize,
}

impl SseBuffer {
    fn push(&mut self, chunk: &[u8]) {
        self.bytes.extend_from_slice(chunk);
    }

    /// Pop the next complete SSE event and return its `data:` payload (empty
    /// string when the event carries no data line). Handles both `\n\n` and
    /// `\r\n\r\n` event separators, splitting at whichever occurs first.
    /// Returns `None` until a complete event is buffered. Only the first
    /// `data:` line of an event is returned; invalid UTF-8 in it decodes
    /// lossily.
    fn next_data(&mut self) -> Option<String> {
        let from = self.searched.saturating_sub(3);
        let found = (from..self.bytes.len()).find_map(|i| {
            let rest = &self.bytes[i..];
            if rest.starts_with(b"\r\n\r\n") {
                Some((i, 4))
            } else if rest.starts_with(b"\n\n") {
                Some((i, 2))
            } else {
                None
            }
        });
        let Some((pos, sep_len)) = found else {
            self.searched = self.bytes.len();
            return None;
        };
        let event_str = String::from_utf8_lossy(&self.bytes[..pos]).into_owned();
        self.bytes.drain(..pos + sep_len);
        self.searched = 0;
        let data = event_str
            .lines()
            .map(|l| l.trim_end_matches('\r'))
            .find(|l| l.starts_with("data: "))
            .map(|l| l[6..].to_string())
            .unwrap_or_default();
        Some(data)
    }
}

#[cfg(test)]
impl From<&[u8]> for SseBuffer {
    fn from(bytes: &[u8]) -> Self {
        Self {
            bytes: bytes.to_vec(),
            searched: 0,
        }
    }
}

/// Whether an SSE data payload is a Google error envelope
/// (`{"error": {...}}`) rather than a content chunk.
fn is_error_payload(data: &str) -> bool {
    serde_json::from_str::<serde_json::Value>(data)
        .map(|v| v.get("error").is_some())
        .unwrap_or(false)
}

/// Non-empty text of a part. Gemini streams empty text parts while thinking;
/// those must be skipped.
fn part_text(part: &GooglePart) -> Option<&str> {
    part.text.as_deref().filter(|t| !t.is_empty())
}

/// Token budget for Gemini 2.x's `thinkingConfig.thinkingBudget` per level.
/// Only reached through [`gemini_thinking_config`], which Vertex AI shares.
pub(crate) fn gemini_thinking_budget(level: ThinkingLevel) -> u32 {
    match level {
        ThinkingLevel::Off => 0,
        ThinkingLevel::Minimal | ThinkingLevel::Low => 1024,
        ThinkingLevel::Medium => 8192,
        // Clamped: 24,576 is the documented thinkingBudget maximum of Gemini
        // 2.5 Flash and Flash-Lite; 2.5 Pro's is 32,768
        // (cloud.google.com/vertex-ai/generative-ai/docs/thinking). There is
        // no per-model table here, so the crate stays within the smallest
        // documented maximum rather than send an out-of-range budget.
        ThinkingLevel::High | ThinkingLevel::XHigh | ThinkingLevel::Max => 24576,
    }
}

/// Which `thinkingConfig` field a Gemini model takes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum GeminiThinkingParam {
    /// `thinkingBudget` — Gemini 2.x, and any id the version rule cannot read
    /// (Gemini 3 still accepts a budget for backward compatibility).
    Budget,
    /// `thinkingLevel` — Gemini 3 and later, restricted to the rungs the
    /// model documents as accepted.
    Level(GeminiRungs),
    /// No `thinkingConfig` at all — Gemini 3+ TTS models, which neither
    /// thinking guide lists as supporting thinking.
    Omit,
}

/// The `thinkingLevel` values a Gemini 3+ model accepts.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum GeminiRungs {
    /// `MINIMAL`, `LOW`, `MEDIUM`, `HIGH` — 3.6 / 3.5 Flash, 3.x Flash-Lite,
    /// 3 Flash.
    All,
    /// `LOW`, `MEDIUM`, `HIGH` — 3.8 / 3.7 Flash, 3.1 Pro, and every 3.x+
    /// model not documented as accepting `MINIMAL`.
    NoMinimal,
    /// `MINIMAL`, `HIGH` — 3.1 Flash Image and 3.1 Flash-Lite Image.
    MinimalHigh,
    /// `HIGH` only — 3 Pro Image.
    HighOnly,
}

/// Reads the Gemini generation off a model id.
///
/// Google's rule is version-based: `thinkingLevel` is "Recommended for Gemini
/// 3 or later models. Use with earlier models results in an error." (REST
/// reference, `ThinkingConfig`). So the version number decides, not a model
/// list. Accepts bare ids (`gemini-3.8-flash`), `models/…` names, Vertex
/// resource paths (`projects/…/publishers/google/models/gemini-3.1-pro-preview`)
/// and `@version` suffixes.
///
/// `MINIMAL` per the Gemini API thinking guide's level table: rejected by 3.8
/// and 3.7 Flash, not supported by 3.1 Pro; accepted by 3.6 / 3.5 Flash, every
/// 3.x Flash-Lite and 3 Flash. Unlisted 3.x+ models (Pro variants, future
/// Flash releases) are treated as not accepting it, so the worst case is
/// `LOW` where `MINIMAL` would have worked, never a 400.
///
/// Image models, per the Vertex AI thinking table ("Supported thinking_level
/// values"): Gemini 3.1 Flash Image and 3.1 Flash-Lite Image take
/// "MINIMAL , HIGH"; Gemini 3 Pro Image takes "HIGH" only. Every 3.x+
/// `-image` id is read that way (`-pro…-image` → `HIGH` only, any other →
/// `MINIMAL`/`HIGH`).
///
/// TTS models (`gemini-3.8-flash-tts`, `gemini-3.8-flash-lite-tts`,
/// `gemini-3.1-flash-tts-preview`) appear in neither guide's list of models
/// that support thinking, so they get no `thinkingConfig`.
///
/// Aliases: `gemini-flash-latest` (3.5 Flash since 2026-05-19) and
/// `gemini-pro-latest` (3 Pro preview since 2026-01-21) point at Gemini 3
/// per the API changelog; they are hot-swapped, so `MINIMAL` is not assumed.
/// `gemini-flash-lite-latest` has no documented target, so it keeps the
/// budget, which every generation accepts.
fn gemini_thinking_param(model: &str) -> GeminiThinkingParam {
    let name = model.rsplit('/').next().unwrap_or(model);
    let name = name.split('@').next().unwrap_or(name).to_ascii_lowercase();
    let Some(rest) = name.strip_prefix("gemini-") else {
        return GeminiThinkingParam::Budget;
    };
    if rest == "flash-latest" || rest == "pro-latest" {
        return GeminiThinkingParam::Level(GeminiRungs::NoMinimal);
    }

    let major_len = rest.bytes().take_while(u8::is_ascii_digit).count();
    let Ok(major) = rest[..major_len].parse::<u32>() else {
        return GeminiThinkingParam::Budget;
    };
    if major < 3 {
        return GeminiThinkingParam::Budget;
    }
    let after_major = &rest[major_len..];
    let (minor, variant) = match after_major.strip_prefix('.') {
        Some(tail) => {
            let len = tail.bytes().take_while(u8::is_ascii_digit).count();
            (tail[..len].parse::<u32>().ok(), &tail[len..])
        }
        None => (Some(0), after_major),
    };

    if variant.contains("-tts") {
        return GeminiThinkingParam::Omit;
    }
    let rungs = if variant.contains("-image") {
        if variant.starts_with("-pro") {
            GeminiRungs::HighOnly
        } else {
            GeminiRungs::MinimalHigh
        }
    } else if major == 3
        && (variant.starts_with("-flash-lite")
            || (variant.starts_with("-flash") && matches!(minor, Some(0 | 5 | 6))))
    {
        GeminiRungs::All
    } else {
        GeminiRungs::NoMinimal
    };
    GeminiThinkingParam::Level(rungs)
}

/// The `thinkingLevel` string for a non-`Off` [`ThinkingLevel`] on a Gemini
/// 3+ model: the requested rung, or the nearest one the model accepts.
///
/// `Minimal` → `LOW` where `MINIMAL` is not accepted. On the `MINIMAL`/`HIGH`
/// image models, `Low` goes down to `MINIMAL` (both are the "minimize latency
/// and cost" end) and `Medium` goes up to `HIGH` (a request for some thinking
/// gets thinking). `Off` never reaches here in practice: it omits
/// `thinkingConfig`, and is mapped like `Minimal` only to keep the match total.
fn gemini_thinking_level(level: ThinkingLevel, rungs: GeminiRungs) -> &'static str {
    use GeminiRungs::*;
    use ThinkingLevel::*;
    match (rungs, level) {
        (HighOnly, _) => "HIGH",
        (MinimalHigh, Off | Minimal | Low) => "MINIMAL",
        (MinimalHigh, _) => "HIGH",
        (All, Off | Minimal) => "MINIMAL",
        (_, Off | Minimal | Low) => "LOW",
        (_, Medium) => "MEDIUM",
        (_, High | XHigh | Max) => "HIGH",
    }
}

/// The `generationConfig.thinkingConfig` object for a request, or `None` to
/// omit it. Shared by the Gemini API and Vertex AI providers so the two
/// cannot drift.
///
/// Exactly one of `thinkingLevel` / `thinkingBudget` is ever sent: Gemini 3
/// rejects a request carrying both, and pre-3 models reject `thinkingLevel`.
/// The choice comes from [`GoogleCompat::thinking_level`] when set, else
/// from the model id (see `gemini_thinking_param`).
///
/// `ThinkingLevel::Off` omits `thinkingConfig` on every model. On Gemini 3
/// that does **not** disable thinking — the model runs at its own default
/// level (3.1 Pro `HIGH`, 3.5–3.8 Flash `MEDIUM`, Flash-Lite `MINIMAL`);
/// `ThinkingLevel::Minimal` is the way to ask for the least thinking.
///
/// [`GoogleCompat::thinking_level`]: super::GoogleCompat::thinking_level
pub(crate) fn gemini_thinking_config(config: &StreamConfig) -> Option<serde_json::Value> {
    if config.thinking_level == ThinkingLevel::Off {
        return None;
    }
    let forced = config
        .model_config
        .as_ref()
        .and_then(|mc| mc.google.as_ref())
        .and_then(|g| g.thinking_level);
    let param = match (forced, gemini_thinking_param(&config.model)) {
        (Some(false), _) => GeminiThinkingParam::Budget,
        (Some(true), GeminiThinkingParam::Budget) => {
            GeminiThinkingParam::Level(GeminiRungs::NoMinimal)
        }
        (_, inferred) => inferred,
    };

    match param {
        // Gemini 2.x: the payload this crate has always sent — budget scales
        // with the level and includeThoughts streams thought summaries back
        // as thought parts.
        GeminiThinkingParam::Budget => Some(serde_json::json!({
            "thinkingBudget": gemini_thinking_budget(config.thinking_level),
            "includeThoughts": true,
        })),
        GeminiThinkingParam::Level(rungs) => Some(serde_json::json!({
            "thinkingLevel": gemini_thinking_level(config.thinking_level, rungs),
            "includeThoughts": true,
        })),
        GeminiThinkingParam::Omit => None,
    }
}

/// Each user content's `functionResponse` parts first, then its other parts,
/// each group in its order. A merged turn of tool results would otherwise
/// interleave a result's images between the responses (`response A, image
/// A, response B`); Gemini matches responses to the calls by count and
/// position, and images among them are a reported cause of its "number of
/// function response parts" 400.
pub(crate) fn function_responses_first(
    mut contents: Vec<serde_json::Value>,
) -> Vec<serde_json::Value> {
    for content in &mut contents {
        if let Some(parts) = content
            .get_mut("parts")
            .and_then(serde_json::Value::as_array_mut)
        {
            let (mut responses, rest): (Vec<_>, Vec<_>) = std::mem::take(parts)
                .into_iter()
                .partition(|part| part.get("functionResponse").is_some());
            responses.extend(rest);
            *parts = responses;
        }
    }
    contents
}

fn build_request_body(config: &StreamConfig) -> serde_json::Value {
    let mut contents: Vec<serde_json::Value> = Vec::new();

    for msg in &config.messages {
        match msg {
            Message::User { content, .. } => {
                let parts = content_to_google_parts(content);
                contents.push(serde_json::json!({
                    "role": "user",
                    "parts": parts,
                }));
            }
            Message::Assistant { content, .. } => {
                let parts = content_to_google_parts(content);
                contents.push(serde_json::json!({
                    "role": "model",
                    "parts": parts,
                }));
            }
            Message::ToolResult {
                tool_call_id,
                tool_name,
                content,
                ..
            } => {
                let text = super::tool_result_text(content);

                let mut fr = serde_json::json!({
                    "name": tool_name,
                    "response": {"result": text},
                });
                if !tool_call_id.is_empty() && !tool_call_id.starts_with("google-fc-") {
                    fr["id"] = serde_json::json!(tool_call_id);
                }
                let mut parts = vec![serde_json::json!({"functionResponse": fr})];

                // Append image parts if present
                for c in content {
                    if let Content::Image { data, mime_type } = c {
                        parts.push(serde_json::json!({
                            "inlineData": {"mimeType": mime_type, "data": data},
                        }));
                    }
                }

                contents.push(serde_json::json!({
                    "role": "user",
                    "parts": parts,
                }));
            }
        }
    }

    // Gemini wants a turn's function responses in one content, matching its
    // function calls — the responses first, then any images and text.
    let contents = function_responses_first(super::merge_same_role(contents, "parts"));
    let mut body = serde_json::json!({
        "contents": contents,
    });

    if !config.system_prompt.is_empty() {
        body["systemInstruction"] = serde_json::json!({
            "parts": [{"text": config.system_prompt}],
        });
    }

    let mut generation_config = serde_json::json!({});
    if let Some(max) = config.max_tokens {
        generation_config["maxOutputTokens"] = serde_json::json!(max);
    }
    if let Some(temp) = config.temperature {
        generation_config["temperature"] = serde_json::json!(temp);
    }
    // Structured outputs: native responseSchema (Gemini's OpenAPI-style
    // schema dialect — pass the caller's schema through as given).
    if let Some(schema) = &config.output_schema {
        generation_config["responseMimeType"] = serde_json::json!("application/json");
        generation_config["responseSchema"] = schema.schema.clone();
    }

    // Thinking: thinkingLevel on Gemini 3+, thinkingBudget on 2.x.
    if let Some(thinking) = gemini_thinking_config(config) {
        generation_config["thinkingConfig"] = thinking;
    }

    if generation_config != serde_json::json!({}) {
        body["generationConfig"] = generation_config;
    }

    if !config.tools.is_empty() {
        let declarations: Vec<serde_json::Value> = config
            .tools
            .iter()
            .map(|t| {
                serde_json::json!({
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.parameters,
                })
            })
            .collect();
        body["tools"] = serde_json::json!([{
            "functionDeclarations": declarations,
        }]);
    }

    body
}

fn content_to_google_parts(content: &[Content]) -> Vec<serde_json::Value> {
    content
        .iter()
        .filter(|c| !matches!(c, Content::Text { text } if text.is_empty()))
        .filter_map(|c| match c {
            Content::Text { text } => Some(serde_json::json!({"text": text})),
            Content::Image { data, mime_type } => Some(serde_json::json!({
                "inlineData": {"mimeType": mime_type, "data": data},
            })),
            Content::ToolCall {
                id,
                name,
                arguments,
                provider_metadata,
            } => {
                let mut fc = serde_json::json!({"name": name, "args": arguments});
                if !id.is_empty() && !id.starts_with("google-fc-") {
                    fc["id"] = serde_json::json!(id);
                }
                let mut part = serde_json::json!({"functionCall": fc});
                if let Some(sig) = provider_metadata
                    .as_ref()
                    .and_then(|m| m.get("thought_signature"))
                    .and_then(|v| v.as_str())
                {
                    part["thoughtSignature"] = serde_json::json!(sig);
                }
                Some(part)
            }
            Content::Thinking { .. } => None,
        })
        .collect()
}

// Google API response types
#[derive(Deserialize)]
struct GoogleChunk {
    #[serde(default)]
    candidates: Option<Vec<GoogleCandidate>>,
    #[serde(default, rename = "promptFeedback")]
    prompt_feedback: Option<GooglePromptFeedback>,
    #[serde(default, rename = "usageMetadata")]
    usage_metadata: Option<GoogleUsageMetadata>,
}

#[derive(Deserialize)]
struct GooglePromptFeedback {
    #[serde(default, rename = "blockReason")]
    block_reason: Option<String>,
}

#[derive(Deserialize)]
struct GoogleCandidate {
    #[serde(default)]
    content: Option<GoogleContent>,
    #[serde(default, rename = "finishReason")]
    finish_reason: Option<String>,
    /// Gemini's explanation for some finish reasons (a malformed call).
    #[serde(default, rename = "finishMessage")]
    finish_message: Option<String>,
}

#[derive(Deserialize)]
struct GoogleContent {
    #[serde(default)]
    parts: Vec<GooglePart>,
}

#[derive(Deserialize)]
struct GooglePart {
    #[serde(default)]
    text: Option<String>,
    /// True when this part is a thought summary (thinkingConfig.includeThoughts).
    #[serde(default)]
    thought: Option<bool>,
    #[serde(default, rename = "functionCall")]
    function_call: Option<GoogleFunctionCall>,
    #[serde(default, rename = "thoughtSignature")]
    thought_signature: Option<String>,
}

#[derive(Deserialize)]
struct GoogleFunctionCall {
    name: String,
    #[serde(default)]
    args: Option<serde_json::Value>,
    #[serde(default)]
    id: Option<String>,
}

#[derive(Deserialize)]
struct GoogleUsageMetadata {
    #[serde(default, rename = "promptTokenCount")]
    prompt_token_count: Option<u64>,
    #[serde(default, rename = "candidatesTokenCount")]
    candidates_token_count: Option<u64>,
    #[serde(default, rename = "totalTokenCount")]
    total_token_count: Option<u64>,
    #[serde(default, rename = "cachedContentTokenCount")]
    cached_content_token_count: Option<u64>,
    #[serde(default, rename = "thoughtsTokenCount")]
    thoughts_token_count: Option<u64>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::provider::{GoogleCompat, ModelConfig};

    /// A model turn whose only content is thinking (not replayed), or the
    /// loop's error message (one empty text block), is not sent as an empty
    /// content (#369).
    #[test]
    fn an_empty_model_turn_is_not_sent() {
        for turn in [
            vec![Content::thinking("thought")],
            vec![Content::Text {
                text: String::new(),
            }],
        ] {
            let mut config = StreamConfig::new("gemini-2.5-pro", "k");
            config.messages = vec![
                Message::user("first"),
                Message::assistant(turn, StopReason::Error, "m", "google", Usage::default()),
                Message::user("again"),
            ];
            let body = build_request_body(&config);
            assert_eq!(
                body["contents"],
                serde_json::json!([{"role": "user", "parts": [{"text": "first"}, {"text": "again"}]}])
            );
        }
    }

    /// A parallel tool call's results go back as one user turn, with the
    /// steering message after them, so roles alternate (#223).
    #[test]
    fn parallel_tool_results_go_back_in_one_turn() {
        let mut config = StreamConfig::new("gemini-2.5-pro", "k");
        config.messages = crate::provider::parallel_tool_conversation();
        let body = build_request_body(&config);
        let contents = body["contents"].as_array().unwrap();
        let roles: Vec<&str> = contents
            .iter()
            .map(|c| c["role"].as_str().unwrap())
            .collect();
        assert_eq!(roles, ["user", "model", "user"], "{body}");
        let parts = contents[2]["parts"].as_array().unwrap();
        assert_eq!(parts[0]["functionResponse"]["name"], "read_file");
        assert_eq!(
            parts[0]["functionResponse"]["response"]["result"],
            "fn a() {}"
        );
        assert_eq!(
            parts[1]["functionResponse"]["response"]["result"],
            "fn b() {}"
        );
        assert_eq!(parts[2]["text"], "also check c.rs");
        assert_eq!(parts.len(), 3);
    }

    /// A result with an image: every function response first, in call
    /// order, then the image and the text that followed.
    #[test]
    fn merged_tool_results_put_the_function_responses_first() {
        let mut config = StreamConfig::new("gemini-2.5-pro", "k");
        config.messages = crate::provider::parallel_tool_conversation();
        if let Message::ToolResult { content, .. } = &mut config.messages[2] {
            content.push(Content::Image {
                data: "AAAA".into(),
                mime_type: "image/png".into(),
            });
        }
        let body = build_request_body(&config);
        let parts = body["contents"][2]["parts"].as_array().unwrap();
        assert_eq!(parts.len(), 4, "{body}");
        assert_eq!(
            parts[0]["functionResponse"]["response"]["result"],
            "fn a() {}"
        );
        assert_eq!(
            parts[1]["functionResponse"]["response"]["result"],
            "fn b() {}"
        );
        assert_eq!(parts[2]["inlineData"]["data"], "AAAA");
        assert_eq!(parts[3]["text"], "also check c.rs");
    }

    #[test]
    fn thinking_blocks_are_dropped_on_replay() {
        // Gemini does not accept echoed thought summaries; signatures ride on
        // functionCall parts instead. Pin the drop so it stays deliberate.
        let mut config = StreamConfig::new("gemini-2.5-pro", "k");
        config.messages = vec![
            Message::user("go"),
            Message::assistant(
                vec![
                    Content::thinking("thought"),
                    Content::Text {
                        text: "answer".into(),
                    },
                ],
                StopReason::Stop,
                "m",
                "google",
                Usage::default(),
            ),
        ];
        let body = build_request_body(&config);
        let parts = body["contents"][1]["parts"].as_array().unwrap();
        assert_eq!(parts.len(), 1, "only the text part is replayed");
        assert_eq!(parts[0]["text"], "answer");
    }

    #[test]
    fn thinking_level_sets_thinking_config() {
        let config = StreamConfig {
            model: "gemini-2.5-pro".into(),
            system_prompt: "".into(),
            messages: vec![Message::user("hi")],
            tools: vec![],
            thinking_level: ThinkingLevel::Medium,
            api_key: "test".into(),
            max_tokens: None,
            temperature: None,
            model_config: None,
            cache_config: CacheConfig::default(),
            output_schema: None,
        };
        let body = build_request_body(&config);
        assert_eq!(
            body["generationConfig"]["thinkingConfig"]["thinkingBudget"],
            8192
        );
        assert_eq!(
            body["generationConfig"]["thinkingConfig"]["includeThoughts"],
            true
        );
    }

    #[test]
    fn thinking_off_omits_thinking_config() {
        let config = StreamConfig {
            model: "gemini-2.5-pro".into(),
            system_prompt: "".into(),
            messages: vec![Message::user("hi")],
            tools: vec![],
            thinking_level: ThinkingLevel::Off,
            api_key: "test".into(),
            max_tokens: None,
            temperature: None,
            model_config: None,
            cache_config: CacheConfig::default(),
            output_schema: None,
        };
        let body = build_request_body(&config);
        assert!(body["generationConfig"]["thinkingConfig"].is_null());
    }

    #[test]
    fn structured_output_sets_response_schema() {
        let config = StreamConfig {
            model: "gemini-2.5-pro".into(),
            system_prompt: "".into(),
            messages: vec![Message::user("Hello")],
            tools: vec![],
            thinking_level: ThinkingLevel::Off,
            api_key: "test".into(),
            max_tokens: None,
            temperature: None,
            model_config: None,
            cache_config: CacheConfig::default(),
            output_schema: Some(crate::provider::OutputSchema::new(
                "structured_output",
                serde_json::json!({"type": "object", "properties": {"x": {"type": "number"}}}),
            )),
        };
        let body = build_request_body(&config);
        assert_eq!(
            body["generationConfig"]["responseMimeType"],
            "application/json"
        );
        assert_eq!(
            body["generationConfig"]["responseSchema"]["properties"]["x"]["type"],
            "number"
        );
    }

    #[test]
    fn test_build_google_request() {
        let config = StreamConfig {
            model: "gemini-2.0-flash".into(),
            system_prompt: "Be helpful".into(),
            messages: vec![Message::user("Hello")],
            tools: vec![],
            thinking_level: ThinkingLevel::Off,
            api_key: "test".into(),
            max_tokens: Some(1024),
            temperature: Some(0.7),
            model_config: None,
            cache_config: CacheConfig::default(),
            output_schema: None,
        };

        let body = build_request_body(&config);
        assert!(body["contents"].is_array());
        assert_eq!(body["contents"][0]["role"], "user");
        assert!(body["systemInstruction"].is_object());
        assert_eq!(body["generationConfig"]["maxOutputTokens"], 1024);
        let temp = body["generationConfig"]["temperature"].as_f64().unwrap();
        assert!((temp - 0.7).abs() < 0.01);
    }

    #[test]
    fn test_content_to_google_parts_text() {
        let content = vec![Content::Text {
            text: "hello".into(),
        }];
        let parts = content_to_google_parts(&content);
        assert_eq!(parts.len(), 1);
        assert_eq!(parts[0]["text"], "hello");
    }

    #[test]
    fn test_content_to_google_parts_filters_empty_text() {
        let content = vec![
            Content::Text { text: "".into() },
            Content::Text {
                text: "hello".into(),
            },
            Content::Text { text: "".into() },
        ];
        let parts = content_to_google_parts(&content);
        assert_eq!(parts.len(), 1);
        assert_eq!(parts[0]["text"], "hello");
    }

    #[test]
    fn test_content_to_google_parts_tool_call() {
        let content = vec![Content::ToolCall {
            id: "tc-1".into(),
            name: "bash".into(),
            arguments: serde_json::json!({"command": "ls"}),
            provider_metadata: None,
        }];
        let parts = content_to_google_parts(&content);
        assert_eq!(parts[0]["functionCall"]["name"], "bash");
    }

    #[test]
    fn test_parse_chunk_with_function_call_and_thought_signature() {
        let data = r#"{"candidates": [{"content": {"parts": [{"functionCall": {"name": "bash", "args": {"command": "echo hi"}, "id": "abc123"}, "thoughtSignature": "SIG_DATA"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}}"#;

        let chunk: GoogleChunk = serde_json::from_str(data).unwrap();
        let candidates = chunk.candidates.unwrap();
        assert_eq!(candidates.len(), 1);

        let parts = &candidates[0].content.as_ref().unwrap().parts;
        assert_eq!(parts.len(), 1);

        let fc = parts[0].function_call.as_ref().unwrap();
        assert_eq!(fc.name, "bash");
        assert_eq!(fc.id.as_deref(), Some("abc123"));
        assert_eq!(fc.args.as_ref().unwrap()["command"], "echo hi");

        assert_eq!(parts[0].thought_signature.as_deref(), Some("SIG_DATA"));
    }

    #[test]
    fn test_parse_chunk_with_empty_text() {
        // Gemini sends empty text parts during thinking -- part_text (used by
        // the streaming loop) must skip them and keep non-empty ones.
        let data = r#"{"candidates": [{"content": {"parts": [{"text": ""}, {"text": "Hello"}], "role": "model"}, "index": 0}]}"#;

        let chunk: GoogleChunk = serde_json::from_str(data).unwrap();
        let candidates = chunk.candidates.unwrap();
        let parts = &candidates[0].content.as_ref().unwrap().parts;
        assert_eq!(part_text(&parts[0]), None, "empty text parts are skipped");
        assert_eq!(part_text(&parts[1]), Some("Hello"));
    }

    #[test]
    fn test_parse_chunk_with_crlf_sse() {
        // Full pipeline: SseBuffer (the production splitter) on a CRLF
        // stream, then chunk parsing.
        let mut buf = SseBuffer::from("data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"Blue\"}], \"role\": \"model\"}, \"finishReason\": \"STOP\", \"index\": 0}]}\r\n\r\n".as_bytes());

        let data = buf.next_data().expect("complete CRLF event");
        assert!(buf.bytes.is_empty(), "event consumed from buffer");
        assert_eq!(buf.next_data(), None);

        let chunk: GoogleChunk = serde_json::from_str(&data).unwrap();
        let candidates = chunk.candidates.unwrap();
        let text = &candidates[0].content.as_ref().unwrap().parts[0].text;
        assert_eq!(text.as_deref(), Some("Blue"));
    }

    #[test]
    fn test_next_sse_data_partial_events_stay_buffered() {
        let mut buf = SseBuffer::from("data: {\"a\":1}\n\ndata: partial".as_bytes());
        assert_eq!(buf.next_data().as_deref(), Some("{\"a\":1}"));
        assert_eq!(buf.next_data(), None, "incomplete event waits");
        assert_eq!(buf.bytes, b"data: partial");
        buf.push(b"\n\n");
        assert_eq!(buf.next_data().as_deref(), Some("partial"));
    }

    #[test]
    fn test_next_sse_data_consumes_events_without_data_lines() {
        // SSE comments/keepalives and leading separators must be CONSUMED
        // (returning Some("")), never None — returning None would wedge the
        // buffer and drop every subsequent event.
        let mut buf = SseBuffer::from(": keepalive\n\n\n\ndata: x\n\n".as_bytes());
        assert_eq!(buf.next_data().as_deref(), Some(""));
        assert_eq!(buf.next_data().as_deref(), Some(""));
        assert_eq!(buf.next_data().as_deref(), Some("x"));
        assert_eq!(buf.next_data(), None);
    }

    #[test]
    fn test_is_error_payload() {
        assert!(is_error_payload(
            r#"{"error": {"code": 429, "status": "RESOURCE_EXHAUSTED"}}"#
        ));
        assert!(!is_error_payload(r#"{"candidates": []}"#));
        assert!(!is_error_payload("not json"));
    }

    #[test]
    fn test_next_sse_data_splits_earliest_separator_first() {
        // A CRLF-separated event earlier in the buffer must split before a
        // later LF separator (the old inline logic preferred the LF match and
        // merged the two events).
        let mut buf = SseBuffer::from("data: one\r\n\r\ndata: two\n\n".as_bytes());
        assert_eq!(buf.next_data().as_deref(), Some("one"));
        assert_eq!(buf.next_data().as_deref(), Some("two"));
        assert_eq!(buf.next_data(), None);
    }

    #[test]
    fn test_sse_buffer_decodes_invalid_utf8_lossily_and_moves_on() {
        // A malformed byte must not wedge the buffer (`None` means "wait
        // for more"): the event decodes with U+FFFD and the next one follows.
        let mut buf = SseBuffer::from(&b"data: {\"a\":\"\xFF\"}\n\ndata: x\n\n"[..]);
        let first = buf.next_data().expect("the malformed event is consumed");
        assert!(first.contains('\u{FFFD}'), "{first:?}");
        assert_eq!(buf.next_data().as_deref(), Some("x"));
        assert!(buf.bytes.is_empty());
    }

    #[test]
    fn test_sse_buffer_joins_a_code_point_and_a_separator_split_across_pushes() {
        // `你` is three bytes; every push ends mid-character or mid-separator.
        let bytes = "data: 你\r\n\r\ndata: 好\n\n".as_bytes();
        let mut buf = SseBuffer::default();
        let mut got = Vec::new();
        for b in bytes {
            buf.push(std::slice::from_ref(b));
            while let Some(data) = buf.next_data() {
                got.push(data);
            }
        }
        assert_eq!(got, ["你", "好"]);
        assert!(buf.bytes.is_empty());
    }

    #[test]
    fn test_sse_buffer_finds_every_event_a_push_completes() {
        // A search that found nothing leaves `searched` at the end; the next
        // push may complete several events, and each must be found — the
        // second one starts before where the first search stopped.
        let mut buf = SseBuffer::default();
        buf.push(b"data: aaaaaaaa");
        assert_eq!(buf.next_data(), None);
        buf.push(b"\n\ndata: b\n\n");
        assert_eq!(buf.next_data().as_deref(), Some("aaaaaaaa"));
        assert_eq!(buf.next_data().as_deref(), Some("b"));
        assert!(buf.bytes.is_empty());
    }

    #[test]
    fn test_sse_buffer_splits_the_same_however_the_stream_is_chunked() {
        // Differential: random streams over separator bytes (and a
        // multi-byte character), cut into random chunks — empty ones
        // included — must split exactly as the whole stream does at once.
        let mut state = 0x2545_F491_4F6C_DD1Du64;
        let mut next = move |n: u64| {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            state % n
        };
        let alphabet: [&[u8]; 6] = [b"\r", b"\n", b"\n", b"a", b"data: ", "好".as_bytes()];
        let drain = |buf: &mut SseBuffer, out: &mut Vec<String>| {
            while let Some(data) = buf.next_data() {
                out.push(data);
            }
        };
        for case in 0..5_000 {
            let mut stream = Vec::new();
            for _ in 0..next(24) {
                stream.extend_from_slice(alphabet[next(6) as usize]);
            }
            let mut whole = Vec::new();
            drain(&mut SseBuffer::from(&stream[..]), &mut whole);

            let mut chunked = Vec::new();
            let mut buf = SseBuffer::default();
            let mut at = 0;
            while at < stream.len() {
                let end = (at + next(8) as usize).min(stream.len());
                buf.push(&stream[at..end]);
                drain(&mut buf, &mut chunked);
                at = end;
            }
            assert_eq!(
                chunked,
                whole,
                "case {case}: {:?}",
                String::from_utf8_lossy(&stream)
            );
        }
    }

    #[test]
    fn test_thought_signature_round_trip() {
        let content = vec![Content::ToolCall {
            id: "abc123".into(),
            name: "bash".into(),
            arguments: serde_json::json!({"command": "echo hi"}),
            provider_metadata: Some(serde_json::json!({"thought_signature": "SIG_DATA"})),
        }];

        let parts = content_to_google_parts(&content);
        assert_eq!(parts.len(), 1);

        assert_eq!(parts[0]["functionCall"]["name"], "bash");
        assert_eq!(parts[0]["functionCall"]["id"], "abc123");
        assert_eq!(parts[0]["functionCall"]["args"]["command"], "echo hi");
        assert_eq!(parts[0]["thoughtSignature"], "SIG_DATA");
    }

    #[test]
    fn test_tool_call_without_thought_signature() {
        // Synthetic IDs (google-fc-*) should not be sent to Gemini
        let content = vec![Content::ToolCall {
            id: "google-fc-0".into(),
            name: "bash".into(),
            arguments: serde_json::json!({"command": "ls"}),
            provider_metadata: None,
        }];

        let parts = content_to_google_parts(&content);
        assert!(parts[0]["functionCall"].get("id").is_none());
        assert!(parts[0].get("thoughtSignature").is_none());
    }

    #[test]
    fn test_function_response_includes_id() {
        let config = StreamConfig {
            model: "gemini-2.5-flash".into(),
            system_prompt: "".into(),
            messages: vec![
                Message::Assistant {
                    content: vec![Content::ToolCall {
                        id: "abc123".into(),
                        name: "bash".into(),
                        arguments: serde_json::json!({"command": "echo hi"}),
                        provider_metadata: None,
                    }],
                    stop_reason: StopReason::ToolUse,
                    model: "test".into(),
                    provider: "test".into(),
                    usage: Usage::default(),
                    timestamp: 0,
                    error_message: None,
                },
                Message::ToolResult {
                    tool_call_id: "abc123".into(),
                    tool_name: "bash".into(),
                    content: vec![Content::Text { text: "hi".into() }],
                    is_error: false,
                    timestamp: 0,
                },
            ],
            tools: vec![],
            thinking_level: ThinkingLevel::Off,
            api_key: "test".into(),
            max_tokens: None,
            temperature: None,
            model_config: None,
            cache_config: CacheConfig::default(),
            output_schema: None,
        };

        let body = build_request_body(&config);
        let msgs = body["contents"].as_array().unwrap();
        let tool_result = &msgs[1]["parts"][0]["functionResponse"];
        assert_eq!(tool_result["name"], "bash");
        assert_eq!(tool_result["id"], "abc123");
        assert_eq!(tool_result["response"]["result"], "hi");
    }

    #[test]
    fn test_function_response_synthetic_id_omitted() {
        let config = StreamConfig {
            model: "gemini-2.5-flash".into(),
            system_prompt: "".into(),
            messages: vec![
                Message::Assistant {
                    content: vec![Content::ToolCall {
                        id: "google-fc-0".into(),
                        name: "bash".into(),
                        arguments: serde_json::json!({"command": "ls"}),
                        provider_metadata: None,
                    }],
                    stop_reason: StopReason::ToolUse,
                    model: "test".into(),
                    provider: "test".into(),
                    usage: Usage::default(),
                    timestamp: 0,
                    error_message: None,
                },
                Message::ToolResult {
                    tool_call_id: "google-fc-0".into(),
                    tool_name: "bash".into(),
                    content: vec![Content::Text {
                        text: "output".into(),
                    }],
                    is_error: false,
                    timestamp: 0,
                },
            ],
            tools: vec![],
            thinking_level: ThinkingLevel::Off,
            api_key: "test".into(),
            max_tokens: None,
            temperature: None,
            model_config: None,
            cache_config: CacheConfig::default(),
            output_schema: None,
        };

        let body = build_request_body(&config);
        let msgs = body["contents"].as_array().unwrap();
        let tool_result = &msgs[1]["parts"][0]["functionResponse"];
        assert!(
            tool_result.get("id").is_none(),
            "Synthetic ID should not be included"
        );
    }

    // --- thinkingLevel (Gemini 3+) vs thinkingBudget (2.x) -----------------

    const ALL_LEVELS: [ThinkingLevel; 7] = [
        ThinkingLevel::Off,
        ThinkingLevel::Minimal,
        ThinkingLevel::Low,
        ThinkingLevel::Medium,
        ThinkingLevel::High,
        ThinkingLevel::XHigh,
        ThinkingLevel::Max,
    ];

    fn thinking_config(model: &str, level: ThinkingLevel) -> serde_json::Value {
        let mut config = StreamConfig::new(model, "k");
        config.messages = vec![Message::user("hi")];
        config.thinking_level = level;
        build_request_body(&config)["generationConfig"]["thinkingConfig"].clone()
    }

    fn thinking_config_with(
        model: &str,
        level: ThinkingLevel,
        compat: GoogleCompat,
    ) -> serde_json::Value {
        let mut config = StreamConfig::new(model, "k");
        config.messages = vec![Message::user("hi")];
        config.thinking_level = level;
        let mut mc = ModelConfig::google(model, "test");
        mc.google = Some(compat);
        config.model_config = Some(mc);
        build_request_body(&config)["generationConfig"]["thinkingConfig"].clone()
    }

    /// The payload every release before this one sent, for every level.
    fn legacy_budget_payload(level: ThinkingLevel) -> serde_json::Value {
        let budget = match level {
            ThinkingLevel::Off => return serde_json::Value::Null,
            ThinkingLevel::Minimal | ThinkingLevel::Low => 1024,
            ThinkingLevel::Medium => 8192,
            _ => 24576,
        };
        serde_json::json!({"thinkingBudget": budget, "includeThoughts": true})
    }

    #[test]
    fn gemini_2x_payload_is_unchanged_budget() {
        // Near-misses: every id here must keep the pre-3 payload byte for byte.
        for model in [
            "gemini-2.5-pro",
            "gemini-2.5-flash",
            "gemini-2.5-flash-lite",
            "gemini-2.5-flash-preview-09-2025",
            "gemini-2.0-flash",
            "models/gemini-2.5-flash",
            "projects/p/locations/us-central1/publishers/google/models/gemini-2.5-pro",
            "gemini-2.5-pro@001",
            "gemini-flash-lite-latest",
            "gemini-exp-1206",
            "gemma-3-27b-it",
        ] {
            for level in ALL_LEVELS {
                let got = thinking_config(model, level);
                assert_eq!(got, legacy_budget_payload(level), "{model} {level:?}");
                assert_eq!(
                    serde_json::to_string(&got).unwrap(),
                    serde_json::to_string(&legacy_budget_payload(level)).unwrap(),
                    "{model} {level:?}: byte-identical"
                );
            }
        }
    }

    #[test]
    fn gemini_3_8_flash_sends_level_and_clamps_minimal_to_low() {
        // 3.8 / 3.7 Flash reject MINIMAL ("Not supported (error)").
        for model in [
            "gemini-3.8-flash",
            "gemini-3.7-flash",
            "gemini-3.8-flash-cyber",
        ] {
            let expect = [
                None,        // Off: omitted, the model runs at its default
                Some("LOW"), // Minimal clamped
                Some("LOW"),
                Some("MEDIUM"),
                Some("HIGH"),
                Some("HIGH"), // XHigh clamped
                Some("HIGH"), // Max clamped
            ];
            for (level, want) in ALL_LEVELS.into_iter().zip(expect) {
                let got = thinking_config(model, level);
                match want {
                    None => assert!(got.is_null(), "{model} {level:?}: {got}"),
                    Some(want) => assert_eq!(
                        got,
                        serde_json::json!({"thinkingLevel": want, "includeThoughts": true}),
                        "{model} {level:?}: level only, never both"
                    ),
                }
            }
        }
    }

    #[test]
    fn gemini_3_flash_lite_and_3_5_flash_pass_minimal_through() {
        for model in [
            "gemini-3.5-flash-lite",
            "gemini-3.1-flash-lite",
            "gemini-3.5-flash",
            "gemini-3.6-flash",
            "gemini-3-flash-preview",
        ] {
            assert_eq!(
                thinking_config(model, ThinkingLevel::Minimal),
                serde_json::json!({"thinkingLevel": "MINIMAL", "includeThoughts": true}),
                "{model}"
            );
            assert_eq!(
                thinking_config(model, ThinkingLevel::Low)["thinkingLevel"],
                "LOW"
            );
        }
    }

    #[test]
    fn gemini_3_pro_has_no_minimal() {
        for model in [
            "gemini-3.1-pro-preview",
            "gemini-3-pro-preview",
            "gemini-pro-latest",
            "gemini-flash-latest", // hot-swapped alias: MINIMAL not assumed
            "gemini-3.8-flash",    // rejects MINIMAL (thinking guide)
            "gemini-4-argon",      // rejects MINIMAL; thinking cannot be off
        ] {
            assert_eq!(
                thinking_config(model, ThinkingLevel::Minimal)["thinkingLevel"],
                "LOW",
                "{model}"
            );
            assert_eq!(
                thinking_config(model, ThinkingLevel::Max)["thinkingLevel"],
                "HIGH",
                "{model}"
            );
        }
    }

    #[test]
    fn off_omits_thinking_config_on_every_gemini_generation() {
        // User decision: Off sends nothing. On Gemini 3 the model then runs
        // at its own default level (it cannot be switched off).
        for model in [
            "gemini-3.8-flash",
            "gemini-3.5-flash-lite",
            "gemini-3.1-pro-preview",
            "gemini-3-flash-preview",
            "gemini-flash-latest",
            "projects/p/locations/global/publishers/google/models/gemini-3.8-flash",
            "gemini-3.1-flash-image",
            "gemini-3-pro-image",
            "gemini-3.8-flash-tts",
            "gemini-2.5-flash",
        ] {
            let mut config = StreamConfig::new(model, "k");
            config.messages = vec![Message::user("hi")];
            config.thinking_level = ThinkingLevel::Off;
            let body = build_request_body(&config);
            assert!(
                body["generationConfig"].get("thinkingConfig").is_none(),
                "{model}: {body}"
            );
        }
        // Forcing thinkingLevel does not bring Off back either.
        assert!(thinking_config_with(
            "my-gemini-alias",
            ThinkingLevel::Off,
            GoogleCompat::force_thinking_level()
        )
        .is_null());
    }

    #[test]
    fn positive_control_off_revert_only_touches_off() {
        // The same 3.x ids still get a thinkingConfig one rung up, so the
        // omission above is Off-specific, not a broken 3.x path.
        for model in [
            "gemini-3.8-flash",
            "gemini-3.5-flash-lite",
            "gemini-3.1-pro-preview",
        ] {
            assert!(
                thinking_config(model, ThinkingLevel::Minimal)
                    .get("thinkingLevel")
                    .is_some(),
                "{model}"
            );
        }
    }

    #[test]
    fn gemini_image_models_get_only_their_accepted_levels() {
        // Vertex thinking table: 3.1 Flash Image and 3.1 Flash-Lite Image take
        // "MINIMAL , HIGH"; 3 Pro Image takes "HIGH" only.
        for model in [
            "gemini-3.1-flash-image",
            "gemini-3.1-flash-image-preview",
            "gemini-3.1-flash-lite-image",
            "projects/p/locations/global/publishers/google/models/gemini-3.1-flash-image",
        ] {
            let expect = [
                None,
                Some("MINIMAL"),
                Some("MINIMAL"), // Low: down to the low-cost end
                Some("HIGH"),    // Medium: up, a request for thinking gets it
                Some("HIGH"),
                Some("HIGH"),
                Some("HIGH"),
            ];
            for (level, want) in ALL_LEVELS.into_iter().zip(expect) {
                let got = thinking_config(model, level);
                match want {
                    None => assert!(got.is_null(), "{model} {level:?}"),
                    Some(want) => assert_eq!(
                        got,
                        serde_json::json!({"thinkingLevel": want, "includeThoughts": true}),
                        "{model} {level:?}"
                    ),
                }
            }
        }
        for model in ["gemini-3-pro-image", "gemini-3-pro-image-preview"] {
            assert!(thinking_config(model, ThinkingLevel::Off).is_null());
            for level in ALL_LEVELS.into_iter().skip(1) {
                assert_eq!(
                    thinking_config(model, level),
                    serde_json::json!({"thinkingLevel": "HIGH", "includeThoughts": true}),
                    "{model} {level:?}"
                );
            }
        }
        // 2.5 Flash Image stays on the legacy budget path, byte for byte.
        for level in ALL_LEVELS {
            assert_eq!(
                thinking_config("gemini-2.5-flash-image", level),
                legacy_budget_payload(level)
            );
        }
    }

    #[test]
    fn gemini_3_tts_models_get_no_thinking_config() {
        for model in [
            "gemini-3.8-flash-tts",
            "gemini-3.8-flash-lite-tts",
            "gemini-3.1-flash-tts-preview",
        ] {
            for level in ALL_LEVELS {
                assert!(thinking_config(model, level).is_null(), "{model} {level:?}");
            }
        }
        // Positive control: the non-TTS sibling does get one.
        assert_eq!(
            thinking_config("gemini-3.8-flash", ThinkingLevel::Low)["thinkingLevel"],
            "LOW"
        );
        // 2.x TTS ids keep the pre-3 payload unchanged.
        assert_eq!(
            thinking_config("gemini-2.5-flash-preview-tts", ThinkingLevel::Low),
            legacy_budget_payload(ThinkingLevel::Low)
        );
    }

    #[test]
    fn gemini_3_ids_are_read_through_prefixes_and_paths() {
        for model in [
            "gemini-3",
            "models/gemini-3.8-flash",
            "GEMINI-3.8-FLASH",
            "projects/p/locations/global/publishers/google/models/gemini-3.8-flash",
            "gemini-3.8-flash@001",
            "gemini-4-flash", // later generations inherit the version rule
        ] {
            let got = thinking_config(model, ThinkingLevel::Medium);
            assert_eq!(
                got,
                serde_json::json!({"thinkingLevel": "MEDIUM", "includeThoughts": true}),
                "{model}"
            );
        }
    }

    #[test]
    fn positive_control_2_5_and_3_x_payloads_differ() {
        // Guards against a helper that always returns the same shape.
        let old = thinking_config("gemini-2.5-flash", ThinkingLevel::High);
        let new = thinking_config("gemini-3.5-flash", ThinkingLevel::High);
        assert_ne!(old, new);
        assert_eq!(old["thinkingBudget"], 24576);
        assert_eq!(new["thinkingLevel"], "HIGH");
    }

    #[test]
    fn google_compat_overrides_the_model_id_rule() {
        // Force level on an id the rule cannot read (a proxy alias).
        let got = thinking_config_with(
            "my-gemini-alias",
            ThinkingLevel::Minimal,
            GoogleCompat::force_thinking_level(),
        );
        assert_eq!(
            got,
            serde_json::json!({"thinkingLevel": "LOW", "includeThoughts": true})
        );
        // Forcing level on a recognised 3.x id keeps its MINIMAL knowledge.
        let got = thinking_config_with(
            "gemini-3.5-flash-lite",
            ThinkingLevel::Minimal,
            GoogleCompat::force_thinking_level(),
        );
        assert_eq!(got["thinkingLevel"], "MINIMAL");
        // Force budget on a 3.x id: the legacy payload, exactly.
        for level in ALL_LEVELS {
            let got = thinking_config_with(
                "gemini-3.8-flash",
                level,
                GoogleCompat::force_thinking_budget(),
            );
            assert_eq!(got, legacy_budget_payload(level), "{level:?}");
        }
        // Default compat == no compat.
        assert_eq!(
            thinking_config_with(
                "gemini-3.8-flash",
                ThinkingLevel::High,
                GoogleCompat::default()
            ),
            thinking_config("gemini-3.8-flash", ThinkingLevel::High)
        );
    }

    #[test]
    fn google_compat_serde_default_and_round_trip() {
        let compat: GoogleCompat = serde_json::from_value(serde_json::json!({})).unwrap();
        assert_eq!(compat, GoogleCompat::default());
        let mut mc = ModelConfig::google("gemini-2.5-flash", "G");
        let before = serde_json::to_value(&mc).unwrap();
        assert!(
            before.get("google").is_none(),
            "None is omitted on serialize"
        );
        mc.google = Some(GoogleCompat::force_thinking_level());
        let back: ModelConfig = serde_json::from_value(serde_json::to_value(&mc).unwrap()).unwrap();
        assert_eq!(back.google, Some(GoogleCompat::force_thinking_level()));
    }

    #[test]
    fn tool_result_sends_every_text_block() {
        // An MCP server may answer in several text blocks (ESLint's: an intro,
        // the findings, a closing note); the model must see all of them.
        let config = StreamConfig {
            model: "gemini-2.5-flash".into(),
            system_prompt: "".into(),
            messages: vec![Message::ToolResult {
                tool_call_id: "call-1".into(),
                tool_name: "lint".into(),
                content: vec![
                    Content::Text {
                        text: "Here are the results:".into(),
                    },
                    Content::Text {
                        text: "[{\"line\":3}]".into(),
                    },
                ],
                is_error: false,
                timestamp: 0,
            }],
            tools: vec![],
            thinking_level: ThinkingLevel::Off,
            api_key: "test".into(),
            max_tokens: None,
            temperature: None,
            model_config: None,
            cache_config: CacheConfig::default(),
            output_schema: None,
        };
        let body = build_request_body(&config);
        assert_eq!(
            body["contents"][0]["parts"][0]["functionResponse"]["response"]["result"],
            "Here are the results:\n[{\"line\":3}]"
        );
    }

    #[test]
    fn tool_result_with_an_image_sends_every_text_block_and_the_image() {
        let config = StreamConfig {
            model: "gemini-2.5-flash".into(),
            system_prompt: "".into(),
            messages: vec![Message::ToolResult {
                tool_call_id: "call-1".into(),
                tool_name: "screenshot".into(),
                content: vec![
                    Content::Text { text: "a".into() },
                    Content::Image {
                        data: "AAAA".into(),
                        mime_type: "image/png".into(),
                    },
                    Content::Text { text: "b".into() },
                ],
                is_error: false,
                timestamp: 0,
            }],
            tools: vec![],
            thinking_level: ThinkingLevel::Off,
            api_key: "test".into(),
            max_tokens: None,
            temperature: None,
            model_config: None,
            cache_config: CacheConfig::default(),
            output_schema: None,
        };
        let body = build_request_body(&config);
        let parts = &body["contents"][0]["parts"];
        assert_eq!(parts[0]["functionResponse"]["response"]["result"], "a\nb");
        assert_eq!(parts[1]["inlineData"]["data"], "AAAA");
        assert_eq!(parts.as_array().unwrap().len(), 2);
    }
}