From 8c09d8c7d365b900448c5e08fbeaa31beeb4dbc8 Mon Sep 17 00:00:00 2001 From: YoVinchen Date: Mon, 6 Apr 2026 11:11:30 +0800 Subject: [PATCH] feat(proxy): add Gemini Native schema, shadow store, transform, and streaming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - gemini_schema: Gemini generateContent request/response type definitions - gemini_shadow: session-scoped shadow store for thinking signature and tool-call state replay across streaming chunks - transform_gemini: bidirectional Anthropic Messages ↔ Gemini Native request/response conversion with thinking block and tool-use support - streaming_gemini: Gemini SSE → Anthropic SSE streaming adapter with incremental thinking/text/tool_use delta emission --- .../src/proxy/providers/gemini_schema.rs | 278 ++++++ .../src/proxy/providers/gemini_shadow.rs | 389 +++++++++ .../src/proxy/providers/streaming_gemini.rs | 619 ++++++++++++++ .../src/proxy/providers/transform_gemini.rs | 788 ++++++++++++++++++ 4 files changed, 2074 insertions(+) create mode 100644 src-tauri/src/proxy/providers/gemini_schema.rs create mode 100644 src-tauri/src/proxy/providers/gemini_shadow.rs create mode 100644 src-tauri/src/proxy/providers/streaming_gemini.rs create mode 100644 src-tauri/src/proxy/providers/transform_gemini.rs diff --git a/src-tauri/src/proxy/providers/gemini_schema.rs b/src-tauri/src/proxy/providers/gemini_schema.rs new file mode 100644 index 000000000..dc07b09b0 --- /dev/null +++ b/src-tauri/src/proxy/providers/gemini_schema.rs @@ -0,0 +1,278 @@ +//! Gemini tool schema helpers. +//! +//! Gemini `FunctionDeclaration` supports two schema channels: +//! - `parameters`: a restricted `Schema` subset +//! - `parametersJsonSchema`: richer JSON Schema via arbitrary JSON `Value` +//! +//! Anthropic tool schemas are closer to JSON Schema, so we choose the richer +//! channel when unsupported `Schema` fields are present. + +use serde_json::{json, Map, Value}; + +#[derive(Debug, Clone, PartialEq)] +pub enum GeminiFunctionParameters { + Schema(Value), + JsonSchema(Value), +} + +pub fn build_gemini_function_parameters(input_schema: Value) -> GeminiFunctionParameters { + let schema = normalize_json_schema(input_schema); + + if requires_parameters_json_schema(&schema) { + GeminiFunctionParameters::JsonSchema(schema) + } else { + GeminiFunctionParameters::Schema(to_gemini_schema(schema)) + } +} + +fn normalize_json_schema(schema: Value) -> Value { + match schema { + Value::Object(mut obj) => { + obj.remove("$schema"); + obj.remove("$id"); + + if let Some(properties) = obj + .get_mut("properties") + .and_then(|value| value.as_object_mut()) + { + for value in properties.values_mut() { + *value = normalize_json_schema(value.clone()); + } + } + + if let Some(items) = obj.get_mut("items") { + *items = normalize_json_schema(items.clone()); + } + + for key in ["anyOf", "oneOf", "allOf", "prefixItems"] { + if let Some(values) = obj.get_mut(key).and_then(|value| value.as_array_mut()) { + for value in values.iter_mut() { + *value = normalize_json_schema(value.clone()); + } + } + } + + for key in ["not", "if", "then", "else", "additionalProperties"] { + if let Some(value) = obj.get_mut(key) { + *value = normalize_json_schema(value.clone()); + } + } + + Value::Object(obj) + } + Value::Array(values) => { + Value::Array(values.into_iter().map(normalize_json_schema).collect()) + } + other => other, + } +} + +fn requires_parameters_json_schema(schema: &Value) -> bool { + match schema { + Value::Object(obj) => object_requires_parameters_json_schema(obj), + Value::Array(values) => values.iter().any(requires_parameters_json_schema), + _ => false, + } +} + +fn object_requires_parameters_json_schema(obj: &Map) -> bool { + for (key, value) in obj { + match key.as_str() { + "type" => { + if value.is_array() { + return true; + } + } + "format" | "title" | "description" | "nullable" | "enum" | "maxItems" | "minItems" + | "required" | "minProperties" | "maxProperties" | "minLength" | "maxLength" + | "pattern" | "example" | "propertyOrdering" | "default" | "minimum" | "maximum" => {} + "properties" => { + let Some(properties) = value.as_object() else { + return true; + }; + if properties.values().any(requires_parameters_json_schema) { + return true; + } + } + "items" => { + if !value.is_object() || requires_parameters_json_schema(value) { + return true; + } + } + "anyOf" => { + let Some(values) = value.as_array() else { + return true; + }; + if values.iter().any(requires_parameters_json_schema) { + return true; + } + } + // JSON Schema keywords that Gemini `parameters` does not accept. + "$ref" + | "$defs" + | "definitions" + | "additionalProperties" + | "unevaluatedProperties" + | "patternProperties" + | "oneOf" + | "allOf" + | "const" + | "not" + | "if" + | "then" + | "else" + | "dependentRequired" + | "dependentSchemas" + | "contains" + | "minContains" + | "maxContains" + | "prefixItems" + | "exclusiveMinimum" + | "exclusiveMaximum" + | "multipleOf" + | "examples" => return true, + // Be conservative for unknown keywords. + _ => return true, + } + } + + false +} + +fn to_gemini_schema(schema: Value) -> Value { + match schema { + Value::Object(obj) => { + let mut result = Map::new(); + + for (key, value) in obj { + match key.as_str() { + "type" | "format" | "title" | "description" | "nullable" | "enum" + | "maxItems" | "minItems" | "required" | "minProperties" | "maxProperties" + | "minLength" | "maxLength" | "pattern" | "example" | "propertyOrdering" + | "default" | "minimum" | "maximum" => { + result.insert(key, value); + } + "properties" => { + if let Some(properties) = value.as_object() { + let converted = properties + .iter() + .map(|(name, property_schema)| { + (name.clone(), to_gemini_schema(property_schema.clone())) + }) + .collect(); + result.insert("properties".to_string(), Value::Object(converted)); + } + } + "items" => { + if value.is_object() { + result.insert("items".to_string(), to_gemini_schema(value)); + } + } + "anyOf" => { + if let Some(values) = value.as_array() { + result.insert( + "anyOf".to_string(), + Value::Array( + values + .iter() + .map(|value| to_gemini_schema(value.clone())) + .collect(), + ), + ); + } + } + _ => {} + } + } + + Value::Object(result) + } + other => other, + } +} + +pub fn build_gemini_function_declaration( + name: &str, + description: Option<&str>, + input_schema: Value, +) -> Value { + let mut declaration = Map::new(); + declaration.insert("name".to_string(), json!(name)); + declaration.insert("description".to_string(), json!(description.unwrap_or(""))); + + match build_gemini_function_parameters(input_schema) { + GeminiFunctionParameters::Schema(schema) => { + declaration.insert("parameters".to_string(), schema); + } + GeminiFunctionParameters::JsonSchema(schema) => { + declaration.insert("parametersJsonSchema".to_string(), schema); + } + } + + Value::Object(declaration) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn uses_schema_for_simple_openapi_subset() { + let schema = json!({ + "type": "object", + "properties": { + "city": { "type": "string", "description": "Target city" } + }, + "required": ["city"] + }); + + let result = build_gemini_function_declaration("weather", Some("Weather lookup"), schema); + + assert!(result.get("parameters").is_some()); + assert!(result.get("parametersJsonSchema").is_none()); + assert_eq!(result["parameters"]["properties"]["city"]["type"], "string"); + } + + #[test] + fn uses_parameters_json_schema_for_additional_properties() { + let schema = json!({ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "city": { "type": "string" } + }, + "required": ["city"], + "additionalProperties": false + }); + + let result = build_gemini_function_declaration("weather", Some("Weather lookup"), schema); + + assert!(result.get("parameters").is_none()); + assert!(result.get("parametersJsonSchema").is_some()); + assert!(result["parametersJsonSchema"].get("$schema").is_none()); + assert_eq!( + result["parametersJsonSchema"]["additionalProperties"], + false + ); + } + + #[test] + fn uses_parameters_json_schema_for_one_of() { + let schema = json!({ + "type": "object", + "properties": { + "target": { + "oneOf": [ + { "type": "string" }, + { "type": "integer" } + ] + } + } + }); + + let result = build_gemini_function_declaration("search", Some("Search"), schema); + + assert!(result.get("parameters").is_none()); + assert!(result.get("parametersJsonSchema").is_some()); + } +} diff --git a/src-tauri/src/proxy/providers/gemini_shadow.rs b/src-tauri/src/proxy/providers/gemini_shadow.rs new file mode 100644 index 000000000..f1f5d6d39 --- /dev/null +++ b/src-tauri/src/proxy/providers/gemini_shadow.rs @@ -0,0 +1,389 @@ +//! Gemini Native shadow state +//! +//! Keeps provider/session-scoped assistant content snapshots and tool call metadata +//! so Gemini thought signatures and tool turns can be replayed without bloating +//! the main proxy files. + +use serde_json::Value; +use std::collections::{HashMap, VecDeque}; +use std::sync::RwLock; + +/// Composite key for a Gemini shadow session. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct GeminiShadowKey { + pub provider_id: String, + pub session_id: String, +} + +impl GeminiShadowKey { + pub fn new(provider_id: impl Into, session_id: impl Into) -> Self { + Self { + provider_id: provider_id.into(), + session_id: session_id.into(), + } + } +} + +/// Gemini function call metadata captured from an assistant turn. +#[derive(Debug, Clone, PartialEq)] +pub struct GeminiToolCallMeta { + pub id: Option, + pub name: String, + pub args: Value, + pub thought_signature: Option, +} + +impl GeminiToolCallMeta { + pub fn new( + id: Option>, + name: impl Into, + args: Value, + thought_signature: Option>, + ) -> Self { + Self { + id: id.map(Into::into), + name: name.into(), + args, + thought_signature: thought_signature.map(Into::into), + } + } +} + +/// Stored assistant turn snapshot. +#[derive(Debug, Clone, PartialEq)] +pub struct GeminiAssistantTurn { + pub assistant_content: Value, + pub tool_calls: Vec, +} + +impl GeminiAssistantTurn { + pub fn new(assistant_content: Value, tool_calls: Vec) -> Self { + Self { + assistant_content, + tool_calls, + } + } +} + +/// Session snapshot returned by read APIs. +#[derive(Debug, Clone, PartialEq)] +pub struct GeminiShadowSessionSnapshot { + pub provider_id: String, + pub session_id: String, + pub turns: Vec, +} + +#[derive(Debug, Clone)] +struct GeminiShadowSession { + turns: VecDeque, +} + +impl GeminiShadowSession { + fn new() -> Self { + Self { + turns: VecDeque::new(), + } + } +} + +#[derive(Debug, Clone)] +struct GeminiShadowInner { + sessions: HashMap, + session_order: VecDeque, +} + +impl GeminiShadowInner { + fn new() -> Self { + Self { + sessions: HashMap::new(), + session_order: VecDeque::new(), + } + } +} + +/// Thread-safe shadow store for Gemini Native replay state. +/// +/// The store is intentionally small and explicit: +/// - sessions are keyed by `(provider_id, session_id)` +/// - each session keeps only a bounded number of recent assistant turns +/// - the oldest session is evicted first when the store is full +#[derive(Debug)] +pub struct GeminiShadowStore { + max_sessions: usize, + max_turns_per_session: usize, + inner: RwLock, +} + +impl Default for GeminiShadowStore { + fn default() -> Self { + Self::with_limits(200, 64) + } +} + +impl GeminiShadowStore { + #[allow(dead_code)] + pub fn new() -> Self { + Self::default() + } + + pub fn with_limits(max_sessions: usize, max_turns_per_session: usize) -> Self { + Self { + max_sessions: max_sessions.max(1), + max_turns_per_session: max_turns_per_session.max(1), + inner: RwLock::new(GeminiShadowInner::new()), + } + } + + /// Record a Gemini assistant turn for later replay. + pub fn record_assistant_turn( + &self, + provider_id: impl Into, + session_id: impl Into, + assistant_content: Value, + tool_calls: Vec, + ) -> GeminiShadowSessionSnapshot { + let key = GeminiShadowKey::new(provider_id, session_id); + let turn = GeminiAssistantTurn::new(assistant_content, tool_calls); + + let mut inner = self.inner.write().expect("gemini shadow lock poisoned"); + Self::touch_session_order(&mut inner.session_order, &key); + + let snapshot = { + let session = inner + .sessions + .entry(key.clone()) + .or_insert_with(GeminiShadowSession::new); + session.turns.push_back(turn); + while session.turns.len() > self.max_turns_per_session { + session.turns.pop_front(); + } + Self::snapshot_session(&key, session) + }; + Self::prune_sessions(&mut inner, self.max_sessions); + snapshot + } + + /// Get the latest assistant content for a provider/session pair. + #[allow(dead_code)] + pub fn latest_assistant_content(&self, provider_id: &str, session_id: &str) -> Option { + self.get_session(provider_id, session_id) + .and_then(|snapshot| { + snapshot + .turns + .last() + .map(|turn| turn.assistant_content.clone()) + }) + } + + /// Get the latest tool calls for a provider/session pair. + #[allow(dead_code)] + pub fn latest_tool_calls( + &self, + provider_id: &str, + session_id: &str, + ) -> Option> { + self.get_session(provider_id, session_id) + .and_then(|snapshot| snapshot.turns.last().map(|turn| turn.tool_calls.clone())) + } + + /// Read a full session snapshot. + pub fn get_session( + &self, + provider_id: &str, + session_id: &str, + ) -> Option { + let key = GeminiShadowKey::new(provider_id, session_id); + let mut inner = self.inner.write().expect("gemini shadow lock poisoned"); + let snapshot = inner + .sessions + .get(&key) + .map(|session| Self::snapshot_session(&key, session)); + if snapshot.is_some() { + Self::touch_session_order(&mut inner.session_order, &key); + } + snapshot + } + + /// Remove a single session from the store. + #[allow(dead_code)] + pub fn clear_session(&self, provider_id: &str, session_id: &str) -> bool { + let key = GeminiShadowKey::new(provider_id, session_id); + let mut inner = self.inner.write().expect("gemini shadow lock poisoned"); + let removed = inner.sessions.remove(&key).is_some(); + if removed { + Self::remove_key_from_order(&mut inner.session_order, &key); + } + removed + } + + /// Remove all sessions for a provider. + #[allow(dead_code)] + pub fn clear_provider(&self, provider_id: &str) -> usize { + let mut inner = self.inner.write().expect("gemini shadow lock poisoned"); + let keys: Vec<_> = inner + .sessions + .keys() + .filter(|key| key.provider_id == provider_id) + .cloned() + .collect(); + for key in &keys { + inner.sessions.remove(key); + Self::remove_key_from_order(&mut inner.session_order, key); + } + keys.len() + } + + /// Number of tracked sessions. + #[allow(dead_code)] + pub fn session_count(&self) -> usize { + self.inner + .read() + .expect("gemini shadow lock poisoned") + .sessions + .len() + } + + fn snapshot_session( + key: &GeminiShadowKey, + session: &GeminiShadowSession, + ) -> GeminiShadowSessionSnapshot { + GeminiShadowSessionSnapshot { + provider_id: key.provider_id.clone(), + session_id: key.session_id.clone(), + turns: session.turns.iter().cloned().collect(), + } + } + + fn touch_session_order(order: &mut VecDeque, key: &GeminiShadowKey) { + if let Some(pos) = order.iter().position(|existing| existing == key) { + order.remove(pos); + } + order.push_back(key.clone()); + } + + #[allow(dead_code)] + fn remove_key_from_order(order: &mut VecDeque, key: &GeminiShadowKey) { + if let Some(pos) = order.iter().position(|existing| existing == key) { + order.remove(pos); + } + } + + fn prune_sessions(inner: &mut GeminiShadowInner, max_sessions: usize) { + while inner.sessions.len() > max_sessions { + let Some(evicted_key) = inner.session_order.pop_front() else { + break; + }; + inner.sessions.remove(&evicted_key); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn record_and_read_latest_turn() { + let store = GeminiShadowStore::with_limits(8, 4); + let snapshot = store.record_assistant_turn( + "provider-a", + "session-1", + json!({"parts": [{"text": "hello", "thoughtSignature": "sig-1"}]}), + vec![GeminiToolCallMeta::new( + Some("call-1"), + "get_weather", + json!({"location": "Tokyo"}), + Some("sig-1"), + )], + ); + + assert_eq!(snapshot.provider_id, "provider-a"); + assert_eq!(snapshot.session_id, "session-1"); + assert_eq!(snapshot.turns.len(), 1); + + let content = store + .latest_assistant_content("provider-a", "session-1") + .expect("content"); + assert_eq!(content["parts"][0]["text"], "hello"); + assert_eq!(content["parts"][0]["thoughtSignature"], "sig-1"); + + let tool_calls = store + .latest_tool_calls("provider-a", "session-1") + .expect("tool calls"); + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0].id.as_deref(), Some("call-1")); + assert_eq!(tool_calls[0].name, "get_weather"); + assert_eq!(tool_calls[0].args["location"], "Tokyo"); + assert_eq!(tool_calls[0].thought_signature.as_deref(), Some("sig-1")); + } + + #[test] + fn sessions_are_isolated_by_provider_and_session_id() { + let store = GeminiShadowStore::with_limits(8, 4); + + store.record_assistant_turn("provider-a", "session-1", json!({"text": "a"}), vec![]); + store.record_assistant_turn("provider-b", "session-1", json!({"text": "b"}), vec![]); + store.record_assistant_turn("provider-a", "session-2", json!({"text": "c"}), vec![]); + + assert_eq!(store.session_count(), 3); + assert_eq!( + store.latest_assistant_content("provider-a", "session-1"), + Some(json!({"text": "a"})) + ); + assert_eq!( + store.latest_assistant_content("provider-b", "session-1"), + Some(json!({"text": "b"})) + ); + assert_eq!( + store.latest_assistant_content("provider-a", "session-2"), + Some(json!({"text": "c"})) + ); + } + + #[test] + fn retains_only_latest_turns_per_session() { + let store = GeminiShadowStore::with_limits(8, 2); + + store.record_assistant_turn("provider-a", "session-1", json!({"idx": 1}), vec![]); + store.record_assistant_turn("provider-a", "session-1", json!({"idx": 2}), vec![]); + store.record_assistant_turn("provider-a", "session-1", json!({"idx": 3}), vec![]); + + let snapshot = store + .get_session("provider-a", "session-1") + .expect("snapshot"); + assert_eq!(snapshot.turns.len(), 2); + assert_eq!(snapshot.turns[0].assistant_content, json!({"idx": 2})); + assert_eq!(snapshot.turns[1].assistant_content, json!({"idx": 3})); + } + + #[test] + fn evicts_oldest_session_when_capacity_is_exceeded() { + let store = GeminiShadowStore::with_limits(2, 2); + + store.record_assistant_turn("provider-a", "session-1", json!({"idx": 1}), vec![]); + store.record_assistant_turn("provider-a", "session-2", json!({"idx": 2}), vec![]); + store.record_assistant_turn("provider-a", "session-3", json!({"idx": 3}), vec![]); + + assert!(store.get_session("provider-a", "session-1").is_none()); + assert!(store.get_session("provider-a", "session-2").is_some()); + assert!(store.get_session("provider-a", "session-3").is_some()); + } + + #[test] + fn clear_session_and_provider_work() { + let store = GeminiShadowStore::with_limits(8, 4); + + store.record_assistant_turn("provider-a", "session-1", json!({"idx": 1}), vec![]); + store.record_assistant_turn("provider-a", "session-2", json!({"idx": 2}), vec![]); + store.record_assistant_turn("provider-b", "session-3", json!({"idx": 3}), vec![]); + + assert!(store.clear_session("provider-a", "session-1")); + assert!(store.get_session("provider-a", "session-1").is_none()); + + let removed = store.clear_provider("provider-a"); + assert_eq!(removed, 1); + assert!(store.get_session("provider-a", "session-2").is_none()); + assert!(store.get_session("provider-b", "session-3").is_some()); + } +} diff --git a/src-tauri/src/proxy/providers/streaming_gemini.rs b/src-tauri/src/proxy/providers/streaming_gemini.rs new file mode 100644 index 000000000..d68662167 --- /dev/null +++ b/src-tauri/src/proxy/providers/streaming_gemini.rs @@ -0,0 +1,619 @@ +//! Gemini Native streaming conversion module. +//! +//! Converts Gemini `streamGenerateContent?alt=sse` chunks into Anthropic-style +//! SSE events for Claude-compatible clients. + +use super::gemini_shadow::{GeminiShadowStore, GeminiToolCallMeta}; +use crate::proxy::sse::{strip_sse_field, take_sse_block}; +use bytes::Bytes; +use futures::stream::{Stream, StreamExt}; +use serde_json::{json, Value}; +use std::collections::HashSet; +use std::sync::Arc; + +fn anthropic_usage_from_gemini(usage: Option<&Value>) -> Value { + let Some(usage) = usage else { + return json!({ + "input_tokens": 0, + "output_tokens": 0 + }); + }; + + let input_tokens = usage + .get("promptTokenCount") + .and_then(|value| value.as_u64()) + .unwrap_or(0); + let total_tokens = usage + .get("totalTokenCount") + .and_then(|value| value.as_u64()) + .unwrap_or(0); + let output_tokens = total_tokens.saturating_sub(input_tokens); + + let mut result = json!({ + "input_tokens": input_tokens, + "output_tokens": output_tokens + }); + + if let Some(cached) = usage + .get("cachedContentTokenCount") + .and_then(|value| value.as_u64()) + { + result["cache_read_input_tokens"] = json!(cached); + } + + result +} + +fn map_finish_reason(reason: Option<&str>, has_tool_use: bool, blocked: bool) -> &'static str { + if blocked { + return "refusal"; + } + + match reason { + Some("MAX_TOKENS") => "max_tokens", + Some("SAFETY") + | Some("RECITATION") + | Some("SPII") + | Some("BLOCKLIST") + | Some("PROHIBITED_CONTENT") => "refusal", + _ if has_tool_use => "tool_use", + _ => "end_turn", + } +} + +fn extract_visible_text(parts: &[Value]) -> String { + parts + .iter() + .filter(|part| part.get("thought").and_then(|value| value.as_bool()) != Some(true)) + .filter_map(|part| part.get("text").and_then(|value| value.as_str())) + .collect::() +} + +fn extract_tool_calls(parts: &[Value]) -> Vec { + parts + .iter() + .filter_map(|part| { + let function_call = part.get("functionCall")?; + Some(GeminiToolCallMeta::new( + function_call.get("id").and_then(|value| value.as_str()), + function_call + .get("name") + .and_then(|value| value.as_str()) + .unwrap_or(""), + function_call + .get("args") + .cloned() + .unwrap_or_else(|| json!({})), + part.get("thoughtSignature") + .or_else(|| part.get("thought_signature")) + .and_then(|value| value.as_str()), + )) + }) + .collect() +} + +fn extract_text_thought_signature(parts: &[Value]) -> Option { + parts + .iter() + .filter(|part| part.get("text").is_some() && part.get("functionCall").is_none()) + .filter_map(|part| { + part.get("thoughtSignature") + .or_else(|| part.get("thought_signature")) + .and_then(|value| value.as_str()) + }) + .next_back() + .map(ToString::to_string) +} + +fn merge_tool_call_snapshots( + tool_call_snapshots: &mut Vec, + incoming: Vec, +) { + for tool_call in incoming { + let existing_index = + tool_call_snapshots + .iter() + .position(|existing| match (&existing.id, &tool_call.id) { + (Some(existing_id), Some(incoming_id)) => existing_id == incoming_id, + _ => existing.name == tool_call.name, + }); + + if let Some(index) = existing_index { + tool_call_snapshots[index] = tool_call; + } else { + tool_call_snapshots.push(tool_call); + } + } +} + +fn build_shadow_assistant_parts( + text: Option<&str>, + text_thought_signature: Option<&str>, + tool_calls: &[GeminiToolCallMeta], +) -> Vec { + let mut parts = Vec::new(); + + if text.filter(|text| !text.is_empty()).is_some() || text_thought_signature.is_some() { + let mut part = json!({ + "text": text.unwrap_or("") + }); + if let Some(signature) = text_thought_signature { + part["thoughtSignature"] = json!(signature); + } + parts.push(part); + } + + for tool_call in tool_calls { + let mut part = json!({ + "functionCall": { + "id": tool_call.id.clone().unwrap_or_default(), + "name": tool_call.name, + "args": tool_call.args + } + }); + + if let Some(signature) = &tool_call.thought_signature { + part["thoughtSignature"] = json!(signature); + } + + parts.push(part); + } + + parts +} + +fn encode_sse(event_name: &str, payload: &Value) -> Bytes { + Bytes::from(format!( + "event: {event_name}\ndata: {}\n\n", + serde_json::to_string(payload).unwrap_or_default() + )) +} + +pub fn create_anthropic_sse_stream_from_gemini( + stream: impl Stream> + Send + 'static, + shadow_store: Option>, + provider_id: Option, + session_id: Option, +) -> impl Stream> + Send { + async_stream::stream! { + let mut buffer = String::new(); + let mut message_id: Option = None; + let mut current_model: Option = None; + let mut has_sent_message_start = false; + let mut accumulated_text = String::new(); + let mut text_block_index: Option = None; + let mut next_content_index: u32 = 0; + let mut open_indices: HashSet = HashSet::new(); + let mut tool_call_snapshots: Vec = Vec::new(); + let mut text_thought_signature: Option = None; + let mut latest_usage: Option = None; + let mut latest_finish_reason: Option = None; + let mut blocked_text: Option = None; + tokio::pin!(stream); + + while let Some(chunk) = stream.next().await { + match chunk { + Ok(bytes) => { + let text = String::from_utf8_lossy(&bytes); + buffer.push_str(&text); + + while let Some(block) = take_sse_block(&mut buffer) { + if block.trim().is_empty() { + continue; + } + + let mut data_lines: Vec = Vec::new(); + for line in block.lines() { + if let Some(data) = strip_sse_field(line, "data") { + data_lines.push(data.to_string()); + } + } + + if data_lines.is_empty() { + continue; + } + + let data = data_lines.join("\n"); + if data.trim() == "[DONE]" { + break; + } + + let chunk_json: Value = match serde_json::from_str(&data) { + Ok(value) => value, + Err(_) => continue, + }; + + if message_id.is_none() { + message_id = chunk_json + .get("responseId") + .and_then(|value| value.as_str()) + .map(ToString::to_string); + } + if current_model.is_none() { + current_model = chunk_json + .get("modelVersion") + .and_then(|value| value.as_str()) + .map(ToString::to_string); + } + if latest_usage.is_none() { + latest_usage = chunk_json.get("usageMetadata").cloned(); + } + + if !has_sent_message_start { + let event = json!({ + "type": "message_start", + "message": { + "id": message_id.clone().unwrap_or_default(), + "type": "message", + "role": "assistant", + "model": current_model.clone().unwrap_or_default(), + "usage": anthropic_usage_from_gemini(chunk_json.get("usageMetadata")) + } + }); + yield Ok(encode_sse("message_start", &event)); + has_sent_message_start = true; + } + + if let Some(reason) = chunk_json + .get("promptFeedback") + .and_then(|value| value.get("blockReason")) + .and_then(|value| value.as_str()) + { + blocked_text = Some(format!("Request blocked by Gemini safety filters: {reason}")); + } + + if let Some(candidate) = chunk_json + .get("candidates") + .and_then(|value| value.as_array()) + .and_then(|value| value.first()) + { + if let Some(reason) = candidate.get("finishReason").and_then(|value| value.as_str()) { + latest_finish_reason = Some(reason.to_string()); + } + if let Some(usage) = chunk_json.get("usageMetadata") { + latest_usage = Some(usage.clone()); + } + if let Some(parts) = candidate + .get("content") + .and_then(|value| value.get("parts")) + .and_then(|value| value.as_array()) + { + if let Some(signature) = extract_text_thought_signature(parts) { + text_thought_signature = Some(signature); + } + merge_tool_call_snapshots( + &mut tool_call_snapshots, + extract_tool_calls(parts), + ); + let visible_text = extract_visible_text(parts); + if !visible_text.is_empty() { + let is_cumulative = visible_text.starts_with(&accumulated_text); + let delta = if is_cumulative { + visible_text[accumulated_text.len()..].to_string() + } else { + visible_text.clone() + }; + + if !delta.is_empty() { + let index = *text_block_index.get_or_insert_with(|| { + let assigned = next_content_index; + next_content_index += 1; + assigned + }); + + if !open_indices.contains(&index) { + let start_event = json!({ + "type": "content_block_start", + "index": index, + "content_block": { + "type": "text", + "text": "" + } + }); + yield Ok(encode_sse("content_block_start", &start_event)); + open_indices.insert(index); + } + + let delta_event = json!({ + "type": "content_block_delta", + "index": index, + "delta": { + "type": "text_delta", + "text": delta + } + }); + yield Ok(encode_sse("content_block_delta", &delta_event)); + if is_cumulative { + accumulated_text = visible_text; + } else { + accumulated_text.push_str(&delta); + } + } + } + } + } + } + } + Err(error) => { + yield Err(std::io::Error::other(error.to_string())); + return; + } + } + } + + if !has_sent_message_start { + let event = json!({ + "type": "message_start", + "message": { + "id": message_id.clone().unwrap_or_default(), + "type": "message", + "role": "assistant", + "model": current_model.clone().unwrap_or_default(), + "usage": anthropic_usage_from_gemini(latest_usage.as_ref()) + } + }); + yield Ok(encode_sse("message_start", &event)); + } + + if accumulated_text.is_empty() { + if let Some(blocked_text) = blocked_text.clone() { + let index = *text_block_index.get_or_insert_with(|| { + let assigned = next_content_index; + next_content_index += 1; + assigned + }); + + if !open_indices.contains(&index) { + let start_event = json!({ + "type": "content_block_start", + "index": index, + "content_block": { + "type": "text", + "text": "" + } + }); + yield Ok(encode_sse("content_block_start", &start_event)); + open_indices.insert(index); + } + + let delta_event = json!({ + "type": "content_block_delta", + "index": index, + "delta": { + "type": "text_delta", + "text": blocked_text + } + }); + yield Ok(encode_sse("content_block_delta", &delta_event)); + } + } + + if let Some(index) = text_block_index { + if open_indices.remove(&index) { + let stop_event = json!({ + "type": "content_block_stop", + "index": index + }); + yield Ok(encode_sse("content_block_stop", &stop_event)); + } + } + + let tool_calls = tool_call_snapshots; + for tool_call in &tool_calls { + let index = next_content_index; + next_content_index += 1; + + let start_event = json!({ + "type": "content_block_start", + "index": index, + "content_block": { + "type": "tool_use", + "id": tool_call.id.clone().unwrap_or_default(), + "name": tool_call.name + } + }); + yield Ok(encode_sse("content_block_start", &start_event)); + + let delta_event = json!({ + "type": "content_block_delta", + "index": index, + "delta": { + "type": "input_json_delta", + "partial_json": serde_json::to_string(&tool_call.args).unwrap_or_else(|_| "{}".to_string()) + } + }); + yield Ok(encode_sse("content_block_delta", &delta_event)); + + let stop_event = json!({ + "type": "content_block_stop", + "index": index + }); + yield Ok(encode_sse("content_block_stop", &stop_event)); + } + + if let (Some(store), Some(provider_id), Some(session_id)) = ( + shadow_store.as_ref(), + provider_id.as_deref(), + session_id.as_deref(), + ) { + let shadow_text = if accumulated_text.is_empty() { + blocked_text.as_deref() + } else { + Some(accumulated_text.as_str()) + }; + let shadow_parts = build_shadow_assistant_parts( + shadow_text, + text_thought_signature.as_deref(), + &tool_calls, + ); + if !shadow_parts.is_empty() { + store.record_assistant_turn( + provider_id, + session_id, + json!({ "parts": shadow_parts }), + tool_calls.clone(), + ); + } + } + + let stop_reason = map_finish_reason( + latest_finish_reason.as_deref(), + !tool_calls.is_empty(), + blocked_text.is_some(), + ); + let usage = anthropic_usage_from_gemini(latest_usage.as_ref()); + let message_delta = json!({ + "type": "message_delta", + "delta": { + "stop_reason": stop_reason, + "stop_sequence": Value::Null + }, + "usage": usage + }); + yield Ok(encode_sse("message_delta", &message_delta)); + + let message_stop = json!({ "type": "message_stop" }); + yield Ok(encode_sse("message_stop", &message_stop)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::proxy::providers::gemini_shadow::GeminiShadowStore; + use crate::proxy::providers::transform_gemini::anthropic_to_gemini_with_shadow; + use std::sync::Arc; + + fn collect_stream_output(chunks: Vec<&str>) -> String { + let owned_chunks: Vec = chunks.into_iter().map(ToString::to_string).collect(); + let stream = futures::stream::iter( + owned_chunks + .into_iter() + .map(|chunk| Ok::(Bytes::from(chunk))), + ); + let converted = create_anthropic_sse_stream_from_gemini(stream, None, None, None); + futures::executor::block_on(async move { + converted + .collect::>() + .await + .into_iter() + .map(|item| String::from_utf8(item.unwrap().to_vec()).unwrap()) + .collect::>() + .join("") + }) + } + + fn collect_stream_output_with_shadow( + chunks: Vec<&str>, + store: Arc, + provider_id: &str, + session_id: &str, + ) -> String { + let owned_chunks: Vec = chunks.into_iter().map(ToString::to_string).collect(); + let stream = futures::stream::iter( + owned_chunks + .into_iter() + .map(|chunk| Ok::(Bytes::from(chunk))), + ); + let converted = create_anthropic_sse_stream_from_gemini( + stream, + Some(store), + Some(provider_id.to_string()), + Some(session_id.to_string()), + ); + futures::executor::block_on(async move { + converted + .collect::>() + .await + .into_iter() + .map(|item| String::from_utf8(item.unwrap().to_vec()).unwrap()) + .collect::>() + .join("") + }) + } + + #[test] + fn converts_text_stream_to_anthropic_sse() { + let output = collect_stream_output(vec![ + "data: {\"responseId\":\"resp_1\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hel\"}]}}],\"usageMetadata\":{\"promptTokenCount\":10,\"totalTokenCount\":13}}\n\n", + "data: {\"responseId\":\"resp_1\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"finishReason\":\"STOP\",\"content\":{\"parts\":[{\"text\":\"Hello\"}]}}],\"usageMetadata\":{\"promptTokenCount\":10,\"totalTokenCount\":15}}\n\n", + ]); + + assert!(output.contains("event: message_start")); + assert!(output.contains("\"type\":\"text_delta\"")); + assert!(output.contains("\"text\":\"Hel\"")); + assert!(output.contains("\"text\":\"lo\"")); + assert!(output.contains("\"stop_reason\":\"end_turn\"")); + assert!(output.contains("event: message_stop")); + } + + #[test] + fn converts_function_call_stream_to_tool_use_events() { + let output = collect_stream_output(vec![ + "data: {\"responseId\":\"resp_2\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"finishReason\":\"STOP\",\"content\":{\"parts\":[{\"functionCall\":{\"id\":\"call_1\",\"name\":\"get_weather\",\"args\":{\"city\":\"Tokyo\"}},\"thoughtSignature\":\"sig-1\"}]}}],\"usageMetadata\":{\"promptTokenCount\":5,\"totalTokenCount\":8}}\n\n", + ]); + + assert!(output.contains("\"type\":\"tool_use\"")); + assert!(output.contains("\"name\":\"get_weather\"")); + assert!(output.contains("\"type\":\"input_json_delta\"")); + assert!(output.contains("\"stop_reason\":\"tool_use\"")); + } + + #[test] + fn converts_crlf_delimited_stream_to_anthropic_sse() { + let output = collect_stream_output(vec![ + "data: {\"responseId\":\"resp_3\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hi\"}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":6}}\r\n\r\n", + "data: {\"responseId\":\"resp_3\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"finishReason\":\"STOP\",\"content\":{\"parts\":[{\"text\":\"Hi there\"}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":9}}\r\n\r\n", + ]); + + assert!(output.contains("event: message_start")); + assert!(output.contains("\"type\":\"text_delta\"")); + assert!(output.contains("\"text\":\"Hi\"")); + assert!(output.contains("\"text\":\" there\"")); + assert!(output.contains("event: message_stop")); + } + + #[test] + fn stores_full_text_for_shadow_replay_across_delta_chunks() { + let store = Arc::new(GeminiShadowStore::with_limits(8, 4)); + let output = collect_stream_output_with_shadow( + vec![ + "data: {\"responseId\":\"resp_4\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hel\"}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":6}}\n\n", + "data: {\"responseId\":\"resp_4\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"finishReason\":\"STOP\",\"content\":{\"parts\":[{\"text\":\"lo\"},{\"text\":\"\",\"thoughtSignature\":\"sig-1\"}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":8}}\n\n", + ], + store.clone(), + "provider-a", + "session-1", + ); + + assert!(output.contains("\"text\":\"Hel\"")); + assert!(output.contains("\"text\":\"lo\"")); + + let shadow = store + .latest_assistant_content("provider-a", "session-1") + .unwrap(); + assert_eq!(shadow["parts"][0]["text"], "Hello"); + assert_eq!(shadow["parts"][0]["thoughtSignature"], "sig-1"); + + let second_turn = anthropic_to_gemini_with_shadow( + json!({ + "messages": [ + { "role": "user", "content": "Hi" }, + { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }] }, + { "role": "user", "content": "Continue" } + ] + }), + Some(store.as_ref()), + Some("provider-a"), + Some("session-1"), + ) + .unwrap(); + + assert_eq!(second_turn["contents"][1]["role"], "model"); + assert_eq!(second_turn["contents"][1]["parts"][0]["text"], "Hello"); + assert_eq!( + second_turn["contents"][1]["parts"][0]["thoughtSignature"], + "sig-1" + ); + } +} diff --git a/src-tauri/src/proxy/providers/transform_gemini.rs b/src-tauri/src/proxy/providers/transform_gemini.rs new file mode 100644 index 000000000..2b0a601f4 --- /dev/null +++ b/src-tauri/src/proxy/providers/transform_gemini.rs @@ -0,0 +1,788 @@ +//! Gemini Native format conversion module. +//! +//! Converts Anthropic Messages requests to Gemini `generateContent` requests, +//! and Gemini `GenerateContentResponse` payloads back to Anthropic Messages +//! responses for Claude-compatible clients. + +use super::gemini_schema::build_gemini_function_declaration; +use super::gemini_shadow::{GeminiAssistantTurn, GeminiShadowStore, GeminiToolCallMeta}; +use crate::proxy::error::ProxyError; +use serde_json::{json, Map, Value}; + +pub fn anthropic_to_gemini(body: Value) -> Result { + anthropic_to_gemini_with_shadow(body, None, None, None) +} + +pub fn anthropic_to_gemini_with_shadow( + body: Value, + shadow_store: Option<&GeminiShadowStore>, + provider_id: Option<&str>, + session_id: Option<&str>, +) -> Result { + let mut result = json!({}); + let shadow_turns = shadow_store + .zip(provider_id) + .zip(session_id) + .and_then(|((store, provider_id), session_id)| store.get_session(provider_id, session_id)) + .map(|snapshot| snapshot.turns) + .unwrap_or_default(); + + if let Some(system) = build_system_instruction(body.get("system"))? { + result["systemInstruction"] = system; + } + + if let Some(messages) = body.get("messages").and_then(|value| value.as_array()) { + result["contents"] = json!(convert_messages_to_contents(messages, &shadow_turns)?); + } + + if let Some(generation_config) = build_generation_config(&body) { + result["generationConfig"] = generation_config; + } + + if let Some(tools) = body.get("tools").and_then(|value| value.as_array()) { + let function_declarations: Vec = tools + .iter() + .filter(|tool| tool.get("type").and_then(|value| value.as_str()) != Some("BatchTool")) + .map(|tool| { + build_gemini_function_declaration( + tool.get("name") + .and_then(|value| value.as_str()) + .unwrap_or(""), + tool.get("description").and_then(|value| value.as_str()), + tool.get("input_schema") + .cloned() + .unwrap_or_else(|| json!({})), + ) + }) + .collect(); + + if !function_declarations.is_empty() { + result["tools"] = json!([{ "functionDeclarations": function_declarations }]); + } + } + + if let Some(tool_config) = map_tool_choice(body.get("tool_choice"))? { + result["toolConfig"] = tool_config; + } + + Ok(result) +} + +pub fn gemini_to_anthropic(body: Value) -> Result { + gemini_to_anthropic_with_shadow(body, None, None, None) +} + +pub fn gemini_to_anthropic_with_shadow( + body: Value, + shadow_store: Option<&GeminiShadowStore>, + provider_id: Option<&str>, + session_id: Option<&str>, +) -> Result { + if let Some(block_reason) = body + .get("promptFeedback") + .and_then(|value| value.get("blockReason")) + .and_then(|value| value.as_str()) + { + let text = format!("Request blocked by Gemini safety filters: {block_reason}"); + return Ok(json!({ + "id": body.get("responseId").and_then(|value| value.as_str()).unwrap_or(""), + "type": "message", + "role": "assistant", + "content": [{ "type": "text", "text": text }], + "model": body.get("modelVersion").and_then(|value| value.as_str()).unwrap_or(""), + "stop_reason": "refusal", + "stop_sequence": Value::Null, + "usage": build_anthropic_usage(body.get("usageMetadata")) + })); + } + + let candidate = body + .get("candidates") + .and_then(|value| value.as_array()) + .and_then(|value| value.first()) + .ok_or_else(|| { + ProxyError::TransformError("No candidates in Gemini response".to_string()) + })?; + + let parts = candidate + .get("content") + .and_then(|value| value.get("parts")) + .and_then(|value| value.as_array()) + .cloned() + .unwrap_or_default(); + + let mut content = Vec::new(); + let mut has_tool_use = false; + + for part in &parts { + if part.get("thought").and_then(|value| value.as_bool()) == Some(true) { + continue; + } + + if let Some(text) = part.get("text").and_then(|value| value.as_str()) { + if !text.is_empty() { + content.push(json!({ + "type": "text", + "text": text + })); + } + continue; + } + + if let Some(function_call) = part.get("functionCall") { + has_tool_use = true; + content.push(json!({ + "type": "tool_use", + "id": function_call.get("id").and_then(|value| value.as_str()).unwrap_or(""), + "name": function_call.get("name").and_then(|value| value.as_str()).unwrap_or(""), + "input": function_call.get("args").cloned().unwrap_or_else(|| json!({})) + })); + } + } + + let stop_reason = map_finish_reason( + candidate + .get("finishReason") + .and_then(|value| value.as_str()), + has_tool_use, + ); + + let anthropic_response = json!({ + "id": body.get("responseId").and_then(|value| value.as_str()).unwrap_or(""), + "type": "message", + "role": "assistant", + "content": content, + "model": body.get("modelVersion").and_then(|value| value.as_str()).unwrap_or(""), + "stop_reason": stop_reason, + "stop_sequence": Value::Null, + "usage": build_anthropic_usage(body.get("usageMetadata")) + }); + + if let (Some(store), Some(provider_id), Some(session_id), Some(content)) = ( + shadow_store, + provider_id, + session_id, + candidate.get("content"), + ) { + store.record_assistant_turn( + provider_id, + session_id, + content.clone(), + extract_tool_call_meta(&parts), + ); + } + + Ok(anthropic_response) +} + +pub fn extract_gemini_model(body: &Value) -> Option<&str> { + body.get("model").and_then(|value| value.as_str()) +} + +fn build_system_instruction(system: Option<&Value>) -> Result, ProxyError> { + let Some(system) = system else { + return Ok(None); + }; + + if let Some(text) = system.as_str() { + if text.is_empty() { + return Ok(None); + } + return Ok(Some(json!({ + "parts": [{ "text": text }] + }))); + } + + let Some(blocks) = system.as_array() else { + return Err(ProxyError::TransformError( + "Anthropic system must be a string or an array".to_string(), + )); + }; + + let texts: Vec<&str> = blocks + .iter() + .filter_map(|block| block.get("text").and_then(|value| value.as_str())) + .filter(|text| !text.is_empty()) + .collect(); + + if texts.is_empty() { + return Ok(None); + } + + Ok(Some(json!({ + "parts": [{ "text": texts.join("\n\n") }] + }))) +} + +fn build_generation_config(body: &Value) -> Option { + let mut config = Map::new(); + + if let Some(value) = body.get("max_tokens") { + config.insert("maxOutputTokens".to_string(), value.clone()); + } + if let Some(value) = body.get("temperature") { + config.insert("temperature".to_string(), value.clone()); + } + if let Some(value) = body.get("top_p") { + config.insert("topP".to_string(), value.clone()); + } + if let Some(value) = body.get("stop_sequences") { + config.insert("stopSequences".to_string(), value.clone()); + } + + if config.is_empty() { + None + } else { + Some(Value::Object(config)) + } +} + +fn convert_messages_to_contents( + messages: &[Value], + shadow_turns: &[GeminiAssistantTurn], +) -> Result, ProxyError> { + let mut contents = Vec::new(); + let mut tool_name_by_id = std::collections::HashMap::::new(); + let total_assistant_messages = messages + .iter() + .filter(|message| message.get("role").and_then(|value| value.as_str()) == Some("assistant")) + .count(); + let shadow_start_index = total_assistant_messages.saturating_sub(shadow_turns.len()); + let mut assistant_seen_index = 0usize; + + for message in messages { + let role = message + .get("role") + .and_then(|value| value.as_str()) + .unwrap_or("user"); + + let gemini_role = if role == "assistant" { "model" } else { "user" }; + + let parts = if role == "assistant" { + let shadow_index = assistant_seen_index + .checked_sub(shadow_start_index) + .filter(|index| *index < shadow_turns.len()); + assistant_seen_index += 1; + + if let Some(index) = shadow_index { + let shadow_turn = &shadow_turns[index]; + merge_tool_names_from_shadow(shadow_turn, &mut tool_name_by_id); + if let Some(parts) = shadow_parts(&shadow_turn.assistant_content) { + parts + } else { + convert_message_content_to_parts( + message.get("content"), + role, + &mut tool_name_by_id, + )? + } + } else { + convert_message_content_to_parts( + message.get("content"), + role, + &mut tool_name_by_id, + )? + } + } else { + convert_message_content_to_parts(message.get("content"), role, &mut tool_name_by_id)? + }; + + contents.push(json!({ + "role": gemini_role, + "parts": parts + })); + } + + Ok(contents) +} + +fn convert_message_content_to_parts( + content: Option<&Value>, + role: &str, + tool_name_by_id: &mut std::collections::HashMap, +) -> Result, ProxyError> { + let Some(content) = content else { + return Ok(Vec::new()); + }; + + if let Some(text) = content.as_str() { + return Ok(vec![json!({ "text": text })]); + } + + let Some(blocks) = content.as_array() else { + return Err(ProxyError::TransformError( + "Anthropic message content must be a string or array".to_string(), + )); + }; + + let mut parts = Vec::new(); + + for block in blocks { + let block_type = block + .get("type") + .and_then(|value| value.as_str()) + .unwrap_or(""); + + match block_type { + "text" => { + if let Some(text) = block.get("text").and_then(|value| value.as_str()) { + parts.push(json!({ "text": text })); + } + } + "image" => { + let source = block.get("source").ok_or_else(|| { + ProxyError::TransformError("Gemini image block missing source".to_string()) + })?; + + let source_type = source + .get("type") + .and_then(|value| value.as_str()) + .unwrap_or(""); + + if source_type != "base64" { + return Err(ProxyError::TransformError(format!( + "Gemini Native only supports base64 image sources, got `{source_type}`" + ))); + } + + parts.push(json!({ + "inlineData": { + "mimeType": source.get("media_type").and_then(|value| value.as_str()).unwrap_or("image/png"), + "data": source.get("data").and_then(|value| value.as_str()).unwrap_or("") + } + })); + } + "document" => { + let source = block.get("source").ok_or_else(|| { + ProxyError::TransformError("Gemini document block missing source".to_string()) + })?; + + let source_type = source + .get("type") + .and_then(|value| value.as_str()) + .unwrap_or(""); + + if source_type != "base64" { + return Err(ProxyError::TransformError(format!( + "Gemini Native only supports base64 document sources, got `{source_type}`" + ))); + } + + parts.push(json!({ + "inlineData": { + "mimeType": source.get("media_type").and_then(|value| value.as_str()).unwrap_or("application/pdf"), + "data": source.get("data").and_then(|value| value.as_str()).unwrap_or("") + } + })); + } + "tool_use" => { + if role != "assistant" { + return Err(ProxyError::TransformError( + "tool_use blocks are only valid in assistant messages".to_string(), + )); + } + + let id = block + .get("id") + .and_then(|value| value.as_str()) + .unwrap_or(""); + let name = block + .get("name") + .and_then(|value| value.as_str()) + .unwrap_or(""); + if !id.is_empty() && !name.is_empty() { + tool_name_by_id.insert(id.to_string(), name.to_string()); + } + + parts.push(json!({ + "functionCall": { + "id": id, + "name": name, + "args": block.get("input").cloned().unwrap_or_else(|| json!({})) + } + })); + } + "tool_result" => { + let tool_use_id = block + .get("tool_use_id") + .and_then(|value| value.as_str()) + .unwrap_or(""); + let name = tool_name_by_id + .get(tool_use_id) + .cloned() + .unwrap_or_default(); + + parts.push(json!({ + "functionResponse": { + "id": tool_use_id, + "name": name, + "response": normalize_tool_result_response(block.get("content")) + } + })); + } + "thinking" | "redacted_thinking" => {} + _ => {} + } + } + + Ok(parts) +} + +fn normalize_tool_result_response(content: Option<&Value>) -> Value { + match content { + Some(Value::String(text)) => json!({ "content": text }), + Some(Value::Array(blocks)) => { + let texts: Vec<&str> = blocks + .iter() + .filter(|block| block.get("type").and_then(|value| value.as_str()) == Some("text")) + .filter_map(|block| block.get("text").and_then(|value| value.as_str())) + .collect(); + + if texts.is_empty() { + json!({ "content": Value::Array(blocks.clone()) }) + } else { + json!({ "content": texts.join("\n") }) + } + } + Some(value) => json!({ "content": value.clone() }), + None => json!({ "content": "" }), + } +} + +fn shadow_parts(content: &Value) -> Option> { + content + .get("parts") + .and_then(|value| value.as_array()) + .cloned() + .or_else(|| content.as_array().cloned()) +} + +fn merge_tool_names_from_shadow( + turn: &GeminiAssistantTurn, + tool_name_by_id: &mut std::collections::HashMap, +) { + for tool_call in &turn.tool_calls { + if let Some(id) = &tool_call.id { + tool_name_by_id.insert(id.clone(), tool_call.name.clone()); + } + } +} + +fn extract_tool_call_meta(parts: &[Value]) -> Vec { + parts + .iter() + .filter_map(|part| { + let function_call = part.get("functionCall")?; + Some(GeminiToolCallMeta::new( + function_call.get("id").and_then(|value| value.as_str()), + function_call + .get("name") + .and_then(|value| value.as_str()) + .unwrap_or(""), + function_call + .get("args") + .cloned() + .unwrap_or_else(|| json!({})), + part.get("thoughtSignature") + .or_else(|| part.get("thought_signature")) + .and_then(|value| value.as_str()), + )) + }) + .collect() +} + +fn map_tool_choice(tool_choice: Option<&Value>) -> Result, ProxyError> { + let Some(tool_choice) = tool_choice else { + return Ok(None); + }; + + match tool_choice { + Value::String(choice) => Ok(match choice.as_str() { + "auto" => Some(json!({ + "functionCallingConfig": { "mode": "AUTO" } + })), + "none" => Some(json!({ + "functionCallingConfig": { "mode": "NONE" } + })), + other => { + return Err(ProxyError::TransformError(format!( + "Unsupported Gemini tool_choice string: {other}" + ))); + } + }), + Value::Object(object) => { + let Some(choice_type) = object.get("type").and_then(|value| value.as_str()) else { + return Ok(None); + }; + + let config = match choice_type { + "auto" => json!({ "mode": "AUTO" }), + "none" => json!({ "mode": "NONE" }), + "any" => json!({ "mode": "ANY" }), + "tool" => { + let name = object + .get("name") + .and_then(|value| value.as_str()) + .unwrap_or(""); + json!({ + "mode": "ANY", + "allowedFunctionNames": [name] + }) + } + other => { + return Err(ProxyError::TransformError(format!( + "Unsupported Gemini tool_choice type: {other}" + ))); + } + }; + + Ok(Some(json!({ "functionCallingConfig": config }))) + } + _ => Ok(None), + } +} + +fn build_anthropic_usage(usage: Option<&Value>) -> Value { + let Some(usage) = usage else { + return json!({ + "input_tokens": 0, + "output_tokens": 0 + }); + }; + + let input_tokens = usage + .get("promptTokenCount") + .and_then(|value| value.as_u64()) + .unwrap_or(0); + let total_tokens = usage + .get("totalTokenCount") + .and_then(|value| value.as_u64()) + .unwrap_or(0); + let output_tokens = total_tokens.saturating_sub(input_tokens); + + let mut result = json!({ + "input_tokens": input_tokens, + "output_tokens": output_tokens + }); + + if let Some(cached) = usage + .get("cachedContentTokenCount") + .and_then(|value| value.as_u64()) + { + result["cache_read_input_tokens"] = json!(cached); + } + + result +} + +fn map_finish_reason(reason: Option<&str>, has_tool_use: bool) -> Value { + let mapped = match reason { + Some("MAX_TOKENS") => Some("max_tokens"), + Some("STOP") | Some("FINISH_REASON_UNSPECIFIED") | None => { + if has_tool_use { + Some("tool_use") + } else { + Some("end_turn") + } + } + Some("SAFETY") + | Some("RECITATION") + | Some("SPII") + | Some("BLOCKLIST") + | Some("PROHIBITED_CONTENT") => Some("refusal"), + Some(other) => { + log::warn!("[Claude/Gemini] Unknown Gemini finishReason `{other}`, using end_turn"); + Some("end_turn") + } + }; + + match mapped { + Some(value) => json!(value), + None => Value::Null, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn anthropic_to_gemini_maps_system_and_messages() { + let input = json!({ + "model": "gemini-2.5-pro", + "max_tokens": 128, + "system": "You are helpful.", + "messages": [ + { "role": "user", "content": "Hello" } + ] + }); + + let result = anthropic_to_gemini(input).unwrap(); + assert_eq!( + result["systemInstruction"]["parts"][0]["text"], + "You are helpful." + ); + assert_eq!(result["contents"][0]["role"], "user"); + assert_eq!(result["contents"][0]["parts"][0]["text"], "Hello"); + assert_eq!(result["generationConfig"]["maxOutputTokens"], 128); + } + + #[test] + fn anthropic_to_gemini_maps_tools_and_tool_results() { + let input = json!({ + "messages": [ + { + "role": "assistant", + "content": [ + { "type": "tool_use", "id": "call_1", "name": "get_weather", "input": { "city": "Tokyo" } } + ] + }, + { + "role": "user", + "content": [ + { "type": "tool_result", "tool_use_id": "call_1", "content": "Sunny" } + ] + } + ], + "tools": [ + { + "name": "get_weather", + "description": "Weather lookup", + "input_schema": { "type": "object", "properties": { "city": { "type": "string" } } } + } + ], + "tool_choice": { "type": "tool", "name": "get_weather" } + }); + + let result = anthropic_to_gemini(input).unwrap(); + assert_eq!( + result["tools"][0]["functionDeclarations"][0]["name"], + "get_weather" + ); + assert!(result["tools"][0]["functionDeclarations"][0] + .get("parameters") + .is_some()); + assert_eq!( + result["contents"][0]["parts"][0]["functionCall"]["name"], + "get_weather" + ); + assert_eq!( + result["contents"][1]["parts"][0]["functionResponse"]["name"], + "get_weather" + ); + assert_eq!( + result["toolConfig"]["functionCallingConfig"]["allowedFunctionNames"][0], + "get_weather" + ); + } + + #[test] + fn anthropic_to_gemini_uses_parameters_json_schema_for_rich_tool_schema() { + let input = json!({ + "tools": [ + { + "name": "search", + "description": "Search data", + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "query": { "type": "string" } + }, + "required": ["query"], + "additionalProperties": false + } + } + ] + }); + + let result = anthropic_to_gemini(input).unwrap(); + let declaration = &result["tools"][0]["functionDeclarations"][0]; + + assert!(declaration.get("parameters").is_none()); + assert!(declaration.get("parametersJsonSchema").is_some()); + assert!(declaration["parametersJsonSchema"].get("$schema").is_none()); + assert_eq!( + declaration["parametersJsonSchema"]["additionalProperties"], + false + ); + } + + #[test] + fn gemini_to_anthropic_maps_text_and_usage() { + let input = json!({ + "responseId": "resp_1", + "modelVersion": "gemini-2.5-pro", + "candidates": [{ + "finishReason": "STOP", + "content": { + "parts": [{ "text": "Hello from Gemini" }] + } + }], + "usageMetadata": { + "promptTokenCount": 12, + "totalTokenCount": 20, + "cachedContentTokenCount": 3 + } + }); + + let result = gemini_to_anthropic(input).unwrap(); + assert_eq!(result["id"], "resp_1"); + assert_eq!(result["content"][0]["type"], "text"); + assert_eq!(result["content"][0]["text"], "Hello from Gemini"); + assert_eq!(result["stop_reason"], "end_turn"); + assert_eq!(result["usage"]["input_tokens"], 12); + assert_eq!(result["usage"]["output_tokens"], 8); + assert_eq!(result["usage"]["cache_read_input_tokens"], 3); + } + + #[test] + fn gemini_to_anthropic_maps_function_calls_to_tool_use() { + let input = json!({ + "responseId": "resp_2", + "modelVersion": "gemini-2.5-pro", + "candidates": [{ + "finishReason": "STOP", + "content": { + "parts": [{ + "functionCall": { + "id": "call_1", + "name": "get_weather", + "args": { "city": "Tokyo" } + } + }] + } + }], + "usageMetadata": { + "promptTokenCount": 10, + "totalTokenCount": 15 + } + }); + + let result = gemini_to_anthropic(input).unwrap(); + assert_eq!(result["content"][0]["type"], "tool_use"); + assert_eq!(result["content"][0]["id"], "call_1"); + assert_eq!(result["stop_reason"], "tool_use"); + } + + #[test] + fn gemini_to_anthropic_maps_blocked_prompt_to_refusal() { + let input = json!({ + "responseId": "resp_3", + "modelVersion": "gemini-2.5-flash", + "promptFeedback": { "blockReason": "SAFETY" }, + "usageMetadata": { + "promptTokenCount": 4, + "totalTokenCount": 4 + } + }); + + let result = gemini_to_anthropic(input).unwrap(); + assert_eq!(result["stop_reason"], "refusal"); + assert_eq!(result["content"][0]["type"], "text"); + assert!(result["content"][0]["text"] + .as_str() + .unwrap() + .contains("SAFETY")); + } +}