bamboo-llm 2026.10.21

//! Google Gemini provider implementation.

mod stream;

pub use stream::{parse_gemini_sse_event, GeminiStreamState};

use crate::providers::common::log_identity;
use async_trait::async_trait;
use reqwest::{
    header::{HeaderMap, HeaderValue, CONTENT_TYPE},
    Client,
};
use serde_json::{json, Value};

use crate::protocol::gemini::GeminiRequest;
use crate::protocol::ToProvider;
use crate::provider::{
    LLMError, LLMProvider, LLMRequestOptions, LLMStream, ProviderVisibleToolFootprint,
    ProviderVisibleToolSegment, ProviderVisibleToolSegmentKind, Result,
};
use crate::providers::common::model_fetcher;
use crate::providers::common::request_overrides;
use crate::types::LLMChunk;
use bamboo_config::{KeywordMaskingConfig, RequestOverridesConfig};
use bamboo_domain::Message;
use bamboo_domain::ReasoningEffort;

pub(crate) fn apply_required_tool_choice(body: &mut Value, required_tool: Option<&str>) {
    if let Some(name) = required_tool {
        body["toolConfig"] = json!({
            "functionCallingConfig": {
                "mode": "ANY",
                "allowedFunctionNames": [name]
            }
        });
    }
}

pub(crate) fn apply_generation_config(
    request: &mut GeminiRequest,
    max_output_tokens: Option<u32>,
    reasoning_effort: Option<ReasoningEffort>,
) {
    if max_output_tokens.is_none() && reasoning_effort.is_none() {
        return;
    }
    let mut generation_config = serde_json::Map::new();
    if let Some(max_tokens) = max_output_tokens {
        generation_config.insert("maxOutputTokens".to_string(), json!(max_tokens));
    }
    if let Some(effort) = reasoning_effort {
        generation_config.insert(
            "thinkingConfig".to_string(),
            json!({ "thinkingLevel": effort.as_str().to_ascii_uppercase() }),
        );
    }
    request.generation_config = Some(Value::Object(generation_config));
}
use bamboo_domain::ToolSchema;

/// Google Gemini API provider.
pub struct GeminiProvider {
    client: Client,
    api_key: String,
    base_url: String,
    default_reasoning_effort: Option<ReasoningEffort>,
    request_overrides: Option<RequestOverridesConfig>,
    masking_config: KeywordMaskingConfig,
}

impl GeminiProvider {
    /// Create a new Gemini provider with an API key.
    pub fn new(api_key: impl Into<String>) -> Self {
        Self {
            client: Client::new(),
            api_key: api_key.into(),
            base_url: "https://generativelanguage.googleapis.com/v1beta".to_string(),
            default_reasoning_effort: None,
            request_overrides: None,
            masking_config: KeywordMaskingConfig::default(),
        }
    }

    /// Configure keyword masking applied as a last-moment scan of every outbound
    /// request body (see [`crate::masking`]).
    pub fn with_masking(mut self, masking_config: KeywordMaskingConfig) -> Self {
        self.masking_config = masking_config;
        self
    }

    /// Overrides the internal HTTP client (e.g., to enable a proxy).
    pub fn with_client(mut self, client: Client) -> Self {
        self.client = client;
        self
    }

    /// Set a custom base URL (e.g., for proxies or alternative endpoints).
    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = url.into();
        self
    }

    /// Configure default reasoning effort for requests sent through this provider.
    pub fn with_reasoning_effort(mut self, effort: Option<ReasoningEffort>) -> Self {
        self.default_reasoning_effort = effort;
        self
    }

    /// Configure request overrides for this provider.
    pub fn with_request_overrides(mut self, overrides: Option<RequestOverridesConfig>) -> Self {
        self.request_overrides = overrides;
        self
    }

    fn build_headers(&self, endpoint: &str, model: Option<&str>) -> Result<HeaderMap> {
        let mut headers = HeaderMap::new();
        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
        // Authenticate via header rather than a `?key=` URL query parameter so
        // the API key never leaks into HTTP/proxy/debug logs (issue #12).
        // A malformed/non-ASCII key surfaces as `LLMError::Auth` instead of
        // being silently dropped (which would send an unauthenticated request
        // and produce a confusing 401) — mirrors the Anthropic sibling.
        headers.insert(
            "x-goog-api-key",
            HeaderValue::from_str(&self.api_key)
                .map_err(|e| LLMError::Auth(format!("Invalid API key: {}", e)))?,
        );
        // `x-goog-api-key` is intentionally overridable: `request_overrides`
        // (applied next) may replace it for per-endpoint/operator auth overrides.
        request_overrides::apply_overrides_to_header_map(
            &mut headers,
            self.request_overrides.as_ref(),
            endpoint,
            model,
        );
        Ok(headers)
    }

    /// Build the `streamGenerateContent` URL for `model`. Auth travels via the
    /// `x-goog-api-key` header (see [`GeminiProvider::build_headers`]), so the
    /// API key never appears in this URL.
    fn stream_url(&self, model: &str) -> String {
        // Gemini returns a JSON array when `alt=sse` is omitted even though the
        // endpoint name is `streamGenerateContent`. The downstream adapter is
        // intentionally an SSE parser, so request the SSE representation
        // explicitly instead of silently producing an empty assistant turn.
        format!(
            "{}/models/{}:streamGenerateContent?alt=sse",
            self.base_url, model
        )
    }

    /// Build the `list_models` URL. Auth travels via the `x-goog-api-key`
    /// header (see [`GeminiProvider::build_headers`]), so the API key never
    /// appears in this URL.
    fn list_models_url(&self) -> String {
        format!("{}/models", self.base_url.trim_end_matches('/'))
    }
}

#[async_trait]
impl LLMProvider for GeminiProvider {
    async fn provider_visible_tool_footprint(
        &self,
        _ir: &crate::prompt_ir::PromptIR,
        tools: &[ToolSchema],
        _model: &str,
        _required_tool: Option<&str>,
    ) -> Result<ProviderVisibleToolFootprint> {
        if tools.is_empty() {
            return Ok(ProviderVisibleToolFootprint::default());
        }
        let projected: Vec<crate::protocol::gemini::GeminiTool> = tools.to_vec().to_provider()?;
        Ok(ProviderVisibleToolFootprint {
            segments: vec![ProviderVisibleToolSegment::from_serializable(
                ProviderVisibleToolSegmentKind::InitialFullDefinition,
                &projected,
            )?],
        })
    }

    async fn chat_stream(
        &self,
        messages: &[Message],
        tools: &[ToolSchema],
        max_output_tokens: Option<u32>,
        model: &str,
    ) -> Result<LLMStream> {
        self.chat_stream_with_options(messages, tools, max_output_tokens, model, None)
            .await
    }

    async fn chat_stream_with_options(
        &self,
        messages: &[Message],
        tools: &[ToolSchema],
        max_output_tokens: Option<u32>,
        model: &str,
        options: Option<&LLMRequestOptions>,
    ) -> Result<LLMStream> {
        crate::image_budget::validate_tool_image_budget(messages).map_err(LLMError::Api)?;
        tracing::debug!("Gemini provider using model_hash={}", log_identity(model));
        let reasoning_effort = options
            .and_then(|o| o.reasoning_effort)
            .or(self.default_reasoning_effort);
        let request_reasoning_effort = options.and_then(|o| o.reasoning_effort);
        let required_tool = crate::provider::required_tool_from_options(options, tools)?;
        let reasoning_source = if request_reasoning_effort.is_some() {
            "request"
        } else if self.default_reasoning_effort.is_some() {
            "provider_default"
        } else {
            "none"
        };
        let request_purpose = options
            .and_then(|o| o.request_purpose.as_deref())
            .unwrap_or("unknown");
        let session_log_id = options
            .and_then(|o| o.session_id.as_deref())
            .unwrap_or("unknown-session");

        // Auth is supplied via the x-goog-api-key header (see build_headers).
        let url = self.stream_url(model);

        let build_request = |effort: Option<ReasoningEffort>| -> Result<GeminiRequest> {
            // Convert messages using the new protocol system
            let messages_vec: Vec<Message> = messages.to_vec();
            let mut request: GeminiRequest = messages_vec.to_provider()?;

            // Add tools if present
            if !tools.is_empty() {
                let tools_vec: Vec<ToolSchema> = tools.to_vec();
                request.tools = Some(tools_vec.to_provider()?);
            }

            apply_generation_config(&mut request, max_output_tokens, effort);

            Ok(request)
        };

        let request = build_request(reasoning_effort)?;
        let mut request_json = serde_json::to_value(&request).map_err(LLMError::Json)?;
        request_overrides::validate_budgeted_overrides(
            self.request_overrides.as_ref(),
            request_overrides::ENDPOINT_STREAM_GENERATE_CONTENT,
            Some(model),
        )?;
        request_overrides::apply_overrides_to_body(
            &mut request_json,
            self.request_overrides.as_ref(),
            request_overrides::ENDPOINT_STREAM_GENERATE_CONTENT,
            Some(model),
        );
        apply_required_tool_choice(&mut request_json, required_tool);
        // Last-moment scan: mask every text value in the fully-assembled body.
        crate::masking::mask_outbound_body(&mut request_json, &self.masking_config);
        tracing::info!(
            "[session_hash={}] Gemini request protocol=streamGenerateContent model_hash={} reasoning_effort={} reasoning_source={} request_reasoning_enabled={} max_output_tokens={} purpose_hash={}",
            log_identity(session_log_id),
            log_identity(model),
            reasoning_effort
                .map(ReasoningEffort::as_str)
                .unwrap_or("none"),
            reasoning_source,
            reasoning_effort.is_some(),
            max_output_tokens
                .map(|tokens| tokens.to_string())
                .unwrap_or_else(|| "none".to_string()),
            log_identity(request_purpose)
        );
        tracing::debug!(
            message_count = messages.len(),
            tool_count = tools.len(),
            "Gemini request prepared"
        );

        let headers = self.build_headers(
            request_overrides::ENDPOINT_STREAM_GENERATE_CONTENT,
            Some(model),
        )?;
        // Retry the initial request on transient failures (issue #18); the
        // returned body is unread, so SSE streaming below is unaffected.
        let response =
            crate::retry::send_provider_request(crate::retry::global(), "Gemini", || {
                self.client
                    .post(&url)
                    .headers(headers.clone())
                    .json(&request_json)
            })
            .await?;

        if !response.status().is_success() {
            let status = response.status();
            let text = response.text().await.map_err(LLMError::Http)?;

            if status == 401 || status == 403 {
                return Err(LLMError::Auth(format!(
                    "Gemini authentication failed: {}. Please check your API key.",
                    text
                )));
            }

            return Err(LLMError::Api(format!(
                "Gemini API error: HTTP {}: {}",
                status, text
            )));
        }

        tracing::debug!("Gemini stream started successfully");

        // Parse SSE stream with Gemini-specific parser
        // Parse SSE stream with Gemini-specific parser. Uses the multi-chunk
        // adapter: a single Gemini event can yield several chunks (a final
        // `usageMetadata` carries both a cache hit and output/thinking usage),
        // and `streamGenerateContent?alt=sse` sends no [DONE], so emitting all
        // chunks from the one event — rather than deferring — guarantees usage
        // is delivered even when the connection closes immediately afterwards.
        let mut state = GeminiStreamState::default();
        let model_for_log = model.to_string();
        let requested_reasoning_for_log = reasoning_effort;
        let mut logged_summary = false;

        let stream = crate::providers::common::sse::llm_stream_from_sse_multi(
            response,
            move |event, data| {
                let chunks = stream::parse_with_provider_usage(&mut state, event, data)?;

                if !logged_summary
                    && (requested_reasoning_for_log.is_some() || state.observed_thinking_signal)
                    && chunks.iter().any(|c| matches!(c, LLMChunk::Done))
                {
                    tracing::info!(
                        "Gemini reasoning summary: model_hash={} requested_effort={} observed_thinking_signal={} thinking_parts_count={} thinking_text_chars={}",
                        log_identity(&model_for_log),
                        requested_reasoning_for_log
                            .map(ReasoningEffort::as_str)
                            .unwrap_or("none"),
                        state.observed_thinking_signal,
                        state.thinking_parts_count,
                        state.thinking_text_chars
                    );
                    logged_summary = true;
                }

                Ok(chunks)
            },
        );

        Ok(stream)
    }

    async fn list_models(&self) -> Result<Vec<String>> {
        let headers = self.build_headers(request_overrides::ENDPOINT_MODELS, None)?;
        let url = self.list_models_url();
        model_fetcher::fetch_model_list(&self.client, &url, headers, "Gemini").await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn tool_footprint_reuses_grouped_gemini_function_declarations() {
        let tools = vec!["lookup", "summarize"]
            .into_iter()
            .map(|name| ToolSchema {
                schema_type: "function".to_string(),
                function: bamboo_domain::FunctionSchema {
                    name: name.to_string(),
                    description: format!("Use {name}"),
                    parameters: json!({"type":"object"}),
                },
            })
            .collect::<Vec<_>>();
        let expected: Vec<crate::protocol::gemini::GeminiTool> =
            tools.clone().to_provider().unwrap();
        let footprint = GeminiProvider::new("k")
            .provider_visible_tool_footprint(
                &crate::prompt_ir::PromptIR::default(),
                &tools,
                "gemini-test",
                None,
            )
            .await
            .unwrap();

        assert_eq!(footprint.segments.len(), 1);
        assert_eq!(
            footprint.segments[0].kind,
            ProviderVisibleToolSegmentKind::InitialFullDefinition
        );
        assert_eq!(
            footprint.segments[0].serialized,
            serde_json::to_string(&expected).unwrap()
        );
        let rendered: Value = serde_json::from_str(&footprint.segments[0].serialized).unwrap();
        assert_eq!(
            rendered[0]["functionDeclarations"]
                .as_array()
                .unwrap()
                .len(),
            2
        );
    }

    #[test]
    fn generation_config_encodes_selected_effort_without_budget_or_capping() {
        for (effort, wire) in [
            (ReasoningEffort::Disabled, "NONE"),
            (ReasoningEffort::Low, "LOW"),
            (ReasoningEffort::Medium, "MEDIUM"),
            (ReasoningEffort::High, "HIGH"),
            (ReasoningEffort::Xhigh, "XHIGH"),
            (ReasoningEffort::Max, "MAX"),
        ] {
            let mut request: GeminiRequest = vec![Message::user("hello")].to_provider().unwrap();
            apply_generation_config(&mut request, Some(1), Some(effort));
            let config = request.generation_config.unwrap();
            assert_eq!(config["thinkingConfig"]["thinkingLevel"], wire);
            assert!(config["thinkingConfig"].get("thinkingBudget").is_none());
            assert_eq!(config["maxOutputTokens"], 1);
        }
    }

    #[test]
    fn required_tool_choice_uses_gemini_any_allowlist_shape() {
        let mut body = json!({
            "generationConfig": {
                "thinkingConfig": {"thinkingBudget": 4096}
            }
        });
        apply_required_tool_choice(&mut body, Some("load_skill"));
        assert_eq!(
            body["toolConfig"],
            json!({
                "functionCallingConfig": {
                    "mode": "ANY",
                    "allowedFunctionNames": ["load_skill"]
                }
            })
        );
        assert_eq!(
            body["generationConfig"]["thinkingConfig"]["thinkingBudget"],
            4096
        );
    }

    #[test]
    fn test_new_provider() {
        let provider = GeminiProvider::new("test_key");
        assert_eq!(provider.api_key, "test_key");
        assert_eq!(
            provider.base_url,
            "https://generativelanguage.googleapis.com/v1beta"
        );
    }

    #[test]
    fn test_with_base_url() {
        let provider =
            GeminiProvider::new("test_key").with_base_url("https://custom.googleapis.com/v1");
        assert_eq!(provider.base_url, "https://custom.googleapis.com/v1");
    }

    #[test]
    fn test_chained_builders() {
        let provider = GeminiProvider::new("test_key").with_base_url("https://custom.api.com");

        assert_eq!(provider.api_key, "test_key");
        assert_eq!(provider.base_url, "https://custom.api.com");
    }

    #[test]
    fn test_url_construction() {
        let provider =
            GeminiProvider::new("my_api_key_123").with_base_url("https://test.api.com/v1beta");

        // The API key must NOT appear in the URL; auth is sent via the
        // x-goog-api-key header (see test_build_headers_sends_api_key).
        // These assertions exercise the SAME helpers the production request
        // paths use (`stream_url` / `list_models_url`), so a regression that
        // reintroduces `?key=` in real code fails here.
        let stream_url = provider.stream_url("gemini-custom");
        assert_eq!(
            stream_url,
            "https://test.api.com/v1beta/models/gemini-custom:streamGenerateContent?alt=sse"
        );
        assert!(
            stream_url.ends_with("?alt=sse"),
            "streamGenerateContent must explicitly request SSE transport"
        );
        assert!(
            !stream_url.contains("key="),
            "API key must not be embedded in the stream URL"
        );

        let models_url = provider.list_models_url();
        assert_eq!(models_url, "https://test.api.com/v1beta/models");
        assert!(
            !models_url.contains("key="),
            "API key must not be embedded in the list_models URL"
        );
    }

    /// Header-based auth: the API key travels in `x-goog-api-key`, never in the
    /// URL query string (issue #12).
    #[test]
    fn test_build_headers_sends_api_key() {
        let provider = GeminiProvider::new("my_api_key_123");

        let stream_headers = provider
            .build_headers(request_overrides::ENDPOINT_STREAM_GENERATE_CONTENT, None)
            .expect("valid key should build headers");
        assert_eq!(
            stream_headers
                .get("x-goog-api-key")
                .expect("x-goog-api-key header should be set on streamGenerateContent requests")
                .to_str()
                .expect("valid header value"),
            "my_api_key_123"
        );

        let models_headers = provider
            .build_headers(request_overrides::ENDPOINT_MODELS, None)
            .expect("valid key should build headers");
        assert_eq!(
            models_headers
                .get("x-goog-api-key")
                .expect("x-goog-api-key header should be set on list_models requests")
                .to_str()
                .expect("valid header value"),
            "my_api_key_123"
        );
    }

    /// A malformed/non-ASCII key (e.g. one carrying a trailing newline) must
    /// surface as `LLMError::Auth` instead of being silently dropped to an
    /// unauthenticated request (which would otherwise produce a confusing 401).
    #[test]
    fn test_build_headers_rejects_invalid_api_key() {
        let provider = GeminiProvider::new("bad\nkey");
        let result =
            provider.build_headers(request_overrides::ENDPOINT_STREAM_GENERATE_CONTENT, None);
        assert!(
            matches!(result, Err(LLMError::Auth(_))),
            "expected LLMError::Auth for invalid key"
        );
    }

    // ========== MODEL REQUIREMENT ARCHITECTURE TESTS ==========
    // These tests ensure the design principle:
    // "Provider must not have a default model field or with_model() method"

    /// Test: GeminiProvider does NOT have a model field
    #[test]
    fn gemini_provider_has_no_model_field() {
        // This test documents the provider structure:
        // pub struct GeminiProvider {
        //     client: Client,
        //     api_key: String,
        //     base_url: String,
        //     // NO model field!
        // }
        //
        // If someone adds a model field, this test should be updated
        // to reflect the architecture change.
        let provider = GeminiProvider::new("test_key");
        // Verify we can access known fields
        assert_eq!(provider.api_key, "test_key");
        assert_eq!(
            provider.base_url,
            "https://generativelanguage.googleapis.com/v1beta"
        );
        // There is NO provider.model field to access
    }

    /// Test: GeminiProvider does NOT have with_model() method
    #[test]
    fn gemini_provider_has_no_with_model_method() {
        let provider = GeminiProvider::new("test_key");

        // Available builder method:
        let provider = provider.with_base_url("https://custom.api.com");

        // There is NO .with_model("gemini-pro") method
        // Model is passed to chat_stream() as a parameter

        assert_eq!(provider.base_url, "https://custom.api.com");
    }
}