codetether-agent 4.0.0

//! Gemini Web provider  drives the Gemini chat UI's undocumented
//! BardChatUi endpoint using browser cookies stored in HashiCorp Vault.
//!
//! This provider reverse-engineers the same HTTP request the Gemini web app
//! sends so no official API key is required.  Authentication is entirely
//! cookie-based (Netscape cookies.txt from Cookie-Editor stored in Vault under
//! `secret/codetether/providers/gemini-web` as the `cookies` key).
//!
//! Supported models (Gemini 3 / 3.1 family):
//! - `gemini-web-fast`       Gemini 3 Fast (mode_id fbb127bbb056c959)
//! - `gemini-web-thinking`   Gemini 3 Thinking (mode_id 5bf011840784117a)
//! - `gemini-web-pro`        Gemini 3.1 Pro (mode_id 9d8ca3786ebdfbea)
//! - `gemini-web-deep-think` Gemini 3 Deep Think (mode_id e6fa609c3fa255c0)
//!
//! The model is selected via the `x-goog-ext-525001261-jspb` request header.

use super::{
    CompletionRequest, CompletionResponse, ContentPart, FinishReason, Message, ModelInfo, Provider,
    Role, StreamChunk, Usage,
};
use anyhow::{Context, Result};
use async_trait::async_trait;
use futures::StreamExt as _;
use regex::Regex;
use reqwest::Client;
use serde_json::{Value, json};
use std::sync::OnceLock;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::Mutex;

const GEMINI_ORIGIN: &str = "https://gemini.google.com";
const GEMINI_ENDPOINT: &str = "https://gemini.google.com/u/1/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate";

/// How long cached session tokens remain valid before re-scraping the home page.
const TOKEN_TTL: Duration = Duration::from_secs(20 * 60);

/// (model_id, mode_id, display_name, context_window)
const MODELS: &[(&str, &str, &str, usize)] = &[
    (
        "gemini-web-fast",
        "fbb127bbb056c959",
        "Gemini 3 Fast",
        1_048_576_usize,
    ),
    (
        "gemini-web-thinking",
        "5bf011840784117a",
        "Gemini 3 Thinking",
        1_048_576_usize,
    ),
    (
        "gemini-web-pro",
        "9d8ca3786ebdfbea",
        "Gemini 3.1 Pro",
        1_048_576_usize,
    ),
    (
        "gemini-web-deep-think",
        "e6fa609c3fa255c0",
        "Gemini 3 Deep Think",
        1_048_576_usize,
    ),
];

#[derive(Clone)]
struct SessionTokens {
    at_token: String,
    f_sid: String,
    bl: String,
}

pub struct GeminiWebProvider {
    client: Client,
    /// Raw contents of cookies.txt (Netscape / Cookie-Editor export format)
    cookies: String,
    /// Cached session tokens with the instant they were fetched.
    token_cache: Mutex<Option<(SessionTokens, Instant)>>,
}

impl std::fmt::Debug for GeminiWebProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GeminiWebProvider")
            .field("cookies", &"[redacted]")
            .finish_non_exhaustive()
    }
}

impl GeminiWebProvider {
    /// Create a new provider from Netscape cookies.txt content.
    pub fn new(cookies: String) -> Result<Self> {
        let client = Client::builder()
            .user_agent(
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) \
                 AppleWebKit/537.36 (KHTML, like Gecko) \
                 Chrome/131.0.0.0 Safari/537.36",
            )
            .timeout(std::time::Duration::from_secs(120))
            .build()
            .context("Failed to build reqwest client for GeminiWebProvider")?;
        Ok(Self {
            client,
            cookies,
            token_cache: Mutex::new(None),
        })
    }

    /// Convert a Netscape cookies.txt (7 tab-separated columns) into a
    /// `Cookie: name=value; ...` header string.
    ///
    /// Standard Netscape format columns:
    ///   0: domain  1: flag  2: path  3: secure  4: expiration  5: name  6: value
    ///
    /// If only 2 columns are present (simple `name\tvalue` format), those are
    /// used directly so the function remains backward-compatible.
    fn cookie_header(&self) -> String {
        self.cookies
            .lines()
            .filter(|l| {
                let t = l.trim();
                !t.is_empty() && !t.starts_with('#')
            })
            .filter_map(|line| {
                let parts: Vec<&str> = line.split('\t').collect();
                let (name, value) = if parts.len() >= 7 {
                    // Standard Netscape format
                    (parts[5].trim(), parts[6].trim())
                } else if parts.len() >= 2 {
                    // Simple name\tvalue format (backward compat)
                    (parts[0].trim(), parts[1].trim())
                } else {
                    return None;
                };
                if name.is_empty() {
                    return None;
                }
                Some(format!("{name}={value}"))
            })
            .collect::<Vec<_>>()
            .join("; ")
    }

    /// GET the Gemini home page and extract the three tokens we need:
    ///   - `at_token`  XSRF token (key `thykhd` or legacy `SNlM0e`)
    ///   - `f_sid`     session ID (`FdrFJe`)
    ///   - `bl`        build label (`cfb2h`)
    ///
    /// Regexes are compiled once and reused across calls via `OnceLock`.
    async fn get_session_tokens(&self) -> Result<SessionTokens> {
        static RE_NEW: OnceLock<Regex> = OnceLock::new();
        static RE_OLD: OnceLock<Regex> = OnceLock::new();
        static RE_BL:  OnceLock<Regex> = OnceLock::new();
        static RE_SID: OnceLock<Regex> = OnceLock::new();

        let re_new = RE_NEW.get_or_init(|| Regex::new(r#""thykhd":"([^"]+)""#).unwrap());
        let re_old = RE_OLD.get_or_init(|| Regex::new(r#""SNlM0e":"([^"]+)""#).unwrap());
        let re_bl  = RE_BL .get_or_init(|| Regex::new(r#""cfb2h":"([^"]+)""#) .unwrap());
        let re_sid = RE_SID.get_or_init(|| Regex::new(r#""FdrFJe":"([^"]+)""#).unwrap());

        let cookie_hdr = self.cookie_header();
        let html = self
            .client
            .get(GEMINI_ORIGIN)
            .header("Cookie", &cookie_hdr)
            .send()
            .await
            .context("Failed to fetch Gemini home page")?
            .text()
            .await
            .context("Failed to read Gemini home page body")?;

        let at_token = if let Some(cap) = re_new.captures(&html) {
            cap[1].to_string()
        } else if let Some(cap) = re_old.captures(&html) {
            cap[1].to_string()
        } else {
            anyhow::bail!(
                "Could not find Gemini at-token (thykhd / SNlM0e)  \
                 cookies may be expired or invalid"
            );
        };

        let bl    = re_bl .captures(&html).map(|c| c[1].to_string()).unwrap_or_default();
        let f_sid = re_sid.captures(&html).map(|c| c[1].to_string()).unwrap_or_default();

        Ok(SessionTokens { at_token, f_sid, bl })
    }

    /// Return session tokens, re-fetching only when the cached copy has
    /// exceeded `TOKEN_TTL` (20 minutes).
    async fn get_or_refresh_tokens(&self) -> Result<SessionTokens> {
        let mut cache = self.token_cache.lock().await;
        if let Some((ref tokens, fetched_at)) = *cache {
            if fetched_at.elapsed() < TOKEN_TTL {
                return Ok(tokens.clone());
            }
        }
        let fresh = self.get_session_tokens().await?;
        *cache = Some((fresh.clone(), Instant::now()));
        Ok(fresh)
    }

    /// Build the `f.req` JSON payload  a two-element outer array whose
    /// second slot is a JSON-encoded 69-element inner array.
    fn build_freq(prompt: &str) -> String {
        let ts = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let mut inner: Vec<Value> = Vec::with_capacity(69);
        inner.push(json!([[prompt, 0, null, null, null, null, 0]])); // [0]
        inner.push(json!(["en"]));                                    // [1]
        inner.push(json!([null, null, null]));                        // [2]
        inner.push(Value::Null);                                      // [3]
        inner.push(Value::Null);                                      // [4]
        inner.push(Value::Null);                                      // [5]
        inner.push(json!([1]));                                       // [6]
        inner.push(json!(1));                                         // [7]
        inner.push(Value::Null);                                      // [8]
        inner.push(json!([1, 0, null, null, null, null, null, 0]));   // [9]
        inner.push(Value::Null);                                      // [10]
        inner.push(Value::Null);                                      // [11]
        inner.push(json!([0]));                                       // [12]
        for _ in 0..40 { inner.push(Value::Null); }                  // [13]-[52]
        inner.push(json!(0));                                         // [53]
        for _ in 0..5  { inner.push(Value::Null); }                  // [54]-[58]
        inner.push(json!("CD1035A5-0E0E-4B68-B744-23C2D8960DF5"));   // [59]
        inner.push(Value::Null);                                      // [60]
        inner.push(json!([]));                                        // [61]
        for _ in 0..4  { inner.push(Value::Null); }                  // [62]-[65]
        inner.push(json!([ts, 0]));                                   // [66]
        inner.push(Value::Null);                                      // [67]
        inner.push(json!(2));                                         // [68]

        debug_assert_eq!(inner.len(), 69, "f.req inner list must be exactly 69 elements");

        let inner_json = serde_json::to_string(&inner).unwrap_or_default();
        serde_json::to_string(&json!([null, inner_json])).unwrap_or_default()
    }

    /// Walk streaming lines and return the text from the *last* parseable
    /// `inner[4][0][1][0]` block.
    ///
    /// Gemini sends cumulative text (each chunk is the full answer so far);
    /// the last valid block is the most complete and most reliable.
    fn extract_text(raw: &str) -> String {
        let mut last = String::new();
        for line in raw.lines() {
            let line = line.trim();
            if line.is_empty() || !line.starts_with('[') { continue; }
            let Ok(outer) = serde_json::from_str::<Value>(line) else { continue };
            let Some(arr) = outer.as_array() else { continue };
            let Some(two) = arr.get(2).and_then(Value::as_str) else { continue };
            let Ok(inner) = serde_json::from_str::<Value>(two) else { continue };
            if let Some(text) = inner
                .get(4).and_then(|v| v.get(0))
                .and_then(|v| v.get(1)).and_then(|v| v.get(0))
                .and_then(Value::as_str)
            {
                last = text.to_string();
            }
        }
        last
    }

    /// Look up the `mode_id` string for a given model identifier.
    fn mode_id(model: &str) -> &'static str {
        MODELS
            .iter()
            .find(|(id, _, _, _)| *id == model)
            .map(|(_, mid, _, _)| *mid)
            .unwrap_or("fbb127bbb056c959")
    }

    /// Build a `RequestBuilder` for the StreamGenerate endpoint.
    async fn build_request(&self, prompt: &str, model: &str) -> Result<reqwest::RequestBuilder> {
        let tokens = self
            .get_or_refresh_tokens()
            .await
            .context("Failed to obtain Gemini session tokens")?;

        let cookie_hdr = self.cookie_header();
        let freq = Self::build_freq(prompt);
        let mode_id = Self::mode_id(model);

        let ext_header = {
            let v: Value = json!([1, null, null, null, mode_id, null, null, 0, [4], null, null, 3]);
            serde_json::to_string(&v).unwrap_or_default()
        };

        let reqid = (SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis()
            % 900_000
            + 100_000)
            .to_string();

        let url = reqwest::Url::parse_with_params(
            GEMINI_ENDPOINT,
            &[
                ("bl",     tokens.bl.as_str()),
                ("f.sid",  tokens.f_sid.as_str()),
                ("hl",     "en"),
                ("_reqid", reqid.as_str()),
                ("rt",     "c"),
            ],
        )
        .context("Failed to build Gemini endpoint URL")?;

        Ok(self
            .client
            .post(url)
            .header("Cookie",                    cookie_hdr)
            .header("X-Same-Domain",             "1")
            .header("Origin",                    GEMINI_ORIGIN)
            .header("Referer",                   format!("{}/app", GEMINI_ORIGIN))
            .header("x-goog-ext-525001261-jspb", ext_header)
            .form(&[("f.req", freq), ("at", tokens.at_token)]))
    }

    /// Core blocking request: POST and collect the complete response.
    async fn ask(&self, prompt: &str, model: &str) -> Result<String> {
        let resp = self
            .build_request(prompt, model)
            .await?
            .send()
            .await
            .context("Failed to send request to Gemini StreamGenerate")?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            anyhow::bail!(
                "Gemini StreamGenerate returned HTTP {}: {}",
                status,
                &body[..body.len().min(500)]
            );
        }

        let body = resp.text().await.context("Failed to read Gemini response body")?;
        let text = Self::extract_text(&body);
        if text.is_empty() {
            anyhow::bail!(
                "No text found in Gemini response  raw body (first 500 chars): {}",
                &body[..body.len().min(500)]
            );
        }
        Ok(text)
    }
}

#[async_trait]
impl Provider for GeminiWebProvider {
    fn name(&self) -> &str { "gemini-web" }

    async fn list_models(&self) -> Result<Vec<ModelInfo>> {
        Ok(MODELS
            .iter()
            .map(|(id, _, label, ctx)| ModelInfo {
                id:                      id.to_string(),
                name:                    label.to_string(),
                provider:                "gemini-web".to_string(),
                context_window:          *ctx,
                max_output_tokens:       Some(65_536),
                supports_vision:         false,
                supports_tools:          false,
                supports_streaming:      true,
                input_cost_per_million:  Some(0.0),
                output_cost_per_million: Some(0.0),
            })
            .collect())
    }

    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse> {
        let prompt = request
            .messages.iter()
            .map(|m| {
                let role = match m.role {
                    Role::System    => "System",
                    Role::User      => "User",
                    Role::Assistant => "Assistant",
                    Role::Tool      => "Tool",
                };
                let text = m.content.iter()
                    .filter_map(|p| match p {
                        ContentPart::Text { text } => Some(text.as_str()),
                        _ => None,
                    })
                    .collect::<Vec<_>>().join("");
                format!("{role}: {text}")
            })
            .collect::<Vec<_>>().join("\n");

        let text = self.ask(&prompt, &request.model).await
            .context("Gemini Web completion failed")?;

        Ok(CompletionResponse {
            message: Message {
                role:    Role::Assistant,
                content: vec![ContentPart::Text { text }],
            },
            usage: Usage {
                prompt_tokens:     0,
                completion_tokens: 0,
                total_tokens:      0,
                cache_read_tokens:  None,
                cache_write_tokens: None,
            },
            finish_reason: FinishReason::Stop,
        })
    }

    async fn complete_stream(
        &self,
        request: CompletionRequest,
    ) -> Result<futures::stream::BoxStream<'static, StreamChunk>> {
        let prompt = request
            .messages.iter()
            .map(|m| {
                let role = match m.role {
                    Role::System    => "System",
                    Role::User      => "User",
                    Role::Assistant => "Assistant",
                    Role::Tool      => "Tool",
                };
                let text = m.content.iter()
                    .filter_map(|p| match p {
                        ContentPart::Text { text } => Some(text.as_str()),
                        _ => None,
                    })
                    .collect::<Vec<_>>().join("");
                format!("{role}: {text}")
            })
            .collect::<Vec<_>>().join("\n");

        let resp = self
            .build_request(&prompt, &request.model)
            .await?
            .send()
            .await
            .context("Failed to send streaming request to Gemini")?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            anyhow::bail!(
                "Gemini StreamGenerate returned HTTP {}: {}",
                status,
                &body[..body.len().min(500)]
            );
        }

        // Process the byte stream and emit text deltas as chunks arrive.
        // Gemini sends cumulative text, so we track prev_len and emit only
        // the newly-arrived portion on each iteration.
        let byte_stream = resp.bytes_stream();
        let (tx, rx) = futures::channel::mpsc::channel::<StreamChunk>(32);

        tokio::spawn(async move {
            futures::pin_mut!(byte_stream);
            let mut buf = String::new();
            let mut prev_len: usize = 0;
            let mut tx = tx;

            while let Some(chunk_result) = byte_stream.next().await {
                let Ok(bytes) = chunk_result else { break };
                let Ok(s) = std::str::from_utf8(&bytes) else { continue };
                buf.push_str(s);

                let current_text = Self::extract_text(&buf);
                if current_text.len() > prev_len {
                    let delta = current_text[prev_len..].to_string();
                    prev_len = current_text.len();
                    if tx.try_send(StreamChunk::Text(delta)).is_err() {
                        return; // receiver dropped
                    }
                }
            }

            // Final pass for any remaining delta in the last chunk
            let final_text = Self::extract_text(&buf);
            if final_text.len() > prev_len {
                let _ = tx.try_send(StreamChunk::Text(final_text[prev_len..].to_string()));
            }
            let _ = tx.try_send(StreamChunk::Done { usage: None });
        });

        let stream = futures::stream::unfold(rx, |mut rx| async {
            use futures::StreamExt as _;
            rx.next().await.map(|chunk| (chunk, rx))
        });

        Ok(Box::pin(stream))
    }
}