feat: completions API improvements, gemini endpoint, response types

2026-02-15 17:08:53 -06:00
parent afa96b88a5
commit ca9f808ee3
8 changed files with 1031 additions and 742 deletions
--- a/src/api/completions.rs
+++ b/src/api/completions.rs
@@ -15,45 +15,77 @@ use super::types::*;
 use super::util::{err_response, now_unix};
 use super::AppState;

+/// Extract a conversation/session ID from a flexible JSON value.
+/// Accepts a plain string or an object with an "id" field.
+fn extract_conversation_id(conv: &Option<serde_json::Value>) -> Option<String> {
+    match conv {
+        Some(serde_json::Value::String(s)) => Some(s.clone()),
+        Some(obj) => obj["id"].as_str().map(|s| s.to_string()),
+        None => None,
+    }
+}
+
+/// System fingerprint for completions responses (derived from crate version at compile time).
+fn system_fingerprint() -> String {
+    format!("fp_{}", env!("CARGO_PKG_VERSION").replace('.', ""))
+}
+
+/// Build a streaming chunk JSON with all required OpenAI fields.
+/// Includes system_fingerprint, service_tier, and logprobs:null in choices.
+fn chunk_json(
+    id: &str, model: &str,
+    choices: serde_json::Value,
+    usage: Option<serde_json::Value>,
+) -> String {
+    let mut obj = serde_json::json!({
+        "id": id,
+        "object": "chat.completion.chunk",
+        "created": now_unix(),
+        "model": model,
+        "system_fingerprint": system_fingerprint(),
+        "service_tier": "default",
+        "choices": choices,
+    });
+    if let Some(u) = usage {
+        obj["usage"] = u;
+    }
+    serde_json::to_string(&obj).unwrap_or_default()
+}
+
+/// Build a single choice for a streaming chunk (delta + finish_reason + logprobs).
+fn chunk_choice(index: u32, delta: serde_json::Value, finish_reason: Option<&str>) -> serde_json::Value {
+    serde_json::json!({
+        "index": index,
+        "delta": delta,
+        "logprobs": serde_json::Value::Null,
+        "finish_reason": finish_reason,
+    })
+}
+
+// ─── Finish reason mapping ───────────────────────────────────────────────────
+
+/// Map Google's finishReason → OpenAI's finish_reason.
+/// Google: STOP, MAX_TOKENS, SAFETY, RECITATION, OTHER, BLOCKLIST, PROHIBITED_CONTENT
+/// OpenAI: stop, length, content_filter, tool_calls (handled separately)
+fn google_to_openai_finish_reason(stop_reason: Option<&str>) -> &'static str {
+    match stop_reason {
+        Some("MAX_TOKENS") => "length",
+        Some("SAFETY") | Some("RECITATION") | Some("BLOCKLIST") | Some("PROHIBITED_CONTENT") => "content_filter",
+        _ => "stop",
+    }
+}
+
 // ─── Input extraction ────────────────────────────────────────────────────────

 /// Extract user text from Chat Completions messages array.
 ///
-/// When tool results are present, builds the full conversation including
-/// tool call results so the model can continue after tool use.
+/// Builds the full conversation context including all messages (system, user,
+/// assistant, tool) so the model has complete history — matching how OpenAI
+/// sends the entire messages array to the model.
 fn extract_chat_input(messages: &[CompletionMessage]) -> String {
-    let has_tool_results = messages.iter().any(|m| m.role == "tool");
-
-    if has_tool_results {
-        // Build full conversation context including tool results
-        return build_conversation_with_tools(messages);
-    }
-
-    // Simple path: no tools, just extract system + last user message
-    let mut system_parts = Vec::new();
-    let mut user_parts = Vec::new();
-
-    for msg in messages {
-        let text = extract_message_text(&msg.content);
-        if text.is_empty() {
-            continue;
-        }
-        match msg.role.as_str() {
-            "system" | "developer" => system_parts.push(text),
-            "user" => user_parts.push(text),
-            _ => {}
-        }
-    }
-
-    let mut result = String::new();
-    if !system_parts.is_empty() {
-        result.push_str(&system_parts.join("\n"));
-        result.push_str("\n\n");
-    }
-    if let Some(last) = user_parts.last() {
-        result.push_str(last);
-    }
-    result.trim().to_string()
+    // Always build the full conversation — we used to only take the last user
+    // message which broke multi-turn conversations via the messages array.
+    build_conversation_with_tools(messages)
 }

 /// Extract text content from a message's content field (string or array).
@@ -179,18 +211,36 @@ pub(crate) async fn handle_completions(
    // Store generation parameters for MITM injection
    {
        use crate::mitm::store::GenerationParams;
+        let (response_mime_type, response_schema) = match body.response_format.as_ref() {
+            Some(rf) => match rf.format_type.as_str() {
+                "json_object" | "json" => (Some("application/json".to_string()), None),
+                "json_schema" => {
+                    let schema = rf.json_schema.as_ref().and_then(|js| js.schema.clone());
+                    (Some("application/json".to_string()), schema)
+                }
+                _ => (None, None),
+            },
+            None => (None, None),
+        };
        let gp = GenerationParams {
            temperature: body.temperature,
            top_p: body.top_p,
            top_k: None, // OpenAI doesn't have top_k
            max_output_tokens: body.max_tokens.or(body.max_completion_tokens),
-            stop_sequences: None, // TODO: body.stop
+            stop_sequences: body.stop.clone().map(|s| s.into_vec()),
            frequency_penalty: body.frequency_penalty,
            presence_penalty: body.presence_penalty,
+            reasoning_effort: body.reasoning_effort.clone(),
+            response_mime_type,
+            response_schema,
+            google_search: body.web_search,
        };
        // Only store if at least one param is set
        if gp.temperature.is_some() || gp.top_p.is_some() || gp.max_output_tokens.is_some()
            || gp.frequency_penalty.is_some() || gp.presence_penalty.is_some()
+            || gp.reasoning_effort.is_some() || gp.stop_sequences.is_some()
+            || gp.response_mime_type.is_some() || gp.response_schema.is_some()
+            || gp.google_search
        {
            state.mitm_store.set_generation_params(gp).await;
        } else {
@@ -216,8 +266,28 @@ pub(crate) async fn handle_completions(
        );
    }

-    // Fresh cascade per request
-    let cascade_id = match state.backend.create_cascade().await {
+    let n = (body.n.max(1)).min(5); // Cap at 5 to prevent abuse
+    if n > 1 && body.stream {
+        warn!("n={n} requested with streaming — streaming only supports n=1, ignoring n");
+    }
+
+    // Session/conversation: reuse cascade if conversation ID provided
+    let session_id_str = extract_conversation_id(&body.conversation);
+
+    // Helper to create a cascade (reuses session or creates fresh)
+    let create_cascade = |state: Arc<AppState>, session_id: Option<String>| async move {
+        if let Some(ref sid) = session_id {
+            state
+                .sessions
+                .get_or_create(Some(sid), || state.backend.create_cascade())
+                .await
+                .map(|sr| sr.cascade_id)
+        } else {
+            state.backend.create_cascade().await
+        }
+    };
+
+    let cascade_id = match create_cascade(Arc::clone(&state), session_id_str.clone()).await {
        Ok(cid) => cid,
        Err(e) => {
            return err_response(
@@ -228,7 +298,7 @@ pub(crate) async fn handle_completions(
        }
    };

-    // Send message
+    // Send message on primary cascade
    state.mitm_store.set_active_cascade(&cascade_id).await;
    match state
        .backend
@@ -275,7 +345,7 @@ pub(crate) async fn handle_completions(
            include_usage,
        )
        .await
-    } else {
+    } else if n <= 1 {
        chat_completions_sync(
            state,
            completion_id,
@@ -284,6 +354,108 @@ pub(crate) async fn handle_completions(
            body.timeout,
        )
        .await
+    } else {
+        // n > 1: fire additional (n-1) parallel cascades
+        let mut extra_cascade_ids = Vec::with_capacity((n - 1) as usize);
+        for _ in 1..n {
+            match state.backend.create_cascade().await {
+                Ok(cid) => {
+                    // Send the same message on each extra cascade
+                    match state.backend.send_message(&cid, &user_text, model.model_enum).await {
+                        Ok((200, _)) => {
+                            let bg = Arc::clone(&state.backend);
+                            let cid2 = cid.clone();
+                            tokio::spawn(async move { let _ = bg.update_annotations(&cid2).await; });
+                            extra_cascade_ids.push(cid);
+                        }
+                        _ => {} // Skip failed cascades
+                    }
+                }
+                Err(_) => {} // Skip failed cascade creation
+            }
+        }
+
+        // Poll all cascades in parallel
+        let mut handles = Vec::with_capacity(n as usize);
+        let all_cascade_ids: Vec<String> = std::iter::once(cascade_id.clone())
+            .chain(extra_cascade_ids)
+            .collect();
+
+        for cid in &all_cascade_ids {
+            let st = Arc::clone(&state);
+            let cid = cid.clone();
+            let timeout = body.timeout;
+            handles.push(tokio::spawn(async move {
+                let result = poll_for_response(&st, &cid, timeout).await;
+                let mitm = match st.mitm_store.take_usage(&cid).await {
+                    Some(u) => Some(u),
+                    None => st.mitm_store.take_usage("_latest").await,
+                };
+                (result, mitm)
+            }));
+        }
+
+        let mut choices = Vec::with_capacity(n as usize);
+        let mut total_prompt = 0u64;
+        let mut total_completion = 0u64;
+        let mut total_cached = 0u64;
+        let mut total_thinking = 0u64;
+
+        for (i, handle) in handles.into_iter().enumerate() {
+            if let Ok((result, mitm)) = handle.await {
+                let finish_reason = google_to_openai_finish_reason(
+                    mitm.as_ref().and_then(|u| u.stop_reason.as_deref()),
+                );
+                let (pt, ct, cached, thinking) = if let Some(ref mu) = mitm {
+                    (mu.input_tokens, mu.output_tokens, mu.cache_read_input_tokens, mu.thinking_output_tokens)
+                } else if let Some(u) = &result.usage {
+                    (u.input_tokens, u.output_tokens, 0, 0)
+                } else {
+                    (0, 0, 0, 0)
+                };
+                total_prompt += pt;
+                total_completion += ct;
+                total_cached += cached;
+                total_thinking += thinking;
+
+                let mut message = serde_json::json!({
+                    "role": "assistant",
+                    "content": result.text,
+                });
+                if let Some(ref thinking_text) = result.thinking {
+                    message["reasoning_content"] = serde_json::json!(thinking_text);
+                }
+
+                choices.push(serde_json::json!({
+                    "index": i,
+                    "message": message,
+                    "logprobs": serde_json::Value::Null,
+                    "finish_reason": finish_reason,
+                }));
+            }
+        }
+
+        Json(serde_json::json!({
+            "id": completion_id,
+            "object": "chat.completion",
+            "created": now_unix(),
+            "model": model_name,
+            "system_fingerprint": system_fingerprint(),
+            "service_tier": "default",
+            "choices": choices,
+            "usage": {
+                "prompt_tokens": total_prompt,
+                "completion_tokens": total_completion,
+                "total_tokens": total_prompt + total_completion,
+                "prompt_tokens_details": {
+                    "cached_tokens": total_cached,
+                },
+                "completion_tokens_details": {
+                    "reasoning_tokens": total_thinking,
+                },
+            },
+        }))
+        .into_response()
    }
 }

@@ -307,21 +479,26 @@ async fn chat_completions_stream(
        state.mitm_store.clear_response_async().await;

        // Initial role chunk
-        yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({
-            "id": completion_id,
-            "object": "chat.completion.chunk",
-            "created": now_unix(),
-            "model": model_name,
-            "choices": [{
-                "index": 0,
-                "delta": {"role": "assistant", "content": ""},
-                "finish_reason": serde_json::Value::Null,
-            }],
-        })).unwrap_or_default()));
+        yield Ok::<_, std::convert::Infallible>(Event::default().data(chunk_json(
+            &completion_id, &model_name,
+            serde_json::json!([chunk_choice(0, serde_json::json!({"role": "assistant", "content": ""}), None)]),
+            None,
+        )));

        let mut keepalive_counter: u64 = 0;
        let mut last_thinking_len: usize = 0;

+        // Helper: build usage JSON from MITM tokens
+        let build_usage = |pt: u64, ct: u64, crt: u64, tt: u64| -> serde_json::Value {
+            serde_json::json!({
+                "prompt_tokens": pt,
+                "completion_tokens": ct,
+                "total_tokens": pt + ct,
+                "prompt_tokens_details": { "cached_tokens": crt },
+                "completion_tokens_details": { "reasoning_tokens": tt },
+            })
+        };
+
        while start.elapsed().as_secs() < timeout {
            // ── Check for MITM-captured function calls FIRST ──
            // This runs independently of LS steps — the MITM captures tool calls
@@ -346,49 +523,28 @@ async fn chat_completions_stream(
                            },
                        }));
                    }
-                    yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                        "id": completion_id,
-                        "object": "chat.completion.chunk",
-                        "created": now_unix(),
-                        "model": model_name,
-                        "choices": [{
-                            "index": 0,
-                            "delta": {"tool_calls": tool_calls},
-                            "finish_reason": serde_json::Value::Null,
-                        }],
-                    })).unwrap_or_default()));
+                    yield Ok(Event::default().data(chunk_json(
+                        &completion_id, &model_name,
+                        serde_json::json!([chunk_choice(0, serde_json::json!({"tool_calls": tool_calls}), None)]),
+                        None,
+                    )));

-                    yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                        "id": completion_id,
-                        "object": "chat.completion.chunk",
-                        "created": now_unix(),
-                        "model": model_name,
-                        "choices": [{
-                            "index": 0,
-                            "delta": {},
-                            "finish_reason": "tool_calls",
-                        }],
-                    })).unwrap_or_default()));
+                    yield Ok(Event::default().data(chunk_json(
+                        &completion_id, &model_name,
+                        serde_json::json!([chunk_choice(0, serde_json::json!({}), Some("tool_calls"))]),
+                        None,
+                    )));
                    if include_usage {
                        let mitm = state.mitm_store.take_usage(&cascade_id).await
                            .or(state.mitm_store.take_usage("_latest").await);
                        let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
                            (u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
                        } else { (0, 0, 0, 0) };
-                        yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                            "id": completion_id,
-                            "object": "chat.completion.chunk",
-                            "created": now_unix(),
-                            "model": model_name,
-                            "choices": [],
-                            "usage": {
-                                "prompt_tokens": pt,
-                                "completion_tokens": ct,
-                                "total_tokens": pt + ct,
-                                "prompt_tokens_details": { "cached_tokens": crt },
-                                "completion_tokens_details": { "reasoning_tokens": tt },
-                            },
-                        })).unwrap_or_default()));
+                        yield Ok(Event::default().data(chunk_json(
+                            &completion_id, &model_name,
+                            serde_json::json!([]),
+                            Some(build_usage(pt, ct, crt, tt)),
+                        )));
                    }
                    yield Ok(Event::default().data("[DONE]"));
                    return;
@@ -413,17 +569,11 @@ async fn chat_completions_stream(
                        let delta = &tc[last_thinking_len..];
                        last_thinking_len = tc.len();

-                        yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                            "id": completion_id,
-                            "object": "chat.completion.chunk",
-                            "created": now_unix(),
-                            "model": model_name,
-                            "choices": [{
-                                "index": 0,
-                                "delta": {"reasoning_content": delta},
-                                "finish_reason": serde_json::Value::Null,
-                            }],
-                        })).unwrap_or_default()));
+                        yield Ok(Event::default().data(chunk_json(
+                            &completion_id, &model_name,
+                            serde_json::json!([chunk_choice(0, serde_json::json!({"reasoning_content": delta}), None)]),
+                            None,
+                        )));
                    }
                }

@@ -436,17 +586,11 @@ async fn chat_completions_stream(
                        };

                        if !delta.is_empty() {
-                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                "id": completion_id,
-                                "object": "chat.completion.chunk",
-                                "created": now_unix(),
-                                "model": model_name,
-                                "choices": [{
-                                    "index": 0,
-                                    "delta": {"content": delta},
-                                    "finish_reason": serde_json::Value::Null,
-                                }],
-                            })).unwrap_or_default()));
+                            yield Ok(Event::default().data(chunk_json(
+                                &completion_id, &model_name,
+                                serde_json::json!([chunk_choice(0, serde_json::json!({"content": delta}), None)]),
+                                None,
+                            )));
                            last_text = text;
                        }
                    }
@@ -454,37 +598,24 @@ async fn chat_completions_stream(
                    // Check if MITM response is complete
                    if state.mitm_store.is_response_complete() && !last_text.is_empty() {
                        debug!("Completions: MITM response complete (bypass), text length={}", last_text.len());
-                        yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                            "id": completion_id,
-                            "object": "chat.completion.chunk",
-                            "created": now_unix(),
-                            "model": model_name,
-                            "choices": [{
-                                "index": 0,
-                                "delta": {},
-                                "finish_reason": "stop",
-                            }],
-                        })).unwrap_or_default()));
+                        // Take usage FIRST so we can read stop_reason for finish_reason
+                        let mitm = state.mitm_store.take_usage(&cascade_id).await
+                            .or(state.mitm_store.take_usage("_latest").await);
+                        let fr = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
+                        yield Ok(Event::default().data(chunk_json(
+                            &completion_id, &model_name,
+                            serde_json::json!([chunk_choice(0, serde_json::json!({}), Some(fr))]),
+                            None,
+                        )));
                        if include_usage {
-                            let mitm = state.mitm_store.take_usage(&cascade_id).await
-                                .or(state.mitm_store.take_usage("_latest").await);
                            let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
                                (u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
                            } else { (0, 0, 0, 0) };
-                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                "id": completion_id,
-                                "object": "chat.completion.chunk",
-                                "created": now_unix(),
-                                "model": model_name,
-                                "choices": [],
-                                "usage": {
-                                    "prompt_tokens": pt,
-                                    "completion_tokens": ct,
-                                    "total_tokens": pt + ct,
-                                    "prompt_tokens_details": { "cached_tokens": crt },
-                                    "completion_tokens_details": { "reasoning_tokens": tt },
-                                },
-                            })).unwrap_or_default()));
+                            yield Ok(Event::default().data(chunk_json(
+                                &completion_id, &model_name,
+                                serde_json::json!([]),
+                                Some(build_usage(pt, ct, crt, tt)),
+                            )));
                        }
                        yield Ok(Event::default().data("[DONE]"));
                        return;
@@ -514,48 +645,27 @@ async fn chat_completions_stream(
                                    },
                                }));
                            }
-                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                "id": completion_id,
-                                "object": "chat.completion.chunk",
-                                "created": now_unix(),
-                                "model": model_name,
-                                "choices": [{
-                                    "index": 0,
-                                    "delta": {"tool_calls": tool_calls},
-                                    "finish_reason": serde_json::Value::Null,
-                                }],
-                            })).unwrap_or_default()));
-                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                "id": completion_id,
-                                "object": "chat.completion.chunk",
-                                "created": now_unix(),
-                                "model": model_name,
-                                "choices": [{
-                                    "index": 0,
-                                    "delta": {},
-                                    "finish_reason": "tool_calls",
-                                }],
-                            })).unwrap_or_default()));
+                            yield Ok(Event::default().data(chunk_json(
+                                &completion_id, &model_name,
+                                serde_json::json!([chunk_choice(0, serde_json::json!({"tool_calls": tool_calls}), None)]),
+                                None,
+                            )));
+                            yield Ok(Event::default().data(chunk_json(
+                                &completion_id, &model_name,
+                                serde_json::json!([chunk_choice(0, serde_json::json!({}), Some("tool_calls"))]),
+                                None,
+                            )));
                            if include_usage {
                                let mitm = state.mitm_store.take_usage(&cascade_id).await
                                    .or(state.mitm_store.take_usage("_latest").await);
                                let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
                                    (u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
                                } else { (0, 0, 0, 0) };
-                                yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                    "id": completion_id,
-                                    "object": "chat.completion.chunk",
-                                    "created": now_unix(),
-                                    "model": model_name,
-                                    "choices": [],
-                                    "usage": {
-                                        "prompt_tokens": pt,
-                                        "completion_tokens": ct,
-                                        "total_tokens": pt + ct,
-                                        "prompt_tokens_details": { "cached_tokens": crt },
-                                        "completion_tokens_details": { "reasoning_tokens": tt },
-                                    },
-                                })).unwrap_or_default()));
+                                yield Ok(Event::default().data(chunk_json(
+                                    &completion_id, &model_name,
+                                    serde_json::json!([]),
+                                    Some(build_usage(pt, ct, crt, tt)),
+                                )));
                            }
                            yield Ok(Event::default().data("[DONE]"));
                            return;
@@ -587,17 +697,11 @@ async fn chat_completions_stream(
                                let delta = &tc[last_thinking_len..];
                                last_thinking_len = tc.len();

-                                yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                    "id": completion_id,
-                                    "object": "chat.completion.chunk",
-                                    "created": now_unix(),
-                                    "model": model_name,
-                                    "choices": [{
-                                        "index": 0,
-                                        "delta": {"reasoning_content": delta},
-                                        "finish_reason": serde_json::Value::Null,
-                                    }],
-                                })).unwrap_or_default()));
+                                yield Ok(Event::default().data(chunk_json(
+                                    &completion_id, &model_name,
+                                    serde_json::json!([chunk_choice(0, serde_json::json!({"reasoning_content": delta}), None)]),
+                                    None,
+                                )));
                            }
                        }

@@ -611,17 +715,11 @@ async fn chat_completions_stream(
                            };

                            if !delta.is_empty() {
-                                yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                    "id": completion_id,
-                                    "object": "chat.completion.chunk",
-                                    "created": now_unix(),
-                                    "model": model_name,
-                                    "choices": [{
-                                        "index": 0,
-                                        "delta": {"content": delta},
-                                        "finish_reason": serde_json::Value::Null,
-                                    }],
-                                })).unwrap_or_default()));
+                                yield Ok(Event::default().data(chunk_json(
+                                    &completion_id, &model_name,
+                                    serde_json::json!([chunk_choice(0, serde_json::json!({"content": delta}), None)]),
+                                    None,
+                                )));
                                last_text = text.to_string();
                            }
                        }
@@ -629,37 +727,23 @@ async fn chat_completions_stream(
                        // Done check: need DONE status AND non-empty text
                        if is_response_done(steps) && !last_text.is_empty() {
                            debug!("Completions stream done, text length={}", last_text.len());
-                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                "id": completion_id,
-                                "object": "chat.completion.chunk",
-                                "created": now_unix(),
-                                "model": model_name,
-                                "choices": [{
-                                    "index": 0,
-                                    "delta": {},
-                                    "finish_reason": "stop",
-                                }],
-                            })).unwrap_or_default()));
+                            let mitm = state.mitm_store.take_usage(&cascade_id).await
+                                .or(state.mitm_store.take_usage("_latest").await);
+                            let fr = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
+                            yield Ok(Event::default().data(chunk_json(
+                                &completion_id, &model_name,
+                                serde_json::json!([chunk_choice(0, serde_json::json!({}), Some(fr))]),
+                                None,
+                            )));
                            if include_usage {
-                                let mitm = state.mitm_store.take_usage(&cascade_id).await
-                                    .or(state.mitm_store.take_usage("_latest").await);
                                let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
                                    (u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
                                } else { (0, 0, 0, 0) };
-                                yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                    "id": completion_id,
-                                    "object": "chat.completion.chunk",
-                                    "created": now_unix(),
-                                    "model": model_name,
-                                    "choices": [],
-                                    "usage": {
-                                        "prompt_tokens": pt,
-                                        "completion_tokens": ct,
-                                        "total_tokens": pt + ct,
-                                        "prompt_tokens_details": { "cached_tokens": crt },
-                                        "completion_tokens_details": { "reasoning_tokens": tt },
-                                    },
-                                })).unwrap_or_default()));
+                                yield Ok(Event::default().data(chunk_json(
+                                    &completion_id, &model_name,
+                                    serde_json::json!([]),
+                                    Some(build_usage(pt, ct, crt, tt)),
+                                )));
                            }
                            yield Ok(Event::default().data("[DONE]"));
                            return;
@@ -673,37 +757,23 @@ async fn chat_completions_stream(
                                    let run_status = td["status"].as_str().unwrap_or("");
                                    if run_status.contains("IDLE") && !last_text.is_empty() {
                                        debug!("Completions IDLE, text length={}", last_text.len());
-                                        yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                            "id": completion_id,
-                                            "object": "chat.completion.chunk",
-                                            "created": now_unix(),
-                                            "model": model_name,
-                                            "choices": [{
-                                                "index": 0,
-                                                "delta": {},
-                                                "finish_reason": "stop",
-                                            }],
-                                        })).unwrap_or_default()));
+                                        let mitm = state.mitm_store.take_usage(&cascade_id).await
+                                            .or(state.mitm_store.take_usage("_latest").await);
+                                        let fr = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
+                                        yield Ok(Event::default().data(chunk_json(
+                                            &completion_id, &model_name,
+                                            serde_json::json!([chunk_choice(0, serde_json::json!({}), Some(fr))]),
+                                            None,
+                                        )));
                                        if include_usage {
-                                            let mitm = state.mitm_store.take_usage(&cascade_id).await
-                                                .or(state.mitm_store.take_usage("_latest").await);
                                            let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
                                                (u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
                                            } else { (0, 0, 0, 0) };
-                                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                                                "id": completion_id,
-                                                "object": "chat.completion.chunk",
-                                                "created": now_unix(),
-                                                "model": model_name,
-                                                "choices": [],
-                                                "usage": {
-                                                    "prompt_tokens": pt,
-                                                    "completion_tokens": ct,
-                                                    "total_tokens": pt + ct,
-                                                    "prompt_tokens_details": { "cached_tokens": crt },
-                                                    "completion_tokens_details": { "reasoning_tokens": tt },
-                                                },
-                                            })).unwrap_or_default()));
+                                            yield Ok(Event::default().data(chunk_json(
+                                                &completion_id, &model_name,
+                                                serde_json::json!([]),
+                                                Some(build_usage(pt, ct, crt, tt)),
+                                            )));
                                        }
                                        yield Ok(Event::default().data("[DONE]"));
                                        return;
@@ -728,37 +798,23 @@ async fn chat_completions_stream(

        // Timeout
        warn!("Completions stream timeout after {}s", timeout);
-        yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-            "id": completion_id,
-            "object": "chat.completion.chunk",
-            "created": now_unix(),
-            "model": model_name,
-            "choices": [{
-                "index": 0,
-                "delta": {"content": if last_text.is_empty() { "[Timeout waiting for response]" } else { "" }},
-                "finish_reason": "stop",
-            }],
-        })).unwrap_or_default()));
+        let mitm = state.mitm_store.take_usage(&cascade_id).await
+            .or(state.mitm_store.take_usage("_latest").await);
+        let fr = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
+        yield Ok(Event::default().data(chunk_json(
+            &completion_id, &model_name,
+            serde_json::json!([chunk_choice(0, serde_json::json!({"content": if last_text.is_empty() { "[Timeout waiting for response]" } else { "" }}), Some(fr))]),
+            None,
+        )));
        if include_usage {
-            let mitm = state.mitm_store.take_usage(&cascade_id).await
-                .or(state.mitm_store.take_usage("_latest").await);
            let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
                (u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
            } else { (0, 0, 0, 0) };
-            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
-                "id": completion_id,
-                "object": "chat.completion.chunk",
-                "created": now_unix(),
-                "model": model_name,
-                "choices": [],
-                "usage": {
-                    "prompt_tokens": pt,
-                    "completion_tokens": ct,
-                    "total_tokens": pt + ct,
-                    "prompt_tokens_details": { "cached_tokens": crt },
-                    "completion_tokens_details": { "reasoning_tokens": tt },
-                },
-            })).unwrap_or_default()));
+            yield Ok(Event::default().data(chunk_json(
+                &completion_id, &model_name,
+                serde_json::json!([]),
+                Some(build_usage(pt, ct, crt, tt)),
+            )));
        }
        yield Ok(Event::default().data("[DONE]"));
    };
@@ -789,7 +845,10 @@ async fn chat_completions_sync(
        Some(u) => Some(u),
        None => state.mitm_store.take_usage("_latest").await,
    };
-    let (prompt_tokens, completion_tokens, cached_tokens, thinking_tokens) = if let Some(mitm_usage) = mitm {
+
+    let finish_reason = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
+
+    let (prompt_tokens, completion_tokens, cached_tokens, thinking_tokens) = if let Some(ref mitm_usage) = mitm {
        (mitm_usage.input_tokens, mitm_usage.output_tokens, mitm_usage.cache_read_input_tokens, mitm_usage.thinking_output_tokens)
    } else if let Some(u) = &result.usage {
        (u.input_tokens, u.output_tokens, 0, 0)
@@ -811,10 +870,13 @@ async fn chat_completions_sync(
        "object": "chat.completion",
        "created": now_unix(),
        "model": model_name,
+        "system_fingerprint": system_fingerprint(),
+        "service_tier": "default",
        "choices": [{
            "index": 0,
            "message": message,
-            "finish_reason": "stop",
+            "logprobs": serde_json::Value::Null,
+            "finish_reason": finish_reason,
        }],
        "usage": {
            "prompt_tokens": prompt_tokens,
--- a/src/api/gemini.rs
+++ b/src/api/gemini.rs
@@ -57,6 +57,14 @@ pub(crate) struct GeminiRequest {
    /// Stop sequences.
    #[serde(default, alias = "stopSequences")]
    pub stop_sequences: Option<Vec<String>>,
+    /// Thinking level for Gemini 3 models: "minimal", "low", "medium", "high".
+    /// Maps directly to thinkingConfig.thinkingLevel.
+    #[serde(default, alias = "thinkingLevel")]
+    pub thinking_level: Option<String>,
+    /// Enable Google Search grounding. See Gemini API docs.
+    /// When true, injects {"googleSearch": {}} into tools via MITM.
+    #[serde(default, alias = "googleSearch")]
+    pub google_search: bool,
 }

 fn default_timeout() -> u64 {
@@ -130,10 +138,33 @@ pub(crate) async fn handle_gemini(
    // Extract user text
    let user_text = match &body.input {
        serde_json::Value::String(s) => s.clone(),
+        serde_json::Value::Array(arr) => {
+            // Support array input: can be strings or {text: "..."} objects
+            let mut parts: Vec<String> = Vec::new();
+            for item in arr {
+                match item {
+                    serde_json::Value::String(s) => parts.push(s.clone()),
+                    serde_json::Value::Object(obj) => {
+                        if let Some(text) = obj.get("text").and_then(|v| v.as_str()) {
+                            parts.push(text.to_string());
+                        }
+                    }
+                    _ => {}
+                }
+            }
+            if parts.is_empty() {
+                return err_response(
+                    StatusCode::BAD_REQUEST,
+                    "Gemini input array contains no text parts".to_string(),
+                    "invalid_request_error",
+                );
+            }
+            parts.join("\n")
+        }
        _ => {
            return err_response(
                StatusCode::BAD_REQUEST,
-                "Gemini endpoint requires input as a string".to_string(),
+                "Gemini endpoint requires input as a string or array of text parts".to_string(),
                "invalid_request_error",
            );
        }
@@ -176,9 +207,14 @@ pub(crate) async fn handle_gemini(
            stop_sequences: body.stop_sequences.clone(),
            frequency_penalty: None,
            presence_penalty: None,
+            reasoning_effort: body.thinking_level.clone(),
+            response_mime_type: None,
+            response_schema: None,
+            google_search: body.google_search,
        };
        if gp.temperature.is_some() || gp.top_p.is_some() || gp.top_k.is_some()
            || gp.max_output_tokens.is_some() || gp.stop_sequences.is_some()
+            || gp.reasoning_effort.is_some() || gp.google_search
        {
            state.mitm_store.set_generation_params(gp).await;
        } else {
@@ -443,6 +479,7 @@ async fn gemini_stream(
                        })
                        .collect();

+                    let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await;
                    yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({
                        "candidates": [{
                            "content": {
@@ -451,6 +488,7 @@ async fn gemini_stream(
                            },
                            "finishReason": "STOP",
                        }],
+                        "usageMetadata": usage_meta,
                        "modelVersion": model_name,
                    })).unwrap_or_default()));
                    yield Ok(Event::default().data("[DONE]"));
@@ -499,7 +537,8 @@ async fn gemini_stream(
                // Check completion
                let complete = state.mitm_store.is_response_complete();
                if complete && !last_text.is_empty() {
-                    // Final chunk with finishReason
+                    // Final chunk with finishReason + usageMetadata
+                    let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await;
                    yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
                        "candidates": [{
                            "content": {
@@ -508,6 +547,7 @@ async fn gemini_stream(
                            },
                            "finishReason": "STOP",
                        }],
+                        "usageMetadata": usage_meta,
                        "modelVersion": model_name,
                    })).unwrap_or_default()));
                    yield Ok(Event::default().data("[DONE]"));
@@ -570,6 +610,7 @@ async fn gemini_stream(

                        // Done check
                        if is_response_done(steps) && !last_text.is_empty() {
+                            let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await;
                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
                                "candidates": [{
                                    "content": {
@@ -578,6 +619,7 @@ async fn gemini_stream(
                                    },
                                    "finishReason": "STOP",
                                }],
+                                "usageMetadata": usage_meta,
                                "modelVersion": model_name,
                            })).unwrap_or_default()));
                            yield Ok(Event::default().data("[DONE]"));
--- a/src/api/responses.rs
+++ b/src/api/responses.rs
@@ -142,12 +142,15 @@ fn build_response_object(data: ResponseData, params: &RequestParams) -> Response
        output: data.output,
        parallel_tool_calls: true,
        previous_response_id: params.previous_response_id.clone(),
-        reasoning: Reasoning::default(),
+        reasoning: Reasoning {
+            effort: params.reasoning_effort.clone(),
+            summary: None,
+        },
        store: params.store,
        temperature: params.temperature,
-        text: TextFormat::default(),
-        tool_choice: "auto",
-        tools: vec![],
+        text: params.text_format.clone(),
+        tool_choice: params.tool_choice.clone(),
+        tools: params.tools.clone(),
        top_p: params.top_p,
        truncation: "disabled",
        usage: data.usage,
@@ -230,6 +233,13 @@ pub(crate) async fn handle_responses(
    }

    // Store client tools in MitmStore for MITM injection
+    // Detect web_search_preview tool (OpenAI spec) → enable Google Search grounding
+    let has_web_search = body.tools.as_ref().map_or(false, |tools| {
+        tools.iter().any(|t| {
+            let t_type = t["type"].as_str().unwrap_or("");
+            t_type == "web_search_preview" || t_type == "web_search"
+        })
+    });
    if let Some(ref tools) = body.tools {
        let gemini_tools = openai_tools_to_gemini(tools);
        if !gemini_tools.is_empty() {
@@ -243,6 +253,28 @@ pub(crate) async fn handle_responses(
    }

    // Store generation parameters for MITM injection
+    // Extract text.format for structured output (json_schema)
+    let (response_mime_type, response_schema, text_format) = if let Some(ref text_val) = body.text {
+        let fmt_type = text_val["format"]["type"].as_str().unwrap_or("text");
+        if fmt_type == "json_schema" {
+            let name = text_val["format"]["name"].as_str().map(|s| s.to_string());
+            let schema = text_val["format"]["schema"].as_object().map(|o| serde_json::Value::Object(o.clone()));
+            let strict = text_val["format"]["strict"].as_bool();
+            let tf = TextFormat {
+                format: TextFormatInner {
+                    format_type: "json_schema".to_string(),
+                    name: name.clone(),
+                    schema: schema.clone(),
+                    strict,
+                },
+            };
+            (Some("application/json".to_string()), schema, tf)
+        } else {
+            (None, None, TextFormat::default())
+        }
+    } else {
+        (None, None, TextFormat::default())
+    };
    {
        use crate::mitm::store::GenerationParams;
        let gp = GenerationParams {
@@ -253,8 +285,15 @@ pub(crate) async fn handle_responses(
            stop_sequences: None,
            frequency_penalty: None,
            presence_penalty: None,
+            reasoning_effort: body.reasoning_effort.clone(),
+            response_mime_type,
+            response_schema,
+            google_search: has_web_search,
        };
-        if gp.temperature.is_some() || gp.top_p.is_some() || gp.max_output_tokens.is_some() {
+        if gp.temperature.is_some() || gp.top_p.is_some() || gp.max_output_tokens.is_some()
+            || gp.reasoning_effort.is_some() || gp.response_mime_type.is_some()
+            || gp.response_schema.is_some() || gp.google_search
+        {
            state.mitm_store.set_generation_params(gp).await;
        } else {
            state.mitm_store.clear_generation_params().await;
@@ -337,6 +376,11 @@ pub(crate) async fn handle_responses(
        previous_response_id: body.previous_response_id.clone(),
        user: body.user.clone(),
        metadata: body.metadata.clone().unwrap_or(serde_json::json!({})),
+        max_tool_calls: body.max_tool_calls,
+        reasoning_effort: body.reasoning_effort.clone(),
+        tool_choice: body.tool_choice.clone().unwrap_or(serde_json::json!("auto")),
+        tools: body.tools.clone().unwrap_or_default(),
+        text_format,
    };

    if body.stream {
@@ -365,6 +409,11 @@ struct RequestParams {
    previous_response_id: Option<String>,
    user: Option<String>,
    metadata: serde_json::Value,
+    max_tool_calls: Option<u32>,
+    reasoning_effort: Option<String>,
+    tool_choice: serde_json::Value,
+    tools: Vec<serde_json::Value>,
+    text_format: TextFormat,
 }

 /// Build Usage from the best available source, and extract thinking text from MITM:
@@ -471,10 +520,15 @@ async fn handle_responses_sync(
        while start.elapsed().as_secs() < timeout {
            // Check for function calls
            let captured = state.mitm_store.take_any_function_calls().await;
-            if let Some(ref calls) = captured {
+            if let Some(ref raw_calls) = captured {
+                let calls: Vec<_> = if let Some(max) = params.max_tool_calls {
+                    raw_calls.iter().take(max as usize).collect()
+                } else {
+                    raw_calls.iter().collect()
+                };
                if !calls.is_empty() {
                    let mut output_items: Vec<serde_json::Value> = Vec::new();
-                    for fc in calls {
+                    for fc in &calls {
                        let call_id = format!(
                            "call_{}",
                            uuid::Uuid::new_v4().to_string().replace('-', "")[..24].to_string()
@@ -567,6 +621,14 @@ async fn handle_responses_sync(
    // Check for captured function calls from MITM (clears the active flag)
    let captured_tool_calls = state.mitm_store.take_any_function_calls().await;

+    // Enforce max_tool_calls limit
+    let captured_tool_calls = captured_tool_calls.map(|mut calls| {
+        if let Some(max) = params.max_tool_calls {
+            calls.truncate(max as usize);
+        }
+        calls
+    });
+
    // If we have captured tool calls, return them as function_call output items
    if let Some(ref calls) = captured_tool_calls {
        info!(
@@ -714,7 +776,12 @@ async fn handle_responses_stream(
            while start.elapsed().as_secs() < timeout {
                // Check for function calls first
                let captured = state.mitm_store.take_any_function_calls().await;
-                if let Some(ref calls) = captured {
+                if let Some(ref raw_calls) = captured {
+                    let calls: Vec<_> = if let Some(max) = params.max_tool_calls {
+                        raw_calls.iter().take(max as usize).collect()
+                    } else {
+                        raw_calls.iter().collect()
+                    };
                    if !calls.is_empty() {
                        let msg_output_index: u32 = if thinking_emitted { 1 } else { 0 };
                        for (i, fc) in calls.iter().enumerate() {
@@ -762,7 +829,7 @@ async fn handle_responses_stream(

                        // Build output for final response
                        let mut output_items: Vec<serde_json::Value> = Vec::new();
-                        for fc in calls {
+                        for fc in &calls {
                            let call_id = format!(
                                "call_{}",
                                uuid::Uuid::new_v4().to_string().replace('-', "")[..24].to_string()
--- a/src/api/search.rs
+++ b/src/api/search.rs
@@ -0,0 +1,288 @@
+//! Pure search endpoint (/v1/search) — triggers Google Search via the LS.
+//!
+//! Sends a minimal prompt to the LS with Google Search grounding enabled,
+//! captures the grounding metadata from the response, and returns the
+//! search results without the model's generated text.
+//!
+//! The LS triggers the actual search request to Google's servers;
+//! we capture the results via MITM, never calling Google directly.
+
+use axum::{
+    extract::State,
+    http::StatusCode,
+    response::{IntoResponse, Json},
+};
+use std::sync::Arc;
+use tracing::{info, warn};
+
+use super::models::{lookup_model, MODELS};
+use super::polling::poll_for_response;
+use super::util::err_response;
+use super::AppState;
+
+/// Search request body.
+#[derive(serde::Deserialize)]
+pub(crate) struct SearchRequest {
+    /// Search query.
+    pub query: String,
+    /// Model to use for grounding. Defaults to gemini-3-flash (cheapest).
+    #[serde(default = "default_search_model")]
+    pub model: String,
+    /// Timeout in seconds.
+    #[serde(default = "default_search_timeout")]
+    pub timeout: u64,
+    /// Conversation/session ID for context reuse.
+    #[serde(default)]
+    pub conversation: Option<String>,
+    /// Max output tokens — keep low since we only want grounding metadata.
+    #[serde(default = "default_search_max_tokens")]
+    pub max_output_tokens: u64,
+}
+
+fn default_search_model() -> String {
+    "gemini-3-flash".to_string()
+}
+
+fn default_search_timeout() -> u64 {
+    30
+}
+
+fn default_search_max_tokens() -> u64 {
+    256
+}
+
+/// GET /v1/search?q=... — quick search via query param.
+pub(crate) async fn handle_search_get(
+    State(state): State<Arc<AppState>>,
+    axum::extract::Query(params): axum::extract::Query<SearchQueryParams>,
+) -> axum::response::Response {
+    info!("GET /v1/search q={}", params.q);
+
+    let body = SearchRequest {
+        query: params.q,
+        model: params.model.unwrap_or_else(default_search_model),
+        timeout: params.timeout.unwrap_or(default_search_timeout()),
+        conversation: None,
+        max_output_tokens: params.max_tokens.unwrap_or(default_search_max_tokens()),
+    };
+
+    do_search(state, body).await
+}
+
+#[derive(serde::Deserialize)]
+pub(crate) struct SearchQueryParams {
+    pub q: String,
+    #[serde(default)]
+    pub model: Option<String>,
+    #[serde(default)]
+    pub timeout: Option<u64>,
+    #[serde(default)]
+    pub max_tokens: Option<u64>,
+}
+
+/// POST /v1/search — full search request.
+pub(crate) async fn handle_search_post(
+    State(state): State<Arc<AppState>>,
+    Json(body): Json<SearchRequest>,
+) -> axum::response::Response {
+    info!("POST /v1/search q={}", body.query);
+    do_search(state, body).await
+}
+
+async fn do_search(state: Arc<AppState>, body: SearchRequest) -> axum::response::Response {
+    let model = match lookup_model(&body.model) {
+        Some(m) => m,
+        None => {
+            let names: Vec<&str> = MODELS.iter().map(|m| m.name).collect();
+            return err_response(
+                StatusCode::BAD_REQUEST,
+                format!("Unknown model: {}. Available: {names:?}", body.model),
+                "invalid_request_error",
+            );
+        }
+    };
+
+    let token = state.backend.oauth_token().await;
+    if token.is_empty() {
+        return err_response(
+            StatusCode::UNAUTHORIZED,
+            "No OAuth token.".into(),
+            "authentication_error",
+        );
+    }
+
+    // Enable Google Search grounding via GenerationParams
+    {
+        use crate::mitm::store::GenerationParams;
+        let gp = GenerationParams {
+            max_output_tokens: Some(body.max_output_tokens),
+            google_search: true,
+            ..Default::default()
+        };
+        state.mitm_store.set_generation_params(gp).await;
+    }
+
+    // Clear any stale tools — we only want googleSearch
+    state.mitm_store.clear_tools().await;
+
+    // Create a prompt that encourages the model to ground its response
+    let search_prompt = format!(
+        "Search the web for the following query and provide a brief summary of the results:\n\n{}",
+        body.query
+    );
+
+    // Session management
+    let session_id_str = body.conversation.clone();
+    let cascade_id = if let Some(ref sid) = session_id_str {
+        match state
+            .sessions
+            .get_or_create(Some(sid), || state.backend.create_cascade())
+            .await
+        {
+            Ok(sr) => sr.cascade_id,
+            Err(e) => {
+                return err_response(
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Failed to create session: {e}"),
+                    "server_error",
+                );
+            }
+        }
+    } else {
+        match state.backend.create_cascade().await {
+            Ok(id) => id,
+            Err(e) => {
+                return err_response(
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Failed to create cascade: {e}"),
+                    "server_error",
+                );
+            }
+        }
+    };
+
+    // Set active cascade for MITM correlation
+    state.mitm_store.set_active_cascade(&cascade_id).await;
+
+    // Send the search message
+    if let Err(e) = state
+        .backend
+        .send_message(&cascade_id, &search_prompt, model.model_enum)
+        .await
+    {
+        state.mitm_store.clear_active_cascade().await;
+        state.mitm_store.clear_generation_params().await;
+        return err_response(
+            StatusCode::INTERNAL_SERVER_ERROR,
+            format!("Failed to send search message: {e}"),
+            "server_error",
+        );
+    }
+
+    // Poll for response
+    let poll_result = poll_for_response(&state, &cascade_id, body.timeout).await;
+
+    // Extract grounding metadata
+    let grounding = state.mitm_store.take_grounding().await;
+
+    // The poll result text contains the model's summary (grounded response)
+    let response_text = if !poll_result.text.is_empty() {
+        poll_result.text.clone()
+    } else {
+        // Fall back to MITM captured text
+        state.mitm_store.take_response_text().await.unwrap_or_default()
+    };
+
+    // Clean up
+    state.mitm_store.clear_active_cascade().await;
+    state.mitm_store.clear_generation_params().await;
+    state.mitm_store.clear_response_async().await;
+
+    // Build the search response
+    let mut response = serde_json::json!({
+        "object": "search_result",
+        "query": body.query,
+        "model": model.name,
+        "summary": response_text,
+    });
+
+    // Include grounding metadata if available
+    if let Some(ref gm) = grounding {
+        // Extract structured search results
+        let mut search_results = Vec::new();
+
+        // groundingChunks → individual web results
+        if let Some(chunks) = gm.get("groundingChunks").and_then(|v| v.as_array()) {
+            for chunk in chunks {
+                if let Some(web) = chunk.get("web") {
+                    search_results.push(serde_json::json!({
+                        "title": web.get("title").and_then(|v| v.as_str()).unwrap_or(""),
+                        "url": web.get("uri").and_then(|v| v.as_str()).unwrap_or(""),
+                    }));
+                }
+            }
+        }
+
+        // groundingSupports → citations with source references
+        let mut citations = Vec::new();
+        if let Some(supports) = gm.get("groundingSupports").and_then(|v| v.as_array()) {
+            for support in supports {
+                let text = support.get("segment")
+                    .and_then(|s| s.get("text"))
+                    .and_then(|v| v.as_str())
+                    .unwrap_or("");
+                let indices: Vec<u64> = support.get("groundingChunkIndices")
+                    .and_then(|v| v.as_array())
+                    .map(|arr| arr.iter().filter_map(|i| i.as_u64()).collect())
+                    .unwrap_or_default();
+                let scores: Vec<f64> = support.get("confidenceScores")
+                    .and_then(|v| v.as_array())
+                    .map(|arr| arr.iter().filter_map(|s| s.as_f64()).collect())
+                    .unwrap_or_default();
+
+                citations.push(serde_json::json!({
+                    "text": text,
+                    "source_indices": indices,
+                    "confidence_scores": scores,
+                }));
+            }
+        }
+
+        // searchEntryPoint → rendered search widget HTML
+        let search_url = gm.get("searchEntryPoint")
+            .and_then(|sep| sep.get("renderedContent"))
+            .and_then(|v| v.as_str());
+
+        // webSearchQueries → the actual queries Google used
+        let queries = gm.get("webSearchQueries")
+            .and_then(|v| v.as_array())
+            .map(|arr| arr.iter().filter_map(|q| q.as_str().map(|s| s.to_string())).collect::<Vec<_>>());
+
+        response["results"] = serde_json::json!(search_results);
+        response["citations"] = serde_json::json!(citations);
+        if let Some(qs) = queries {
+            response["search_queries"] = serde_json::json!(qs);
+        }
+        if let Some(url) = search_url {
+            response["search_widget_html"] = serde_json::json!(url);
+        }
+
+        // Include raw grounding metadata for advanced consumers
+        response["grounding_metadata"] = gm.clone();
+    } else {
+        response["results"] = serde_json::json!([]);
+        response["citations"] = serde_json::json!([]);
+        warn!("Search completed but no grounding metadata captured — model may not have grounded");
+    }
+
+    // Include usage if available
+    if let Some(ref u) = poll_result.usage {
+        response["usage"] = serde_json::json!({
+            "input_tokens": u.input_tokens,
+            "output_tokens": u.output_tokens,
+            "total_tokens": u.input_tokens + u.output_tokens,
+        });
+    }
+
+    Json(response).into_response()
+}
--- a/src/api/types.rs
+++ b/src/api/types.rs
@@ -38,6 +38,17 @@ pub(crate) struct ResponsesRequest {
    /// Tool choice: "auto", "required", "none", or {"type":"function","function":{"name":"X"}}.
    #[serde(default)]
    pub tool_choice: Option<serde_json::Value>,
+    /// Reasoning effort — forwarded as thinkingConfig.thinkingLevel to Google.
+    /// Values: "low", "medium", "high".
+    #[serde(default)]
+    pub reasoning_effort: Option<String>,
+    /// Maximum number of tool calls allowed per response.
+    #[serde(default)]
+    pub max_tool_calls: Option<u32>,
+    /// Text output format — {format: {type: "json_schema", name: "...", schema: {...}}}.
+    /// When json_schema, injects responseMimeType + responseSchema via MITM.
+    #[serde(default)]
+    pub text: Option<serde_json::Value>,
 }

 /// Stream options for Chat Completions (controls usage emission in final chunk).
@@ -85,6 +96,80 @@ pub(crate) struct CompletionRequest {
    /// Presence penalty — forwarded to Google via MITM.
    #[serde(default)]
    pub presence_penalty: Option<f64>,
+    /// Reasoning effort — forwarded as thinkingConfig.thinkingLevel to Google.
+    /// Values: "low", "medium", "high".
+    #[serde(default)]
+    pub reasoning_effort: Option<String>,
+    /// Stop sequences — forwarded as generationConfig.stopSequences to Google.
+    /// Up to 4 sequences where the API will stop generating further tokens.
+    #[serde(default)]
+    pub stop: Option<StopSequence>,
+    /// Response format — {"type": "json_object"} or {"type": "json_schema", "json_schema": {...}}.
+    /// Injected as responseMimeType (+ responseSchema) in generationConfig via MITM.
+    #[serde(default)]
+    pub response_format: Option<ResponseFormat>,
+    /// Session/conversation ID for multi-turn reuse (custom extension).
+    #[serde(default)]
+    pub conversation: Option<serde_json::Value>,
+    /// Metadata — accepted and ignored (no upstream equivalent).
+    #[serde(default)]
+    pub metadata: Option<serde_json::Value>,
+    /// Number of completions to generate. Each uses a separate cascade (costs N× quota).
+    /// Defaults to 1. Only supported in sync mode; streaming always uses n=1.
+    #[serde(default = "default_n")]
+    pub n: u32,
+    /// Enable Google Search grounding. When true, the model can search the web
+    /// and responses include grounding metadata with search results/citations.
+    #[serde(default)]
+    pub web_search: bool,
+}
+
+fn default_n() -> u32 { 1 }
+
+/// Stop sequence can be a single string or array of strings (OpenAI accepts both).
+#[derive(Deserialize, Clone)]
+#[serde(untagged)]
+pub(crate) enum StopSequence {
+    Single(String),
+    Multiple(Vec<String>),
+}
+
+impl StopSequence {
+    pub fn into_vec(self) -> Vec<String> {
+        match self {
+            StopSequence::Single(s) => vec![s],
+            StopSequence::Multiple(v) => v,
+        }
+    }
+}
+
+/// Response format for structured output.
+/// Supports:
+/// - `{"type": "json_object"}` — JSON mode (responseMimeType only)
+/// - `{"type": "json_schema", "json_schema": {"name": "...", "schema": {...}}}` — structured output (responseMimeType + responseSchema)
+/// - `{"type": "text"}` — plain text (default, no injection)
+#[derive(Deserialize, Clone)]
+pub(crate) struct ResponseFormat {
+    #[serde(rename = "type")]
+    pub format_type: String,
+    /// JSON schema definition for structured output.
+    /// Only used when format_type is "json_schema".
+    #[serde(default)]
+    pub json_schema: Option<JsonSchemaFormat>,
+}
+
+/// JSON schema structured output format.
+#[derive(Deserialize, Clone)]
+pub(crate) struct JsonSchemaFormat {
+    /// Schema name (for client identification).
+    #[serde(default)]
+    pub name: Option<String>,
+    /// The actual JSON schema object — forwarded as Gemini's responseSchema.
+    #[serde(default)]
+    pub schema: Option<serde_json::Value>,
+    /// Whether to enable strict schema adherence.
+    #[serde(default)]
+    pub strict: Option<bool>,
 }

 #[derive(Deserialize)]
@@ -132,7 +217,7 @@ pub(crate) struct ResponsesResponse {
    pub store: bool,
    pub temperature: f64,
    pub text: TextFormat,
-    pub tool_choice: &'static str,
+    pub tool_choice: serde_json::Value,
    pub tools: Vec<serde_json::Value>,
    pub top_p: f64,
    pub truncation: &'static str,
@@ -180,7 +265,16 @@ pub(crate) struct TextFormat {
 #[derive(Serialize, Clone)]
 pub(crate) struct TextFormatInner {
    #[serde(rename = "type")]
-    pub format_type: &'static str,
+    pub format_type: String,
+    /// JSON schema — present when format_type is "json_schema".
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub name: Option<String>,
+    /// The actual schema object.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub schema: Option<serde_json::Value>,
+    /// Whether strict mode was requested.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub strict: Option<bool>,
 }

 impl Usage {
@@ -220,7 +314,10 @@ impl Default for TextFormat {
    fn default() -> Self {
        Self {
            format: TextFormatInner {
-                format_type: "text",
+                format_type: "text".to_string(),
+                name: None,
+                schema: None,
+                strict: None,
            },
        }
    }
--- a/src/mitm/proxy.rs
+++ b/src/mitm/proxy.rs
@@ -755,7 +755,7 @@ async fn handle_http_over_tls(
                            info!("MITM: stored {} function call(s) from initial body", streaming_acc.function_calls.len());
                        }

-                        // Capture response + thinking text directly into MitmStore
+                        // Capture response + thinking text + grounding directly into MitmStore
                        if bypass_ls {
                            if !streaming_acc.response_text.is_empty() {
                                store.set_response_text(&streaming_acc.response_text).await;
@@ -763,6 +763,9 @@ async fn handle_http_over_tls(
                            if !streaming_acc.thinking_text.is_empty() {
                                store.set_thinking_text(&streaming_acc.thinking_text).await;
                            }
+                            if let Some(ref gm) = streaming_acc.grounding_metadata {
+                                store.set_grounding(gm.clone()).await;
+                            }
                            if streaming_acc.is_complete {
                                store.mark_response_complete();
                            }
@@ -827,7 +830,7 @@ async fn handle_http_over_tls(
                    info!("MITM: stored {} function call(s) from body chunk", streaming_acc.function_calls.len());
                }

-                // Capture response + thinking text directly into MitmStore
+                // Capture response + thinking text + grounding directly into MitmStore
                if bypass_ls {
                    if !streaming_acc.response_text.is_empty() {
                        store.set_response_text(&streaming_acc.response_text).await;
@@ -835,6 +838,9 @@ async fn handle_http_over_tls(
                    if !streaming_acc.thinking_text.is_empty() {
                        store.set_thinking_text(&streaming_acc.thinking_text).await;
                    }
+                    if let Some(ref gm) = streaming_acc.grounding_metadata {
+                        store.set_grounding(gm.clone()).await;
+                    }
                    if streaming_acc.is_complete {
                        store.mark_response_complete();
                    }
@@ -883,6 +889,10 @@ async fn handle_http_over_tls(

        // Capture usage data
        if is_streaming_response {
+            // Store grounding metadata before consuming the accumulator
+            if let Some(ref gm) = streaming_acc.grounding_metadata {
+                store.set_grounding(gm.clone()).await;
+            }
            if streaming_acc.is_complete || streaming_acc.output_tokens > 0 {
                // Function calls are stored immediately when detected (above),
                // so no need to store them again here.