From ab4846513551dac0d30d8caf3f5dcf09af5a67b8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 06:01:24 -0700 Subject: [PATCH] Studio: add Gemini provider with web_search, code_execution, prompt caching, and Nano Banana image generation (#5720) * Studio: add Gemini provider with web_search, code_execution, prompt caching, and Nano Banana image generation Wires Google's native Gemini API into Studio's external-provider stack so users can pick gemini-2.5-pro / gemini-2.5-flash / gemini-2.5-flash-image (Nano Banana) alongside the existing OpenAI / Anthropic / OpenRouter providers. Gemini does not speak OpenAI Chat Completions on its primary endpoint; the new `_stream_gemini` async generator translates between the two shapes the same way `_stream_anthropic` handles the Messages API. Backend: - New `_stream_gemini` translator in external_provider.py. Converts OpenAI messages -> Gemini `contents` + `systemInstruction`; maps generationConfig (temperature / topP / topK / maxOutputTokens); forwards `tools: [{googleSearch: {}}]` for web_search and `{codeExecution: {}}` for code_execution; passes `cachedContent` through for prompt caching; sets `responseModalities=[TEXT, IMAGE]` for Nano Banana image generation. - Translates streamed `GenerateContentResponse` SSE frames back into OpenAI chat.completion.chunk frames (text deltas, function_call -> tool_calls deltas, inlineData -> image_b64 tool_end envelope, usage chunk before [DONE]). - Registry entry switched to native base URL `https://generativelanguage.googleapis.com/v1beta` with `openai_compatible: False` and the `x-goog-api-key` auth header. Model lineup curated to current 2.5 / 2.0 family + Nano Banana. Frontend: - Provider-capability matrix: Gemini supports temperature, top_p, top_k, presence_penalty (matches generationConfig); min_p / repetition_penalty hidden because the API does not accept them. - `providerSupportsBuiltinWebSearch` / `providerSupportsBuiltinCodeExecution` / `providerSupportsBuiltinImageGeneration` extended for Gemini. - Prompt caching toggle now also lit on Gemini. Tests: - 21 new tests in `test_gemini_provider.py` using httpx.MockTransport. Cover request body shape conversion, URL/header wiring, web_search forwarded as googleSearch, function-call translation both directions, prompt caching passthrough, image generation emitting image_b64, grounded-search citations -> tool_end, finish_reason mapping, and vision data URL -> inlineData translation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: forward presence_penalty to Gemini and recover function name from tool_call_id Two follow-up fixes for the Gemini provider: * Thread presence_penalty into _stream_gemini and set generationConfig.presencePenalty when non-zero. The OpenAI-side capability matrix already exposes the slider for Gemini, so the value was being collected and silently dropped on the way out. * When an OpenAI role=tool message omits 'name' and only carries 'tool_call_id', recover the function name from the matching functionCall on the prior assistant turn. Gemini 400s on an empty functionResponse name. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: surface Gemini code execution parts as code_execution tool events The Gemini stream parser only handled text/functionCall/inlineData parts, so when the user toggled the Code pill on a Gemini model the sandbox output (executableCode + codeExecutionResult parts) was dropped on the floor while adjacent text reached the UI. Reviewers flagged this as the headline feature being silently broken. Translate both parts into the existing code_execution tool envelope that CodeExecutionToolUI already consumes for OpenAI / Anthropic: * executableCode -> tool_start with kind=code_execution and the source code under arguments.code. We mint a tool_call_id and stash it so the matching result block can pair to it. * codeExecutionResult -> tool_end on that id with the stdout under result. Non-OK outcomes (OUTCOME_FAILED / OUTCOME_DEADLINE_EXCEEDED) are prefixed onto the text so the failure is visible. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: native Gemini model catalog, function-call ids, and honest cache claim Three follow-ups to the Gemini provider PR after the codex pass: * list_models() now translates Gemini's native /v1beta/models payload ({models[{name, baseModelId, displayName, supportedGenerationMethods}]}) into the OpenAI-compatible shape Studio expects. Without this the picker stayed empty for Gemini and fell back to hardcoded defaults. Embedding-only models are filtered out. * Forward the OpenAI tool_call id into Gemini's functionCall.id and mirror it onto functionResponse.id. Two parallel calls to the same function name can now be paired unambiguously on the follow-up turn. * Drop Gemini from the prompt-caching capability set. The wire flow requires a separate cachedContents POST first and the boolean Studio emits today is a no-op; the toggle should not advertise a feature it cannot apply. Leaves a pointer to the docs for the eventual two-step orchestration. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: distinct tool_calls index per emitted Gemini function call Codex flagged that the Gemini stream parser hardcoded tool_calls[0].index to 0 on every emitted functionCall. OpenAI reassemblers key tool_calls by index when joining deltas, so two parallel function calls in one assistant turn collapsed onto a single slot and the second call's arguments overwrote the first. Track the running count via len(emitted_function_call_ids) - 1 and emit it as the per-call index. The dedupe guard above (skip when fc_id already in the set) means the index is monotonic and stable for the lifetime of the stream. Regression test asserts [0, 1] across two parallel calls in one candidate parts list. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: surface Gemini 3.5/3.1/3 + Nano Banana 2/Pro and plumb thinking budget `gemini-2.0-flash` / `gemini-2.0-flash-exp` were retired by Google in 2026 (`/v1beta/models/gemini-2.0-flash:streamGenerateContent` returns HTTP 404 "no longer available to new users"), and the picker had nothing past the 2.x family. Verified against the live ListModels catalog: drop the retired ids from `default_models` + allowlist and surface the chat-capable 3.5 / 3.1 / 3 families plus the Nano Banana image trio. Also plumb `enable_thinking` / `reasoning_effort` into Gemini's `generationConfig.thinkingConfig`. Without this, Gemini 3.5 Flash, gemini-pro-latest, and the 3.x previews silently spend the caller's `max_tokens` budget on hidden "thoughts" before emitting any visible answer -- the chat shows a truncated stub like "The capital of" and streams stop. Mapping: - enable_thinking=False / reasoning_effort=none -> thinkingBudget=0 (Flash tier; Pro tier coerces to a small positive budget because the API 400s on 0 with "This model only works in thinking mode") - minimal/low/medium/high -> 512/2048/8192/24576 budget tokens - max/xhigh -> -1 (dynamic) - default (neither knob set) -> thinkingConfig omitted, model decides Frontend `getExternalReasoningCapabilities` now surfaces a `reasoning_effort` picker for every Gemini chat id (Pro tier hides the "none" option; image-tier ids stay knob-less). Adds 6 unit tests covering Flash/Pro effort mapping, the off-toggle coercion on Pro, default omission, and the nano-banana-pro-preview alias routing through the image modalities path. 28 -> 34 tests in `test_gemini_provider.py`, all green; full backend suite still passes (1459/1460; the unrelated test_help_output flake is pre-existing and not in any file this PR touches). Live verification against generativelanguage.googleapis.com on 2026-05-24 with `_stream_gemini` directly: text gemini-3.5-flash single PASS multi PASS text gemini-3.1-pro-preview single PASS multi PASS text gemini-3.1-flash-lite single PASS multi PASS text gemini-3-pro-preview single PASS multi PASS text gemini-3-flash-preview single PASS multi PASS text gemini-2.5-pro single PASS multi PASS text gemini-2.5-flash single PASS multi PASS text gemini-2.5-flash-lite single PASS multi PASS text gemini-flash-latest single PASS multi PASS text gemini-flash-lite-latest single PASS multi PASS text gemini-pro-latest single PASS multi PASS image gemini-2.5-flash-image PASS (1082 KB png returned) image gemini-3.1-flash-image-preview PASS (Nano Banana 2) image gemini-3-pro-image-preview PASS (Nano Banana Pro) tool web_search PASS tool code_execution PASS -> 16/16 e2e through the actual ExternalProviderClient code path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten Gemini provider after review (PR #5720) Fixes a batch of bugs surfaced by a second-pass review on top of the 3.5/3.1/3 + Nano Banana 2/Pro additions in c6724dbd. Backend (external_provider.py): - Constructor normalises legacy /v1beta/openai base URLs to /v1beta so Gemini providers saved before the native switch keep working without a manual re-config. - Skip thinkingConfig, googleSearch, and codeExecution on image-tier models (-image / nano-banana). The image responseModalities path is mutually exclusive with text-tool wiring and stale UI state would otherwise 400 the turn. - _PRO_THINKING_PREFIXES now includes gemini-3.5-pro and uses anchored prefix matching (exact id or "-...") so the image-tier gemini-3-pro-image-preview cannot accidentally match the pro guard. - Gemini 3 functionCall thoughtSignature is round-tripped through the tool_calls envelope via extra_content.google.thought_signature on emit, and replayed as a sibling of functionCall on the next request. - finishReason swaps STOP -> tool_calls when any functionCall was emitted on the same turn so OAI clients trigger tool execution (matches the OpenAI Chat Completions contract). - usageMetadata.thoughtsTokenCount is rolled into output_tokens and surfaced on output_tokens_details.reasoning_tokens so total_tokens reflects the full billable spend instead of dropping the hidden reasoning slice. Registry (providers.py): - Drop gemini-3-pro-preview from default_models. Google shut it down on 2026-03-09 and auto-redirects to gemini-3.1-pro-preview; we surface the canonical id only. - Add model_id_deny_exact = ("gemini-3-pro-preview",) so the live ListModels fetch does not re-surface the redirect alias. Route schema (models/inference.py): - enable_prompt_caching widened to Optional[Union[bool, str]] so the /v1/chat/completions caller can pass a Gemini cachedContent resource name (e.g. cachedContents/abc123). Without this widening _stream_gemini s string cachedContent passthrough was unreachable from the public route (bool_parsing 422). stream_chat_completion signature mirrors. Frontend (provider-capabilities.ts, chat-page.tsx, chat-adapter.ts): - providerSupportsBuiltinImageGeneration now also recognises nano-banana ids (nano-banana-pro-preview was hidden from the image pill before). - providerSupportsBuiltinWebSearch takes the model id so Gemini image models hide the Search pill (mirrors the backend skip). - providerSupportsBuiltinCodeExecution uses the same isGeminiImageModel guard for nano-banana ids. - GEMINI_THINKING_PRO_PREFIXES gains gemini-3.5-pro; gemini-3-pro tightened to gemini-3-pro-preview to avoid the image-id overlap. - Updated 3 callers of providerSupportsBuiltinWebSearch to thread the selected model id through. Tests (test_gemini_provider.py): 34 -> 42, all green - test_image_models_skip_thinking_config - test_image_models_drop_text_only_tools - test_gemini_35_pro_recognized_as_pro_thinking - test_legacy_openai_base_url_normalized - test_finish_reason_swaps_to_tool_calls_when_function_call_emitted - test_thought_signature_round_trips_into_gemini_function_call - test_thought_signature_emitted_in_tool_call_delta - test_usage_chunk_includes_thoughts_tokens Verification: - Backend pytest 1518/1519 passing (one unrelated Qwen3.5 flash-attn test fails on main as well; nothing in this PR touches that path). - Frontend npx tsc -b clean. - Live e2e 16/16 against generativelanguage.googleapis.com through the patched _stream_gemini code path (all 11 chat models single + multi turn, all 3 image models returned image bytes, web_search and code_execution tools both emit the expected envelope). - Live /api/providers/models against the patched backend surfaces 16 ids (gemini-3-pro-preview correctly filtered via deny_exact). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: address second-pass review findings on Gemini (PR #5720) Round-2 reviewer.py flagged a phantom web_search card on image turns (12/12 reviewers), route-layer stripping of tool_calls / tool_call_id / name, an over-narrow image-mode tool guard, and silent safety blocks. This patch fixes all four. Backend (external_provider.py): - web_search_active is now derived from the outbound tools_array (whether googleSearch was actually forwarded), not the raw enabled_tools intent. Image-mode turns dropped the tool above so the inbound stream no longer emits a phantom "search complete" tool_start / tool_end on those turns. - text_tools_allowed now uses is_image_model (covers both `-image` / `nano-banana` picker models AND text models that requested `image_generation` via enabled_tools). Verified against the live Gemini API which rejects both googleSearch and codeExecution alongside responseModalities=["TEXT","IMAGE"] with explicit 400s ("Search as tool is not enabled for this model", "Code execution is not enabled for this model"). - promptFeedback.blockReason is surfaced as a 400 content-filter error chunk instead of returning an empty successful assistant response. The streaming loop closes the response before exiting. Route (routes/inference.py): - _build_external_messages now propagates tool_calls (assistant), tool_call_id, and name (tool result) through every code path (string content, multimodal content, non-vision fallback). Without this Gemini 3 function-call round trips lost their thoughtSignature + tool_call_id at the route boundary, and functionResponse.name arrived empty on the second turn. - Assistant messages with content=None and tool_calls populated are preserved as a synthetic empty-string content turn so the Gemini translator can rebuild the functionCall part. Tests (test_gemini_provider.py): 42 -> 45, all green - test_image_models_suppress_phantom_web_search_card - test_image_generation_tool_drops_text_tools - test_prompt_feedback_block_reason_surfaces_as_error Verification: - Backend pytest 1736 / 1736 (the two pre-existing unrelated fails on main, test_help_output and Qwen3.5 flash-attn pin, are skipped). - Frontend npx tsc -b clean. - Live e2e 16/16 against generativelanguage.googleapis.com: 11 chat models single + multi turn, 3 image models returning image bytes, web_search and code_execution both PASS. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix third-pass Gemini findings (PR #5720) Round 3 review follow-ups: Backend (studio/backend/core/inference/external_provider.py): - Close response AND aiter_lines iterator in a finally so normal, prompt-block, and cancellation exits all clean up (eliminates the RuntimeWarning about aclose never being awaited). - Pair the synthetic web_search tool_start with a tool_end on the promptFeedback.blockReason path so the UI does not leave a stuck "searching..." spinner after the error toast. - Preserve native id and thoughtSignature on executableCode and codeExecutionResult tool events under google.native_part, and pair the tool_end on the code-exec id so multi-turn code-execution replays do not lose Gemini-required history. - Carry part-level thoughtSignature on text deltas via delta.extra_content.google.thought_signature and on inline image tool_end via google.thought_signature so Gemini 3 image editing and tool turns round-trip the signature on the next request. - Guess remote image_url MIME from the URL path so PNG / WebP / GIF inputs are not silently relabeled as JPEG. - Roll usageMetadata.toolUsePromptTokenCount into translated input tokens and surface thoughtsTokenCount as completion_tokens_details.reasoning_tokens in _build_usage_chunk. - Only normalize the Google-hosted /v1beta/openai legacy base URL; custom proxies whose paths happen to end in /openai are left untouched. - Forward ChatCompletionRequest.tools and tool_choice through stream_chat_completion into _stream_gemini, translating to tools[].functionDeclarations and toolConfig.functionCallingConfig. Frontend: - chat-adapter: when Gemini image-generation is enabled for the turn, also disable Search and Code so the request, builder, and active pills agree with what the backend actually sends (the backend already strips text tools when image_generation is in enabled_tools). - chat-adapter: consume OpenAI-shape delta.tool_calls chunks so Gemini function-call deltas without text surface as tool-call parts. - shared-composer: disable Search and Code pills while Gemini image mode is active so the UI matches the request. Tests (studio/backend/tests/test_gemini_provider.py): adds coverage for proxy base-url gating, remote image MIME inference, toolUsePromptTokenCount, reasoning_tokens propagation, prompt-block web_search tool_end pairing, native code-exec id/thoughtSignature metadata, inline image thoughtSignature, text-chunk extra_content, OpenAI tools/tool_choice translation, and image-model tool drop. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: Gemini 3 thinkingLevel + image-model Search grounding (PR #5720) Gemini 3.x migrated to a string `thinkingConfig.thinkingLevel` (MINIMAL/LOW/MEDIUM/HIGH) and rejects `thinkingBudget`+`thinkingLevel` in the same request. Gemini 3 also cannot turn thinking fully off, so the lowest position is "minimal" (Flash) or "low" (Pro rejects "minimal"). - external_provider._stream_gemini: split thinking translation by family. Gemini 3.x (3 / 3.1 / 3.5 + gemini-pro-latest / gemini-flash-latest / gemini-flash-lite-latest) emits thinkingConfig.thinkingLevel; effort none/off coerces to "low" on Pro and "minimal" on Flash. Gemini 2.5 stays on thinkingBudget. - external_provider._stream_gemini: allow `tools: [{googleSearch: {}}]` on the Gemini 3 image family (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, nano-banana-pro). Google's docs document Search grounding on these. codeExecution stays blocked on image mode (still mutually exclusive with responseModalities). - provider-capabilities.ts: mirror the Gemini 3 effort ladders in resolveGeminiReasoningCapabilities (Pro: low/medium/high; Flash: minimal/low/medium/high; 2.5 Flash keeps the off-position). - provider-capabilities.ts: providerSupportsBuiltinWebSearch now returns true on the documented Gemini 3 image models so the pill is reachable; older image ids (gemini-2.5-flash-image) still hide. Tests: splits the existing thinkingBudget cases by family (Gemini 3 checks thinkingLevel; Gemini 2.5 keeps thinkingBudget), adds positive googleSearch coverage for Gemini 3 image models and negative googleSearch coverage for legacy image models. References: - https://ai.google.dev/gemini-api/docs/thinking - https://ai.google.dev/gemini-api/docs/gemini-3 - https://ai.google.dev/gemini-api/docs/models/gemini-3-pro-image-preview * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: attach Gemini code_execution inline images to the code card (PR #5720) When a text Gemini turn wires codeExecution and the sandbox produces a matplotlib plot, the inline image part ships right after the codeExecutionResult. Previously this surfaced as a separate empty image_generation card. Track the most recent code_execution tool_call_id + result text and, when an inline image follows with code_execution active, emit a second tool_end on the same id that appends the image as a data: URI under the `__IMAGES__:` marker the chat-adapter already understands. Image-picker turns (`-image` / `nano-banana`) keep the standalone image_generation envelope so Nano Banana outputs render the same way. Tests: covers the merged code-execution card emission with no standalone image_generation event when code_execution is the active tool. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix fourth-pass Gemini findings (PR #5720) Round 4 review follow-ups: Backend: - `_is_openai_compatible` + `_auth_headers` detect Gemini connections pointed at a custom OpenAI-compatible proxy (non-Google host whose path ends in `/openai`) and route them through the OpenAI-compat surface with `Authorization: Bearer ...` instead of the native `_stream_gemini` translator + `x-goog-api-key`. Google-hosted Gemini keeps the native dispatch path it migrated to in this PR. - `_stream_gemini` thinkingLevel handling for Gemini 3 Pro now coerces both "minimal" and "medium" effort to "low" / "high" respectively (Pro tier only accepts low/high per https://ai.google.dev/gemini-api/docs/thinking). - `providers.py` `default_models` restores the advertised `gemini-3.5-pro` and the rolling `gemini-pro-latest` / `gemini-flash-latest` / `gemini-flash-lite-latest` aliases that the allowlist already admits. Frontend: - chat-adapter: lean on `providerSupportsBuiltinWebSearch` (which already encodes the Gemini 3 image-model Search allowance) instead of blanket-disabling Search whenever Gemini image mode is active. Code execution stays blocked because Gemini image mode rejects it. - shared-composer: mirror the same gate -- only the Code pill is unconditionally disabled in Gemini image mode; the Search pill is driven by `supportsBuiltinWebSearch`. - provider-capabilities: Gemini 3 Pro reasoning levels now expose only "low" and "high" (no Medium pill) to match the API. Tests: covers the Gemini 3 Pro medium / minimal coercion, the custom proxy OAI-compat dispatch + Authorization Bearer auth, and the native-vs-proxy detection. Also closes the mocked httpx.AsyncClient inside the test event loop so the Python 3.13 `aclose was never awaited` warning no longer fires. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix fifth-pass Gemini findings (PR #5720) Round 5 review follow-ups: Backend: - `_is_openai_compatible` + `_auth_headers` now treat ANY non-Google Gemini base URL as OpenAI-compat (LiteLLM / custom OAI gateways / OpenAI-compat vLLM routers), not just paths ending in `/openai`. Pre-existing saved Gemini proxies on `/v1` keep working. - Gemini 3 thinkingLevel coercion narrowed to the documented inconsistencies: only "minimal" is coerced to "low" on Pro tier. "medium" passes through (Gemini 3.1 Pro accepts it per https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro). - `_stream_gemini` only flips `responseModalities=[TEXT,IMAGE]` when the selected model is image-capable. A stale `enabled_tools=["image_generation"]` on a text model is silently dropped instead of producing an invalid Gemini request. - `_stream_gemini` validates the model id against `[A-Za-z0-9._-]+` before URL interpolation so a model like `../cachedContents/x` cannot redirect the request to an unintended endpoint with the configured API key attached. - Empty-text Gemini parts that still carry `thoughtSignature` emit a content-free delta with `extra_content.google.thought_signature` so Gemini 3 turns that end with a signature-only fragment do not lose the replay state. - ConnectError / ReadTimeout / generic HTTPError paths in `_stream_gemini` now close the synthetic web_search tool_start with a matching tool_end before the error chunk so the UI does not leave a stuck "searching..." card on transport failure. - `providers.py` default_models drop the non-existent `gemini-3.5-pro` (Google launched only `gemini-3.5-flash` at I/O 2026; Pro tier remains `gemini-3.1-pro-preview`). - `routes/inference.py` only forwards `payload.top_k` when the caller explicitly set it on the request (Pydantic `model_fields_set`). Omitted top_k stays omitted, restoring the pre-PR behavior where Gemini uses its server default. - `ChatCompletionRequest.enable_prompt_caching` adds a `mode="before"` validator that coerces the canonical string literals "true"/"false" back to bool so historical opt-out callers keep working after the field widened to `Union[bool, str]` for Gemini cache resource names. Frontend: - `providerSupportsBuiltinWebSearch` / Code / Image now accept the saved connection `baseUrl` and return false for custom OAI-compat Gemini proxies. Backend skips `_stream_gemini` for those bases, so native tool envelopes never reach them; hiding the pills keeps the request, builder, and UI consistent. - `provider-capabilities.ts` Gemini 3 Pro effort ladder restores `["low", "medium", "high"]` to match Google's documented levels. - Call sites in `chat-page.tsx` and `chat-adapter.ts` pass through `provider.baseUrl` so the proxy gate fires. Tests: covers Gemini 3 Pro medium pass-through, custom proxy dispatch on `/v1` and `/openai` bases, path-traversal model id rejection, top_k omission when not explicit, text-model image_generation drop, empty-text + thoughtSignature surfacing, and enable_prompt_caching string coercion. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix sixth-pass Gemini findings (PR #5720) Round 6 review follow-ups: Frontend: - chat-adapter `delta.tool_calls` accumulates fragments by `id` / `index` instead of pushing a new tool-call card per chunk. The standard OpenAI Chat Completions stream contract sends `id`/`name` on the first chunk and partial `function.arguments` on subsequent chunks; our previous handler parsed each fragment as a standalone tool call. Local llama.cpp and OAI-compat providers that stream fragments now reassemble into a single function-call part. - chat-adapter also preserves `extra_content` on streamed tool-call deltas so Gemini 3 `thoughtSignature` survives to the next turn. - provider-capabilities Gemini 3 Pro restores "medium" in the reasoning-effort ladder (Google's official Gemini API thinking doc lists low/medium/high for Gemini 3.1 Pro; my earlier round 4 coercion was wrong). - provider-capabilities orders `gemini-2.5-flash-lite` ahead of the broader `gemini-2.5-flash` prefix so Flash-Lite falls into the "no native thinking knob" branch as documented. * Studio: round-trip Gemini tool_calls and tool results (PR #5720) Recurring round 3-6 P1: the chat-adapter renders Gemini function-call parts and code-execution events but `toOpenAIMessage` only serialized text + image content, so the next turn lost the assistant `tool_calls[]` (including Gemini 3's required `extra_content.google.thought_signature`) and the matching `role="tool"` result. Gemini 3 multi-turn function calling and code execution failed validation on the second turn. Frontend: - types/api.ts widens OpenAIChatMessage to permit `role="tool"`, `tool_calls`, `tool_call_id`, `name`, and `content: null`. Adds OpenAIToolCallPart with `extra_content` for the Gemini round-trip. - chat-adapter: new `toOpenAIMessages` expands an assistant turn with tool-call parts into [assistant w/ tool_calls + extra_content, role=tool result, ...]. tool result content is JSON-serialized so the backend translator can rebuild Gemini's `functionResponse` shape. - chat-adapter outbound history now uses `flatMap(toOpenAIMessages)` so each assistant tool-call round-trips through the standard OAI shape the backend's `_stream_gemini` already understands. * Studio: replay Gemini code_execution and image native parts on history (PR #5720) Multi-turn Gemini history previously lost the native executableCode, codeExecutionResult, and inlineData parts because the outbound translator regenerated a generic functionCall for every assistant tool_call. Stow the native dict on tool_end (frontend) and replay it verbatim with thoughtSignature (backend) so follow-up turns preserve the prior execution and image generation state. Skip role="tool" fan-out for server-side builtin tools so Gemini does not 400 on a functionResponse with no matching user-declared function. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: complete Gemini built-in tool replay round-trip (PR #5720) Round 7 follow-up to the multi-turn native-part work. Three asymmetric storage/consume gaps remained between the backend translator and the chat adapter, so realistic Gemini follow-up turns degraded to generic functionCalls instead of native history. - Frontend collectAssistantToolCalls now drops web_search outright, drops code_execution / image_generation when the native part is missing, and promotes args.google to extra_content.google so the backend native_part replay branch actually fires. - Backend image_generation tool_end now emits google.native_part with the inlineData (mimeType + base64) and thoughtSignature so the follow-up image-edit turn can replay the prior image as a native Gemini model part. - Backend code-execution plot tool_end now stows google.native_part with the inlineData so the merged code-exec card can round-trip executableCode + codeExecutionResult + inlineData on the same id. - Added regression tests for image-gen native-part replay and the code-exec plot native_part stow. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 8 Gemini follow-ups (PR #5720) - Text-part thoughtSignature: stow on the assistant message during streaming and replay onto the last text part on the next turn so Gemini 3 strict function-calling does not reject history. - Function declarations: recursively strip Gemini-unsupported OpenAPI keys (additionalProperties, $schema, $defs, strict, etc.) so OpenAI strict tools stop 400ing as INVALID_ARGUMENT on Gemini. - OpenAI-compat fallback: forward tools/tool_choice so custom Gemini proxies (LiteLLM, gateways) keep function-calling. - enable_prompt_caching: cover the Pydantic v1 legacy off/on/f/n/t/y string set so explicit opt-outs stay opt-out (Gemini was sending cachedContent: "off" otherwise). - Frontend collectAssistantToolCalls / collectToolResultMessages: use google.native_part + result presence to disambiguate provider builtins from same-named user-declared functions. - Added regression tests for text-signature replay and schema sanitization. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 9 Gemini follow-ups (PR #5720) Two round-9 convergent finds across the 12 reviewers: - Server-side web_search was leaking onto the next turn as a fake user functionCall/functionResponse. The previous heuristic (skip builtin only when no native_part AND no result) let it through because the synthetic tool card has a non-empty result string. Always skip web_search by name on both serializers, accept that a user-declared function literally named "web_search" must use a different name. - Assistant `extra_content` was dropped by ChatMessage validation before _stream_gemini could replay text-part thought signatures. Add the field to ChatMessage and forward it through _build_external_messages so the multi-turn signature path actually carries data. Includes a regression test for the ChatMessage round-trip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 10 Gemini follow-ups (PR #5720) Three convergent round-10 reviewer findings closed: - Tag synthetic provider-side builtins with `args._server_tool=True` via a central helper that runs in every `_emit_tool_event` / `_emit_synthetic_tool_event` path. The frontend filter now skips on that marker instead of on the public tool name, so local llama.cpp `web_search` and OpenAI function tools literally named `web_search` / `code_execution` / `image_generation` round-trip cleanly while Gemini grounding / hosted code-exec / hosted image cards stay skipped. - Gate Gemini image-mode (responseModalities=[TEXT,IMAGE]) on the Images pill (enabled_tools containing `image_generation`). Selecting an image-capable model with the pill off no longer forces image output the UI says is disabled. - Frontend missing-key guard now exempts custom Gemini OAI-compat proxies (LiteLLM, gateways) the same way the backend already does, so a saved Gemini connection on `http://localhost:4000/v1` with no API key stops being blocked. Existing tests updated to pass `enabled_tools=["image_generation"]` on image-mode capture paths. * Studio: round 11 Gemini follow-ups (PR #5720) Four round-11 findings closed: - Kimi _stream_kimi_web_search's local _synthetic_chunk helper now runs through _stamp_server_tool_marker so Kimi search history is not replayed as a fake user functionCall on the next turn (was an asymmetric miss after the round-10 tagging work). - OpenAI Responses path (/v1/responses for gpt-5.x) forwards caller-supplied tools / tool_choice, translating the Chat Completions function-tool shape into the Responses native shape. Without this, standard OpenAI tools silently dropped on Responses-routed traffic. - Decoupled the Gemini image-tier model-id guards (text-tool / thinking strip) from the Images pill flip (responseModalities=[TEXT,IMAGE]). gemini-2.5-flash-image with Search/Code on and the Images pill OFF no longer forwards googleSearch + thinkingConfig (Gemini 400s on those for legacy image ids). - Gemini-only extra_content is now forwarded by _build_external_messages only when provider_type=="gemini" so Google's thought_signature does not leak into OpenAI / Mistral / Kimi / OpenRouter request bodies as an unknown field. Added a regression test for the image-tier strict-guard split and extended the extra_content test to cover the non-Gemini suppression. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 12 Gemini follow-ups (PR #5720) Three round-12 convergent findings closed: - extra_content leak to custom Gemini OAI-compat proxies (8/12 reviewers). _build_external_messages now gates extra_content on the native generativelanguage.googleapis.com host, not just provider_type=="gemini", so LiteLLM / custom gateways routed through /chat/completions do not get an unknown top-level field. - OpenAI Responses function-tool round-trip (5/12 reviewers). I added user `tools` forwarding in round 11 but did not parse the matching response.output_item.done items of type=function_call. The parser now translates them into Chat Completions delta.tool_calls and the terminal chunk reports finish_reason="tool_calls" when the model invoked a user function. - Image-tier model with Images pill OFF (2/12). Google's image models default to text+image when responseModalities is omitted, so the previous fix silently still billed image output. Force responseModalities=["TEXT"] when the Images pill is off and the selected model is image-capable. Updated the two pre-existing tests that pinned the synthetic-tool arguments shape to include the new `_server_tool: True` marker, and added a regression test for the Responses function-call output translation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 13 Gemini/Responses follow-ups (PR #5720) Three round-13 convergent findings closed: - OpenAI Responses function_call indices: my round-12 translator hardcoded every emitted tool_calls[*].index to 0, so parallel function calls collapsed for index-keyed clients. Track and increment function_call_index per emit (mirrors the Gemini branch's distinct-index pattern). 10/12 reviewers flagged. - _SERVER_SIDE_BUILTIN_TOOL_NAMES now includes web_fetch so Anthropic-hosted web_fetch cards carry the _server_tool marker and the frontend history serializer doesn't replay them as fake user functions. 4 reviewers flagged. - OpenAI Responses follow-up tool results now serialize as Responses-shape function_call / function_call_output items keyed by call_id, instead of Chat Completions role="tool" content. Skips assistant tool_calls tagged with _server_tool so hosted builtins don't round-trip as user functions. 2 reviewers flagged. Updated the Anthropic code_execution and web_fetch test argument pins to include the new _server_tool marker, and added two regression tests (distinct indices on parallel function_call, function_call_output round-trip). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 14 Gemini follow-ups (PR #5720) Three round-14 findings closed: - Remote `image_url` translation (5 reviewers convergent). Public HTTPS image URLs can't be sent as `fileData.fileUri` -- Gemini reserves that path for Files API URIs and YouTube. Fetch the bytes server-side and inline them as base64 `inlineData`, mirroring the pre-PR OpenAI-compat behaviour. YouTube URLs and generativelanguage.googleapis.com/v1beta/files/* stay as `fileData`. - Nullable JSON Schema type arrays. OpenAI strict tools commonly use `"type": ["string", "null"]`; the Gemini sanitizer now flattens that to `"type": "string", "nullable": true` so strict function tools stop 400ing. - Parallel functionResponses now ride on one user content block with multiple `functionResponse` parts, matching Google's parallel tool docs. Consecutive `role="tool"` messages merge into the previous user turn instead of splitting into separate Gemini user turns. Three regression tests added (remote URL fetch + inline, Files API / YouTube fileData preservation, schema nullable flattening, parallel-tool grouping). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: SSRF harden Gemini remote image fetch (PR #5720) Round 15 convergent finding (12/12 reviewers). My round-14 fix to download user-controlled image URLs for inlineData inlining was an SSRF / data-exfiltration path: no scheme check, no private-host guard, no size cap, no Content-Type validation, redirects could bounce to internal services, and the full URL was logged. Replace the inline fetch with `_safe_fetch_image_for_gemini`: - Require https:// (reject http, file, data, ftp, etc). - Resolve the hostname via socket.getaddrinfo and reject if ANY resolved address is private / loopback / link-local / multicast / reserved / unspecified (covers 127.0.0.0/8, 10/8, 172.16/12, 192.168/16, ::1, 169.254/16 metadata, RFC 6890). - Block IP-literal URLs that resolve into those same ranges. - Cap response body at 10 MB (Content-Length pre-check + streamed byte counter). - Require Content-Type to start with `image/`. - Disable redirect following so a 302 to a private host can't slip past the address check. - Use a short 15s timeout and a tiny connection pool dedicated to these fetches. - Log only the host name + error class -- no full URL, no signed querystring leak. If the guard rejects, the image part is silently dropped (instead of forwarding raw bytes or a fileData fallback). Files API URIs and YouTube URLs still ride as `fileData.fileUri` unchanged. Tests: replaced the live-fetch test with a `_safe_fetch_image_for_gemini` monkeypatch, added four new SSRF-guard tests (non-https rejected, loopback / private IP literals rejected, hostnames that resolve to private IPs rejected). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 16 Gemini follow-ups (PR #5720) - IP-pinned image fetch (`_safe_fetch_image_for_gemini`): reuse the validated-once-then-pin pattern from `tools._fetch_page_text` via `asyncio.to_thread`, so DNS rebinding between validation and the HTTP connect cannot redirect us at a private/metadata address. Catch malformed-bracketed IPv6 urlparse errors. Follow up to 4 redirect hops with per-hop SSRF re-validation. - Replace contains-substring detection of Gemini Files API + YouTube URLs with parsed scheme/host/path checks, so attacker URLs like `https://evil.example/path/youtube.com/x.png` no longer skip the safe-fetch path and serialize as `fileData.fileUri`. - `_build_external_messages`: strip per-tool-call `extra_content` for non-native-Gemini providers; the Gemini-only `thought_signature` payload was leaking through `tool_calls[]` into /chat/completions on OpenAI, Anthropic, and custom Gemini OAI-compat gateways. - `_server_tool` marker now gated on the function name being one of the canonical builtin names (`web_search`, `web_fetch`, `code_execution`, `image_generation`) AND the marker being set, so a user function whose schema happens to define an `_server_tool` field is no longer dropped. Frontend filter mirrors the same gate, plus a backward-compat fallback for pre-PR persisted server-tool cards (no marker) routed via name + native_part / web-tool heuristic. - Gemini schema sanitizer collapses `anyOf: [{X}, {"type":"null"}]` to `{X, "nullable": true}` so Optional[X] tool args from OpenAI/Pydantic schemas no longer 400 the Gemini request. - Frontend tool-result serializer emits `{"result":""}` for empty string outputs so the ChatMessage validator does not reject `role="tool"` with empty content. - Coerce `medium` thinkingLevel to `high` for legacy `gemini-3-pro*` / `gemini-3-pro-preview*` (only low/high documented; shut down 2026-03-09); 3.1+ Pro still passes through. - Hide Gemini native thinking ladder on custom OAI-compat Gemini gateways by routing `getExternalReasoningCapabilities` through `isGeminiCustomOpenAICompatBase(baseUrl)`; thread baseUrl through all four call sites. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 17 Gemini follow-ups (PR #5720) - Frontend `collectAssistantToolCalls` and `collectToolResultMessages` no longer drop unmarked `web_search` / `web_fetch` cards by name alone: a user-defined function with one of those names must round-trip. Pre-PR persisted `code_execution` / `image_generation` cards still get filtered via a shape heuristic (kind/command/code/ prompt fields) instead of bare name. - `_build_external_messages._filter_tool_calls` now drops marked server-side builtin `tool_calls` entirely for non-native-Gemini providers, not just their `extra_content`. An assistant turn whose only payload was a marked builtin is dropped completely so the receiving provider does not see an orphan tool_call. - `_stream_anthropic` translates OpenAI top-level `tool_calls` into Anthropic native `{type:"tool_use", id, name, input}` content blocks, and translates `role="tool"` follow-ups into `role:"user"` messages carrying a `tool_result` block. Anthropic's native Messages API rejects the OpenAI shapes. - `_safe_fetch_image_for_gemini_sync` factors URL validation through `_safe_parse_https`, so malformed `port` access (e.g. `https://host:bad/x.png`) and malformed redirect targets (e.g. a 302 to `https://[bad/x.png`) drop the image instead of raising mid- request. - `tool_choice="none"` now disables hosted builtins (Gemini googleSearch / codeExecution and OpenAI Responses web_search / shell / image_generation), not just user function declarations. - Schema sanitizer handles multi-type `anyOf` with null (`Union[str, int, None]`): keep the slim non-null anyOf and add `nullable: true` so Gemini does not reject `{"type":"null"}`. - Image fetch falls back to the caller-provided MIME (guessed from URL extension) when the server omits Content-Type instead of dropping the image as `non-image content-type=`. - Per-request aggregate caps on remote image inlining (8 images, 20MB total) so a single chat request cannot force unbounded backend downloads. - Frontend exposes the reasoning ladder for `gemini-2.5-flash-lite` (`none/minimal/low/medium/high/max`) so the UI can drive the thinkingBudget the backend already supports. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 18 Gemini follow-ups (PR #5720) - `tool_choice="none"` now opts out of hosted builtin tools on every provider path, not just Gemini and OpenAI Responses. Anthropic web_search / web_fetch / code_execution, Kimi `$web_search` early return, and OpenRouter `plugins:[{id:"web"}]` are all gated on `tool_choice_disabled`. Passing `enabled_tools=[...]` with `tool_choice="none"` no longer triggers provider-side search / code execution for any provider. - `_stream_anthropic` accepts `tool_choice` and threads it through; the dispatcher in `stream_chat_completion` forwards it. - Frontend `isServerSideBuiltinToolPart` simplified to drop only on (marker) OR (canonical name + native_part). The previous shape heuristic on `args.kind`/`args.command`/`args.code`/`args.prompt` dropped real user-declared `code_execution` / `image_generation` functions. Pre-PR persisted hosted cards lacking the marker now leak to non-native providers on switch -- preferred to silently deleting legitimate function-call history. - Backend `_is_marked_server_builtin_tool_call` and the OpenAI Responses translator's matching filter accept BOTH `_server_tool` marker AND `args.google.native_part` as durable provider-side signals so Gemini code_execution / image_generation cards are still dropped on a provider switch. - Per-request remote image count cap now counts ATTEMPTS, not just successful inlines, so 100 failing/slow URLs cannot each consume the 15s fetch timeout. Data: URL images now share the same count and byte caps as fetched remote URLs. - OpenAI Responses translator tracks skipped server-builtin `function_call` ids and drops their matching `role="tool"` follow-ups, preventing orphan `function_call_output` items in the outbound body. - Gemini schema sanitizer preserves multi-type unions with null: `{"type":["string","integer","null"]}` becomes `anyOf:[{string},{integer}] + nullable:true` instead of being flattened to the first non-null type. - Gemini model id validation moved to the top of `_stream_gemini` so an invalid model id rejects the request before any remote image fetch / message translation side effect. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 19 Gemini follow-ups (PR #5720) - `_build_external_messages` now skips an empty assistant turn when `_filter_tool_calls` drops every synthetic builtin tool_call (was guarded only on the `content is None` branch; the string-content and list-content branches still forwarded `{"role":"assistant","content":""}` which several providers reject). Also tracks the dropped server-builtin tool_call ids and skips the matching `role="tool"` follow-ups so the receiving provider does not see an orphan tool_result. - OpenRouter `web_search_active` (the synthetic tool_start / tool_end emitter) is now also gated on `tool_choice_disabled` so a request with `tool_choice="none"` does not surface a fake web_search card in the chat UI even though the plugin was correctly stripped from the outbound body. - `_stream_anthropic` translates an OpenAI role="tool" with list content (`content=[{"type":"text","text":"..."}]`) into a native `tool_result` block on a user message; previously only the string-content shape was translated, so list-content tool results were forwarded as invalid `role:"tool"` messages. - Gemini `data:` URL image_url parts now require an `image/*` MIME type; a `data:text/html;base64,...` is dropped instead of being forwarded as `inlineData.mimeType="text/html"` (Gemini rejects the malformed image part). Symmetric with the fetched-remote image fetch path that already rejects non-image Content-Type. - YouTube `fileData.fileUri` now declares `video/mp4` as the mimeType instead of `image/jpeg` guessed from the URL path. The YouTube/fileData input is the documented Gemini video path; the guessed image MIME made valid YouTube inputs malformed. - OpenAI Responses translator preserves `response.output` ordering on assistant turns that emitted both text and a function_call: assistant text is now serialized BEFORE the function_call item so the subsequent function_call_output (the matching role=tool follow-up) lands in the right position. Previously the order was function_call -> assistant text -> function_call_output, which can confuse multi-turn function-calling flows. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: round 20 Gemini follow-ups (PR #5720) Convergent reviewer findings from round 20: - tool_choice="none" no longer flips responseModalities=[TEXT,IMAGE] on image-tier Gemini models. Forced-function tool_choice (e.g. {type:function, function:{name:lookup}}) also drops hosted Search / code execution from the Gemini body so the caller's pinned user function is not silently joined by hosted builtins. - Gemini code-execution thoughtSignature replay now uses an ordered parts list (native_part.parts[]) so per-part signatures stay attached to the exact part Gemini emitted. The previous merged shape fanned one top-level thoughtSignature across executableCode + codeExecutionResult + inlineData and tripped Gemini 3 strict validators. Backward-compat fallback keeps pre-round-21 persisted history working: a legacy native_part with a single subpart still replays the signature on that subpart; merged legacy objects pin the signature to executableCode only. - Remote-image fetch threads the remaining per-request byte budget into _safe_fetch_image_for_gemini, so over-budget URLs are refused via Content-Length pre-check / short read instead of fully downloaded then discarded after the aggregate cap check. - Gemini role=tool with OpenAI list-form content ([{type:text,text:result}]) now flattens text parts before building functionResponse.response.result; previously the parts arrived as the result value instead of the actual tool output. - Frontend chat-adapter merges native_part by concatenating parts lists (preserving per-part thoughtSignature). Wire types expose enable_prompt_caching as boolean|string (Gemini cached-content name) and OpenAIChatDelta now carries tool_calls and extra_content. - Test test_openrouter_no_synthetic_web_search_event_on_tool_choice_none reads _toolEvent from the top-level SSE payload so a backend regression cannot mask the assertion. Adds 7 regression tests covering image_generation gate, forced-function gate, native_part list replay, legacy fallback, list-content functionResponse flattening, fetch byte-budget threading, and wire types. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply forced-function tool_choice gate to Anthropic, OpenRouter, Kimi Previously only the Gemini path treated `tool_choice={"type":"function", "function":{"name":...}}` as a hosted-tool opt-out. Anthropic, OpenRouter, and Kimi still attached hosted web_search / web_fetch / code_execution when the caller explicitly pinned a user function plus `enabled_tools=[...]`. That contradicts the explicit function pin and bills the caller for unwanted server-side calls. Mirror the Gemini gate symmetrically: - Anthropic web_search / web_fetch / code_execution - OpenRouter `plugins:[{id:"web"}]` + the synthetic web_search SSE event the same path emits at stream close - Kimi `_stream_kimi_web_search` dispatch Adds 4 regression tests: - test_anthropic_forced_function_tool_choice_drops_hosted_tools - test_openrouter_forced_function_tool_choice_drops_web_plugin - test_kimi_forced_function_tool_choice_skips_web_search_helper - test_openrouter_no_synthetic_web_search_event_on_forced_function_tool_choice All 146 existing backend tests still pass. * Strip Gemini-only synthetic tool history on local-GGUF dispatch After a Gemini chat that ran code_execution / image_generation, switching the same thread to a local GGUF model used to forward the synthetic provider-side tool_calls (tagged with `args._server_tool` or carrying a Gemini `args.google.native_part` payload) and the message-level `extra_content` to llama-server. The receiving backend has no tool declaration for those names and no use for Gemini thoughtSignature metadata; in the worst case it can produce an orphan tool_call_id and a confused continuation. Add `_strip_provider_synthetic_tool_history()` and wire it through the two local message builders: - `_openai_messages_for_passthrough` (OAI-compat passthrough) - `_openai_messages_for_gguf_chat` (standard GGUF chat path) Real user-function `tool_calls` and their matching `role="tool"` replies survive unchanged; only synthetic provider-side cards and Gemini-only `extra_content` are stripped. If the synthetic call was the assistant turn's only payload, the now-empty turn is dropped too so llama-server does not reject the request. Adds 2 regression tests: - test_strip_provider_synthetic_tool_history_drops_synthetic_only - test_strip_provider_synthetic_tool_history_drops_empty_assistant 142 existing backend tests still pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Disable Search/Code composer pills for Gemini image-tier models For external Gemini image-tier models (gemini-2.5-flash-image, gemini-3.x-image-preview, etc.), the backend unconditionally strips code_execution and strips web_search on older image ids. Search is still allowed on Gemini 3.x Pro/Flash image models, which supportsBuiltinWebSearch already encodes per model. Before this commit the composer pill gates were: searchDisabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch) codeDisabled = !modelLoaded || !(supportsTools || supportsBuiltinCodeExecution) || imageModeDisablesCode `supportsTools` here is a local-runtime fallback that becomes true when any tool-capable local model has been loaded in the session. With a local tool-capable runtime active, switching the chat to an external Gemini image-tier model used to leave Search/Code clickable, even though the backend will silently drop the tool on the wire. Detect "external provider is Gemini AND the model is image-tier" (via supportsBuiltinImageGeneration) and gate the two pills strictly on the provider's own builtin support in that case. Non-Gemini paths and non-image Gemini models keep the supportsTools fallback unchanged. * Apply forced-function tool_choice gate to OpenAI Responses path Round 22 added the gate for Gemini / Anthropic / OpenRouter / Kimi but missed the OpenAI Responses translator. When a caller pinned a user function via `tool_choice={"type":"function","function":{"name":...}}` plus `enabled_tools=["web_search","code_execution","image_generation"]`, the Responses body still attached `{"type":"web_search"}`, `{"type":"shell"}`, and `{"type":"image_generation"}` server tools. The function pin should suppress those for the same privacy + billing reason the other provider paths now do. Compute `_responses_tool_choice_forced_function` next to `_responses_tool_choice_none` and gate each hosted-tool append on `_responses_hosted_builtins_allowed = not none and not forced_function`. The fix has to be applied in TWO places: the initial body builder and `_build_body()` (called by the container-expiry retry path). User function declarations still flow through so the pin has something to target, and the Responses-shape `{type:"function", name:"..."}` `tool_choice` is forwarded unchanged. Adds regression test `test_openai_responses_forced_function_tool_choice_drops_hosted_tools`. All 166 existing backend tests across Gemini + Responses + image-gen + code-exec suites still pass. * Round 24 P1s: SSRF shared-address gap + extra_content text-only leak + custom-Gemini model list Three convergent P1s from round 24 review: 1. SSRF: the shared SSRF validator in `tools._validate_and_resolve_host` used a denylist (is_private / loopback / link_local / multicast / reserved / unspecified). Python classifies shared address space (100.64.0.0/10 carrier-grade NAT, plus 240.0.0.0/4, benchmarking ranges, etc.) with `is_private=False` AND `is_global=False`. The new Gemini server-side image fetcher therefore accepts URLs whose hostname resolves to 100.64.0.1 in cloud/VPC deployments. Add `not ip.is_global` as the primary gate -- a single source of truth that covers every current and future non-global range. 2. _strip_provider_synthetic_tool_history previously only stripped message-level `extra_content` when the assistant turn had tool_calls. A plain text Gemini reply carrying `extra_content.google.thought_signature` flowed through to llama-server when the thread was switched to a local GGUF backend. Always strip message-level `extra_content` on assistant turns. 3. routes/providers.list_provider_models applied Gemini's native `model_id_allowlist` regex to every Gemini provider, including custom OAI-compatible bases (LiteLLM, deployment gateways). IDs like `google/gemini-2.5-flash` and team-prefixed deployment aliases got filtered out even though the chat-dispatch path now routes them via the OpenAI-compatible client. Skip registry-level model-id filters when the configured Gemini base_url host is not the canonical `generativelanguage.googleapis.com`, mirroring the chat-dispatch gate. Three regression tests added: - test_validate_and_resolve_host_blocks_shared_address_space - test_strip_provider_synthetic_tool_history_drops_text_only_extra_content - test_gemini_custom_oai_compat_base_skips_native_allowlist * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Round 25 P1s: skip synthetic server-tool replay + inline $ref/$defs into Gemini schema Two convergent reviewer findings on the native Gemini path: 1. _stream_gemini's tool_calls replay loop falls through to a generic functionCall emission whenever it sees an assistant tool_call. Marked server-side builtin cards (web_search / web_fetch tagged with _server_tool or args.google.native_part) hit that fallthrough with no replayable native_part, which produces an outbound functionCall whose name is not a declared user function. The Gemini turn 400s on the undeclared name. Guard the loop to drop those entries instead, while keeping the existing code_execution / image_generation native-part replay branch intact. 2. _sanitize_gemini_schema uses a strict allowlist that drops local $ref / $defs references. Pydantic-generated tool schemas hoist nested object shapes into $defs and reference them via {"$ref": "#/$defs/X"}, so a property like address: {"$ref": "#/$defs/Address"} collapsed to {} on the wire and the model lost the nested fields, types, and required keys. Resolve local #/... pointers against the schema root and inline the referenced subtree, with local siblings overriding the reference (normal JSON Schema composition) and a seen-ref guard for self-referential schemas. Added regression coverage: - test_gemini_native_skips_synthetic_server_builtin_replay - test_function_declarations_inline_local_refs_into_gemini_schema - test_function_declarations_inline_local_refs_in_anyof_and_items - test_function_declarations_self_referential_schema_terminates All 145 Gemini provider tests pass; touched provider regression set (OpenAI Responses, code execution, image generation, Anthropic code execution, Anthropic web_fetch) also 43/43 green. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Round 26 P1s: drop orphan Gemini functionResponse + Anthropic /messages synthetic-history strip Reviewer round 26 surfaced two convergent asymmetric-fix bugs. 1. _stream_gemini drops a synthetic server-tool tool_call (web_search / web_fetch tagged _server_tool) and also replays code_execution / image_generation tool_calls as Gemini-native executableCode / codeExecutionResult / inlineData parts. The matching role="tool" follow-up was still falling through to the generic functionResponse branch, producing either an orphan functionResponse (synthetic case) or a duplicate response pointing at a name with no functionDeclarations entry (native-part case). Both forms 400 the next Gemini turn. Track skipped + native-replayed tool_call_ids in _gemini_skip_tool_result_ids and short-circuit the role="tool" branch on a match. 2. The Anthropic-compatible local /v1/messages route only called _drop_empty_assistant_sentinels on the OpenAI-translated history, while the sibling /v1/chat/completions and GGUF passthrough builders chain that with _strip_provider_synthetic_tool_history. An Anthropic caller replaying a prior provider-side tool_use therefore forwarded fake builtin tool history straight into local llama-server. Apply the same strip on the Anthropic route after the anthropic_messages_to_openai conversion. Regression coverage added: - test_gemini_native_skips_orphan_function_response_for_dropped_builtin - test_gemini_native_skips_orphan_function_response_for_native_part_replay Gemini suite 147/147; touched provider regression set 43/43. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Round 27 P1s: native_part location fallback + Gemini image request budget for base64 Two convergent reviewer findings on the native Gemini path. 1. _stream_gemini's synthetic-builtin detector at lines 3519-3524 recognizes args.google.native_part as a server-tool marker, but _native_part was only loaded from tc.extra_content.google.native_part. A direct OpenAI-compatible API caller or imported third-party thread round-trips the payload through function.arguments because tool_calls[].extra_content is not in the OpenAI spec. The round-25 guard then saw a synthetic builtin with no _native_part and dropped the entire assistant turn, so the next native Gemini request lost the prior executableCode / inlineData / codeExecutionResult context. Fall back to args.google.native_part when extra_content path is missing, mirroring what the synthetic detector already accepts. 2. _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES capped DECODED bytes at 20MB. Gemini receives images base64-encoded inside JSON, and base64 inflates payload size by ~4/3. With 20MB decoded the actual JSON body is ~26.7MB plus prompt overhead, well over Gemini's ~20MB request limit. Drop the decoded cap to 14MB so realistic multi- image turns stay safely under 20MB encoded. Added regression test test_gemini_native_part_falls_back_to_args_google covering an OpenAI-compat-shaped image_generation tool_call whose native_part lives only in function.arguments. Gemini suite 148/148. * Fix TS build errors from main merge: restore imageParts + refusal return [] + cast image-edit ref Three errors in chat-adapter.ts surfaced by the frontend tsc step after merging main into feat/gemini-provider: 1. The Anthropic refusal early-return used main's but toOpenAIMessages returns SerializedMessage[]; flip to . 2. Restore -- the line was lost when removing main's conflict block from the function body. 3. selectedImageEditReference splice was inserting OpenAIChatMessage into a SerializedMessage[] array; the shapes differ on tool_calls.id nullability. Cast the reference message through unknown -- it carries no tool_calls, so the runtime payload is structurally compatible. Reproduced locally with `tsc -b --pretty false` (now passes). Build also failing in the in-repo `npm run build` step on PR CI; this commit unblocks all 12 failing UI/API workflows. * Tighten verbose comments in external_provider.py + chat-adapter.ts Compress multi-line explanatory comments in the Gemini translator and the chat adapter without changing any behaviour. All 148 Gemini provider tests still pass; tsc --noEmit clean. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .../core/inference/external_provider.py | 3035 +++++++-- studio/backend/core/inference/providers.py | 69 +- studio/backend/core/inference/tools.py | 11 +- studio/backend/models/inference.py | 45 +- studio/backend/routes/inference.py | 356 +- studio/backend/routes/providers.py | 55 +- .../tests/test_anthropic_code_execution.py | 8 +- .../backend/tests/test_anthropic_web_fetch.py | 7 +- studio/backend/tests/test_gemini_provider.py | 5501 +++++++++++++++++ .../tests/test_openai_code_execution.py | 10 +- .../tests/test_openai_image_generation.py | 3 + .../test_openai_responses_translation.py | 248 + .../src/components/assistant-ui/thread.tsx | 1 + .../src/features/chat/api/chat-adapter.ts | 783 ++- .../frontend/src/features/chat/chat-page.tsx | 10 +- .../src/features/chat/external-providers.ts | 7 + .../features/chat/provider-capabilities.ts | 245 +- .../src/features/chat/shared-composer.tsx | 39 +- .../frontend/src/features/chat/types/api.ts | 53 +- 19 files changed, 9874 insertions(+), 612 deletions(-) create mode 100644 studio/backend/tests/test_gemini_provider.py diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 8f34bb23fc..8a1edd608b 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -8,10 +8,12 @@ Most registry providers expose OpenAI-compatible /v1/chat/completions endpoints; Anthropic uses native Messages API with translation in this client. """ +import base64 import json as _json +import mimetypes import re import time -from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional +from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional, Union from urllib.parse import urlparse import httpx @@ -506,6 +508,250 @@ def _apply_mistral_reasoning_controls( _http_client = httpx.AsyncClient() +# Cap per-image fetch well below Gemini's ~20 MB total request budget. +_GEMINI_REMOTE_IMAGE_MAX_BYTES = 10 * 1024 * 1024 +_GEMINI_REMOTE_IMAGE_TIMEOUT_S = 15.0 + + +def _safe_fetch_image_for_gemini_sync( + url: str, + fallback_mime: str, + max_bytes: int = _GEMINI_REMOTE_IMAGE_MAX_BYTES, +) -> Optional[tuple[str, str]]: + """Synchronous IP-pinned HTTPS image fetch with SSRF guards. + + Uses the same pinned-IP + SNI pattern as `tools._fetch_page_text` so + DNS rebinding between validation and the actual connection cannot + redirect us to a private/metadata address. Follows up to 4 hops, + re-validating each redirect target. Returns (mime, base64) or None. + + `max_bytes` is clamped to the per-image cap and additionally lets + the caller pass the remaining per-request budget so an over-budget + URL is rejected via Content-Length (or read short-circuit) instead + of being fully downloaded then discarded after the fact. + """ + import urllib.error + import urllib.request + from urllib.parse import urljoin, urlunparse + + # Refuse upfront if the per-request budget is already spent. + _byte_limit = min(max(0, int(max_bytes)), _GEMINI_REMOTE_IMAGE_MAX_BYTES) + if _byte_limit <= 0: + return None + + # Share tools.py's pinned-IP hardening: validate-once-then-pin. + from .tools import ( + _NoRedirect, + _SNIHTTPSHandler, + _validate_and_resolve_host, + ) + + def _safe_parse_https(raw_url: str) -> Optional[tuple[Any, str, int]]: + """Validate https + hostname + port. Returns (parsed, host, port) or + None. Handles malformed-port and malformed-bracketed-IPv6 URLs that + would otherwise raise ValueError mid-build. + """ + try: + parsed_url = urlparse(raw_url) + host_value = parsed_url.hostname + port_value = parsed_url.port or 443 + except (ValueError, UnicodeError) as _err: + logger.info( + "Gemini image fetch: refusing malformed url err=%s", + type(_err).__name__, + ) + return None + scheme_value = (parsed_url.scheme or "").lower() + if scheme_value != "https": + logger.info( + "Gemini image fetch: refusing non-https scheme=%s", + scheme_value, + ) + return None + if not host_value: + logger.info("Gemini image fetch: refusing url with no hostname") + return None + return parsed_url, host_value, port_value + + parsed_info = _safe_parse_https(url) + if parsed_info is None: + return None + parsed, current_host, current_port = parsed_info + current_url = url + ok, reason, pinned_ip = _validate_and_resolve_host(current_host, current_port) + if not ok: + logger.warning( + "Gemini image fetch: refusing host=%s reason=%s", + current_host, + reason, + ) + return None + + for _hop in range(4): + # Pin to validated IP; SNI + cert still use hostname via _SNIHTTPSHandler. + cp_info = _safe_parse_https(current_url) + if cp_info is None: + return None + cp, _cp_host, _cp_port = cp_info + ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip + ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str + pinned_url = urlunparse(cp._replace(netloc = ip_netloc)) + + opener = urllib.request.build_opener( + _NoRedirect, + _SNIHTTPSHandler(current_host), + ) + req = urllib.request.Request( + pinned_url, + headers = {"Host": current_host}, + method = "GET", + ) + + try: + resp = opener.open(req, timeout = _GEMINI_REMOTE_IMAGE_TIMEOUT_S) + except urllib.error.HTTPError as e: + if e.code not in (301, 302, 303, 307, 308): + logger.info( + "Gemini image fetch: status=%d host=%s", + e.code, + current_host, + ) + return None + location = e.headers.get("Location") + if not location: + return None + try: + current_url = urljoin(current_url, location) + except (ValueError, UnicodeError) as _err: + logger.info( + "Gemini image fetch: refusing malformed redirect err=%s", + type(_err).__name__, + ) + return None + rp_info = _safe_parse_https(current_url) + if rp_info is None: + return None + _rp, current_host, current_port = rp_info + ok2, reason2, pinned_ip = _validate_and_resolve_host( + current_host, current_port + ) + if not ok2: + logger.warning( + "Gemini image fetch: refusing redirect host=%s reason=%s", + current_host, + reason2, + ) + return None + continue + except (urllib.error.URLError, OSError) as _err: + logger.warning( + "Gemini image fetch failed host=%s err=%s", + current_host, + type(_err).__name__, + ) + return None + + with resp: + status = getattr(resp, "status", None) or resp.getcode() + if status != 200: + logger.info( + "Gemini image fetch: status=%s host=%s", status, current_host + ) + return None + _hdr_mime = ( + (resp.headers.get("content-type") or "").split(";")[0].strip().lower() + ) + # Declared non-image MIME is a refusal; missing MIME falls back to caller's. + if _hdr_mime and not _hdr_mime.startswith("image/"): + logger.info( + "Gemini image fetch: non-image content-type=%s host=%s", + _hdr_mime, + current_host, + ) + return None + _final_mime_pre = _hdr_mime if _hdr_mime else fallback_mime + if not isinstance(_final_mime_pre, str) or not _final_mime_pre.startswith( + "image/" + ): + logger.info( + "Gemini image fetch: missing content-type and no image fallback host=%s", + current_host, + ) + return None + _hdr_len = resp.headers.get("content-length") + if _hdr_len and _hdr_len.isdigit() and int(_hdr_len) > _byte_limit: + logger.info( + "Gemini image fetch: declared %s bytes exceeds cap=%s host=%s", + _hdr_len, + _byte_limit, + current_host, + ) + return None + # Read cap+1 to detect oversize without buffering unbounded data. + raw = resp.read(_byte_limit + 1) + if len(raw) > _byte_limit: + logger.info( + "Gemini image fetch: streamed bytes exceed cap=%s host=%s", + _byte_limit, + current_host, + ) + return None + return _final_mime_pre, base64.b64encode(raw).decode("ascii") + + logger.info("Gemini image fetch: too many redirects host=%s", current_host) + return None + + +async def _safe_fetch_image_for_gemini( + url: str, + fallback_mime: str, + max_bytes: int = _GEMINI_REMOTE_IMAGE_MAX_BYTES, +) -> Optional[tuple[str, str]]: + """Async wrapper running the IP-pinned fetch on a worker thread. + + SSRF guards (https only, pinned IP, per-hop redirect re-check, size + cap, image/* content-type) live in the sync helper. `max_bytes` + carries the remaining per-request budget so over-budget URLs are + rejected up front. + """ + import asyncio + + return await asyncio.to_thread( + _safe_fetch_image_for_gemini_sync, url, fallback_mime, max_bytes + ) + + +# Synthetic-tool names stamped onto outbound _toolEvent.arguments so the +# frontend can distinguish provider-side cards from real user-declared +# tools of the same name. Mirrored on the TS side. +_SERVER_SIDE_BUILTIN_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "code_execution", "image_generation"} +) + + +def _stamp_server_tool_marker(payload: dict[str, Any]) -> None: + """Tag synthetic provider-side tool events so the frontend can + distinguish them from real user-declared / local function tools of + the same name. The marker rides on `arguments._server_tool` and is + only added for known server-side builtin names; user-supplied + tool calls echoed back through these helpers (e.g. Kimi + `$web_search`) keep their existing shape because we keep this scoped + to the canonical builtin names. + """ + if not isinstance(payload, dict): + return + if payload.get("type") != "tool_start": + return + name = payload.get("tool_name") + if not isinstance(name, str) or name not in _SERVER_SIDE_BUILTIN_TOOL_NAMES: + return + args = payload.get("arguments") + if not isinstance(args, dict): + args = {} + payload["arguments"] = args + args["_server_tool"] = True + + def _build_kimi_tool_end( synthetic_chunk_fn: Any, tool_call_id: str, @@ -546,16 +792,23 @@ class ExternalProviderClient: ): self.provider_type = provider_type self.base_url = base_url.rstrip("/") + # Strip a legacy `/openai` suffix from Google-hosted bases so + # configs saved before the native switch still route correctly. + # Custom proxy paths ending in `/openai` are left untouched. + if self.provider_type == "gemini": + _parsed_base = urlparse(self.base_url) + if ( + (_parsed_base.hostname or "").lower() + == "generativelanguage.googleapis.com" + and _parsed_base.path.rstrip("/") == "/v1beta/openai" + ): + self.base_url = self.base_url[: -len("/openai")] self.api_key = api_key self._timeout = httpx.Timeout(timeout, connect = 10.0) - # Separate timeout for SSE streams: reasoning-heavy providers - # (Anthropic Opus 4.7 with adaptive thinking, OpenAI gpt-5.x via - # /v1/responses) can pause for tens of seconds between bytes - # while the model is internally thinking. httpx's read timeout is - # the *gap* between successive reads, not a wall clock — so - # disabling it lets long thinks complete without cutting the - # stream prematurely. connect/write/pool keep the 10s / 120s - # bounds so genuine network failures still surface. + # Disable read timeout on SSE streams: reasoning-heavy models + # pause tens of seconds between bytes while thinking, and httpx's + # read timeout is the per-byte gap, not wall clock. connect/write + # bounds still surface real network failures. self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = None) def _auth_headers(self) -> dict[str, str]: @@ -566,6 +819,14 @@ class ExternalProviderClient: auth_header = provider_info.get("auth_header", "Authorization") auth_prefix = provider_info.get("auth_prefix", "Bearer ") + # Non-Google Gemini bases (LiteLLM, custom gateways) use OAI-compat + # Bearer auth, not Google's x-goog-api-key. Override the registry default. + if self.provider_type == "gemini": + _host = (urlparse(self.base_url).hostname or "").lower() + if _host != "generativelanguage.googleapis.com": + auth_header = "Authorization" + auth_prefix = "Bearer " + headers = {"Content-Type": "application/json"} # Skip auth header when api_key is empty (optional for local providers); # httpx rejects an empty `Bearer ` value as "Illegal header value". @@ -580,6 +841,12 @@ class ExternalProviderClient: from core.inference.providers import get_provider_info info = get_provider_info(self.provider_type) or {} + # Google-hosted Gemini uses the native translator; non-Google + # bases stay on OAI-compat so LiteLLM / custom proxies still work. + if self.provider_type == "gemini": + _host = (urlparse(self.base_url).hostname or "").lower() + if _host != "generativelanguage.googleapis.com": + return True return info.get("openai_compatible", True) async def stream_chat_completion( @@ -594,11 +861,13 @@ class ExternalProviderClient: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, enabled_tools: Optional[list[str]] = None, - enable_prompt_caching: Optional[bool] = None, + enable_prompt_caching: Optional[Union[bool, str]] = None, openai_code_exec_container_id: Optional[str] = None, anthropic_code_exec_container_id: Optional[str] = None, prompt_cache_ttl: Optional[str] = None, compaction_threshold: Optional[int] = None, + tools: Optional[list[dict[str, Any]]] = None, + tool_choice: Optional[Any] = None, fast_mode: Optional[bool] = None, stream: bool = True, ) -> AsyncGenerator[str, None]: @@ -616,7 +885,35 @@ class ExternalProviderClient: ``fast_mode`` only applies to Anthropic Opus 4.6 / 4.7 (silently dropped elsewhere); adds the beta header and ``speed: "fast"``. """ + # tool_choice="none" hard-disables hosted/builtin tools across + # every provider so enabled_tools cannot accidentally bill or leak. + tool_choice_disabled = ( + isinstance(tool_choice, str) and tool_choice.strip().lower() == "none" + ) + if not self._is_openai_compatible(): + # Gemini speaks its own native REST shape (contents/parts); + # `_stream_gemini` translates request/response into the OpenAI + # Chat Completions chunk format the rest of Studio expects. + # API reference: https://ai.google.dev/gemini-api/docs + if self.provider_type == "gemini": + async for line in self._stream_gemini( + messages, + model, + temperature, + top_p, + max_tokens, + top_k, + presence_penalty, + enabled_tools, + enable_prompt_caching, + enable_thinking, + reasoning_effort, + tools, + tool_choice, + ): + yield line + return async for line in self._stream_anthropic( messages, model, @@ -631,6 +928,7 @@ class ExternalProviderClient: anthropic_code_exec_container_id, prompt_cache_ttl, compaction_threshold, + tool_choice, fast_mode = fast_mode, ): yield line @@ -654,20 +952,25 @@ class ExternalProviderClient: enable_prompt_caching, openai_code_exec_container_id, compaction_threshold, + tools, + tool_choice, ): yield line return - # Kimi's $web_search is a builtin_function that requires a client - # round-trip: the first call returns a tool_calls envelope with - # function.arguments populated; the caller echoes those arguments - # back as a role=tool message; the second call streams the final - # answer with the search incorporated. The doc also mandates - # disabling thinking while $web_search is active. Route to a - # dedicated helper so the default OAI-compat path stays single-pass. - # https://platform.kimi.ai/docs/guide/use-web-search + # Kimi $web_search needs a 2-call round-trip + thinking off; route + # to a helper. Forced-function tool_choice suppresses it. + # https://platform.kimi.ai/docs/guide/use-web-search + _kimi_tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) if ( self.provider_type == "kimi" + and not tool_choice_disabled + and not _kimi_tool_choice_forced_function and enabled_tools and "web_search" in enabled_tools ): @@ -694,28 +997,18 @@ class ExternalProviderClient: else: body["max_tokens"] = max_tokens - # Strip body fields a provider's registry entry declares unusable — - # reasoning-class models that lock these to fixed defaults (e.g. - # Kimi k2.5/k2.6 only accept temperature=1, top_p=1) 400 otherwise. - # The frontend capability map already hides the matching sliders; - # this is the matching guard for the pydantic default that the - # route layer would otherwise still fill in. + # Drop fields the registry flags as unusable so reasoning-class + # models with fixed defaults (Kimi k2.6 etc) don't 400 on pydantic + # default values that the route layer still fills in. from core.inference.providers import get_provider_info provider_info = get_provider_info(self.provider_type) or {} for field in provider_info.get("body_omit", ()): body.pop(field, None) - # Kimi (kimi-k2.6, kimi-k2-thinking) accepts a boolean thinking toggle - # via a top-level `thinking` field (the docs show it nested under - # extra_body, but that is an OpenAI Python SDK convention; on the - # wire it merges into the request body). - # - kimi-k2.6 defaults to thinking enabled; clients can pass - # {"type": "disabled"} to suppress it. - # - kimi-k2-thinking is always on; we never send disabled there. - # `keep: all` retains every thinking chunk through the stream, which - # is what we need so our frontend can wrap reasoning_content into - # the chat reasoning panel. + # Kimi thinking is a top-level body field. kimi-k2-thinking is + # always on (ignore the toggle); kimi-k2.6 defaults on, can be + # disabled. `keep: all` preserves every chunk for the UI panel. if self.provider_type == "kimi" and enable_thinking is not None: if model == "kimi-k2-thinking": # Always on; ignore client toggle to avoid an API-level reject. @@ -736,17 +1029,9 @@ class ExternalProviderClient: tpl_kw["enable_thinking"] = bool(enable_thinking) body["chat_template_kwargs"] = tpl_kw - # OpenRouter exposes a unified `reasoning` parameter on every - # chat-completion request — the gateway routes it to whichever - # underlying model actually supports reasoning, and silently - # no-ops for ones that don't. Documented at - # https://openrouter.ai/docs/guides/best-practices/reasoning-tokens - # Shape: `reasoning: {enabled?: bool, effort?: low|medium|high, - # max_tokens?: N, exclude?: bool}` with effort and max_tokens - # mutually exclusive. We forward either an effort level (when - # the user picked one) or a bare {enabled: true}. A small set of - # known routes rejects explicit disable with 400 ("Reasoning is - # mandatory for this endpoint ..."), so only those omit "off". + # OpenRouter's unified `reasoning` field gates per-model thinking. + # Some routes (`*_MANDATORY_REASONING_MODELS`) 400 on explicit off. + # https://openrouter.ai/docs/guides/best-practices/reasoning-tokens if self.provider_type == "openrouter": normalized_or_model = model.strip().lower() if reasoning_effort in ("low", "medium", "high"): @@ -759,17 +1044,22 @@ class ExternalProviderClient: else: body["reasoning"] = {"enabled": False} - # OpenRouter web-search plugin — universal shape that works - # for every model id, including the `openrouter/free` and - # `openrouter/auto` meta-routers. Documented at - # https://openrouter.ai/docs/guides/features/plugins/web-search - # The `:online` model-suffix shortcut is "exactly equivalent - # to" this plugin per the same doc, but only works on - # concrete model ids — meta-routers reject the suffix. - # `plugins: [{id: "web"}]` works everywhere, no model id - # rewrite needed, and idempotent if some future call site - # adds the entry first. - if enabled_tools and "web_search" in enabled_tools: + # OpenRouter web plugin works on every model id including + # meta-routers (unlike the `:online` suffix). Forced-function + # tool_choice suppresses it, matching Gemini/Anthropic. + # https://openrouter.ai/docs/guides/features/plugins/web-search + _or_tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) + if ( + not tool_choice_disabled + and not _or_tool_choice_forced_function + and enabled_tools + and "web_search" in enabled_tools + ): plugins = list(body.get("plugins") or []) if not any( isinstance(p, dict) and p.get("id") == "web" for p in plugins @@ -781,6 +1071,15 @@ class ExternalProviderClient: body.get("model"), ) + # Forward OpenAI-style function tools / tool_choice on every + # OAI-compat route (incl. custom Gemini OpenAI proxies like + # LiteLLM). Without this, callers that wire user-defined tools + # silently lose function-calling on non-native providers. + if tools: + body["tools"] = tools + if tool_choice is not None: + body["tool_choice"] = tool_choice + url = f"{self.base_url}/chat/completions" logger.info( "Proxying chat completion to %s (provider=%s, model=%s)", @@ -816,30 +1115,22 @@ class ExternalProviderClient: ) return - # NOTE: manual __anext__ loop instead of `async for` is intentional. - # On Python 3.13 + httpcore 1.0.x, `async for` auto-calls aclose() on - # early exit (break/return/GeneratorExit) BEFORE our finally block runs. - # That propagates GeneratorExit into PoolByteStream.__aiter__() while it - # calls `await self.aclose()` inside `with AsyncShieldCancellation()`, - # triggering "RuntimeError: async generator ignored GeneratorExit". - # Fix: call response.aclose() FIRST (sets PoolByteStream._closed=True), - # then lines_gen.aclose() is a no-op and GeneratorExit re-raises cleanly. + # Manual __anext__ (not `async for`) so we can close + # the response BEFORE lines_gen, avoiding the httpcore + # 1.0 GeneratorExit -> RuntimeError path on Python 3.13. lines_gen = response.aiter_lines().__aiter__() - # Best-effort diagnostics for the default OAI-compat path. Without - # this, OpenRouter mid-stream errors (200 OK + error event in the - # SSE body) and OpenRouter-router model selection were invisible - # in the backend logs — the user only saw "Provider returned - # error" in the UI with no trail on the server side. + # Diagnostic counters for the OAI-compat path; surfaces + # OpenRouter mid-stream errors that would otherwise be + # invisible server-side. event_counts: dict[str, int] = {} chosen_model: Optional[str] = None - # Web-search tool-card synthesis for OpenRouter. The gateway - # doesn't emit structured web_search_call events — citations - # come back as `annotations` of type=url_citation on delta / - # message objects. Mirror the OpenAI/Anthropic UX by yielding - # a synthetic tool_start at stream open and tool_end at - # stream close with the collected citation list. + # OpenRouter has no web_search_call events — citations + # arrive as url_citation annotations. Synthesise a + # tool_start/tool_end pair to match the OpenAI/Anthropic UX. web_search_active = ( self.provider_type == "openrouter" + and not tool_choice_disabled + and not _or_tool_choice_forced_function and bool(enabled_tools) and "web_search" in (enabled_tools or []) ) @@ -849,6 +1140,7 @@ class ExternalProviderClient: web_search_tool_ended = False def _emit_synthetic_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": f"chatcmpl-{self.provider_type}-synthetic", "object": "chat.completion.chunk", @@ -1104,6 +1396,7 @@ class ExternalProviderClient: synthetic_id = f"chatcmpl-{self.provider_type}-synthetic" def _synthetic_chunk(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": synthetic_id, "object": "chat.completion.chunk", @@ -1441,6 +1734,7 @@ class ExternalProviderClient: anthropic_code_exec_container_id: Optional[str] = None, prompt_cache_ttl: Optional[str] = None, compaction_threshold: Optional[int] = None, + tool_choice: Optional[Any] = None, *, fast_mode: Optional[bool] = None, ) -> AsyncGenerator[str, None]: @@ -1471,6 +1765,42 @@ class ExternalProviderClient: continue content = msg.get("content") + # OpenAI role="tool" with list content -> Anthropic native + # tool_result block on a user message. Translating in the + # string-content branch only (below) leaves the list-content + # form forwarded as an invalid `role:"tool"` message that + # Anthropic rejects. Handle both upfront. + if msg.get("role") == "tool": + _tr_id = msg.get("tool_call_id") or "" + if isinstance(content, list): + _flat_parts: list[str] = [] + for part in content: + if ( + isinstance(part, dict) + and part.get("type") == "text" + and part.get("text") + ): + _flat_parts.append(str(part["text"])) + _flat_result = "".join(_flat_parts) + elif content is None: + _flat_result = "" + elif isinstance(content, str): + _flat_result = content + else: + _flat_result = _json.dumps(content) + filtered.append( + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": _tr_id, + "content": _flat_result, + } + ], + } + ) + continue if isinstance(content, list): # Translate OpenAI multimodal parts -> Anthropic native shapes. # - `image_url` -> `{type:"image", source:...}` @@ -1583,6 +1913,37 @@ class ExternalProviderClient: if title: doc_block["title"] = title anthropic_parts.append(doc_block) + # Assistant tool_calls -> Anthropic tool_use blocks + # appended to the same message. Anthropic native + # Messages API does not accept OpenAI's top-level + # `tool_calls` field; the call lives inside a content + # block with `{type:"tool_use", id, name, input}`. + if msg.get("role") == "assistant" and isinstance( + msg.get("tool_calls"), list + ): + for _tc in msg["tool_calls"]: + if not isinstance(_tc, dict): + continue + _fn = _tc.get("function") or {} + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _raw = _fn.get("arguments") or "{}" + try: + _input = ( + _json.loads(_raw) if isinstance(_raw, str) else _raw + ) + except Exception: + _input = {"_raw": _raw} + if not isinstance(_input, dict): + _input = {"value": _input} + anthropic_parts.append( + { + "type": "tool_use", + "id": _tc.get("id") or f"toolu_{time.time_ns()}", + "name": _fn["name"], + "input": _input, + } + ) # Skip whole-message append when nothing usable survived. # An empty content array (e.g. user dropped only an unparseable # `input_document`) would 400 the Anthropic API with @@ -1590,6 +1951,72 @@ class ExternalProviderClient: if anthropic_parts: filtered.append({"role": msg["role"], "content": anthropic_parts}) else: + # role="tool" follow-up -> Anthropic native tool_result + # block on a `user` message. The OpenAI shape + # (role=tool, content=string, tool_call_id) is not a + # valid Anthropic role. + if msg.get("role") == "tool": + _tr_id = msg.get("tool_call_id") or "" + _tr_content = msg.get("content") + if _tr_content is None: + _tr_content = "" + filtered.append( + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": _tr_id, + "content": ( + _tr_content + if isinstance(_tr_content, str) + else _json.dumps(_tr_content) + ), + } + ], + } + ) + continue + # Assistant turn whose content is a plain string but + # also carries OpenAI `tool_calls`: convert into a + # content-array message with a text block + tool_use + # blocks. Without this, the top-level tool_calls leaks + # through unchanged. + if ( + msg.get("role") == "assistant" + and isinstance(msg.get("tool_calls"), list) + and msg["tool_calls"] + ): + _text_content = msg.get("content") + _blocks: list[dict[str, Any]] = [] + if isinstance(_text_content, str) and _text_content: + _blocks.append({"type": "text", "text": _text_content}) + for _tc in msg["tool_calls"]: + if not isinstance(_tc, dict): + continue + _fn = _tc.get("function") or {} + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _raw = _fn.get("arguments") or "{}" + try: + _input = ( + _json.loads(_raw) if isinstance(_raw, str) else _raw + ) + except Exception: + _input = {"_raw": _raw} + if not isinstance(_input, dict): + _input = {"value": _input} + _blocks.append( + { + "type": "tool_use", + "id": _tc.get("id") or f"toolu_{time.time_ns()}", + "name": _fn["name"], + "input": _input, + } + ) + if _blocks: + filtered.append({"role": "assistant", "content": _blocks}) + continue filtered.append(msg) # Claude 4.7 family removed temperature / top_p / top_k entirely. @@ -1616,34 +2043,16 @@ class ExternalProviderClient: # same as True here (callers that don't set the flag still get # caching). Pass False explicitly to opt out. prompt_caching_enabled = enable_prompt_caching is not False - # Anthropic accepts an optional `ttl` on each cache_control marker - # (default is the 5m ephemeral pool; set "1h" to land in the 1h - # pool instead). Per the prompt-caching docs, 1h cache writes are - # billed at 2x base input vs 1.25x for 5m, but reads are 0.1x for - # both. The 1h pool is the right pick when conversations span - # multiple short bursts more than 5 minutes apart -- the read - # discount makes up for the 1.6x write premium after a single - # additional hit. Anything other than the known TTL strings is - # dropped to avoid sending a malformed marker. - # - # The `extended-cache-ttl-2025-04-11` beta header that originally - # gated 1h TTL has been promoted to GA: as of 2026-05 the live - # API accepts `ttl: "1h"` without any beta opt-in. Verified - # against api.anthropic.com on claude-opus-4-7 (status 200 + - # `ephemeral_1h_input_tokens` populated). The test below pins - # the contract by asserting the header is NOT on the wire so a - # future regression that reintroduces the gate would surface - # before users see a 400. + # Optional 1h cache TTL is GA as of 2026-05 (no beta header). 1h + # writes are 2x vs 5m's 1.25x but reads are 0.1x for both, so 1h + # wins after a single extra hit. Unknown TTL strings drop. cache_marker: dict[str, Any] = {"type": "ephemeral"} if prompt_cache_ttl in ("5m", "1h"): cache_marker["ttl"] = prompt_cache_ttl if system: if prompt_caching_enabled: - # System block is the most stable prefix across turns, so - # it gets its own breakpoint. Skipped when system is - # empty — there's nothing to cache, and an empty marker - # is a no-op. + # System is the most stable cross-turn prefix; own breakpoint. body["system"] = [ { "type": "text", @@ -1749,19 +2158,29 @@ class ExternalProviderClient: if body.get("max_tokens", 0) <= budget_tokens: body["max_tokens"] = budget_tokens + 1024 - # Anthropic server-side web_search — see - # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool - # The tool type is date-pinned per model family. Newer Opus / - # Sonnet 4.6 + 4.7 accept `web_search_20260209` with dynamic - # filtering (Claude writes code to filter results before they - # reach context); everything else uses `web_search_20250305`. - # `_anthropic_web_search_version` picks the right one. Anthropic - # dispatches search calls server-side, returning server_tool_use - # + web_search_tool_result blocks in the SSE stream, plus - # url-citation annotations on text deltas. We translate all of - # that into our local _toolEvent shape so the chat UI renders - # web_search exactly like OpenAI's path. - if enabled_tools and "web_search" in enabled_tools: + # tool_choice="none" or pinned-function suppresses hosted tools + # so a stale UI toggle can't fire server-side search/code-exec. + _anthropic_tool_choice_disabled = ( + isinstance(tool_choice, str) and tool_choice.strip().lower() == "none" + ) + _anthropic_tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) + _anthropic_hosted_builtins_allowed = ( + not _anthropic_tool_choice_disabled + and not _anthropic_tool_choice_forced_function + ) + + # Anthropic web_search (date-pinned per model family). + # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool + if ( + _anthropic_hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + ): anthropic_tools = list(body.get("tools") or []) anthropic_tools.append( { @@ -1772,14 +2191,13 @@ class ExternalProviderClient: ) body["tools"] = anthropic_tools - # Anthropic server-side web_fetch reads a single URL (text/PDF) - # and returns a `web_fetch_tool_result` document block. Opt in - # via `enabled_tools=["web_fetch"]`; no beta header required. - # `_anthropic_web_fetch_version` picks `web_fetch_20260209` - # (dynamic filtering) for Opus 4.6/4.7 + Sonnet 4.6, falling - # back to `web_fetch_20250910` elsewhere; mismatched variants - # return 400 so the per-model picker is required. - web_fetch_enabled = bool(enabled_tools and "web_fetch" in enabled_tools) + # Anthropic web_fetch: only URLs already in conversation. Date-pinned. + # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-fetch-tool + web_fetch_enabled = bool( + _anthropic_hosted_builtins_allowed + and enabled_tools + and "web_fetch" in enabled_tools + ) if web_fetch_enabled: anthropic_tools = list(body.get("tools") or []) anthropic_tools.append( @@ -1792,24 +2210,13 @@ class ExternalProviderClient: body["tools"] = anthropic_tools # Anthropic server-side code execution — see - # https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool - # The tool type is date-pinned per model family. - # `_anthropic_code_execution_version` picks `code_execution_20260120` - # for Opus 4.5+ / Sonnet 4.5+ / Opus 4.7 / Sonnet 4.6 (adds REPL - # state persistence + programmatic tool calling) and falls back - # to `code_execution_20250825` everywhere else. Both versions - # run Python + bash + str_replace file edits inside a 5 GB - # sandboxed container per request, with no internet access, and - # both are unlocked by the same `code-execution-2025-08-25` - # `anthropic-beta` header set further down. On the SSE stream - # Anthropic emits two sub-tool names -- `bash_code_execution` - # and `text_editor_code_execution` -- wrapped in the standard - # server_tool_use / *_tool_result block shape. - # v1 wires the tool only; file uploads (container_upload - # content blocks and generated-file retrieval via the Files - # API) are a deliberate follow-up. + # https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool + # Date-pinned tool type per model; both unlock via the same + # `code-execution-2025-08-25` beta header set below. code_execution_enabled = bool( - enabled_tools and "code_execution" in enabled_tools + _anthropic_hosted_builtins_allowed + and enabled_tools + and "code_execution" in enabled_tools ) if code_execution_enabled: anthropic_tools = list(body.get("tools") or []) @@ -1820,32 +2227,14 @@ class ExternalProviderClient: } ) body["tools"] = anthropic_tools - # Reuse the prior turn's container so filesystem state - # (files written, packages installed, variables set) - # persists across turns of the same thread. Anthropic - # exposes the container id on the Message object's - # top-level `container.id`; on the SSE stream we latch it - # off `message_start.message.container.id` further down - # and emit a `container_ready` _toolEvent so the chat - # adapter persists it on the thread record. A stale id - # (container expired / not found) surfaces as a 4xx - # below, where we emit `container_invalidated` and let - # the next turn fall back to auto-create. + # Reuse the thread's prior container so filesystem state + # persists. Stale ids 4xx and clear via container_invalidated. if anthropic_code_exec_container_id: body["container"] = anthropic_code_exec_container_id - # Server-side context compaction — see - # https://platform.claude.com/docs/en/build-with-claude/compaction - # Beta as of `compact-2026-01-12`. When `compaction_threshold` is - # provided AND the model accepts compaction (Opus 4.6+ / 4.7, - # Sonnet 4.6, Mythos preview), attach - # `context_management.edits[{type:"compact_20260112", trigger: - # {type:"input_tokens", value:N}}]` to the body. Anthropic runs - # the compaction step server-side once the rendered prompt - # crosses the threshold and replies with a top-level - # `context_management` block plus `usage.iterations[]` so we can - # account per-iteration. Below-min thresholds get clamped up to - # 50K so the request doesn't 400. + # Server-side compaction (beta `compact-2026-01-12`). Clamps + # below-min thresholds to 50K so the request doesn't 400. + # https://platform.claude.com/docs/en/build-with-claude/compaction compaction_active = ( compaction_threshold is not None and compaction_threshold > 0 @@ -1878,10 +2267,8 @@ class ExternalProviderClient: url = f"{self.base_url}/messages" completion_id = f"chatcmpl-anthropic-{model.replace('/', '-')}" - # Log the outgoing config keys (not the messages themselves) so we - # can prove which thinking/effort fields actually reached the wire. - # If Anthropic skips reasoning despite a configured effort, this - # tells us whether we sent the field or dropped it on the floor. + # Log outgoing config keys (not messages) to prove which thinking / + # effort fields actually reached the wire. logger.info( "Anthropic request shape (model=%s, has_thinking=%s, thinking=%s, " "output_config=%s, temperature=%s, has_top_p=%s, has_top_k=%s, " @@ -1896,16 +2283,10 @@ class ExternalProviderClient: body.get("max_tokens"), ) - # Translate Anthropic stop reasons onto the OpenAI chat-completions - # `finish_reason` vocabulary. `pause_turn` maps to None so the - # adapter does NOT emit a finish_reason chunk: pause_turn means - # Claude paused a long server-tool turn (web_search / web_fetch) - # and will continue once the user (or our retry) sends back the - # partial assistant message. Forwarding it as "stop" makes the - # OpenAI client think the answer is done and truncates the - # rendered message. `refusal` maps to "content_filter" as the - # nearest semantic match. See - # https://platform.claude.com/docs/en/api/messages#response-stop-reason + # Anthropic stop_reason -> OpenAI finish_reason. `pause_turn` + # maps to None so the UI doesn't treat a paused server-tool turn + # as final. `refusal` -> "content_filter" (closest match). + # https://platform.claude.com/docs/en/api/messages#response-stop-reason _finish_reason_map: dict[str, Optional[str]] = { "end_turn": "stop", "max_tokens": "length", @@ -1918,11 +2299,7 @@ class ExternalProviderClient: logger.info("Proxying Anthropic Messages API to %s (model=%s)", url, model) request_headers = self._auth_headers() - # Anthropic accepts comma-separated beta features in a single - # `anthropic-beta` header. Merge our flags onto whatever the - # registry's extra_headers contributed (currently nothing on - # the beta axis, just anthropic-version) so future betas - # added at the registry level keep working. + # Merge new beta flags onto whatever the registry contributed. existing_beta = request_headers.get("anthropic-beta", "").strip() beta_parts = ( [p.strip() for p in existing_beta.split(",") if p.strip()] @@ -1988,27 +2365,14 @@ class ExternalProviderClient: # "no thinking content" — distinguishes "Anthropic never sent # thinking_delta" from "frontend didn't render the chunks". event_counts: dict[str, int] = {} - # web_search state. Anthropic emits the query inside an - # `input_json_delta` stream on a `server_tool_use` content - # block, then a separate `web_search_tool_result` block - # with the URL list. Unlike OpenAI we get per-call results - # directly, so each tool card carries its own citations. - # `current_server_tool_use`: {id, name, partial_json_buffer} - # `current_result_block`: {tool_use_id, results} - # Both go to None when the matching content_block_stop fires. + # web_search state. Query streams via input_json_delta + # on a server_tool_use block; results land in a separate + # web_search_tool_result block. Per-call citations. current_server_tool_use: Optional[dict[str, Any]] = None current_result_block: Optional[dict[str, Any]] = None web_search_calls: dict[str, dict[str, Any]] = {} - # code_execution state. Anthropic's - # `code_execution_20250825` tool emits the same - # server_tool_use → *_tool_result block shape as - # web_search, but the server_tool_use carries one of - # two sub-tool names (`bash_code_execution` or - # `text_editor_code_execution`) and the result block - # type matches (`bash_code_execution_tool_result` / - # `text_editor_code_execution_tool_result`). Kept - # parallel to web_search state so the two paths don't - # collide when both pills are on in the same turn. + # code_execution state (bash / text_editor sub-tools); + # kept parallel to web_search so concurrent pills don't collide. current_code_exec_use: Optional[dict[str, Any]] = None current_code_exec_result: Optional[dict[str, Any]] = None code_execution_calls: dict[str, dict[str, Any]] = {} @@ -2078,6 +2442,7 @@ class ExternalProviderClient: return f"data: {_json.dumps(chunk)}" def _emit_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -2335,18 +2700,8 @@ class ExternalProviderClient: "inner": inner if isinstance(inner, dict) else {}, } elif block_type == "compaction": - # Server-side compaction emits a `compaction` - # content block on the assistant message. - # Anthropic may include the summary text on - # this start event AND/OR stream it via - # text_delta events on the same block. See - # https://platform.claude.com/docs/en/build-with-claude/compaction - # Capture either form; finalize and emit - # on content_block_stop. The chat-adapter - # persists the block onto the assistant - # message so the next turn's request - # carries it back -- Anthropic then skips - # re-compaction from scratch. + # Summary may arrive on start AND/OR via + # text_delta. Capture both; emit on stop. seed = content_block.get("content") or "" current_compaction = { "content": seed if isinstance(seed, str) else "", @@ -2356,12 +2711,7 @@ class ExternalProviderClient: delta = event.get("delta", {}) delta_type = delta.get("type") if delta_type == "thinking_delta": - # Anthropic streams extended-thinking content as - # thinking_delta events on a separate content - # block. Wrap as inline ... so - # the frontend's parseAssistantContent lifts it - # into the reasoning panel — same pattern as - # the OpenAI Responses path. + # Wrap as ... for parseAssistantContent. thinking_text = delta.get("thinking", "") if thinking_text: if not thinking_open: @@ -2631,19 +2981,9 @@ class ExternalProviderClient: delta_usage = event.get("usage") if isinstance(delta_usage, dict): last_usage.update(delta_usage) - # When a fresh compaction has run, Anthropic - # publishes per-iteration token counts in - # `usage.iterations[]`. The top-level - # input_tokens / output_tokens only cover the - # `message` iteration, NOT the compaction - # passes — billing has to sum the whole - # array. See - # https://platform.claude.com/docs/en/build-with-claude/compaction - # Fold the compaction iterations into - # `compaction_input_tokens` / `compaction_output_tokens` - # so the cost surface can add them without - # re-walking the array (and so the closing - # log line names the figures). + # Compaction iterations aren't in top-level + # input/output_tokens; fold them into + # compaction_{input,output}_tokens for billing. iterations = delta_usage.get("iterations") if isinstance(iterations, list): c_in = 0 @@ -2880,6 +3220,1722 @@ class ExternalProviderClient: self.provider_type, ) + async def _stream_gemini( + self, + messages: list[dict[str, Any]], + model: str, + temperature: float, + top_p: float, + max_tokens: Optional[int], + top_k: Optional[int] = None, + presence_penalty: float = 0.0, + enabled_tools: Optional[list[str]] = None, + enable_prompt_caching: Optional[Any] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + tools: Optional[list[dict[str, Any]]] = None, + tool_choice: Optional[Any] = None, + ) -> AsyncGenerator[str, None]: + """ + Call Google's native Gemini API and translate its streaming + ``streamGenerateContent`` response into OpenAI Chat Completions + chunk format. + + Gemini does NOT speak the OpenAI Chat Completions contract on + its primary endpoint. The wire shape is: + + POST /v1beta/models/{model}:streamGenerateContent?alt=sse + { + "contents": [{"role": "user|model", "parts": [{"text": "..."}]}], + "systemInstruction": {"parts": [{"text": "..."}]}, + "generationConfig": {"temperature": 0.7, "topP": 0.95, "topK": 40, + "maxOutputTokens": 1024}, + "tools": [{"googleSearch": {}}, {"codeExecution": {}}], + "cachedContent": "" // optional, see caching docs + } + + Streamed responses are SSE frames carrying partial + ``GenerateContentResponse`` objects: + + {"candidates": [{"content": {"parts": [{"text": "Hello"}]}, + "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 7, "candidatesTokenCount": 3}} + + Image generation uses the same endpoint with model + ``gemini-2.5-flash-image`` (also called Nano Banana); the + response carries an ``inlineData`` part with the base64 PNG + bytes and a ``mimeType``. We surface that through the same + ``tool_start`` / ``tool_end`` ``image_b64`` envelope the OpenAI + image_generation path uses, so the chat UI renders the image + inline with no extra plumbing. + + References: + - https://ai.google.dev/gemini-api/docs/text-generation + - https://ai.google.dev/gemini-api/docs/function-calling + - https://ai.google.dev/gemini-api/docs/grounding + - https://ai.google.dev/gemini-api/docs/caching + - https://ai.google.dev/gemini-api/docs/image-generation + """ + import json as _json + + # Validate the user-controlled model id BEFORE any message + # translation. A model like `../cachedContents/x` is path- + # traversal that lands in `/v1beta/cachedContents/...`; rejecting + # it here also avoids triggering user-controlled outbound fetches + # (remote image_url inlining) on a request we'll error out + # anyway. Documented catalog ids match `[A-Za-z0-9._-]+`. + if not re.fullmatch(r"[A-Za-z0-9._-]+", model): + yield _error_sse_line( + 400, + f"Invalid Gemini model id: {model!r}", + self.provider_type, + ) + return + + # Translate OpenAI messages -> Gemini contents. system role + # promotes to top-level systemInstruction. + system_text_parts: list[str] = [] + contents: list[dict[str, Any]] = [] + # OpenAI may drop `name` from role="tool" follow-ups. Remember + # prior function names so functionResponse isn't sent name-less + # (Gemini 400s on empty names). + tool_call_names: dict[str, str] = {} + # tool_call_ids whose assistant card was dropped (synthetic + # builtin) or already replayed as native parts. Their role="tool" + # follow-up must be skipped to avoid orphan/duplicate responses. + _gemini_skip_tool_result_ids: set[str] = set() + # Per-request image caps. The byte cap counts DECODED bytes; we + # set it to ~14 MB because base64 expansion + prompt overhead + # must fit Gemini's ~20 MB request limit. + _GEMINI_REMOTE_IMAGE_MAX_COUNT = 8 + _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES = 14 * 1024 * 1024 + _remote_image_count = 0 + _remote_image_total_bytes = 0 + for msg in messages: + role = msg.get("role") + content = msg.get("content", "") + if role == "system": + if isinstance(content, str): + if content: + system_text_parts.append(content) + elif isinstance(content, list): + for part in content: + if ( + isinstance(part, dict) + and part.get("type") == "text" + and part.get("text") + ): + system_text_parts.append(part["text"]) + continue + # Map OpenAI roles to Gemini's two-role contract. + gemini_role = "model" if role == "assistant" else "user" + parts: list[dict[str, Any]] = [] + if isinstance(content, str): + if content: + parts.append({"text": content}) + elif isinstance(content, list): + for part in content: + if not isinstance(part, dict): + continue + ptype = part.get("type") + if ptype == "text": + text = part.get("text", "") + if text: + parts.append({"text": text}) + elif ptype == "image_url": + url = part.get("image_url", {}).get("url", "") + if url.startswith("data:"): + header, _, b64data = url.partition(",") + media_type = ( + header.split(";")[0] + .replace("data:", "") + .strip() + .lower() + or "image/jpeg" + ) + # Symmetry with the fetched remote image + # path, which already rejects non-image + # Content-Type. A `data:text/html;base64,...` + # URL otherwise lands as Gemini inlineData + # with mimeType="text/html" and 400s the + # whole request. + if not media_type.startswith("image/"): + logger.info( + "Gemini inlineData: refusing non-image data URL media_type=%s", + media_type, + ) + elif b64data: + # data: URLs share the same caps as fetched + # URLs so inline payloads don't bypass them. + _data_approx_bytes = (len(b64data) * 3) // 4 + if ( + _remote_image_count + >= _GEMINI_REMOTE_IMAGE_MAX_COUNT + ): + logger.info( + "Gemini inlineData: per-request count cap %d reached, dropping image", + _GEMINI_REMOTE_IMAGE_MAX_COUNT, + ) + elif ( + _remote_image_total_bytes + _data_approx_bytes + > _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES + ): + logger.info( + "Gemini inlineData: per-request byte cap reached, dropping image", + ) + else: + _remote_image_count += 1 + _remote_image_total_bytes += _data_approx_bytes + parts.append( + { + "inlineData": { + "mimeType": media_type, + "data": b64data, + } + } + ) + elif url: + # fileData.fileUri only accepts Files-API URIs + # and YouTube; everything else must be downloaded + # and inlined. Parse fields explicitly so + # attacker URLs like https://evil.com/youtube.com/x + # aren't misclassified as YouTube. + try: + _parsed_image_url = urlparse(url) + except (ValueError, UnicodeError): + _parsed_image_url = None + if _parsed_image_url is None: + _img_scheme = "" + _img_host = "" + _img_path = "" + else: + _img_scheme = (_parsed_image_url.scheme or "").lower() + _img_host = (_parsed_image_url.hostname or "").lower() + _img_path = _parsed_image_url.path or "" + _is_native_uri = ( + _img_scheme == "https" + and _img_host == "generativelanguage.googleapis.com" + and _img_path.startswith("/v1beta/files/") + ) + _is_youtube = _img_scheme == "https" and ( + _img_host == "youtu.be" + or _img_host == "youtube.com" + or _img_host.endswith(".youtube.com") + ) + _guessed, _ = mimetypes.guess_type(_img_path) + _media_type = ( + _guessed + if isinstance(_guessed, str) + and _guessed.startswith("image/") + else "image/jpeg" + ) + if _is_youtube: + # YouTube URIs must use video/mp4; the + # default image/jpeg yields a 400. + parts.append( + { + "fileData": { + "fileUri": url, + "mimeType": "video/mp4", + } + } + ) + elif _is_native_uri: + parts.append( + { + "fileData": { + "fileUri": url, + "mimeType": _media_type, + } + } + ) + elif _remote_image_count >= _GEMINI_REMOTE_IMAGE_MAX_COUNT: + logger.info( + "Gemini image fetch: per-request count cap %d reached, dropping image", + _GEMINI_REMOTE_IMAGE_MAX_COUNT, + ) + else: + # Refuse pre-fetch when the per-request + # byte budget is spent; pass the remainder + # so over-budget URLs reject on Content-Length. + _remaining_bytes = ( + _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES + - _remote_image_total_bytes + ) + if _remaining_bytes <= 0: + logger.info( + "Gemini image fetch: per-request byte cap already reached, dropping image", + ) + else: + # Count attempts before awaiting so + # slow URLs don't each burn the timeout. + _remote_image_count += 1 + _fetched = await _safe_fetch_image_for_gemini( + url, + _media_type, + max_bytes = _remaining_bytes, + ) + if _fetched is not None: + _final_mime, _b64 = _fetched + # base64 expands ~4/3 — recover bytes from len(_b64). + _approx_bytes = (len(_b64) * 3) // 4 + if ( + _remote_image_total_bytes + _approx_bytes + > _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES + ): + logger.info( + "Gemini image fetch: per-request byte cap reached, dropping image", + ) + else: + _remote_image_total_bytes += _approx_bytes + parts.append( + { + "inlineData": { + "mimeType": _final_mime, + "data": _b64, + } + } + ) + # Gemini 3 strict function-calling requires text-part + # thoughtSignatures to be replayed on history; the frontend + # stows the latest one as + # extra_content.google.thought_signature on the assistant + # message and we pin it onto the last text part here. + if role == "assistant" and parts: + _msg_extra = msg.get("extra_content") if isinstance(msg, dict) else None + if isinstance(_msg_extra, dict): + _msg_g = _msg_extra.get("google") or {} + if isinstance(_msg_g, dict): + _msg_sig = _msg_g.get("thought_signature") or _msg_g.get( + "thoughtSignature" + ) + if isinstance(_msg_sig, str) and _msg_sig: + for _idx in range(len(parts) - 1, -1, -1): + if "text" in parts[_idx]: + parts[_idx] = { + **parts[_idx], + "thoughtSignature": _msg_sig, + } + break + # Translate OpenAI tool_calls into Gemini functionCall parts. + # code_execution / image_generation replay their native parts + # (executableCode / codeExecutionResult / inlineData) stowed + # on extra_content.google.native_part. + tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if isinstance(tool_calls, list): + for tc in tool_calls: + if not isinstance(tc, dict): + continue + fn = tc.get("function") or {} + if not isinstance(fn, dict): + continue + args_raw = fn.get("arguments") or "{}" + if isinstance(args_raw, str): + try: + args = _json.loads(args_raw) + except Exception: + args = {"_raw": args_raw} + elif isinstance(args_raw, dict): + args = args_raw + else: + args = {} + fn_name = fn.get("name", "") + tc_id = tc.get("id") + if fn_name and isinstance(tc_id, str) and tc_id: + tool_call_names[tc_id] = fn_name + + # Replay native Gemini code_execution / image_generation parts + # from extra_content.google.native_part, with fallback to + # args.google.native_part for OAI-compat round-trips. + _extra = tc.get("extra_content") + _native_part = None + _google_extra: dict[str, Any] = {} + if isinstance(_extra, dict): + _ge = _extra.get("google") or {} + if isinstance(_ge, dict): + _google_extra = _ge + _native_part = _ge.get("native_part") + if _native_part is None and isinstance(args, dict): + _args_google = args.get("google") + if isinstance(_args_google, dict): + _args_np = _args_google.get("native_part") + if isinstance(_args_np, dict): + _native_part = _args_np + if not _google_extra: + _google_extra = _args_google + + # Synthetic builtin cards (web_search/web_fetch) must + # not become fake functionCalls; drop them. Native + # code_execution / image_generation replay below. + _name_lc = fn_name.lower() if isinstance(fn_name, str) else "" + _is_synthetic_server_builtin = ( + _name_lc + in ( + "web_search", + "web_fetch", + "code_execution", + "image_generation", + ) + and isinstance(args, dict) + and ( + args.get("_server_tool") is True + or isinstance( + (args.get("google") or {}).get("native_part"), dict + ) + ) + ) + if _is_synthetic_server_builtin and not ( + _name_lc in ("code_execution", "image_generation") + and isinstance(_native_part, dict) + ): + # No replayable Gemini native part -- skip + # entirely rather than send a fake functionCall. + # Also remember this tool_call_id so a matching + # role="tool" follow-up does not become an + # orphan functionResponse below. + if isinstance(tc_id, str) and tc_id: + _gemini_skip_tool_result_ids.add(tc_id) + tool_call_names.pop(tc_id, None) + continue + if fn_name in ("code_execution", "image_generation") and isinstance( + _native_part, dict + ): + # code_execution/image_generation history is + # replayed as native parts; the matching + # role="tool" must be skipped or Gemini sees a + # functionResponse with no declared function + # name and 400s the turn. + if isinstance(tc_id, str) and tc_id: + _gemini_skip_tool_result_ids.add(tc_id) + # New shape: `native_part.parts` is an ordered list + # of full part wrappers, each carrying its own + # `thoughtSignature`. This preserves Gemini 3's + # strict per-part replay requirement when the + # frontend has merged executableCode + + # codeExecutionResult + inlineData into the same + # tool-call card. + _native_parts_list = _native_part.get("parts") + if isinstance(_native_parts_list, list): + for _entry in _native_parts_list: + if isinstance(_entry, dict): + parts.append(_entry) + continue + # Legacy single-object native_part: fan the shared + # thoughtSignature only when one subpart exists; + # for code+result, prefer executableCode and drop + # the signature elsewhere. + _legacy_sig = _native_part.get( + "thoughtSignature" + ) or _native_part.get("thought_signature") + _legacy_subparts = [ + _k + for _k in ( + "executableCode", + "codeExecutionResult", + "inlineData", + ) + if isinstance(_native_part.get(_k), dict) + ] + for _native_key in ( + "executableCode", + "codeExecutionResult", + "inlineData", + ): + _sub = _native_part.get(_native_key) + if not isinstance(_sub, dict): + continue + _replay_part: dict[str, Any] = {_native_key: _sub} + if isinstance(_legacy_sig, str) and _legacy_sig: + if len(_legacy_subparts) == 1: + _replay_part["thoughtSignature"] = _legacy_sig + elif _native_key == "executableCode": + _replay_part["thoughtSignature"] = _legacy_sig + parts.append(_replay_part) + continue + + # Forward the OpenAI tool_call id into Gemini's + # functionCall.id so a follow-up turn that issues + # multiple calls to the same function (different + # args, same name) can be disambiguated on the + # response side. Gemini accepts the field per + # https://ai.google.dev/gemini-api/docs/function-calling. + function_call_part: dict[str, Any] = { + "name": fn_name, + "args": args, + } + if isinstance(tc_id, str) and tc_id: + function_call_part["id"] = tc_id + # Gemini 3 function-calling requires the prior + # thoughtSignature to be echoed back as a sibling + # of the functionCall part. The translator stows + # it on the assistant tool_call via + # `extra_content.google.thought_signature` (see + # the inbound emit below). + fc_part: dict[str, Any] = {"functionCall": function_call_part} + sig = _google_extra.get("thought_signature") or _google_extra.get( + "thoughtSignature" + ) + if isinstance(sig, str) and sig: + fc_part["thoughtSignature"] = sig + parts.append(fc_part) + if role == "tool": + # If the matching assistant-side tool_call was either + # dropped (synthetic server-tool with no native part) + # or already replayed as Gemini-native parts + # (code_execution/image_generation native_part), drop + # the follow-up too. Emitting it as a functionResponse + # would be orphaned or duplicate the native result. + _tc_id_for_skip = msg.get("tool_call_id") + if ( + isinstance(_tc_id_for_skip, str) + and _tc_id_for_skip in _gemini_skip_tool_result_ids + ): + continue + # OpenAI's role="tool" follow-up carries the function + # result. Gemini's matching shape is a role="user" turn + # with a functionResponse part. When the caller dropped + # ``name``, recover it from the matching assistant + # tool_call so Gemini doesn't 400 on an empty name. + tool_name = msg.get("name") or msg.get("tool_name") or "" + if not tool_name: + tc_id = msg.get("tool_call_id") + if isinstance(tc_id, str) and tc_id in tool_call_names: + tool_name = tool_call_names[tc_id] + response_payload: Any + if isinstance(content, list): + # OpenAI tool messages may carry list-form content + # (`[{"type":"text","text":"..."}]`). Forwarding the + # content-part objects verbatim into Gemini's + # `functionResponse.response.result` yields + # `result:[{"type":"text","text":"..."}]` instead of + # the actual tool output text; flatten text parts so + # the result mirrors the string-content path. + _flat_parts: list[str] = [] + for _cpart in content: + if ( + isinstance(_cpart, dict) + and _cpart.get("type") == "text" + and isinstance(_cpart.get("text"), str) + ): + _flat_parts.append(_cpart["text"]) + _flat_text = "".join(_flat_parts) + try: + response_payload = _json.loads(_flat_text) + except Exception: + response_payload = {"result": _flat_text} + elif isinstance(content, str): + try: + response_payload = _json.loads(content) + except Exception: + response_payload = {"result": content} + else: + response_payload = content or {} + function_response_part: dict[str, Any] = { + "name": tool_name, + "response": ( + response_payload + if isinstance(response_payload, dict) + else {"result": response_payload} + ), + } + # Mirror tool_call_id onto functionResponse.id so + # Gemini can match the result to the originating + # functionCall when multiple parallel calls were made. + tc_id = msg.get("tool_call_id") + if isinstance(tc_id, str) and tc_id: + function_response_part["id"] = tc_id + parts = [{"functionResponse": function_response_part}] + gemini_role = "user" + if parts: + # Gemini expects parallel functionResponses (multiple + # OpenAI role="tool" messages in a row) to ride on a + # single user content with multiple functionResponse + # parts -- the docs show parallel responses grouped + # together in the next turn. Merge consecutive + # functionResponse-only user blocks so realistic + # parallel tool loops round-trip correctly. + if ( + role == "tool" + and contents + and contents[-1].get("role") == "user" + and all( + isinstance(p, dict) and "functionResponse" in p + for p in (contents[-1].get("parts") or []) + ) + ): + contents[-1]["parts"].extend(parts) + else: + contents.append({"role": gemini_role, "parts": parts}) + + body: dict[str, Any] = {"contents": contents} + if system_text_parts: + body["systemInstruction"] = { + "parts": [{"text": "\n\n".join(system_text_parts)}] + } + + # Generation config -- temperature / topP / topK / maxOutputTokens + # map straight across. The frontend capability matrix restricts + # the sliders the UI exposes for Gemini to this set. + gen_config: dict[str, Any] = {} + if temperature is not None: + gen_config["temperature"] = temperature + if top_p is not None: + gen_config["topP"] = top_p + if top_k is not None and top_k > 0: + gen_config["topK"] = top_k + # Gemini accepts ``presencePenalty`` on generationConfig with the + # same sign convention as the OpenAI knob (positive discourages + # repetition). Forward when the caller bothers to set it. + if presence_penalty: + gen_config["presencePenalty"] = presence_penalty + if max_tokens is not None: + gen_config["maxOutputTokens"] = max_tokens + + # Nano Banana image generation. Gemini only accepts + # `responseModalities: ["TEXT","IMAGE"]` on the image-capable + # model family (id contains `-image` or `nano-banana`). Text- + # only models such as `gemini-2.5-flash` 400 on the same body, + # so only force image mode when the selected model actually + # supports it -- a stale `enabled_tools=["image_generation"]` + # on a text model is silently treated as a regular turn. + # https://ai.google.dev/gemini-api/docs/image-generation + model_lc = model.lower() + is_image_picker_model = "-image" in model_lc or "nano-banana" in model_lc + # tool_choice="none" / forced-function tool_choice must also + # suppress the implicit image-generation hosted tool. Otherwise + # an explicit OpenAI-style opt-out (or an explicit user-function + # pin) still flips `responseModalities=["TEXT","IMAGE"]` on + # image-tier models and bills for image output. + _tool_choice_disabled = ( + isinstance(tool_choice, str) and tool_choice.strip().lower() == "none" + ) + _tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) + _hosted_builtins_allowed = ( + not _tool_choice_disabled and not _tool_choice_forced_function + ) + # Image-tier model IDs reject text-only tools (code_execution, + # user functions) and thinkingConfig regardless of whether the + # Images pill is on -- those are model-level constraints + # documented by Google. The pill only controls whether we ask + # Gemini to actually emit image output via + # `responseModalities: ["TEXT","IMAGE"]`. Decoupling the two + # avoids the case where Images is off + Code/Search is on + # forwards `tools: [{codeExecution: {}}]` plus + # `thinkingConfig` to an image model and 400s. + image_tool_requested = bool( + _hosted_builtins_allowed + and enabled_tools + and "image_generation" in enabled_tools + ) + # Strict tool / thinking strip uses the model-id check. + is_image_model_strict = is_image_picker_model + # The actual modality flip only happens when the user opted in. + is_image_model = is_image_picker_model and image_tool_requested + if is_image_model: + gen_config["responseModalities"] = ["TEXT", "IMAGE"] + elif is_image_picker_model: + # Force TEXT-only so an image-capable model with Images OFF + # doesn't still bill for image output. + gen_config["responseModalities"] = ["TEXT"] + + # Thinking control. Gemini 3 uses thinkingLevel (str), 2.5 uses + # thinkingBudget (int). Gemini 3 has no full-off; minimum is + # "minimal" on Flash, "low" on Pro. + # https://ai.google.dev/gemini-api/docs/thinking + _GEMINI3_THINKING_PREFIXES = ( + "gemini-3.5-", + "gemini-3.1-", + "gemini-3-", + "gemini-pro-latest", + "gemini-flash-latest", + "gemini-flash-lite-latest", + ) + _GEMINI3_PRO_PREFIXES = ( + "gemini-3.5-pro", + "gemini-3.1-pro", + "gemini-3-pro", + "gemini-pro-latest", + ) + _PRO_THINKING_PREFIXES = ("gemini-2.5-pro",) + is_gemini3_thinking = any( + model_lc.startswith(p) for p in _GEMINI3_THINKING_PREFIXES + ) + is_gemini3_pro = any(model_lc.startswith(p) for p in _GEMINI3_PRO_PREFIXES) + _is_pro_thinking_only = any( + model_lc == p or model_lc.startswith(p + "-") + for p in _PRO_THINKING_PREFIXES + ) + effort_lc = (reasoning_effort or "").strip().lower() + if not is_image_model_strict and is_gemini3_thinking: + # Gemini 3.x thinkingLevel matrix: + # 3.1+ Pro: low/medium/high + # 3 Pro: low/high (deprecated 2026-03-09) + # 3.x Flash*: minimal/low/medium/high + # Coerce minimal->low on Pro; medium->high on legacy 3-Pro. + _G3_LEVELS = {"minimal", "low", "medium", "high"} + level: Optional[str] = None + if effort_lc in ("none", "off"): + level = "low" if is_gemini3_pro else "minimal" + elif effort_lc == "max": + level = "high" + elif effort_lc in _G3_LEVELS: + # Coerce legacy 3-Pro (low/high only) inputs. + _is_legacy_gemini3_pro = model_lc.startswith( + ("gemini-3-pro-preview", "gemini-3-pro") + ) and not model_lc.startswith(("gemini-3.1-pro", "gemini-3.5-pro")) + if is_gemini3_pro and effort_lc == "minimal": + level = "low" + elif _is_legacy_gemini3_pro and effort_lc == "medium": + level = "high" + else: + level = effort_lc + elif enable_thinking is True: + level = "high" + elif enable_thinking is False: + level = "low" if is_gemini3_pro else "minimal" + if level is not None: + gen_config["thinkingConfig"] = {"thinkingLevel": level} + elif not is_image_model_strict: + # Gemini 2.5 / older: thinkingBudget int. Effort -> budget + # mirrors the OpenAI minimal/low/medium/high ladder so the + # existing frontend picker maps cleanly. + # NOTE: gemini-2.5-flash-lite rejects positive budgets below + # 512 with HTTP 400, so minimal=512 sits at that floor. + _EFFORT_TO_BUDGET: dict[str, int] = { + "minimal": 512, + "low": 2048, + "medium": 8192, + "high": 24576, + "xhigh": -1, + "max": -1, + } + thinking_budget: Optional[int] = None + if effort_lc == "none" or enable_thinking is False: + # Pro-tier 2.5 rejects budget=0 (400 "only works in + # thinking mode"), so coerce to a small positive value. + thinking_budget = 128 if _is_pro_thinking_only else 0 + elif effort_lc in _EFFORT_TO_BUDGET: + thinking_budget = _EFFORT_TO_BUDGET[effort_lc] + elif enable_thinking is True: + thinking_budget = -1 + if thinking_budget is not None: + gen_config["thinkingConfig"] = { + "thinkingBudget": thinking_budget, + } + + if gen_config: + body["generationConfig"] = gen_config + + # Hosted tools: googleSearch (grounding) and codeExecution. + # Image-mode rejects codeExecution; only Gemini 3 image models + # accept googleSearch. + # https://ai.google.dev/gemini-api/docs/grounding + # https://ai.google.dev/gemini-api/docs/code-execution + def _gemini_image_model_allows_google_search(_m: str) -> bool: + return ( + _m.startswith("gemini-3-pro-image") + or _m.startswith("gemini-3.1-flash-image") + or _m.startswith("nano-banana-pro") + or _m.startswith("nano-banana-2") + ) + + google_search_allowed = ( + not is_image_model_strict + or _gemini_image_model_allows_google_search(model_lc) + ) + code_execution_allowed = not is_image_model_strict + text_tools_allowed = not is_image_model_strict + # tool_choice="none" / forced-function suppresses hosted builtins + # too, matching the Anthropic / OpenRouter gates. + tools_array: list[dict[str, Any]] = [] + if ( + _hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + and google_search_allowed + ): + tools_array.append({"googleSearch": {}}) + if ( + _hosted_builtins_allowed + and enabled_tools + and "code_execution" in enabled_tools + and code_execution_allowed + ): + tools_array.append({"codeExecution": {}}) + # OpenAI-style function declarations -> Gemini functionDeclarations. + # https://ai.google.dev/gemini-api/docs/function-calling#step_1 + # Gemini's Schema accepts only the OpenAPI 3.0 subset documented + # at https://ai.google.dev/api/caching#Schema; OpenAI's strict + # tool definitions routinely include `additionalProperties`, + # `$schema`, `$defs`, `strict`, `examples`, and similar keys + # which 400 the request as INVALID_ARGUMENT. Strip them + # recursively before forwarding. + _GEMINI_ALLOWED_SCHEMA_KEYS = frozenset( + { + "type", + "format", + "title", + "description", + "nullable", + "enum", + "maxItems", + "minItems", + "properties", + "required", + "minProperties", + "maxProperties", + "items", + "minimum", + "maximum", + "minLength", + "maxLength", + "pattern", + "default", + "anyOf", + "propertyOrdering", + } + ) + + def _resolve_local_schema_ref( + root: Optional[dict[str, Any]], ref: str + ) -> Optional[Any]: + # Walk a `#/foo/bar` JSON pointer against the schema root. + # Returns None if the pointer doesn't resolve to a dict, so + # the caller can fall back to the unresolved node. + if not isinstance(root, dict) or not isinstance(ref, str): + return None + if not ref.startswith("#/"): + return None + node: Any = root + for raw_part in ref[2:].split("/"): + if not raw_part: + continue + part = raw_part.replace("~1", "/").replace("~0", "~") + if not isinstance(node, dict) or part not in node: + return None + node = node[part] + return node + + def _sanitize_gemini_schema( + node: Any, + root: Optional[dict[str, Any]] = None, + _seen_refs: Optional[frozenset[str]] = None, + ) -> Any: + # Recursively filter to Gemini's OpenAPI 3.0 subset. At a + # Schema-keyword dict layer we drop keys not in the + # allowlist; under `properties` the keys are user-defined + # field names and the values are themselves Schemas; under + # `items` / `anyOf` the values are also Schemas. + # OpenAI strict tools commonly use JSON Schema's + # `"type": ["string", "null"]` form for nullable fields; + # Gemini's OpenAPI Schema uses `"type": "string"` plus + # `"nullable": true`. Translate that here. + if root is None and isinstance(node, dict): + root = node + if _seen_refs is None: + _seen_refs = frozenset() + if isinstance(node, dict): + # Pydantic / OpenAI strict tools commonly hoist nested + # object schemas into `$defs` and reference them via + # `{"$ref": "#/$defs/Address"}`. Gemini's OpenAPI subset + # has no $ref and drops anything not in the allowlist, + # so the referenced shape would vanish if we didn't + # inline it here. Recurse into the resolved target with + # local siblings overriding the reference (normal JSON + # Schema composition), guarding against ref cycles. + _ref = node.get("$ref") + if isinstance(_ref, str): + if _ref in _seen_refs: + return {} + _target = _resolve_local_schema_ref(root, _ref) + if isinstance(_target, dict): + _merged = { + **_target, + **{k: v for k, v in node.items() if k != "$ref"}, + } + return _sanitize_gemini_schema( + _merged, root, _seen_refs | {_ref} + ) + cleaned: dict[str, Any] = {} + _nullable_from_union = False + _flattened_type: Optional[str] = None + _union_any_of: Optional[list[dict[str, Any]]] = None + _raw_type = node.get("type") + if isinstance(_raw_type, list): + _non_null = [t for t in _raw_type if t != "null"] + if len(_non_null) < len(_raw_type): + _nullable_from_union = True + if len(_non_null) == 1: + _flattened_type = _non_null[0] + elif len(_non_null) > 1: + # Preserve multi-type unions as anyOf; flattening + # to the first non-null type silently drops the + # other branches and changes the tool contract. + _union_any_of = [ + {"type": _t} for _t in _non_null if isinstance(_t, str) + ] + for _k, _v in node.items(): + if _k == "type" and isinstance(_v, list): + # Handled below via _flattened_type. + continue + if _k not in _GEMINI_ALLOWED_SCHEMA_KEYS: + continue + if _k == "properties" and isinstance(_v, dict): + cleaned[_k] = { + _name: _sanitize_gemini_schema(_subschema, root, _seen_refs) + for _name, _subschema in _v.items() + } + elif _k == "items": + cleaned[_k] = _sanitize_gemini_schema(_v, root, _seen_refs) + elif _k == "anyOf" and isinstance(_v, list): + # Optional[X] / Union[A, B, None]: Pydantic emits + # `anyOf: [..., {"type":"null"}]`. Gemini's + # OpenAPI subset rejects `"type": "null"` inside + # anyOf, so drop the null variant and surface it + # via `nullable: true`. If exactly one non-null + # branch remains, collapse it inline; otherwise + # keep the slim anyOf and mark the field + # nullable. + _saw_null = any( + isinstance(_entry, dict) and _entry.get("type") == "null" + for _entry in _v + ) + _non_null_entries = [ + _entry + for _entry in _v + if not ( + isinstance(_entry, dict) + and _entry.get("type") == "null" + ) + ] + if len(_non_null_entries) == 1 and _saw_null: + _inner = _sanitize_gemini_schema( + _non_null_entries[0], root, _seen_refs + ) + if isinstance(_inner, dict): + for _ik, _iv in _inner.items(): + cleaned.setdefault(_ik, _iv) + cleaned.setdefault("nullable", True) + else: + cleaned[_k] = [ + _sanitize_gemini_schema(_entry, root, _seen_refs) + for _entry in _non_null_entries + ] + if _saw_null: + cleaned.setdefault("nullable", True) + elif _k in ("required", "enum", "propertyOrdering"): + # Lists of plain strings; copy verbatim. + cleaned[_k] = _v + else: + cleaned[_k] = _v + if _union_any_of is not None and "anyOf" not in cleaned: + cleaned["anyOf"] = [ + _sanitize_gemini_schema(_s, root, _seen_refs) + for _s in _union_any_of + ] + elif _flattened_type is not None: + cleaned["type"] = _flattened_type + if _nullable_from_union and "nullable" not in cleaned: + cleaned["nullable"] = True + return cleaned + return node + + function_declarations: list[dict[str, Any]] = [] + if tools and text_tools_allowed and not _tool_choice_disabled: + for _tool in tools: + if not isinstance(_tool, dict) or _tool.get("type") != "function": + continue + _fn = _tool.get("function") + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _decl: dict[str, Any] = { + "name": _fn["name"], + "description": _fn.get("description") or "", + } + _params = _fn.get("parameters") + if isinstance(_params, dict): + _decl["parameters"] = _sanitize_gemini_schema(_params) + function_declarations.append(_decl) + if function_declarations: + tools_array.append({"functionDeclarations": function_declarations}) + if tools_array: + body["tools"] = tools_array + # Tool-choice mapping: OpenAI "auto"/"none"/"required"/{name=...} + # -> Gemini toolConfig.functionCallingConfig.mode + allowedFunctionNames. + if tool_choice is not None and function_declarations and text_tools_allowed: + _mode: Optional[str] = None + _allowed: Optional[list[str]] = None + if isinstance(tool_choice, str): + _tc_lc = tool_choice.strip().lower() + if _tc_lc == "auto": + _mode = "AUTO" + elif _tc_lc == "none": + _mode = "NONE" + elif _tc_lc in ("required", "any"): + _mode = "ANY" + elif ( + isinstance(tool_choice, dict) and tool_choice.get("type") == "function" + ): + _fn_pick = tool_choice.get("function") or {} + _name = _fn_pick.get("name") if isinstance(_fn_pick, dict) else None + if isinstance(_name, str) and _name: + _mode = "ANY" + _allowed = [_name] + if _mode is not None: + _fcc: dict[str, Any] = {"mode": _mode} + if _allowed: + _fcc["allowedFunctionNames"] = _allowed + body["toolConfig"] = {"functionCallingConfig": _fcc} + + # Prompt caching. The Gemini caching contract is "create a + # CachedContent resource, then pass its name on + # `cachedContent`". The cache itself is created out of band by + # the caller via POST /cachedContents; here we forward an + # explicit cache id when the dispatcher hands us one (a string + # value on enable_prompt_caching means "use this cache name"). + # https://ai.google.dev/gemini-api/docs/caching + if isinstance(enable_prompt_caching, str) and enable_prompt_caching: + body["cachedContent"] = enable_prompt_caching + + # Model id is already validated at the top of _stream_gemini so + # we never reach a path-traversed URL segment here. + url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse" + completion_id = f"chatcmpl-gemini-{model.replace('/', '-')}" + + logger.info( + "Proxying Gemini streamGenerateContent to %s (model=%s, " + "tools=%s, image=%s)", + url, + model, + [list(t.keys())[0] for t in tools_array] if tools_array else [], + is_image_model, + ) + + def _emit_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": None, + } + ], + "_toolEvent": payload, + } + return f"data: {_json.dumps(chunk)}" + + def _text_chunk( + text: str, extra_content: Optional[dict[str, Any]] = None + ) -> str: + delta: dict[str, Any] = {"content": text} + if extra_content: + delta["extra_content"] = extra_content + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": delta, + "finish_reason": None, + } + ], + } + return f"data: {_json.dumps(chunk)}" + + def _gemini_part_extra(part: dict[str, Any]) -> Optional[dict[str, Any]]: + """Return ``{"google": {"thought_signature": ...}}`` when the + Gemini stream part carries a `thoughtSignature` we need to + replay on a follow-up turn (Gemini 3 image editing + tool + contexts both require an exact signature echo).""" + sig = part.get("thoughtSignature") or part.get("thought_signature") + if isinstance(sig, str) and sig: + return {"google": {"thought_signature": sig}} + return None + + # Gemini finish reasons -> OpenAI vocabulary. Reference: + # https://ai.google.dev/api/rest/v1beta/Candidate#FinishReason + _finish_reason_map: dict[str, Optional[str]] = { + "STOP": "stop", + "MAX_TOKENS": "length", + "SAFETY": "content_filter", + "RECITATION": "content_filter", + "PROHIBITED_CONTENT": "content_filter", + "BLOCKLIST": "content_filter", + "MALFORMED_FUNCTION_CALL": "stop", + "OTHER": "stop", + "FINISH_REASON_UNSPECIFIED": None, + } + + last_usage: Optional[dict[str, Any]] = None + emitted_function_call_ids: set[str] = set() + # True once any Gemini functionCall part has been emitted so the + # final finish_reason swaps STOP -> tool_calls (matches the + # OpenAI Chat Completions contract; an OAI client that sees a + # tool_calls delta followed by finish_reason="stop" never + # executes the tool). + emitted_any_function_call = False + # web_search_active drives the tool_start / tool_end envelope. + # Track on whether `googleSearch` was actually forwarded above, + # not the raw caller intent -- image-mode requests filter the + # tool out, and emitting a phantom "search complete" card on a + # turn where Gemini was never told to search confuses the UI. + web_search_active = any("googleSearch" in t for t in tools_array) + web_search_tool_id = "gemini_web_search" + web_search_tool_started = False + web_search_tool_ended = False + web_search_citations: list[dict[str, str]] = [] + # Tracks the tool_call_id minted on the most recent + # executableCode part so the matching codeExecutionResult can + # close out the same envelope. None between rounds. + gemini_code_exec_pending_id: Optional[str] = None + # The most recently emitted code_execution id + result text. Kept + # *after* the tool_end so a following inline image (matplotlib + # plot rendered by codeExecution) can attach to the same card + # via a `__IMAGES__:` marker instead of spawning a separate + # image_generation event. + last_code_exec_tool_id: Optional[str] = None + last_code_exec_result_text: str = "" + + try: + async with _http_client.stream( + "POST", + url, + json = body, + headers = self._auth_headers(), + timeout = self._stream_timeout, + ) as response: + if response.status_code != 200: + error_body = await response.aread() + error_text = error_body.decode("utf-8", errors = "replace") + logger.error( + "Gemini returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + return + + if web_search_active: + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": web_search_tool_id, + "arguments": {}, + } + ) + web_search_tool_started = True + + # NOTE: same manual __anext__ loop pattern as the other + # streaming helpers (see stream_chat_completion for the + # Python 3.13 + httpcore 1.0.x GeneratorExit ordering). + lines_gen = response.aiter_lines().__aiter__() + final_finish_reason: Optional[str] = None + try: + while True: + try: + line = await lines_gen.__anext__() + except StopAsyncIteration: + break + if not line.strip(): + continue + if not line.startswith("data:"): + continue + data_str = line[len("data:") :].strip() + if not data_str or data_str == "[DONE]": + continue + try: + event = _json.loads(data_str) + except Exception: + logger.warning( + "Gemini: failed to parse SSE chunk: %s", + data_str[:200], + ) + continue + if not isinstance(event, dict): + continue + + # Latch usageMetadata across deltas -- the final + # fragment carries the complete totals. + usage_meta = event.get("usageMetadata") + if isinstance(usage_meta, dict): + last_usage = usage_meta + + # Prompt-level safety block: Gemini ships zero + # candidates plus a `promptFeedback.blockReason` + # (e.g. SAFETY). The downstream OAI client would + # otherwise see an empty successful assistant + # response. Surface as a content_filter error + # event so the UI can render the block reason. + prompt_feedback = event.get("promptFeedback") + if isinstance(prompt_feedback, dict) and prompt_feedback.get( + "blockReason" + ): + block_reason = str(prompt_feedback.get("blockReason")) + # Close out the synthetic web_search start so + # the UI does not show a spinner stuck on + # "searching..." after the error toast lands. + if ( + web_search_active + and web_search_tool_started + and not web_search_tool_ended + ): + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": ( + "(search aborted: Gemini blocked " + f"prompt: {block_reason})" + ), + } + ) + web_search_tool_ended = True + yield _error_sse_line( + 400, + f"Gemini blocked prompt: {block_reason}", + self.provider_type, + ) + return + + candidates = event.get("candidates") or [] + if not isinstance(candidates, list): + continue + for cand in candidates: + if not isinstance(cand, dict): + continue + # Citations / grounding metadata. + # `groundingMetadata.groundingChunks[].web` + # carries `uri` + `title`. Collect for the + # tool_end emission at stream close. + gm = cand.get("groundingMetadata") + if isinstance(gm, dict) and web_search_active: + chunks_list = gm.get("groundingChunks") or [] + if isinstance(chunks_list, list): + for ch in chunks_list: + if not isinstance(ch, dict): + continue + web = ch.get("web") or {} + if not isinstance(web, dict): + continue + u = web.get("uri") or "" + if not u or not isinstance(u, str): + continue + if any( + c["url"] == u for c in web_search_citations + ): + continue + web_search_citations.append( + { + "url": u, + "title": (web.get("title") or u), + "snippet": "", + } + ) + + content_obj = cand.get("content") or {} + parts = ( + content_obj.get("parts") + if isinstance(content_obj, dict) + else None + ) + if isinstance(parts, list): + for part in parts: + if not isinstance(part, dict): + continue + # Text delta. Stow part-level + # `thoughtSignature` on the delta so + # Gemini 3 turns that need an exact + # signature echo round-trip cleanly. + text = part.get("text") + _part_extra = _gemini_part_extra(part) + if isinstance(text, str) and text: + yield _text_chunk( + text, + extra_content = _part_extra, + ) + elif _part_extra is not None and not any( + k in part + for k in ( + "functionCall", + "executableCode", + "codeExecutionResult", + "inlineData", + ) + ): + # Empty-content part carrying a + # thoughtSignature: emit an empty delta + # so the signature is preserved. + yield _text_chunk( + "", + extra_content = _part_extra, + ) + # functionCall -> OpenAI tool_calls + # delta envelope. + fc = part.get("functionCall") + if isinstance(fc, dict): + fc_name = fc.get("name") or "" + fc_args = fc.get("args") or {} + fc_id = ( + fc.get("id") + or f"call_{fc_name}_{time.time_ns()}" + ) + if fc_id in emitted_function_call_ids: + continue + emitted_function_call_ids.add(fc_id) + # Each distinct functionCall in an + # assistant turn needs its own + # tool_calls[*].index. Consumers + # that reassemble tool_calls by + # index collapse all calls onto + # the same slot when this is + # hardcoded to 0, breaking + # parallel/multi-tool turns. + tc_index = len(emitted_function_call_ids) - 1 + tool_call_delta: dict[str, Any] = { + "index": tc_index, + "id": fc_id, + "type": "function", + "function": { + "name": fc_name, + "arguments": _json.dumps(fc_args), + }, + } + # Gemini 3 function-calling: the + # part-level `thoughtSignature` + # must be echoed back on the + # next turn or the model rejects + # the tool-result envelope. Stow + # it on `extra_content.google` + # so the frontend can persist it + # and our outbound translator + # (below) can replay it. + thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + if isinstance(thought_sig, str) and thought_sig: + tool_call_delta["extra_content"] = { + "google": { + "thought_signature": thought_sig, + } + } + emitted_any_function_call = True + tool_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [tool_call_delta] + }, + "finish_reason": None, + } + ], + } + yield f"data: {_json.dumps(tool_chunk)}" + # executableCode + codeExecutionResult + # parts surface as the standard + # code_execution tool_start/tool_end + # envelope (same shape OpenAI and + # Anthropic emit) so the chat + # adapter can render Gemini sandbox + # output through CodeExecutionToolUI. + # https://ai.google.dev/gemini-api/docs/code-execution + exec_code = part.get("executableCode") + if isinstance(exec_code, dict): + code_str = exec_code.get("code") or "" + if code_str: + code_tool_id = ( + exec_code.get("id") + or f"gemini_code_exec_{time.time_ns()}" + ) + gemini_code_exec_pending_id = code_tool_id + # Stow the raw Gemini part so + # follow-up turns can replay + # the native `executableCode` + # (Gemini rejects a generic + # functionCall echo for code + # execution history). + _exec_thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + # Per-part thoughtSignature stays + # bound to its own part (Gemini 3 + # rejects shared signatures). + _exec_part_entry: dict[str, Any] = { + "executableCode": exec_code, + } + if ( + isinstance(_exec_thought_sig, str) + and _exec_thought_sig + ): + _exec_part_entry["thoughtSignature"] = ( + _exec_thought_sig + ) + _exec_native: dict[str, Any] = { + "parts": [_exec_part_entry], + } + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": code_tool_id, + "arguments": { + "kind": "code_execution", + "language": ( + ( + exec_code.get( + "language" + ) + or "PYTHON" + ).lower() + ), + "code": code_str, + "google": { + "native_part": _exec_native, + }, + }, + } + ) + exec_result = part.get("codeExecutionResult") + if isinstance(exec_result, dict): + outcome = exec_result.get("outcome") or "" + output = exec_result.get("output") or "" + # Gemini returns + # OUTCOME_OK / OUTCOME_FAILED / + # OUTCOME_DEADLINE_EXCEEDED. Treat + # non-OK outcomes as stderr so the + # UI surfaces the error. + if outcome and outcome != "OUTCOME_OK": + result_text = ( + f"[{outcome}]\n{output}".rstrip() + ) + else: + result_text = output + # Pair tool_end with the most recent + # executableCode tool_start; fall back + # to exec_result.id then a fresh id. + pair_id = ( + gemini_code_exec_pending_id + or exec_result.get("id") + or f"gemini_code_exec_{time.time_ns()}" + ) + if gemini_code_exec_pending_id is None: + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": pair_id, + "arguments": { + "kind": "code_execution", + "code": "", + }, + } + ) + _result_thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + _result_part_entry: dict[str, Any] = { + "codeExecutionResult": exec_result, + } + if ( + isinstance(_result_thought_sig, str) + and _result_thought_sig + ): + _result_part_entry["thoughtSignature"] = ( + _result_thought_sig + ) + _result_native: dict[str, Any] = { + "parts": [_result_part_entry], + } + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": pair_id, + "result": result_text, + "google": { + "native_part": _result_native, + }, + } + ) + last_code_exec_tool_id = pair_id + last_code_exec_result_text = result_text + gemini_code_exec_pending_id = None + # inlineData: either a Nano Banana + # generation (own card) or a sandbox + # plot attached to the code_execution + # card via the __IMAGES__: marker. + inline = part.get("inlineData") + if isinstance(inline, dict): + b64 = inline.get("data") or "" + mime = inline.get("mimeType") or "image/png" + if b64: + image_uri = f"data:{mime};base64,{b64}" + attached_to_code_exec = ( + not is_image_model + and last_code_exec_tool_id is not None + and bool(enabled_tools) + and "code_execution" + in (enabled_tools or []) + ) + if attached_to_code_exec: + updated_result = ( + last_code_exec_result_text + + "\n__IMAGES__:" + + _json.dumps([image_uri]) + ) + # Stow inlineData so a follow-up + # turn can replay the plot with + # its per-part thoughtSignature. + _plot_thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + _plot_part_entry: dict[str, Any] = { + "inlineData": { + "mimeType": mime, + "data": b64, + }, + } + if ( + isinstance(_plot_thought_sig, str) + and _plot_thought_sig + ): + _plot_part_entry[ + "thoughtSignature" + ] = _plot_thought_sig + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": ( + last_code_exec_tool_id + ), + "result": updated_result, + "google": { + "native_part": { + "parts": [ + _plot_part_entry + ], + }, + }, + } + ) + last_code_exec_result_text = ( + updated_result + ) + else: + img_id = f"img_{time.time_ns()}" + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "image_generation", + "tool_call_id": img_id, + "arguments": { + "kind": "image", + "prompt": "", + }, + } + ) + # Gemini 3 image edit needs + # the prior thoughtSignature + # echoed on the inline image part. + _img_thought_sig = part.get( + "thoughtSignature" + ) or part.get("thought_signature") + _img_tool_end: dict[str, Any] = { + "type": "tool_end", + "tool_call_id": img_id, + "result": "", + "image_b64": b64, + "image_mime": mime, + } + # Stow inlineData so multi-turn + # edits replay the original + # image as native history. + _img_part_entry: dict[str, Any] = { + "inlineData": { + "mimeType": mime, + "data": b64, + }, + } + if ( + isinstance(_img_thought_sig, str) + and _img_thought_sig + ): + _img_part_entry[ + "thoughtSignature" + ] = _img_thought_sig + _img_native: dict[str, Any] = { + "parts": [_img_part_entry], + } + _img_google: dict[str, Any] = { + "native_part": _img_native, + } + if ( + isinstance(_img_thought_sig, str) + and _img_thought_sig + ): + _img_google["thought_signature"] = ( + _img_thought_sig + ) + _img_tool_end["google"] = _img_google + yield _emit_tool_event(_img_tool_end) + finish_reason = cand.get("finishReason") + if isinstance(finish_reason, str): + mapped = _finish_reason_map.get(finish_reason, "stop") + if mapped is not None: + final_finish_reason = mapped + + # End-of-stream emission order: web_search tool_end + # (with citations) -> finish_reason chunk -> usage + # chunk -> [DONE]. Matches the Anthropic / OpenAI + # helpers' contract so the frontend handler does + # not need provider-specific ordering knowledge. + if ( + web_search_active + and web_search_tool_started + and not web_search_tool_ended + ): + blocks: list[str] = [] + for cit in web_search_citations: + line_out = f"Title: {cit['title']}\nURL: {cit['url']}" + if cit.get("snippet"): + line_out += f"\nSnippet: {cit['snippet']}" + blocks.append(line_out) + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": ( + "\n---\n".join(blocks) + if blocks + else "(search complete)" + ), + } + ) + web_search_tool_ended = True + + if final_finish_reason: + # OpenAI clients trigger tool execution when + # finish_reason="tool_calls". Gemini emits + # "STOP" even when the turn was a pure + # functionCall request, so override after the + # fact to match the OAI contract. + if emitted_any_function_call and final_finish_reason == "stop": + final_finish_reason = "tool_calls" + finish_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": final_finish_reason, + } + ], + } + yield f"data: {_json.dumps(finish_chunk)}" + + # Map Gemini usageMetadata onto OpenAI include_usage. + # thoughtsTokenCount is billed output too — fold it in + # so cost calculators don't undercount. + if isinstance(last_usage, dict): + thought_tokens = last_usage.get("thoughtsTokenCount") or 0 + candidate_tokens = last_usage.get("candidatesTokenCount") or 0 + prompt_tokens = last_usage.get("promptTokenCount") or 0 + # Gemini bills tool-call prompt slices separately + # via `toolUsePromptTokenCount`. Fold into input + # so total_tokens does not undercount tool turns. + tool_use_prompt_tokens = ( + last_usage.get("toolUsePromptTokenCount") or 0 + ) + translated_usage = { + "input_tokens": prompt_tokens + tool_use_prompt_tokens, + "output_tokens": candidate_tokens + thought_tokens, + "input_tokens_details": { + "cached_tokens": ( + last_usage.get("cachedContentTokenCount") or 0 + ), + "tool_use_prompt_tokens": tool_use_prompt_tokens, + }, + "output_tokens_details": { + "reasoning_tokens": thought_tokens, + }, + } + usage_line = _build_usage_chunk( + completion_id, "openai", translated_usage + ) + if usage_line: + yield usage_line + + yield "data: [DONE]" + finally: + # Close response first so lines_gen.aclose() becomes + # a no-op (avoids the httpcore 1.0 GeneratorExit + # path and the aclose-never-awaited RuntimeWarning). + await response.aclose() + await lines_gen.aclose() + + except httpx.ConnectError as exc: + logger.error("Connection error to %s: %s", self.provider_type, exc) + if web_search_tool_started and not web_search_tool_ended: + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": f"(search aborted: connection error: {exc})", + } + ) + web_search_tool_ended = True + yield _error_sse_line( + 502, + f"Failed to connect to {self.provider_type}: {exc}", + self.provider_type, + ) + except httpx.ReadTimeout as exc: + logger.error("Read timeout from %s: %s", self.provider_type, exc) + if web_search_tool_started and not web_search_tool_ended: + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": "(search aborted: read timeout)", + } + ) + web_search_tool_ended = True + yield _error_sse_line( + 504, + f"Timeout waiting for {self.provider_type} response", + self.provider_type, + ) + except httpx.HTTPError as exc: + logger.error("HTTP error from %s: %s", self.provider_type, exc) + if web_search_tool_started and not web_search_tool_ended: + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": web_search_tool_id, + "result": f"(search aborted: transport error: {exc})", + } + ) + web_search_tool_ended = True + yield _error_sse_line( + 502, + f"Error communicating with {self.provider_type}: {exc}", + self.provider_type, + ) + async def _stream_openai_responses( self, messages: list[dict[str, Any]], @@ -2893,6 +4949,8 @@ class ExternalProviderClient: enable_prompt_caching: Optional[bool] = None, openai_code_exec_container_id: Optional[str] = None, compaction_threshold: Optional[int] = None, + tools: Optional[list[dict[str, Any]]] = None, + tool_choice: Optional[Any] = None, ) -> AsyncGenerator[str, None]: """ Call OpenAI's /v1/responses endpoint and translate its SSE stream back @@ -2916,6 +4974,12 @@ class ExternalProviderClient: # translate user/assistant messages into the Responses input shape. instructions_parts: list[str] = [] input_items: list[dict[str, Any]] = [] + # When we drop a server-side builtin `function_call` here, the + # matching `role="tool"` follow-up must also be dropped -- + # otherwise the outbound body contains an orphan + # `function_call_output` with no matching `function_call`, which + # OpenAI Responses can reject or mis-associate. + skipped_server_builtin_call_ids: set[str] = set() openai_replay_items: list[dict[str, Any]] = [] previous_response_id: Optional[str] = None for msg in messages: @@ -2932,6 +4996,126 @@ class ExternalProviderClient: instructions_parts.append(part["text"]) continue + # OpenAI Responses uses item-shape history for function + # calling: assistant turns that invoked user tools must + # serialize each call as a `function_call` input item, and + # each role="tool" follow-up as a `function_call_output` + # item keyed by the matching `call_id`. Without this the + # second turn after a function call sends Chat Completions + # shape and Responses 400s the request. + if role == "tool": + _call_id = msg.get("tool_call_id") or "" + # If the matching assistant `function_call` was a + # server-side builtin we already dropped, drop the + # follow-up too to avoid emitting an orphan + # `function_call_output`. + if _call_id and _call_id in skipped_server_builtin_call_ids: + continue + if isinstance(content, list): + _flat_parts: list[str] = [] + for part in content: + if part.get("type") == "text" and part.get("text"): + _flat_parts.append(part["text"]) + _output_text = "".join(_flat_parts) + else: + _output_text = content if isinstance(content, str) else "" + if _call_id: + input_items.append( + { + "type": "function_call_output", + "call_id": _call_id, + "output": _output_text, + } + ) + continue + + # Assistant turns that returned tool_calls translate each + # call as a `function_call` item (carrying name + JSON + # arguments + call_id). Skip builtin server-side cards + # (canonical builtin name + `args._server_tool` marker) + # which never round-trip as user functions. We require both + # checks so a user function literally named `_server_tool` + # in its argument schema is not dropped. + _tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if role == "assistant" and isinstance(_tool_calls, list): + # Preserve the prior `response.output` ordering: the + # model's text precedes its function_call items, and + # the matching role=tool follow-up arrives AFTER the + # call. Without this guard, history replay puts + # function_call -> assistant text -> function_call_output, + # which can put the tool output after an unrelated + # assistant message and confuse multi-turn function + # calling. + if isinstance(content, str) and content: + input_items.append({"role": "assistant", "content": content}) + elif isinstance(content, list): + _asst_parts: list[dict[str, Any]] = [] + for _part in content: + if not isinstance(_part, dict): + continue + _pt = _part.get("type") + if _pt == "text" and _part.get("text"): + _asst_parts.append( + { + "type": "input_text", + "text": _part.get("text", ""), + } + ) + elif _pt == "image_url": + _u = _part.get("image_url", {}).get("url", "") + if _u: + _asst_parts.append( + {"type": "input_image", "image_url": _u} + ) + if _asst_parts: + input_items.append( + {"role": "assistant", "content": _asst_parts} + ) + + for _tc in _tool_calls: + if not isinstance(_tc, dict): + continue + _fn = _tc.get("function") or {} + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _args_raw = _fn.get("arguments") or "" + if not isinstance(_args_raw, str): + try: + _args_raw = _json.dumps(_args_raw) + except Exception: + _args_raw = "" + _fn_name_lc = (_fn.get("name") or "").lower() + _is_server_builtin = False + if _fn_name_lc in _SERVER_SIDE_BUILTIN_TOOL_NAMES: + try: + _args_obj = _json.loads(_args_raw) if _args_raw else {} + except Exception: + _args_obj = None + if isinstance(_args_obj, dict): + if _args_obj.get("_server_tool") is True: + _is_server_builtin = True + else: + _g = _args_obj.get("google") + if isinstance(_g, dict) and isinstance( + _g.get("native_part"), dict + ): + _is_server_builtin = True + _call_id_out = _tc.get("id") or f"call_{time.time_ns()}" + if _is_server_builtin: + skipped_server_builtin_call_ids.add(_call_id_out) + continue + input_items.append( + { + "type": "function_call", + "call_id": _call_id_out, + "name": _fn["name"], + "arguments": _args_raw, + } + ) + # Assistant text already emitted above (in order) so we + # don't fall through to the generic content branches. + continue + if isinstance(content, str): input_items.append({"role": role, "content": content}) continue @@ -3121,44 +5305,15 @@ class ExternalProviderClient: if max_tokens is not None: body["max_output_tokens"] = max_tokens - # Prompt caching on /v1/responses is automatic and free, but the - # default in-memory policy only survives ~5-10 min of inactivity - # (up to ~1 hr). Opt into the 24-hour retention policy so a chat - # left idle overnight still hits the cache on the next turn. - # Pricing is identical to in_memory per OpenAI's docs. - # - # Gated on the base URL because ollama / llama.cpp / "custom" - # presets all collapse to provider_type="openai" in - # toExternalBackendProviderType, so they also land in this - # helper. Those servers expose /v1/responses-shaped routes in - # some configurations but don't implement - # prompt_cache_retention — sending the field unconditionally - # would 400 them. Match the public OpenAI host strictly so the - # field only goes to OpenAI cloud. Studio's openai model picker - # is registry-scoped to gpt-5.x / o3 / gpt-4.5, all of which - # accept this parameter (gpt-5.5+ already defaults to "24h" and - # rejects "in_memory", so it's a safe no-op there). - # OpenAI-family cloud: api.openai.com OR Azure OpenAI Foundry - # (*.openai.azure.com). Both expose the same Responses-API - # extensions used below -- prompt_cache_retention, - # context_management compaction, container shell tool -- so - # treat them uniformly. Non-cloud OpenAI-compatible servers - # (ollama / llama.cpp / vLLM / "custom" preset) hit /v1/responses - # without these extensions and would 400 on the unknown body - # fields, so they intentionally fall outside this gate. + # Opt into 24h prompt-cache retention (free, vs the default + # ~5-10 min). Gated on the OpenAI cloud host because ollama / + # llama.cpp / "custom" presets reach this code path too and + # would 400 on the unknown field. if is_openai_cloud and enable_prompt_caching is not False: body["prompt_cache_retention"] = "24h" - # OpenAI server-side context compaction — see - # https://developers.openai.com/api/docs/guides/compaction - # When `compaction_threshold` is provided on a cloud OpenAI - # request, attach `context_management: [{type:"compaction", - # compact_threshold:N}]` so the API runs server-side - # compaction when the rendered prompt crosses the threshold. - # No beta header is required; no dated version pin. The field - # is silently dropped for non-cloud backends because ollama / - # llama.cpp / "custom" presets land in this helper and would - # 400 on an unknown body field. + # Server-side context compaction (OpenAI cloud only). + # https://developers.openai.com/api/docs/guides/compaction if ( is_openai_cloud and compaction_threshold is not None @@ -3171,55 +5326,92 @@ class ExternalProviderClient: } ] - # OpenAI server-side tools — see - # https://developers.openai.com/api/docs/guides/tools - # https://developers.openai.com/api/docs/guides/tools-shell - # The frontend's Search/Code buttons map to the unified - # enabled_tools shorthand; translate that into the Responses-API - # tool schema. Other built-in tools (file_search, - # code_interpreter, image_generation, computer_use_preview) can - # be added with the same pattern when we surface their toggles. + # Map enabled_tools onto Responses-API server tools (cloud only; + # local OAI-compat backends 400 on these). + # https://developers.openai.com/api/docs/guides/tools code_execution_enabled_openai = bool( enabled_tools and "code_execution" in enabled_tools and is_openai_cloud ) - # OpenAI's image_generation tool is a Responses-API server tool. - # See https://developers.openai.com/api/docs/guides/tools-image-generation - # The model picks size / quality / background server-side and - # delegates rendering to a gpt-image-* family model; the result - # comes back inline as an `image_generation_call` output item - # with a base64 image. Available on every gpt-5.x family member - # plus gpt-4.1 / gpt-4o / o3 per the docs; restrict to cloud - # OpenAI because the local llama.cpp / ollama backends don't - # implement it and would 400. - image_generation_enabled_openai = image_generation_requested + image_generation_enabled_openai = bool( + enabled_tools and "image_generation" in enabled_tools and is_openai_cloud + ) def _openai_image_generation_tool() -> dict[str, Any]: tool: dict[str, Any] = {"type": "image_generation"} if image_generation_has_reference: - # OpenAI's Responses image tool defaults to `auto`. For - # Studio's explicit follow-up edit flow, force edit mode so - # the provider uses the previous response / call id as image - # context instead of treating the text as a fresh generation. + # Force edit mode so the prior call id is used as context. tool["action"] = "edit" return tool - if enabled_tools: - tools_array: list[dict[str, Any]] = [] - if "web_search" in enabled_tools: + # Translate Chat-Completions function tools into the Responses + # function-tool shape (flattened name/description/parameters). + responses_user_function_tools: list[dict[str, Any]] = [] + if tools: + for _tool in tools: + if not isinstance(_tool, dict) or _tool.get("type") != "function": + continue + _fn = _tool.get("function") + if not isinstance(_fn, dict) or not _fn.get("name"): + continue + _entry: dict[str, Any] = { + "type": "function", + "name": _fn["name"], + } + if _fn.get("description"): + _entry["description"] = _fn["description"] + if isinstance(_fn.get("parameters"), dict): + _entry["parameters"] = _fn["parameters"] + responses_user_function_tools.append(_entry) + + # Translate tool_choice into the Responses shape. + _responses_tc_string: Optional[str] = None + if isinstance(tool_choice, str): + _tc_lc = tool_choice.strip().lower() + if _tc_lc in ("auto", "none", "required"): + _responses_tc_string = _tc_lc + responses_tool_choice: Optional[Any] = None + _has_responses_tools = bool(enabled_tools or responses_user_function_tools) + if _responses_tc_string is not None and _has_responses_tools: + responses_tool_choice = _responses_tc_string + elif ( + tool_choice is not None + and responses_user_function_tools + and isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + ): + _fn_pick = tool_choice.get("function") or {} + _name = _fn_pick.get("name") if isinstance(_fn_pick, dict) else None + if isinstance(_name, str) and _name: + responses_tool_choice = {"type": "function", "name": _name} + + _responses_tool_choice_none = _responses_tc_string == "none" + # A pinned user function suppresses hosted builtins (privacy + + # billing), matching the Gemini / Anthropic / OpenRouter gates. + _responses_tool_choice_forced_function = ( + isinstance(tool_choice, dict) + and tool_choice.get("type") == "function" + and isinstance(tool_choice.get("function"), dict) + and bool(tool_choice["function"].get("name")) + ) + _responses_hosted_builtins_allowed = ( + not _responses_tool_choice_none + and not _responses_tool_choice_forced_function + ) + + if ( + enabled_tools or responses_user_function_tools + ) and not _responses_tool_choice_none: + tools_array: list[dict[str, Any]] = list(responses_user_function_tools) + if ( + _responses_hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + ): tools_array.append({"type": "web_search"}) - if code_execution_enabled_openai: - # `container_auto` lets OpenAI auto-create a fresh - # container per request; we capture the resulting - # container_id off the SSE stream and the chat-adapter - # persists it onto the thread record. Subsequent turns - # in the same thread pass it back as - # `openai_code_exec_container_id`, which we translate to - # `container_reference` here so the model sees - # filesystem state from prior turns. Container expires - # after ~20 min of inactivity per OpenAI's default - # policy — a stale id 400s, the chat-adapter clears it - # via container_invalidated, and the next turn falls - # back to auto-create. + if _responses_hosted_builtins_allowed and code_execution_enabled_openai: + # Reuse the thread's container so filesystem state + # persists; auto-create when there isn't one yet. Stale + # ids 400 and are cleared via container_invalidated. shell_env: dict[str, Any] if openai_code_exec_container_id: shell_env = { @@ -3229,10 +5421,12 @@ class ExternalProviderClient: else: shell_env = {"type": "container_auto"} tools_array.append({"type": "shell", "environment": shell_env}) - if image_generation_enabled_openai: + if _responses_hosted_builtins_allowed and image_generation_enabled_openai: tools_array.append(_openai_image_generation_tool()) if tools_array: body["tools"] = tools_array + if responses_tool_choice is not None: + body["tool_choice"] = responses_tool_choice url = f"{self.base_url}/responses" completion_id = f"chatcmpl-openai-{model.replace('/', '-')}" @@ -3246,11 +5440,19 @@ class ExternalProviderClient: first attempt. """ attempt_body = dict(body) - if enabled_tools: - tools_array_attempt: list[dict[str, Any]] = [] - if "web_search" in enabled_tools: + if ( + enabled_tools or responses_user_function_tools + ) and not _responses_tool_choice_none: + tools_array_attempt: list[dict[str, Any]] = list( + responses_user_function_tools + ) + if ( + _responses_hosted_builtins_allowed + and enabled_tools + and "web_search" in enabled_tools + ): tools_array_attempt.append({"type": "web_search"}) - if code_execution_enabled_openai: + if _responses_hosted_builtins_allowed and code_execution_enabled_openai: if container_id_for_this_attempt: env_attempt: dict[str, Any] = { "type": "container_reference", @@ -3261,12 +5463,17 @@ class ExternalProviderClient: tools_array_attempt.append( {"type": "shell", "environment": env_attempt} ) - if image_generation_enabled_openai: + if ( + _responses_hosted_builtins_allowed + and image_generation_enabled_openai + ): tools_array_attempt.append(_openai_image_generation_tool()) if tools_array_attempt: attempt_body["tools"] = tools_array_attempt else: attempt_body.pop("tools", None) + if responses_tool_choice is not None: + attempt_body["tool_choice"] = responses_tool_choice return attempt_body def _is_openai_container_expired_error(error_text: str) -> bool: @@ -3328,52 +5535,26 @@ class ExternalProviderClient: done_emitted = False reasoning_open = False reasoning_emitted = False - # Latched from response.completed / response.incomplete so - # the final log can surface input_tokens_details.cached_tokens — - # the field that proves prompt_cache_retention="24h" is - # actually hitting OpenAI's cache instead of recomputing - # the prefix every turn. + # Per-call function-tool indexing; distinct slots so + # parallel calls don't collide on delta.tool_calls[].index. + saw_function_call = False + function_call_index = 0 + # Latched from response.completed/incomplete; surfaces + # input_tokens_details.cached_tokens to prove cache hits. last_usage: Optional[dict[str, Any]] = None - # Per-call state for OpenAI's server-side web_search tool. Mapped - # back into our local _toolEvent shape so the existing chat-UI - # renderer surfaces web_search the same way it does for local - # tool calls: a "Searching…" tool-call card, then a `tool_end` - # carrying citations formatted as - # Title: …\nURL: …\nSnippet: …\n---\n… - # blocks (which the frontend's parseSourcesFromResult lifts - # into source content parts at end of stream). - # web_search_calls preserves insertion order so we can apply - # the aggregated citation list onto the *last* call's - # tool_end — that's the one the frontend's source-pill - # extraction reads (parseSourcesFromResult flatMaps every - # web_search result, so a single non-empty result is enough - # to surface all sources at message tail). - # OpenAI emits url_citation annotations on text deltas, not - # per call — there's no wire field linking a citation back - # to a specific search invocation. Hence the shared list. - # web_search_calls: { item_id -> {query} } + # web_search state. Citations are emitted on text deltas + # (not per call), so the aggregate list is shared and + # applied to the LAST web_search tool_end (parseSourcesFromResult + # flatmaps every call, one non-empty is enough). web_search_calls: dict[str, dict[str, Any]] = {} all_url_citations: list[dict[str, Any]] = [] - # Shell-tool (code execution) state. OpenAI emits - # `shell_call` items (model requesting a command list) - # paired with `shell_call_output` items (execution - # results). We mirror the Anthropic code-execution UX - # by emitting one `_toolEvent` tool_start per - # shell_call and one tool_end per shell_call_output; - # they're linked via `shell_call_output.call_id` - # matching `shell_call.id`. Items are independent of - # web_search (different keyed map). - # shell_calls: { call_id -> {commands, output} } + # shell_calls (code execution): { call_id -> {commands, output} }. + # shell_call <-> shell_call_output match by call_id; emit + # tool_start/tool_end like the Anthropic UX. shell_calls: dict[str, dict[str, Any]] = {} - # Container id captured from the response stream. When - # it differs from the inbound id, emit a synthetic - # `container_ready` _toolEvent so the frontend can - # persist it onto the thread record for the next turn. - # Where OpenAI surfaces it is documented loosely; we - # probe two known fields (response.container_id on - # response.completed, item.environment.container_id on - # shell_call output items) and latch the first one we - # see. + # Container id latched from response.container_id or + # item.environment.container_id; emit container_ready + # when it differs from the inbound id. latched_container_id: Optional[str] = None container_id_emitted = False current_openai_response_id: Optional[str] = None @@ -3453,6 +5634,7 @@ class ExternalProviderClient: return rendered def _emit_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -3775,17 +5957,9 @@ class ExternalProviderClient: f"ws_{len(web_search_calls)}" ) web_search_calls.setdefault(item_id, {"query": ""}) - # Shell-tool: register the call eagerly so - # the matching shell_call_output can link - # back even if `done` arrives out of order. - # Also probe for container_id on the - # environment field — when container_auto - # auto-creates one, this is the first place - # the new id might surface (OpenAI doesn't - # promise this in docs, but the field is - # cheap to scan and lets us emit - # container_ready earlier than - # response.completed). + # Register shell_call eagerly so out-of-order + # output links back. Probe env.container_id + # to emit container_ready before response.completed. if ( isinstance(item, dict) and item.get("type") == "shell_call" @@ -3978,24 +6152,9 @@ class ExternalProviderClient: } ) elif item.get("type") == "image_generation_call": - # OpenAI's image_generation tool returns - # a single output item with the base64 - # PNG/WebP/JPEG on `result` (sometimes - # `b64_json` depending on output_format). - # `revised_prompt` is what the gpt-image - # backbone actually used after refinement - # of the assistant's request. Emit - # tool_start + tool_end so the chat card - # renders the prompt + the generated - # image inline. The frontend chat-adapter - # decides how to render the base64 blob - # (likely an ) - # based on the `kind: "image"` hint we - # set on tool_start arguments. - # `time_ns()` (nanoseconds) instead of - # millisecond resolution so synthesised - # ids stay unique even when two image - # generations resolve in the same ms. + # Base64 image on `result` (or `b64_json`), + # `revised_prompt` for the rewritten prompt. + # ns-resolution id so concurrent gens stay unique. raw_item_id = item.get("id") item_id = raw_item_id or f"img_{time.time_ns()}" prompt_in = ( @@ -4034,6 +6193,54 @@ class ExternalProviderClient: "prompt": prompt_in, } ) + elif item.get("type") == "function_call": + # Translate to Chat-Completions delta.tool_calls. + # https://platform.openai.com/docs/guides/function-calling?api-mode=responses + fn_call_id = ( + item.get("call_id") + or item.get("id") + or f"call_{time.time_ns()}" + ) + fn_name = item.get("name") or "" + fn_args = item.get("arguments") or "" + if not isinstance(fn_args, str): + try: + fn_args = _json.dumps(fn_args) + except Exception: + fn_args = "" + _tc_index = function_call_index + function_call_index += 1 + yield ( + "data: " + + _json.dumps( + { + "id": completion_id, + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": _tc_index, + "id": fn_call_id, + "type": "function", + "function": { + "name": fn_name, + "arguments": ( + fn_args + ), + }, + } + ], + }, + "finish_reason": None, + } + ], + } + ) + ) + saw_function_call = True elif ( isinstance(event_type, str) @@ -4167,7 +6374,11 @@ class ExternalProviderClient: { "index": 0, "delta": {}, - "finish_reason": "stop", + "finish_reason": ( + "tool_calls" + if saw_function_call + else "stop" + ), } ], } @@ -4448,11 +6659,70 @@ class ExternalProviderClient: models = [model for model in raw_models if isinstance(model, dict)] if not models and self.provider_type == "ollama": models = await self._list_ollama_native_models() + # Gemini's native /v1beta/models returns + # {"models": [{"name": "models/gemini-2.5-flash", ...}]} + # -- repackage into the OpenAI-compatible shape the rest + # of Studio expects so dynamic model discovery works. + if not models and self.provider_type == "gemini": + models = self._parse_gemini_models(data) return models except httpx.HTTPError as exc: logger.error("Failed to list models from %s: %s", self.provider_type, exc) raise + @staticmethod + def _parse_gemini_models(payload: Any) -> list[dict[str, Any]]: + """Translate Gemini's native /v1beta/models payload to OpenAI shape. + + Native response: + {"models": [{"name": "models/gemini-2.5-flash", + "baseModelId": "gemini-2.5-flash", + "displayName": "Gemini 2.5 Flash", + "supportedGenerationMethods": [...]}]} + + We only keep entries that advertise + ``generateContent`` / ``streamGenerateContent`` so the picker + does not surface embedding-only models the chat path can't + drive. + """ + if not isinstance(payload, dict): + return [] + entries = payload.get("models") or [] + if not isinstance(entries, list): + return [] + out: list[dict[str, Any]] = [] + for entry in entries: + if not isinstance(entry, dict): + continue + methods = entry.get("supportedGenerationMethods") or [] + if ( + isinstance(methods, list) + and methods + and not any( + m in methods for m in ("generateContent", "streamGenerateContent") + ) + ): + continue + base_id = entry.get("baseModelId") + name = entry.get("name") or "" + # ``name`` arrives as ``"models/gemini-2.5-flash"``; the + # chat path uses the bare id. + short_id = ( + base_id + if isinstance(base_id, str) and base_id + else (name.split("/", 1)[1] if "/" in name else name) + ) + if not short_id: + continue + out.append( + { + "id": short_id, + "owned_by": "google", + "display_name": entry.get("displayName") or short_id, + } + ) + return out + async def _list_ollama_native_models(self) -> list[dict[str, Any]]: """Fallback when Ollama's /v1/models returns an empty or null catalog.""" root = self.base_url.removesuffix("/v1").rstrip("/") @@ -4757,6 +7027,17 @@ def _build_usage_chunk( "total_tokens": prompt_tokens + completion_tokens, "prompt_tokens_details": {"cached_tokens": cached}, } + # Surface OpenAI Responses / Gemini reasoning-token detail. The + # caller pre-populates last_usage["output_tokens_details"] with + # at least {"reasoning_tokens": ...}; mirror it into the OAI + # `completion_tokens_details` shape so SDKs can render the + # hidden-thoughts slice. + out_details = last_usage.get("output_tokens_details") + if isinstance(out_details, dict) and out_details: + usage_block["completion_tokens_details"] = { + "reasoning_tokens": out_details.get("reasoning_tokens") or 0, + } + usage_block["output_tokens_details"] = out_details chunk = { "id": completion_id, diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index fef9ba3e12..785f1dec3b 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -65,28 +65,77 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { }, "gemini": { "display_name": "Google Gemini", - "base_url": "https://generativelanguage.googleapis.com/v1beta/openai", - # Curated lineup — Google's /v1beta/openai/models returns dozens - # of historical / experimental / embedding ids. Cap to the current - # 3.x family plus the rolling `*-latest` aliases. + # Native Gemini REST endpoint -- the Gemini API does NOT speak + # OpenAI Chat Completions on this base. Requests/responses are + # translated in `_stream_gemini` in external_provider.py. + # API reference: https://ai.google.dev/gemini-api/docs + "base_url": "https://generativelanguage.googleapis.com/v1beta", + # Curated lineup -- the live ListModels response returns dozens + # of historical / experimental / embedding ids. Cap to the + # current chat-capable Gemini families (3.5 / 3.1 / 3 Flash / + # 2.5) plus the Nano Banana image trio and the rolling + # `*-latest` aliases. Excluded on purpose: + # - `gemini-2.0-flash*` (Google retired 2026-06-01; 404 on use) + # - `gemini-3-pro-preview` (shut down 2026-03-09; auto-redirects + # to `gemini-3.1-pro-preview` per Google's deprecation notice, + # so we surface 3.1 directly and skip the redirect). + # The allowlist below blocks the retired ids from re-appearing + # via the live ListModels fetch. Verified against the live + # `/v1beta/models` catalog 2026-05-24. "default_models": [ "gemini-3.1-pro-preview", + "gemini-3.5-flash", "gemini-3.1-flash-lite", "gemini-3-flash-preview", "gemini-pro-latest", "gemini-flash-latest", "gemini-flash-lite-latest", + "gemini-2.5-pro", + "gemini-2.5-flash", + "gemini-2.5-flash-lite", + "gemini-3-pro-image-preview", + "gemini-3.1-flash-image-preview", + "gemini-2.5-flash-image", ], "supports_streaming": True, "supports_vision": True, "supports_tool_calling": True, - "auth_header": "Authorization", - "auth_prefix": "Bearer ", - "notes": "OpenAI-compatible endpoint. API key from https://aistudio.google.com/apikey.", + # The native API takes the API key on the `x-goog-api-key` + # header. An empty `auth_prefix` ensures we send the bare key. + "auth_header": "x-goog-api-key", + "auth_prefix": "", + "openai_compatible": False, + "notes": ( + "Native Gemini API. Translation lives in _stream_gemini. " + "API key from https://aistudio.google.com/apikey. " + "See https://ai.google.dev/gemini-api/docs for endpoint shapes." + ), + # Even after the regex match, drop ids that Google still + # returns from ListModels but routes via implicit redirect. + # gemini-3-pro-preview was shut down 2026-03-09 and is + # auto-aliased to gemini-3.1-pro-preview; we surface the + # canonical id only so users do not see two cards for the + # same underlying model. + "model_id_deny_exact": ("gemini-3-pro-preview",), + # Matches the chat-capable 3.5 / 3.1 / 3 / 2.5 families plus the + # rolling *-latest aliases (which Google rolls forward as new + # generations ship). Image-tier ids (`-image`, `-image-preview`, + # `nano-banana-pro-preview`) flow through the Nano Banana + # `responseModalities` path in `_stream_gemini`. Retired 2.0 + # ids ARE NOT in this regex on purpose -- Google's ListModels + # would otherwise re-surface them and they 404 on use. "model_id_allowlist": re.compile( - r"^(gemini-3\.1-flash-lite|gemini-3-flash-preview|" - r"gemini-3\.1-pro-preview|gemini-pro-latest|" - r"gemini-flash-latest|gemini-flash-lite-latest)$" + r"^(" + r"gemini-3\.5-(?:flash|pro)(?:-preview)?|" + r"gemini-3\.1-(?:flash|pro|flash-lite)(?:-preview)?(?:-customtools)?|" + r"gemini-3\.1-flash-image-preview|" + r"gemini-3-(?:flash|pro)(?:-preview)?|" + r"gemini-3-pro-image-preview|" + r"nano-banana-pro-preview|" + r"gemini-2\.5-pro|gemini-2\.5-flash|gemini-2\.5-flash-lite|" + r"gemini-2\.5-flash-image|" + r"gemini-pro-latest|gemini-flash-latest|gemini-flash-lite-latest" + r")$" ), }, "deepseek": { diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 0e9cce7c3e..35d5903826 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -632,8 +632,17 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str for *_, sockaddr in infos: ip = ipaddress.ip_address(sockaddr[0]) + # `not ip.is_global` rejects every category the denylist below + # also rejects PLUS shared address space (100.64.0.0/10 carrier- + # grade NAT) and benchmarking/documentation/exchange ranges that + # Python classifies with `is_private=False` and `is_global=False` + # (see https://docs.python.org/3/library/ipaddress.html#ipaddress.IPv4Address.is_global). + # The explicit predicates after it give human-readable categories + # in the error message, but a single non-global check is the + # source of truth and prevents future ranges from leaking. if ( - ip.is_private + not ip.is_global + or ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_multicast diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0af9425fdc..b58ca66440 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -581,6 +581,14 @@ class ChatMessage(BaseModel): None, description = "OpenAI tool-result messages: name of the tool whose result this is.", ) + extra_content: Optional[dict] = Field( + None, + description = ( + "Provider-specific extra fields the translator may read. " + "Gemini reads `extra_content.google.thought_signature` " + "from assistant messages to replay text-part signatures." + ), + ) @model_validator(mode = "after") def _validate_role_shape(self) -> "ChatMessage": @@ -752,17 +760,42 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] Override base URL for the external provider.", ) - enable_prompt_caching: Optional[bool] = Field( + enable_prompt_caching: Optional[Union[bool, str]] = Field( None, description = ( "[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, " - "attaches cache_control={type:ephemeral} to the system block so the " - "static prefix is reused across turns. On OpenAI cloud, caching is " - "automatic for prompts >=1024 tokens and this flag is informational. " - "Ignored for every other provider (mistral, gemini, kimi, openrouter, " - "vllm, local, etc.). Treated as enabled when omitted." + "boolean true attaches cache_control={type:ephemeral} to the system " + "block so the static prefix is reused across turns. On OpenAI cloud, " + "caching is automatic for prompts >=1024 tokens and the boolean is " + "informational. On Gemini, pass a string cache resource name such " + "as `cachedContents/abc123` to attach `cachedContent` on the native " + "request (boolean true is a no-op on Gemini because creating the " + "cache requires a separate POST /cachedContents call). Ignored for " + "every other provider. Treated as enabled when omitted." ), ) + + @field_validator("enable_prompt_caching", mode = "before") + @classmethod + def _coerce_enable_prompt_caching(cls, value: Any) -> Any: + """Preserve the pre-PR coercion: the field used to be Optional[bool], + so callers historically sent JSON strings `"true"` / `"false"` and + Pydantic v1 coerced them. Widening to Optional[Union[bool, str]] for + Gemini cache resource names lets `"false"` slip through as a truthy + string. Coerce the canonical bool literals back so explicit opt-outs + stay opt-out.""" + if isinstance(value, str): + lowered = value.strip().lower() + # Match Pydantic v1's BooleanField coercion table (yes/y/on/t/1 + # and no/n/off/f/0) so opt-outs that used to parse still parse. + # Anything else is preserved as a string for Gemini's + # cachedContent resource path. + if lowered in ("true", "t", "1", "yes", "y", "on"): + return True + if lowered in ("false", "f", "0", "no", "n", "off"): + return False + return value + prompt_cache_ttl: Optional[str] = Field( None, description = ( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a156f2397c..92498b8a9a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1705,6 +1705,7 @@ def _build_external_messages( messages: list, supports_vision: bool, provider_type: Optional[str] = None, + base_url: Optional[str] = None, ) -> list[dict]: """ Convert ChatMessage list to OpenAI-compatible dicts for external providers. @@ -1732,14 +1733,171 @@ def _build_external_messages( document_provider = provider_type in _INPUT_DOCUMENT_PROVIDERS anthropic = provider_type == "anthropic" openai = provider_type == "openai" + # `extra_content` is a Gemini-specific carrier for the assistant's + # text-part `thoughtSignature` round-trip on the native + # streamGenerateContent endpoint. Custom Gemini OpenAI-compatible + # gateways (LiteLLM etc.) route through /chat/completions where + # the field is unknown and can be rejected -- gate strictly on the + # Google-hosted Gemini base. + _native_gemini = False + if provider_type == "gemini" and base_url: + try: + from urllib.parse import urlparse as _urlparse + + _host = (_urlparse(base_url).hostname or "").lower() + _native_gemini = _host == "generativelanguage.googleapis.com" + except Exception: + _native_gemini = False + emit_extra_content = _native_gemini + + _SERVER_BUILTIN_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "code_execution", "image_generation"} + ) + + def _is_marked_server_builtin_tool_call(tc: Any) -> bool: + """Return True iff `tc` is a synthetic provider-side tool card + with one of the canonical builtin names and either: + - the new `args._server_tool` marker stamped by the backend, or + - a Gemini `args.google.native_part` payload (durable replay + signal for code_execution / image_generation that predates + the marker). + Such cards must not be forwarded to non-native providers + because they are not real user functions and the receiving API + will reject the orphan tool history. Real user functions with + these names normally have neither signal. + """ + if not isinstance(tc, dict): + return False + fn = tc.get("function") + if not isinstance(fn, dict): + return False + name = (fn.get("name") or "").lower() + if name not in _SERVER_BUILTIN_TOOL_NAMES: + return False + raw_args = fn.get("arguments") or "" + try: + args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args + except Exception: + return False + if not isinstance(args, dict): + return False + if args.get("_server_tool") is True: + return True + google = args.get("google") + return isinstance(google, dict) and isinstance(google.get("native_part"), dict) + + # When we drop a server-side builtin tool_call here, the matching + # `role="tool"` follow-up must also be dropped from the outbound + # history -- otherwise the provider receives an orphan + # tool_call_id with no matching assistant call, which OpenAI + # Responses and Anthropic both reject. + dropped_server_builtin_tool_call_ids: set[str] = set() + + def _filter_tool_calls(tool_calls: Any) -> Optional[list]: + """Sanitize assistant `tool_calls` for non-native-Gemini providers. + + Two concerns: + 1. `tool_calls[i].extra_content` carries Gemini-only + thoughtSignature metadata; strip it for providers that + cannot parse the unknown key. + 2. Marked server-side builtin cards (`_server_tool: true` on + a canonical builtin name, or a Gemini `native_part` + payload) are provider-internal Studio tool cards from a + prior native Gemini turn; forwarding them to OpenAI / + Anthropic / custom OAI-compat gateways sends an orphan + `tool_calls` entry (no matching tool declaration, often + no matching `role="tool"` reply) that can be rejected. + We record the dropped call_ids so the matching role=tool + message is also skipped below. + Native Gemini keeps both untouched so the native translator can + replay them via `native_part`. + """ + if not tool_calls: + return None + if not isinstance(tool_calls, list): + return tool_calls + if emit_extra_content: + return tool_calls + cleaned: list = [] + for _tc in tool_calls: + if _is_marked_server_builtin_tool_call(_tc): + _tc_id = _tc.get("id") if isinstance(_tc, dict) else None + if isinstance(_tc_id, str) and _tc_id: + dropped_server_builtin_tool_call_ids.add(_tc_id) + continue + if not isinstance(_tc, dict): + cleaned.append(_tc) + continue + if "extra_content" not in _tc: + cleaned.append(_tc) + continue + _stripped = {k: v for k, v in _tc.items() if k != "extra_content"} + cleaned.append(_stripped) + return cleaned + result = [] for msg in messages: + # Drop role=tool messages whose matching server-builtin + # tool_call was already filtered above. Forwarding an orphan + # tool_result with no matching tool_call would be rejected by + # OpenAI Responses and Anthropic. + if ( + msg.role == "tool" + and isinstance(msg.tool_call_id, str) + and msg.tool_call_id in dropped_server_builtin_tool_call_ids + ): + continue if isinstance(msg.content, str): - # Skip assistant messages with empty content (some providers reject them) - if msg.role == "assistant" and not msg.content.strip(): + # Drop bare assistant messages with no content AND no + # tool_calls (some providers reject empty assistant turns). + # Preserve assistant turns whose only payload is tool_calls + # so multi-turn function-call loops round-trip. + if ( + msg.role == "assistant" + and not msg.content.strip() + and not msg.tool_calls + ): continue - result.append({"role": msg.role, "content": msg.content}) - elif isinstance(msg.content, list): + out: dict[str, Any] = {"role": msg.role, "content": msg.content} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + out["tool_calls"] = _tcs + elif not msg.content.strip(): + # Every tool_call was a synthetic provider-side + # card and was dropped; the assistant turn would + # be an empty `{"role":"assistant","content":""}` + # which some providers reject. Skip it entirely. + continue + if msg.role == "tool": + if msg.tool_call_id: + out["tool_call_id"] = msg.tool_call_id + if msg.name: + out["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + out["extra_content"] = msg.extra_content + result.append(out) + continue + # Assistant messages with content=None but populated tool_calls + # are valid (post-tool-call assistant turn). Forward them so the + # provider helper can rebuild the functionCall part. + if msg.content is None and msg.role == "assistant" and msg.tool_calls: + _filtered_tcs = _filter_tool_calls(msg.tool_calls) + if not _filtered_tcs: + # Every tool_call on this turn was provider-side + # synthetic and dropped; skipping the whole message + # avoids forwarding an empty assistant turn. + continue + _assistant_only: dict[str, Any] = { + "role": "assistant", + "content": "", + "tool_calls": _filtered_tcs, + } + if emit_extra_content and msg.extra_content: + _assistant_only["extra_content"] = msg.extra_content + result.append(_assistant_only) + continue + if isinstance(msg.content, list): if supports_vision: parts = [] for part in msg.content: @@ -1797,9 +1955,27 @@ def _build_external_messages( # provider would 400 on the unknown part, so # gate by provider_type. parts.append({"type": "compaction", "content": part.content}) - if msg.role == "assistant" and not parts: + entry: dict[str, Any] = {"role": msg.role, "content": parts} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + entry["tool_calls"] = _tcs + elif not parts: + # All tool_calls were synthetic and dropped, + # and no preserved content parts survived. + # Skip rather than forward an empty assistant + # turn that downstream providers reject. + continue + elif msg.role == "assistant" and not parts: continue - result.append({"role": msg.role, "content": parts}) + if msg.role == "tool": + if msg.tool_call_id: + entry["tool_call_id"] = msg.tool_call_id + if msg.name: + entry["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + entry["extra_content"] = msg.extra_content + result.append(entry) else: # Non-vision provider: strip images / documents, keep # text, optionally keep compaction (Anthropic only -- @@ -1835,9 +2011,32 @@ def _build_external_messages( if len(preserved) == 1 and preserved[0]["type"] == "text": # Single text part collapses back to a string for # providers that don't accept content arrays. - result.append({"role": msg.role, "content": preserved[0]["text"]}) + entry = {"role": msg.role, "content": preserved[0]["text"]} else: - result.append({"role": msg.role, "content": preserved}) + entry = {"role": msg.role, "content": preserved} + if msg.role == "assistant" and msg.tool_calls: + _tcs = _filter_tool_calls(msg.tool_calls) + if _tcs: + entry["tool_calls"] = _tcs + else: + # All tool_calls were synthetic and dropped; + # skip if there's no surviving content either. + _entry_content = entry.get("content") + _has_text = ( + isinstance(_entry_content, str) and _entry_content.strip() + ) or ( + isinstance(_entry_content, list) and len(_entry_content) > 0 + ) + if not _has_text: + continue + if msg.role == "tool": + if msg.tool_call_id: + entry["tool_call_id"] = msg.tool_call_id + if msg.name: + entry["name"] = msg.name + if emit_extra_content and msg.role == "assistant" and msg.extra_content: + entry["extra_content"] = msg.extra_content + result.append(entry) return result @@ -1912,6 +2111,7 @@ async def _proxy_to_external_provider( payload.messages, _supports_vision, provider_type = provider_type, + base_url = base_url, ) client = ExternalProviderClient( @@ -1920,6 +2120,14 @@ async def _proxy_to_external_provider( api_key = api_key, ) + # `top_k` defaults to 20 in ChatCompletionRequest because the local + # inference path expects an int, but the external-provider path + # should treat "field omitted from JSON" as "use provider default" + # so callers that send only model/messages do not silently get + # different sampling than before this PR. Pydantic's + # `model_fields_set` tracks explicit-vs-default per request. + _top_k_explicit = payload.top_k if "top_k" in payload.model_fields_set else None + async def _stream(): gen = client.stream_chat_completion( messages = chat_messages, @@ -1928,7 +2136,7 @@ async def _proxy_to_external_provider( top_p = payload.top_p, max_tokens = payload.max_tokens, presence_penalty = payload.presence_penalty, - top_k = payload.top_k, + top_k = _top_k_explicit, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, enabled_tools = payload.enabled_tools, @@ -1937,6 +2145,8 @@ async def _proxy_to_external_provider( anthropic_code_exec_container_id = payload.anthropic_code_exec_container_id, prompt_cache_ttl = payload.prompt_cache_ttl, compaction_threshold = payload.compaction_threshold, + tools = payload.tools, + tool_choice = payload.tool_choice, fast_mode = payload.fast_mode, stream = payload.stream, ) @@ -4480,7 +4690,17 @@ async def anthropic_messages( [m.model_dump() for m in payload.messages], payload.system, ) - openai_messages = _drop_empty_assistant_sentinels(openai_messages) + # Strip synthetic provider-side builtin tool history (web_search, + # web_fetch, code_execution, image_generation cards tagged with + # _server_tool or extra_content.google.native_part) before handing + # off to local llama-server. The local /v1/chat/completions and + # GGUF passthrough builders apply the same strip; without it an + # Anthropic /v1/messages caller replaying a prior provider-side + # tool_use forwards fake builtin tool history to a backend that + # has no matching function declarations. + openai_messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels(openai_messages) + ) # Enforce vision guard + re-encode embedded images to PNG so the # Anthropic endpoint matches the behavior of /v1/chat/completions. @@ -5271,6 +5491,110 @@ def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]: return out +_LOCAL_SERVER_BUILTIN_TOOL_NAMES = frozenset( + {"web_search", "web_fetch", "code_execution", "image_generation"} +) + + +def _strip_provider_synthetic_tool_history(messages: list[dict]) -> list[dict]: + """Drop synthetic provider-side tool_calls + matching role=tool replies + on the local-backend (llama-server / GGUF) dispatch path. + + A Gemini chat that ran code_execution / image_generation persists the + server-side tool card into thread history as an assistant tool_calls + entry tagged with ``args._server_tool`` (or a Gemini + ``args.google.native_part`` payload) plus a follow-up role=tool reply. + When the user switches the SAME thread to a local GGUF model, those + synthetic tool_calls are not real user functions, llama-server has no + matching declaration, and Gemini-only ``extra_content`` / + ``native_part`` payloads are meaningless. Forward only ordinary user + function calls; strip the matched role=tool replies too so the + backend does not see an orphan tool_call_id. + """ + dropped_ids: set[str] = set() + sanitized_assistant: list[dict] = [] + for m in messages: + if m.get("role") != "assistant": + sanitized_assistant.append(m) + continue + tool_calls = m.get("tool_calls") + if not isinstance(tool_calls, list) or not tool_calls: + # Plain text Gemini reply: still strip message-level + # `extra_content` (carries `google.thought_signature` replay + # metadata) so a text-only Gemini turn switched to a local + # GGUF backend does not leak Gemini-only fields to + # llama-server. ChatMessage previously did not have + # `extra_content`, so the field was implicitly dropped -- + # round-22 added it to ChatMessage, which is what made this + # leak possible. + if "extra_content" in m: + m = {k: v for k, v in m.items() if k != "extra_content"} + sanitized_assistant.append(m) + continue + cleaned: list[dict] = [] + for tc in tool_calls: + if not isinstance(tc, dict): + cleaned.append(tc) + continue + fn = tc.get("function") + name = "" + if isinstance(fn, dict): + name = (fn.get("name") or "").lower() + if name in _LOCAL_SERVER_BUILTIN_TOOL_NAMES: + raw_args = fn.get("arguments") if isinstance(fn, dict) else None + args_obj: Any = None + if isinstance(raw_args, str): + try: + args_obj = json.loads(raw_args) if raw_args else None + except Exception: + args_obj = None + elif isinstance(raw_args, dict): + args_obj = raw_args + is_synthetic = False + if isinstance(args_obj, dict): + if args_obj.get("_server_tool") is True: + is_synthetic = True + google = args_obj.get("google") + if isinstance(google, dict) and isinstance( + google.get("native_part"), dict + ): + is_synthetic = True + if is_synthetic: + tc_id = tc.get("id") + if isinstance(tc_id, str) and tc_id: + dropped_ids.add(tc_id) + continue + # Strip Gemini-only `extra_content` on real user tool_calls + # too — llama-server has no use for it and may pass it + # through to the model unchanged. + if "extra_content" in tc: + tc = {k: v for k, v in tc.items() if k != "extra_content"} + cleaned.append(tc) + # Drop top-level message-level `extra_content` (Gemini + # thoughtSignature replay metadata) on local dispatch. + m_clean = {k: v for k, v in m.items() if k != "extra_content"} + if cleaned: + m_clean["tool_calls"] = cleaned + else: + m_clean.pop("tool_calls", None) + if not m_clean.get("content") and not m_clean.get("tool_calls"): + continue # assistant turn now empty, drop + sanitized_assistant.append(m_clean) + + if not dropped_ids: + return sanitized_assistant + out: list[dict] = [] + for m in sanitized_assistant: + if ( + m.get("role") == "tool" + and isinstance(m.get("tool_call_id"), str) + and m["tool_call_id"] in dropped_ids + ): + continue + out.append(m) + return out + + def _openai_messages_for_passthrough(payload) -> list[dict]: """Build OpenAI-format message dicts for the /v1/chat/completions passthrough path. @@ -5287,8 +5611,10 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: ``image_url`` content part so vision + function-calling requests work transparently. """ - messages = _drop_empty_assistant_sentinels( - [m.model_dump(exclude_none = True) for m in payload.messages] + messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels( + [m.model_dump(exclude_none = True) for m in payload.messages] + ) ) if not payload.image_base64: @@ -5337,8 +5663,10 @@ def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict] all per-turn ``image_url`` parts so multi-image chat history keeps each image attached to its original turn. """ - messages = _drop_empty_assistant_sentinels( - [m.model_dump(exclude_none = True) for m in payload.messages] + messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels( + [m.model_dump(exclude_none = True) for m in payload.messages] + ) ) has_message_image = any( isinstance(msg.get("content"), list) diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 2bb1de5366..5d4bd46e62 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -318,22 +318,45 @@ async def list_provider_models( try: models = await client.list_models() - allow_prefixes = info.get("model_id_allow_prefixes") - if allow_prefixes is not None: - prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p)) - if prefix_tuple: - models = [m for m in models if m.get("id", "").startswith(prefix_tuple)] - allowlist = info.get("model_id_allowlist") - if allowlist is not None: - models = [m for m in models if allowlist.match(m.get("id", ""))] - deny_exact = info.get("model_id_deny_exact") - if deny_exact is not None: - deny_ids = {str(m) for m in deny_exact if str(m)} - if deny_ids: - models = [m for m in models if m.get("id", "") not in deny_ids] - denylist = info.get("model_id_denylist") - if denylist is not None: - models = [m for m in models if not denylist.search(m.get("id", ""))] + # Registry-level model-id filters are scoped to the canonical + # native Gemini base. A custom Gemini OAI-compatible proxy + # (LiteLLM, deployment gateway) returns IDs like + # `google/gemini-2.5-flash`, `gemini/gemini-2.5-flash`, or + # team-prefixed deployment aliases; the native allowlist regex + # would strip those out and leave the picker empty even though + # the chat path now routes them via the OAI-compatible + # dispatcher (the same gate ExternalProviderClient applies for + # request building). Match the host check here so the model + # list and chat dispatch agree on what counts as "native". + apply_registry_model_filters = True + if payload.provider_type == "gemini": + try: + from urllib.parse import urlparse as _urlparse + + _host = (_urlparse(base_url).hostname or "").lower() + except Exception: + _host = "" + apply_registry_model_filters = _host == "generativelanguage.googleapis.com" + + if apply_registry_model_filters: + allow_prefixes = info.get("model_id_allow_prefixes") + if allow_prefixes is not None: + prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p)) + if prefix_tuple: + models = [ + m for m in models if m.get("id", "").startswith(prefix_tuple) + ] + allowlist = info.get("model_id_allowlist") + if allowlist is not None: + models = [m for m in models if allowlist.match(m.get("id", ""))] + deny_exact = info.get("model_id_deny_exact") + if deny_exact is not None: + deny_ids = {str(m) for m in deny_exact if str(m)} + if deny_ids: + models = [m for m in models if m.get("id", "") not in deny_ids] + denylist = info.get("model_id_denylist") + if denylist is not None: + models = [m for m in models if not denylist.search(m.get("id", ""))] # Apply an optional cap after filtering so registry entries with a # large remote catalog (e.g. HF Inference Providers) can stay # picker-sized. No popularity sort happens server-side, so this is diff --git a/studio/backend/tests/test_anthropic_code_execution.py b/studio/backend/tests/test_anthropic_code_execution.py index 7f6fe58329..5c88437d17 100644 --- a/studio/backend/tests/test_anthropic_code_execution.py +++ b/studio/backend/tests/test_anthropic_code_execution.py @@ -275,7 +275,13 @@ def test_bash_code_execution_emits_tool_start_and_end(monkeypatch): assert start["type"] == "tool_start" assert start["tool_name"] == "code_execution" assert start["tool_call_id"] == "srvtoolu_1" - assert start["arguments"] == {"kind": "bash", "command": "ls -la"} + # `_server_tool: True` marks this as a provider-side synthetic + # tool card for the frontend's history serializer. + assert start["arguments"] == { + "kind": "bash", + "command": "ls -la", + "_server_tool": True, + } assert end["type"] == "tool_end" assert end["tool_call_id"] == "srvtoolu_1" diff --git a/studio/backend/tests/test_anthropic_web_fetch.py b/studio/backend/tests/test_anthropic_web_fetch.py index da10d679eb..88a922e7eb 100644 --- a/studio/backend/tests/test_anthropic_web_fetch.py +++ b/studio/backend/tests/test_anthropic_web_fetch.py @@ -271,7 +271,12 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch): assert start["type"] == "tool_start" assert start["tool_name"] == "web_fetch" assert start["tool_call_id"] == "srvtoolu_wf1" - assert start["arguments"] == {"url": "https://example.com/article"} + # `_server_tool: True` marks this as a provider-side synthetic + # tool card for the frontend's history serializer. + assert start["arguments"] == { + "url": "https://example.com/article", + "_server_tool": True, + } assert end["type"] == "tool_end" assert end["tool_call_id"] == "srvtoolu_wf1" # The source pill uses Title / URL / snippet as parseSourcesFromResult expects. diff --git a/studio/backend/tests/test_gemini_provider.py b/studio/backend/tests/test_gemini_provider.py new file mode 100644 index 0000000000..4dc97302e2 --- /dev/null +++ b/studio/backend/tests/test_gemini_provider.py @@ -0,0 +1,5501 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the native Gemini API translation layer. + +Gemini does NOT speak OpenAI Chat Completions on its primary endpoint +(`streamGenerateContent`). `_stream_gemini` in +`core/inference/external_provider.py` translates between the two shapes: + + Request: + OpenAI messages [{role, content}] + -> Gemini contents [{role, parts: [{text}|{inlineData}|{functionCall}|...]}] + + systemInstruction.parts[].text for role=system messages + + generationConfig.{temperature,topP,topK,maxOutputTokens} + + tools[{googleSearch:{}}] for web_search + + tools[{codeExecution:{}}] for code_execution + + responseModalities=[TEXT,IMAGE] for Nano Banana (gemini-2.5-flash-image) + + cachedContent for prompt caching + + Response: + Gemini SSE chunks { candidates:[{content:{parts:[...]}, finishReason}], + usageMetadata:{promptTokenCount, candidatesTokenCount} } + -> OpenAI chat.completion.chunk frames + (delta.content for text, delta.tool_calls for functionCall, + _toolEvent for image_b64/web_search, usage block before [DONE]) + +These tests pin the outbound body shape AND the inbound translation +using httpx.MockTransport (no live network). Mirrors the structure of +test_anthropic_cache_ttl.py and test_openai_image_generation.py. +""" + +import asyncio +import base64 +import json + +import httpx +import pytest + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +_active_mock_clients: list[httpx.AsyncClient] = [] + + +def _drive(coro): + # Create a fresh loop per drive so tests don't share asyncio state. + # Close mocked clients + shutdown async-generators inside this loop + # so Python 3.13 doesn't emit the + # `Response.aiter_*.aclose was never awaited` warning on GC. + loop = asyncio.new_event_loop() + try: + result = loop.run_until_complete(coro) + while _active_mock_clients: + mc = _active_mock_clients.pop() + loop.run_until_complete(mc.aclose()) + return result + finally: + try: + loop.run_until_complete(loop.shutdown_asyncgens()) + finally: + loop.close() + + +def _make_gemini_client( + base_url: str = "https://generativelanguage.googleapis.com/v1beta", +) -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "gemini", + base_url = base_url, + api_key = "AIza-test-key", + ) + + +def _mock_http(monkeypatch, handler): + mock_client = httpx.AsyncClient(transport = httpx.MockTransport(handler)) + monkeypatch.setattr(ep_mod, "_http_client", mock_client) + # `_drive` will aclose this at the end of the run inside the same + # event loop so we do not leak an unawaited aclose() coroutine. + _active_mock_clients.append(mock_client) + + +def _gemini_sse(events: list[dict]) -> bytes: + """Encode a list of dicts as Gemini-style SSE frames (`data:` lines).""" + chunks: list[str] = [] + for event in events: + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def _capture_body(monkeypatch, **kwargs) -> dict: + """Drive a single stream and return the captured outbound request body.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + captured["url"] = str(request.url) + captured["method"] = request.method + # Minimal valid Gemini stream so the helper can complete. + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + messages = kwargs.pop("messages", [{"role": "user", "content": "hi"}]) + model = kwargs.pop("model", "gemini-2.5-flash") + temperature = kwargs.pop("temperature", 0.7) + top_p = kwargs.pop("top_p", 0.95) + max_tokens = kwargs.pop("max_tokens", 64) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = messages, + model = model, + temperature = temperature, + top_p = top_p, + max_tokens = max_tokens, + **kwargs, + ): + pass + await client.close() + + _drive(run()) + return captured + + +def _collect(monkeypatch, sse_events, **kwargs) -> list[str]: + """Drive a stream with a custom set of SSE events and return raw lines.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _gemini_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + messages = kwargs.pop("messages", [{"role": "user", "content": "hi"}]) + model = kwargs.pop("model", "gemini-2.5-flash") + temperature = kwargs.pop("temperature", 0.7) + top_p = kwargs.pop("top_p", 0.95) + max_tokens = kwargs.pop("max_tokens", 64) + + out: list[str] = [] + + async def run(): + client = _make_gemini_client() + async for line in client.stream_chat_completion( + messages = messages, + model = model, + temperature = temperature, + top_p = top_p, + max_tokens = max_tokens, + **kwargs, + ): + out.append(line) + await client.close() + + _drive(run()) + return out + + +def _parse_chunks(lines: list[str]) -> list[dict]: + out: list[dict] = [] + for raw in lines: + if not raw.startswith("data:"): + continue + payload = raw[len("data:") :].strip() + if not payload or payload == "[DONE]": + continue + try: + out.append(json.loads(payload)) + except json.JSONDecodeError: + continue + return out + + +# ── request body translation ───────────────────────────────────────── + + +def test_request_body_uses_contents_and_parts_shape(monkeypatch): + """OpenAI messages must be translated to Gemini's `contents` shape.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + {"role": "user", "content": "Follow up"}, + ], + ) + body = captured["body"] + # system -> systemInstruction + assert body["systemInstruction"] == {"parts": [{"text": "Be brief."}]}, body + # user/assistant -> contents with role user/model + assert body["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]}, + {"role": "model", "parts": [{"text": "Hi there"}]}, + {"role": "user", "parts": [{"text": "Follow up"}]}, + ], body["contents"] + # generationConfig fields map across with Google's casing. + gc = body["generationConfig"] + assert gc["temperature"] == 0.7 + assert gc["topP"] == 0.95 + assert gc["maxOutputTokens"] == 64 + + +def test_request_url_targets_stream_generate_content(monkeypatch): + """Helper must POST to /v1beta/models/{model}:streamGenerateContent?alt=sse.""" + captured = _capture_body(monkeypatch, model = "gemini-2.5-pro") + url = captured["url"] + assert ":streamGenerateContent" in url, url + assert "alt=sse" in url, url + assert "/v1beta/models/gemini-2.5-pro" in url, url + assert captured["method"] == "POST" + + +def test_request_auth_header_uses_x_goog_api_key(monkeypatch): + """API key must be sent on `x-goog-api-key`, not Authorization.""" + captured = _capture_body(monkeypatch) + hdrs = captured["headers"] + assert hdrs.get("x-goog-api-key") == "AIza-test-key", hdrs + assert "authorization" not in {k.lower() for k in hdrs}, hdrs + + +def test_top_k_forwarded_only_when_positive(monkeypatch): + """top_k is opt-in; only positive integers reach the wire.""" + captured = _capture_body(monkeypatch, top_k = 40) + assert captured["body"]["generationConfig"]["topK"] == 40 + + captured = _capture_body(monkeypatch, top_k = 0) + assert "topK" not in captured["body"]["generationConfig"] + + +def test_presence_penalty_forwarded_to_generation_config(monkeypatch): + """A non-zero presence_penalty reaches generationConfig.presencePenalty.""" + captured = _capture_body(monkeypatch, presence_penalty = 0.7) + assert captured["body"]["generationConfig"]["presencePenalty"] == 0.7 + + # And the default of zero is omitted, matching top_k semantics. + captured = _capture_body(monkeypatch, presence_penalty = 0.0) + assert "presencePenalty" not in captured["body"]["generationConfig"] + + +# ── thinkingConfig translation ──────────────────────────────────────── + + +def test_gemini25_flash_thinking_disabled_sets_budget_zero(monkeypatch): + """Gemini 2.5 Flash still uses thinkingBudget; 0 = off.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingBudget": 0}, tc + + +def test_gemini3_flash_thinking_disabled_uses_minimal_level(monkeypatch): + """Gemini 3 Flash migrated to thinkingLevel; "off" maps to minimal + (Gemini 3 cannot turn thinking fully off).""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "minimal"}, tc + + +def test_gemini25_pro_thinking_disabled_uses_small_budget(monkeypatch): + """Gemini 2.5 Pro 400s on thinkingBudget=0 ("only works in thinking + mode"); coerce to a small positive budget.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-pro", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc is not None and tc.get("thinkingBudget", 0) > 0, tc + + +def test_gemini3_pro_thinking_disabled_uses_low_level(monkeypatch): + """Gemini 3 Pro uses thinkingLevel and rejects 'minimal' (Pro tier), + so 'off' coerces to 'low' (lowest the API accepts).""" + for model in ( + "gemini-3.1-pro-preview", + "gemini-3-pro-preview", + "gemini-3.5-pro", + "gemini-pro-latest", + ): + captured = _capture_body( + monkeypatch, + model = model, + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, (model, tc) + + +def test_gemini25_flash_effort_levels_map_to_budgets(monkeypatch): + """Gemini 2.5 Flash retains the integer thinkingBudget ladder.""" + cases = { + "minimal": 512, + "low": 2048, + "medium": 8192, + "high": 24576, + "max": -1, + "xhigh": -1, + } + for effort, expected in cases.items(): + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + reasoning_effort = effort, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingBudget": expected}, (effort, tc) + + +def test_gemini3_flash_effort_levels_map_to_thinking_level(monkeypatch): + """Gemini 3 Flash thinkingLevel ladder: minimal/low/medium/high.""" + cases = { + "minimal": "minimal", + "low": "low", + "medium": "medium", + "high": "high", + "max": "high", + } + for effort, expected in cases.items(): + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + reasoning_effort = effort, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": expected}, (effort, tc) + + +def test_gemini3_pro_passes_medium_through(monkeypatch): + """Gemini 3.1+ Pro accepts thinkingLevel="medium" per + https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro; + forward as-is (medium is the documented mid-tier on Gemini 3.1).""" + for model in ( + "gemini-3.1-pro-preview", + "gemini-pro-latest", + ): + captured = _capture_body( + monkeypatch, + model = model, + reasoning_effort = "medium", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "medium"}, (model, tc) + + +def test_gemini3_pro_minimal_effort_coerces_to_low(monkeypatch): + """Gemini 3 Pro rejects thinkingLevel="minimal"; coerce to "low".""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.1-pro-preview", + reasoning_effort = "minimal", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, tc + + +def test_gemini3_flash_effort_none_maps_to_minimal(monkeypatch): + """reasoning_effort='none' on Gemini 3 Flash -> thinkingLevel=minimal.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-flash", + reasoning_effort = "none", + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "minimal"}, tc + + +def test_thinking_default_omits_thinking_config(monkeypatch): + """When neither knob is supplied, thinkingConfig is omitted entirely + (Google's server-side default applies).""" + captured = _capture_body(monkeypatch, model = "gemini-3.5-flash") + gc = captured["body"]["generationConfig"] + assert "thinkingConfig" not in gc, gc + + +def test_nano_banana_alias_routes_through_image_modalities(monkeypatch): + """`nano-banana-pro-preview` is an alias for the Pro image model and + must set responseModalities=[TEXT,IMAGE] when the Images pill is on + (enabled_tools includes "image_generation").""" + captured = _capture_body( + monkeypatch, + model = "nano-banana-pro-preview", + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert gc.get("responseModalities") == ["TEXT", "IMAGE"], gc + + +def test_image_capable_model_without_image_pill_stays_text_only(monkeypatch): + """When the Images pill is off (enabled_tools has no + image_generation), an image-capable model id (gemini-2.5-flash-image) + must force responseModalities=["TEXT"]. Google's image models + default to text+image when responseModalities is omitted, so + omitting it would silently bill image output the UI says is + disabled.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = [], + ) + gc = captured["body"]["generationConfig"] + assert gc.get("responseModalities") == ["TEXT"], gc + + +def test_image_models_skip_thinking_config(monkeypatch): + """Image-tier ids do not benefit from a visible thinking knob and + must NOT forward thinkingConfig even when stale UI state still + sends `reasoning_effort` or `enable_thinking=False`.""" + for model in ( + "gemini-2.5-flash-image", + "gemini-3.1-flash-image-preview", + "gemini-3-pro-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + reasoning_effort = "high", + enable_thinking = False, + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert "thinkingConfig" not in gc, (model, gc) + + +def test_image_models_drop_code_execution(monkeypatch): + """All image-tier ids reject `tools: [{codeExecution: {}}]`; drop + silently. (Gemini 3 image models DO accept googleSearch -- see + test_gemini3_image_models_allow_google_search; older image models + drop everything.)""" + for model in ( + "gemini-2.5-flash-image", + "gemini-3.1-flash-image-preview", + "gemini-3-pro-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + enabled_tools = ["image_generation", "code_execution"], + ) + tools_arr = captured["body"].get("tools") or [] + names = [list(t.keys())[0] for t in tools_arr] + assert "codeExecution" not in names, (model, tools_arr) + + +def test_gemini_35_pro_uses_thinking_level(monkeypatch): + """`gemini-3.5-pro` is part of the Gemini 3 family and uses + thinkingLevel (not thinkingBudget). "Off" maps to "low" because Pro + tier rejects "minimal".""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.5-pro", + enable_thinking = False, + ) + tc = captured["body"]["generationConfig"].get("thinkingConfig") + assert tc == {"thinkingLevel": "low"}, tc + + +def test_gemini3_image_models_allow_google_search(monkeypatch): + """Google documents Search grounding on the Gemini 3 image family + (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + nano-banana-pro). codeExecution stays blocked on image mode.""" + for model in ( + "gemini-3-pro-image-preview", + "gemini-3.1-flash-image-preview", + "nano-banana-pro-preview", + ): + captured = _capture_body( + monkeypatch, + model = model, + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + tools_arr = captured["body"].get("tools") or [] + names = [list(t.keys())[0] for t in tools_arr] + assert "googleSearch" in names, (model, tools_arr) + assert "codeExecution" not in names, (model, tools_arr) + + +def test_legacy_image_models_block_google_search(monkeypatch): + """Older Gemini image ids (gemini-2.5-flash-image) still 400 on + `tools: [{googleSearch: {}}]`; backend keeps stripping it.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + assert "tools" not in captured["body"], captured["body"].get("tools") + + +def test_legacy_openai_base_url_normalized(monkeypatch): + """Saved Gemini providers carrying the legacy `/v1beta/openai` base + (from the pre-PR OpenAI-compat plumbing) now point at the native + endpoint without the user re-saving the connection.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://generativelanguage.googleapis.com/v1beta" + + +def test_finish_reason_swaps_to_tool_calls_when_function_call_emitted(monkeypatch): + """Gemini emits finishReason="STOP" even for pure functionCall turns; + surface as `tool_calls` so OAI clients trigger tool execution.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"functionCall": {"name": "lookup", "args": {"k": "v"}}} + ], + }, + "finishReason": "STOP", + } + ] + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + finish_chunks = [ + c for c in chunks if c.get("choices", [{}])[0].get("finish_reason") is not None + ] + assert finish_chunks, chunks + assert finish_chunks[-1]["choices"][0]["finish_reason"] == "tool_calls", chunks + + +def test_thought_signature_round_trips_into_gemini_function_call(monkeypatch): + """An assistant tool_call carrying `extra_content.google.thought_signature` + must echo the value back as a sibling of the Gemini functionCall part.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "lookup x"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "extra_content": {"google": {"thought_signature": "SIG-ABC"}}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_0", + "name": "lookup", + "content": "{}", + }, + ], + ) + contents = captured["body"]["contents"] + fc_turn = next((c for c in contents if c["role"] == "model"), None) + assert fc_turn is not None, contents + fc_part = next( + (p for p in fc_turn["parts"] if "functionCall" in p), + None, + ) + assert fc_part is not None, fc_turn + assert fc_part.get("thoughtSignature") == "SIG-ABC", fc_part + + +def test_thought_signature_emitted_in_tool_call_delta(monkeypatch): + """A Gemini functionCall part with `thoughtSignature` must surface + that signature on the outbound OpenAI tool_calls delta via + `extra_content.google.thought_signature`.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "lookup", + "args": {"k": "v"}, + "id": "call_xyz", + }, + "thoughtSignature": "SIG-FROM-GEMINI", + } + ], + }, + "finishReason": "STOP", + } + ] + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + deltas = [ + tc + for c in chunks + for tc in (c.get("choices", [{}])[0].get("delta", {}) or {}).get( + "tool_calls", [] + ) + ] + assert deltas, chunks + sig = deltas[0].get("extra_content", {}).get("google", {}).get("thought_signature") + assert sig == "SIG-FROM-GEMINI", deltas + + +def test_image_models_suppress_phantom_web_search_card(monkeypatch): + """When the image guard filters googleSearch out of the outbound + request, the inbound stream must NOT emit web_search tool_start / + tool_end (otherwise the UI shows a misleading 'Search complete' + card on a turn where Gemini never actually searched).""" + sse = [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "drawn"}]}, + "finishReason": "STOP", + } + ] + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation", "web_search", "code_execution"], + ) + chunks = _parse_chunks(lines) + tool_evs = [ + ev + for c in chunks + for ev in [c.get("_toolEvent")] + if isinstance(ev, dict) and ev.get("tool_name") == "web_search" + ] + assert tool_evs == [], tool_evs + + +def test_image_generation_tool_on_image_model_drops_text_tools(monkeypatch): + """`enabled_tools=["image_generation", "web_search", "code_execution"]` + on a Gemini IMAGE model flips responseModalities to TEXT+IMAGE; in + that mode codeExecution must NOT be forwarded (Gemini rejects text + code tools alongside image responseModalities). Older image + families also drop googleSearch.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = [ + "image_generation", + "web_search", + "code_execution", + ], + ) + assert "tools" not in captured["body"], captured["body"] + assert captured["body"]["generationConfig"].get("responseModalities") == [ + "TEXT", + "IMAGE", + ] + + +def test_prompt_feedback_block_reason_surfaces_as_error(monkeypatch): + """`promptFeedback.blockReason` with zero candidates must produce + an error chunk, not a silent empty assistant reply.""" + sse = [ + { + "promptFeedback": {"blockReason": "SAFETY"}, + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + error_chunks = [c for c in chunks if "error" in c] + assert error_chunks, chunks + assert "SAFETY" in ( + error_chunks[0].get("error", {}).get("message") or "" + ), error_chunks + + +def test_usage_chunk_includes_thoughts_tokens(monkeypatch): + """`thoughtsTokenCount` is the hidden-reasoning slice of output; + roll it into `output_tokens` AND surface it on + `output_tokens_details.reasoning_tokens` so total_tokens reflects + the full billable spend.""" + sse = [ + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 20, + "totalTokenCount": 35, + }, + } + ] + chunks = _parse_chunks(_collect(monkeypatch, sse)) + usage_chunk = next((c for c in chunks if isinstance(c.get("usage"), dict)), None) + assert usage_chunk is not None, chunks + usage = usage_chunk["usage"] + assert usage.get("prompt_tokens") == 10, usage + # candidates 5 + thoughts 20 = 25 output tokens; total = 35. + assert usage.get("completion_tokens") == 25, usage + assert usage.get("total_tokens") == 35, usage + + +# ── web_search forwarded as googleSearch tool ──────────────────────── + + +def test_web_search_forwarded_as_google_search_tool(monkeypatch): + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search"], + ) + tools = captured["body"].get("tools") or [] + assert {"googleSearch": {}} in tools, tools + + +def test_code_execution_forwarded_as_code_execution_tool(monkeypatch): + captured = _capture_body( + monkeypatch, + enabled_tools = ["code_execution"], + ) + tools = captured["body"].get("tools") or [] + assert {"codeExecution": {}} in tools, tools + + +def test_omitted_tools_leaves_body_untouched(monkeypatch): + captured = _capture_body(monkeypatch, enabled_tools = []) + assert "tools" not in captured["body"], captured["body"] + + +# ── prompt caching passthrough ─────────────────────────────────────── + + +def test_cached_content_pass_through(monkeypatch): + """A string cache id on enable_prompt_caching is forwarded verbatim.""" + cache_name = "cachedContents/abc123" + captured = _capture_body( + monkeypatch, + enable_prompt_caching = cache_name, + ) + assert captured["body"].get("cachedContent") == cache_name + + +def test_boolean_caching_does_not_set_cached_content(monkeypatch): + """Studio's existing True/False signals shouldn't fabricate a cache id.""" + captured = _capture_body(monkeypatch, enable_prompt_caching = True) + assert "cachedContent" not in captured["body"] + + +# ── image generation: request modalities + response translation ────── + + +def test_image_model_sets_response_modalities(monkeypatch): + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + ) + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] + + +def test_image_generation_tool_sets_response_modalities_on_image_model(monkeypatch): + """`enabled_tools=["image_generation"]` flips responseModalities + only when the selected model is image-capable; otherwise the + request stays plain text (text-only models 400 on + responseModalities).""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + ) + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] + + +def test_image_response_emits_image_b64_tool_event(monkeypatch): + """`inlineData` parts become a tool_end with image_b64 + image_mime.""" + fake_b64 = base64.b64encode(b"PNG-BYTES").decode() + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": fake_b64, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 0, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + image_starts = [e for e in starts if e.get("tool_name") == "image_generation"] + image_ends = [e for e in ends if e.get("image_b64")] + assert len(image_starts) == 1, tool_events + assert len(image_ends) == 1, tool_events + assert image_ends[0]["image_b64"] == fake_b64 + assert image_ends[0]["image_mime"] == "image/png" + + +# ── function calling round-trips both directions ───────────────────── + + +def test_function_call_response_translates_to_tool_calls_delta(monkeypatch): + """Gemini `functionCall` parts become OpenAI `tool_calls` delta chunks.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "get_weather", + "args": {"location": "Paris"}, + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 12, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + tool_call_chunks = [ + c + for c in chunks + if "_toolEvent" not in c + and any( + (isinstance(ch.get("delta"), dict) and "tool_calls" in ch["delta"]) + for ch in c.get("choices", []) + ) + ] + assert len(tool_call_chunks) == 1, chunks + tc = tool_call_chunks[0]["choices"][0]["delta"]["tool_calls"][0] + assert tc["function"]["name"] == "get_weather" + args = json.loads(tc["function"]["arguments"]) + assert args == {"location": "Paris"} + + +def test_tool_message_translates_to_function_response_part(monkeypatch): + """role=tool follow-ups are rewritten to functionResponse parts.""" + messages = [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "Paris"}), + }, + } + ], + }, + { + "role": "tool", + "name": "get_weather", + "content": json.dumps({"temp_c": 18, "summary": "Sunny"}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + # Last turn must be a functionResponse part (Gemini wraps it as a + # role=user turn carrying the result). + last = contents[-1] + assert last["role"] == "user", last + fr = last["parts"][0].get("functionResponse") + assert fr is not None, last + assert fr["name"] == "get_weather" + assert fr["response"] == {"temp_c": 18, "summary": "Sunny"} + # And the assistant turn carries the original functionCall so the + # model sees the round-trip context. + assistant_turn = [c for c in contents if c["role"] == "model"][0] + fc_part = next( + (p for p in assistant_turn["parts"] if "functionCall" in p), + None, + ) + assert fc_part is not None, assistant_turn + assert fc_part["functionCall"]["name"] == "get_weather" + assert fc_part["functionCall"]["args"] == {"location": "Paris"} + + +def test_parallel_function_calls_get_distinct_tool_call_indices(monkeypatch): + """Each emitted functionCall in one assistant turn needs its own + tool_calls[*].index. Hardcoding index=0 collapses parallel calls + onto a single slot in OpenAI-style reassemblers.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "id": "call_alpha", + "name": "search", + "args": {"q": "alpha"}, + } + }, + { + "functionCall": { + "id": "call_beta", + "name": "search", + "args": {"q": "beta"}, + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + tool_call_chunks = [ + c + for c in chunks + if "_toolEvent" not in c + and any( + (isinstance(ch.get("delta"), dict) and "tool_calls" in ch["delta"]) + for ch in c.get("choices", []) + ) + ] + assert len(tool_call_chunks) == 2, tool_call_chunks + indices = [ + c["choices"][0]["delta"]["tool_calls"][0]["index"] for c in tool_call_chunks + ] + assert indices == [0, 1], indices + + +def test_function_call_ids_forwarded_into_gemini_function_call_part(monkeypatch): + """OpenAI tool_call id rides functionCall.id so parallel calls disambiguate.""" + messages = [ + {"role": "user", "content": "x"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_alpha", + "type": "function", + "function": { + "name": "search", + "arguments": json.dumps({"q": "a"}), + }, + }, + { + "id": "call_beta", + "type": "function", + "function": { + "name": "search", + "arguments": json.dumps({"q": "b"}), + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_alpha", + "content": json.dumps({"hits": ["A"]}), + }, + { + "role": "tool", + "tool_call_id": "call_beta", + "content": json.dumps({"hits": ["B"]}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + assistant_parts = next(c for c in contents if c["role"] == "model")["parts"] + call_ids = [p["functionCall"]["id"] for p in assistant_parts if "functionCall" in p] + assert call_ids == ["call_alpha", "call_beta"], assistant_parts + response_ids = [ + p["functionResponse"]["id"] + for c in contents + for p in c["parts"] + if "functionResponse" in p + ] + assert response_ids == ["call_alpha", "call_beta"], contents + + +def test_parse_gemini_models_translates_native_catalog(): + """Gemini's native /v1beta/models payload becomes OpenAI-shape entries.""" + payload = { + "models": [ + { + "name": "models/gemini-2.5-flash", + "baseModelId": "gemini-2.5-flash", + "displayName": "Gemini 2.5 Flash", + "supportedGenerationMethods": [ + "generateContent", + "streamGenerateContent", + ], + }, + { + "name": "models/embedding-001", + "supportedGenerationMethods": ["embedContent"], + }, + { + "name": "models/gemini-2.5-pro", + }, + ] + } + out = ExternalProviderClient._parse_gemini_models(payload) + ids = [m["id"] for m in out] + assert "gemini-2.5-flash" in ids + assert "gemini-2.5-pro" in ids + assert "embedding-001" not in ids + flash = next(m for m in out if m["id"] == "gemini-2.5-flash") + assert flash["display_name"] == "Gemini 2.5 Flash" + assert flash["owned_by"] == "google" + + +def test_code_execution_parts_translate_to_code_execution_tool_events(monkeypatch): + """executableCode + codeExecutionResult parts emit code_execution events.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "print(2+2)", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "4\n", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect(monkeypatch, sse, enabled_tools = ["code_execution"]) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + code_starts = [ + e + for e in tool_events + if e.get("type") == "tool_start" and e.get("tool_name") == "code_execution" + ] + code_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and "4" in str(e.get("result", "")) + ] + assert len(code_starts) == 1, tool_events + assert code_starts[0]["arguments"]["code"] == "print(2+2)" + assert code_starts[0]["arguments"]["language"] == "python" + assert len(code_ends) == 1, tool_events + # tool_start and tool_end must share the same tool_call_id so the + # frontend pairs them onto a single CodeExecutionToolUI block. + assert code_starts[0]["tool_call_id"] == code_ends[0]["tool_call_id"] + + +def test_code_execution_failure_outcome_surfaces_in_result(monkeypatch): + """OUTCOME_FAILED is prefixed onto the result text so the UI shows it.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "1/0", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_FAILED", + "output": "ZeroDivisionError", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 2, + }, + } + ] + lines = _collect(monkeypatch, sse, enabled_tools = ["code_execution"]) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + result_text = next( + (e["result"] for e in tool_events if e.get("type") == "tool_end"), + "", + ) + assert "OUTCOME_FAILED" in result_text + assert "ZeroDivisionError" in result_text + + +def test_tool_message_recovers_name_from_tool_call_id(monkeypatch): + """When name is omitted, recover it from the matching tool_call_id.""" + messages = [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_xyz", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "Paris"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_xyz", + "content": json.dumps({"temp_c": 18}), + }, + ] + captured = _capture_body(monkeypatch, messages = messages) + contents = captured["body"]["contents"] + last = contents[-1] + fr = last["parts"][0].get("functionResponse") + assert fr is not None, last + assert ( + fr["name"] == "get_weather" + ), "name should fall back to the prior tool_call's function name" + + +# ── usage chunk surfaces promptTokenCount / candidatesTokenCount ───── + + +def test_usage_chunk_translates_gemini_token_counts(monkeypatch): + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1234, + "candidatesTokenCount": 56, + "cachedContentTokenCount": 1000, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("choices") == [] and "usage" in c] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["prompt_tokens"] == 1234 + assert usage["completion_tokens"] == 56 + assert usage["total_tokens"] == 1290 + assert usage["prompt_tokens_details"]["cached_tokens"] == 1000 + + +# ── multimodal: vision image -> inlineData ─────────────────────────── + + +def test_vision_data_url_translates_to_inline_data(monkeypatch): + fake = base64.b64encode(b"JPGBYTES").decode() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{fake}", + }, + }, + ], + } + ] + captured = _capture_body(monkeypatch, messages = messages) + parts = captured["body"]["contents"][0]["parts"] + inline_parts = [p for p in parts if "inlineData" in p] + assert len(inline_parts) == 1, parts + assert inline_parts[0]["inlineData"] == { + "mimeType": "image/jpeg", + "data": fake, + } + + +# ── finish reason mapping ──────────────────────────────────────────── + + +@pytest.mark.parametrize( + "gemini_reason, openai_reason", + [ + ("STOP", "stop"), + ("MAX_TOKENS", "length"), + ("SAFETY", "content_filter"), + ("PROHIBITED_CONTENT", "content_filter"), + ], +) +def test_finish_reason_translation(monkeypatch, gemini_reason, openai_reason): + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "x"}], + }, + "finishReason": gemini_reason, + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + finish_chunks = [ + c for c in chunks if any(ch.get("finish_reason") for ch in c.get("choices", [])) + ] + assert any( + ch["choices"][0]["finish_reason"] == openai_reason for ch in finish_chunks + ), finish_chunks + + +# ── grounding citations surface as web_search tool_end ─────────────── + + +def test_grounding_metadata_surfaces_as_tool_end_citations(monkeypatch): + """`groundingMetadata.groundingChunks[].web` -> tool_end result block.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Answer with sources."}], + }, + "groundingMetadata": { + "groundingChunks": [ + { + "web": { + "uri": "https://example.com/a", + "title": "Example A", + } + }, + { + "web": { + "uri": "https://example.com/b", + "title": "Example B", + } + }, + ] + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 7, + "candidatesTokenCount": 3, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["web_search"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + web_search_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "gemini_web_search" + ] + assert len(web_search_ends) == 1, tool_events + result = web_search_ends[0]["result"] + assert "https://example.com/a" in result + assert "https://example.com/b" in result + assert "Example A" in result + assert "Example B" in result + + +# ── round 3 review follow-ups ───────────────────────────────────────── + + +def test_custom_gemini_proxy_base_url_not_rewritten(): + """Only the Google-hosted /v1beta/openai base is normalized; a + custom gateway whose path ends in /openai must be left alone.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://proxy.example.com/team/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://proxy.example.com/team/openai" + + +def test_custom_gemini_proxy_uses_openai_dispatch(): + """Any non-Google Gemini base (LiteLLM, custom OpenAI-compat + routers) must route through the OpenAI-compatible forwarder, not + the native translator. Auth uses Authorization: Bearer ..., not + x-goog-api-key.""" + for base in ( + "https://proxy.example.com/team/openai", + "https://proxy.example.com/v1", + "https://litellm.internal.example/v1", + ): + client = ExternalProviderClient( + provider_type = "gemini", + base_url = base, + api_key = "AIza-test-key", + ) + assert client._is_openai_compatible() is True, base + headers = client._auth_headers() + assert "x-goog-api-key" not in {k.lower() for k in headers}, ( + base, + headers, + ) + assert headers["Authorization"] == "Bearer AIza-test-key", ( + base, + headers, + ) + + +def test_google_hosted_gemini_still_uses_native_dispatch(): + """Google-hosted Gemini keeps native dispatch + x-goog-api-key auth.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + api_key = "AIza-test-key", + ) + assert client._is_openai_compatible() is False + headers = client._auth_headers() + assert headers.get("x-goog-api-key") == "AIza-test-key", headers + + +def test_invalid_gemini_model_id_rejected_before_request(monkeypatch): + """Path-traversal model ids must be rejected before the URL is + interpolated so the configured API key isn't sent to unintended + Gemini endpoints.""" + + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response( + 200, + content = _gemini_sse([]), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + out: list[str] = [] + + async def run(): + client = _make_gemini_client() + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "../cachedContents/leak", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + ): + out.append(line) + await client.close() + + _drive(run()) + # No outbound request should have been issued. + assert captured == [], captured + error_lines = [line for line in out if '"error"' in line] + assert error_lines, out + + +def test_top_k_omitted_when_not_explicit_default_for_gemini(monkeypatch): + """top_k=None means "use provider default"; helper must not emit + `topK` in generationConfig when the caller didn't pass it.""" + captured = _capture_body(monkeypatch, top_k = None) + assert "topK" not in captured["body"]["generationConfig"], captured["body"] + + +def test_text_model_image_generation_tool_silently_dropped(monkeypatch): + """A stale `enabled_tools=["image_generation"]` on a text-only + Gemini model (e.g. gemini-2.5-flash) must NOT switch the request + into image mode -- Google's API 400s on responseModalities for + text models.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash", + enabled_tools = ["image_generation"], + ) + gc = captured["body"]["generationConfig"] + assert "responseModalities" not in gc, gc + + +def test_empty_text_part_with_thought_signature_emits_extra_content( + monkeypatch, +): + """Gemini 3 can ship a content-free fragment whose only payload is + `thoughtSignature`. The translator must still surface that signature + on a delta.extra_content envelope so the next turn can replay it.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"text": "answer"}, + {"thoughtSignature": "SIG-FINAL"}, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 2, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + extra_carriers = [ + c + for c in chunks + if c.get("choices") + and c["choices"][0]["delta"].get("extra_content") + == {"google": {"thought_signature": "SIG-FINAL"}} + ] + assert extra_carriers, chunks + + +def test_enable_prompt_caching_false_string_coerces_to_bool(): + """Pre-PR the field was Optional[bool]; widening to Union[bool,str] + must preserve historical coercion so callers sending `"false"` + still opt out of caching.""" + from models.inference import ChatCompletionRequest + + msg = {"role": "user", "content": "hi"} + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "false", + } + ) + assert req.enable_prompt_caching is False, req.enable_prompt_caching + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "true", + } + ) + assert req.enable_prompt_caching is True + + # An actual cache resource name passes through untouched. + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [msg], + "enable_prompt_caching": "cachedContents/abc123", + } + ) + assert req.enable_prompt_caching == "cachedContents/abc123" + + +def test_legacy_google_openai_base_url_is_rewritten(): + """The Google-hosted /v1beta/openai legacy base IS still rewritten.""" + client = ExternalProviderClient( + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta/openai", + api_key = "AIza-test-key", + ) + assert client.base_url == "https://generativelanguage.googleapis.com/v1beta" + + +def test_remote_image_url_downloads_and_inlines_as_base64(monkeypatch): + """Round 14: arbitrary public HTTPS image URLs cannot be sent as + Gemini fileData (that path is reserved for Files API URIs and + YouTube). The translator must fetch the bytes server-side and + inline them as base64 inlineData.""" + image_bytes = b"FAKEPNGBYTES" + + async def fake_fetch(url, fallback_mime, max_bytes = None): + assert url == "https://cdn.example.com/diagram.png" + return ("image/png", base64.b64encode(image_bytes).decode("ascii")) + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": { + "url": "https://cdn.example.com/diagram.png", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + inline = next((p for p in parts if "inlineData" in p), None) + assert inline is not None, parts + assert inline["inlineData"]["mimeType"] == "image/png" + assert inline["inlineData"]["data"] == base64.b64encode(image_bytes).decode() + assert not any("fileData" in p for p in parts), parts + + +def test_remote_image_url_dropped_when_fetch_returns_none(monkeypatch): + """Round 15: if the SSRF guard rejects the URL (private host, + non-https, oversize, non-image), the helper returns None and the + image part is silently dropped instead of forwarding raw bytes + or a fileData fallback.""" + + async def fake_fetch_reject(url, fallback_mime, max_bytes = None): + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch_reject) + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": {"url": "http://10.0.0.5/private.png"}, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + assert not any("inlineData" in p for p in parts), parts + assert not any("fileData" in p for p in parts), parts + + +def test_safe_fetch_image_rejects_non_https(): + """SSRF guard: only https URLs may be fetched.""" + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini("http://cdn.example.com/x.png", "image/png") + ) + assert res is None + + +def test_safe_fetch_image_rejects_loopback_ip_literal(): + """SSRF guard: refuse loopback / private IP literals before any + network call.""" + for url in ( + "https://127.0.0.1/x.png", + "https://[::1]/x.png", + "https://169.254.169.254/latest/meta-data", + "https://10.0.0.5/x.png", + "https://192.168.1.1/x.png", + ): + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini(url, "image/png") + ) + assert res is None, url + + +def test_safe_fetch_image_rejects_resolved_private_host(monkeypatch): + """SSRF guard: if a hostname resolves to a private IP, refuse.""" + import socket + + def fake_getaddrinfo(host, *_args, **_kwargs): + return [(socket.AF_INET, None, None, "", ("10.0.0.5", 0))] + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + res = asyncio.new_event_loop().run_until_complete( + ep_mod._safe_fetch_image_for_gemini( + "https://internal.example/x.png", "image/png" + ) + ) + assert res is None + + +def test_youtube_and_files_api_uris_stay_as_file_data(monkeypatch): + """Round 14: YouTube URLs and generativelanguage.googleapis.com + Files API URIs are the documented `fileData.fileUri` paths and + must NOT be downloaded; arbitrary public URLs do get fetched.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "explain"}, + { + "type": "image_url", + "image_url": { + "url": "https://www.youtube.com/watch?v=abc123", + }, + }, + { + "type": "image_url", + "image_url": { + "url": "https://generativelanguage.googleapis.com/v1beta/files/abc", + }, + }, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + parts = captured["body"]["contents"][-1]["parts"] + file_uris = [p["fileData"]["fileUri"] for p in parts if "fileData" in p] + assert "https://www.youtube.com/watch?v=abc123" in file_uris, parts + assert ( + "https://generativelanguage.googleapis.com/v1beta/files/abc" in file_uris + ), parts + + +def test_tool_use_prompt_tokens_added_to_input_tokens(monkeypatch): + """`toolUsePromptTokenCount` must roll into the OpenAI prompt + total -- otherwise tool turns silently undercount input tokens.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "result"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "toolUsePromptTokenCount": 100, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 2, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("usage")] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["prompt_tokens"] == 110, usage + assert usage["completion_tokens"] == 7, usage + assert usage["total_tokens"] == 117, usage + assert usage["completion_tokens_details"]["reasoning_tokens"] == 2, usage + + +def test_usage_chunk_reasoning_tokens_surfaced(monkeypatch): + """thoughtsTokenCount must surface as completion_tokens_details. + reasoning_tokens in the emitted OpenAI usage chunk.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 5, + "thoughtsTokenCount": 20, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + usage_chunks = [c for c in chunks if c.get("usage")] + assert len(usage_chunks) == 1, chunks + usage = usage_chunks[0]["usage"] + assert usage["completion_tokens"] == 25, usage + assert usage["completion_tokens_details"]["reasoning_tokens"] == 20, usage + + +def test_prompt_block_pairs_web_search_tool_end(monkeypatch): + """When `promptFeedback.blockReason` triggers after the synthetic + web_search tool_start, the helper must emit a matching tool_end so + the UI does not leave a "searching..." spinner stuck on screen.""" + sse = [ + {"promptFeedback": {"blockReason": "SAFETY"}}, + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["web_search"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + assert len(starts) == 1, tool_events + assert len(ends) == 1, tool_events + assert ends[0]["tool_call_id"] == "gemini_web_search" + assert "aborted" in ends[0]["result"] + error_chunks = [c for c in chunks if c.get("error")] + assert error_chunks, chunks + + +def test_code_execution_tool_events_stow_native_part(monkeypatch): + """executableCode / codeExecutionResult must round-trip native ids + and thoughtSignature in google.native_part so follow-up turns can + replay Gemini's required history shape.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(1+1)", + }, + "thoughtSignature": "SIG-CODE", + }, + { + "codeExecutionResult": { + "id": "result_a", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + starts = [e for e in tool_events if e.get("type") == "tool_start"] + ends = [e for e in tool_events if e.get("type") == "tool_end"] + code_start = next( + (e for e in starts if e.get("tool_name") == "code_execution"), + None, + ) + code_end = next(iter(ends), None) + assert code_start is not None, starts + assert code_start["tool_call_id"] == "code_a", code_start + native = code_start["arguments"]["google"]["native_part"] + # Round 21: native_part now uses an ordered `parts` list so per-part + # `thoughtSignature` survives a frontend merge of executableCode + + # codeExecutionResult into one tool-call card. + start_parts = native["parts"] + assert start_parts[0]["executableCode"]["id"] == "code_a" + assert start_parts[0]["thoughtSignature"] == "SIG-CODE" + assert code_end is not None, ends + assert code_end["tool_call_id"] == "code_a", code_end + native_end = code_end["google"]["native_part"] + end_parts = native_end["parts"] + assert end_parts[0]["codeExecutionResult"]["id"] == "result_a" + + +def test_inline_image_tool_end_carries_thought_signature(monkeypatch): + """Inline image parts with thoughtSignature must persist it on the + emitted tool_end so Gemini 3 image editing can echo it back.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": base64.b64encode(b"PNG").decode(), + }, + "thoughtSignature": "SIG-IMG", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 4, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + model = "gemini-2.5-flash-image", + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + image_ends = [ + e for e in tool_events if e.get("type") == "tool_end" and e.get("image_b64") + ] + assert image_ends, tool_events + assert image_ends[0]["google"]["thought_signature"] == "SIG-IMG" + # Multi-turn image edit must replay the original inlineData part with + # its thoughtSignature; the outbound translator reads + # google.native_part.parts[].inlineData, so stow it on the tool_end + # too. Round 21 changed native_part to an ordered parts list so a + # per-part signature stays attached to inlineData only. + native = image_ends[0]["google"]["native_part"] + image_parts = native["parts"] + assert image_parts[0]["inlineData"]["mimeType"] == "image/png" + assert image_parts[0]["inlineData"]["data"] == base64.b64encode(b"PNG").decode() + assert image_parts[0]["thoughtSignature"] == "SIG-IMG" + + +def test_code_execution_plot_attaches_inline_image_native_part(monkeypatch): + """A code_execution turn that returns a matplotlib plot must stow + the plot's inlineData on the secondary tool_end so the follow-up + turn can replay the image alongside executableCode and + codeExecutionResult.""" + plot_data = base64.b64encode(b"PLOT").decode() + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "plt.plot([0,1])", + }, + }, + { + "codeExecutionResult": { + "id": "result_a", + "outcome": "OUTCOME_OK", + "output": "", + }, + }, + { + "inlineData": { + "mimeType": "image/png", + "data": plot_data, + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + code_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_a" + ] + # Two tool_end events on the same id: one for codeExecutionResult, + # one merging in the inlineData plot. The plot one must carry the + # native inlineData under google.native_part so the frontend + # tool_end merge union joins it with the prior executableCode and + # codeExecutionResult parts on the same card. + assert len(code_ends) == 2, code_ends + image_end = next( + (e for e in code_ends if "__IMAGES__:" in (e.get("result") or "")), + None, + ) + assert image_end is not None, code_ends + native = image_end["google"]["native_part"] + plot_parts = native["parts"] + assert plot_parts[0]["inlineData"]["mimeType"] == "image/png" + assert plot_parts[0]["inlineData"]["data"] == plot_data + + +def test_text_chunk_carries_thought_signature(monkeypatch): + """Text parts with thoughtSignature surface it on delta.extra_content + so frontend persistence can replay it on the follow-up turn.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "hello", + "thoughtSignature": "SIG-TEXT", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 2, + "candidatesTokenCount": 1, + }, + } + ] + lines = _collect(monkeypatch, sse) + chunks = _parse_chunks(lines) + text_chunks = [ + c + for c in chunks + if c.get("choices") and c["choices"][0]["delta"].get("content") == "hello" + ] + assert text_chunks, chunks + extra = text_chunks[0]["choices"][0]["delta"].get("extra_content") + assert extra == {"google": {"thought_signature": "SIG-TEXT"}}, text_chunks + + +def test_openai_tools_translated_into_function_declarations(monkeypatch): + """Standard ChatCompletionRequest.tools must be forwarded into + Gemini's tools[].functionDeclarations envelope.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + }, + "required": ["city"], + }, + }, + } + ], + tool_choice = {"type": "function", "function": {"name": "get_weather"}}, + ) + tools_arr = captured["body"].get("tools") or [] + fn_decls = [t for t in tools_arr if "functionDeclarations" in t] + assert fn_decls, captured["body"] + decls = fn_decls[0]["functionDeclarations"] + assert decls[0]["name"] == "get_weather" + assert decls[0]["parameters"]["properties"]["city"]["type"] == "string" + tool_config = captured["body"].get("toolConfig") + assert tool_config is not None, captured["body"] + fcc = tool_config["functionCallingConfig"] + assert fcc["mode"] == "ANY" + assert fcc["allowedFunctionNames"] == ["get_weather"] + + +def test_tool_choice_auto_maps_to_function_calling_mode_auto(monkeypatch): + """tool_choice="auto" maps to toolConfig.functionCallingConfig.mode.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": {"name": "noop", "parameters": {"type": "object"}}, + } + ], + tool_choice = "auto", + ) + fcc = captured["body"]["toolConfig"]["functionCallingConfig"] + assert fcc["mode"] == "AUTO" + assert "allowedFunctionNames" not in fcc + + +def test_code_exec_inline_image_attaches_to_code_execution_card(monkeypatch): + """A codeExecution sandbox plot (matplotlib) ships as an inline + image part right after the codeExecutionResult. Instead of spawning + a separate empty image_generation card, attach to the same + code_execution tool_end via the `__IMAGES__:` marker the chat + adapter already understands.""" + sse = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "executableCode": { + "id": "code_plot", + "language": "PYTHON", + "code": "import matplotlib.pyplot as plt; plt.plot([1,2,3]); plt.savefig('out.png')", + }, + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "saved", + }, + }, + { + "inlineData": { + "mimeType": "image/png", + "data": base64.b64encode(b"PNGDATA").decode(), + }, + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 4, + }, + } + ] + lines = _collect( + monkeypatch, + sse, + enabled_tools = ["code_execution"], + ) + chunks = _parse_chunks(lines) + tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c] + # No standalone image_generation card should have been emitted. + image_starts = [ + e + for e in tool_events + if e.get("type") == "tool_start" and e.get("tool_name") == "image_generation" + ] + assert not image_starts, tool_events + # The code_execution tool_end should now carry the inline image + # via the `__IMAGES__:` marker. + code_ends = [ + e + for e in tool_events + if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_plot" + ] + assert code_ends, tool_events + final_result = code_ends[-1]["result"] + assert "__IMAGES__:" in final_result, code_ends + assert "data:image/png;base64," in final_result, code_ends + + +def test_code_execution_tool_call_replays_native_executable_code(monkeypatch): + """An assistant tool_call with toolName=code_execution and + extra_content.google.native_part containing the originally-emitted + `executableCode` + `codeExecutionResult` must round-trip as native + Gemini parts (not a generic functionCall) on the next turn.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "compute 2+2"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "code_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(2+2)", + }, + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "4\n", + }, + "thoughtSignature": "SIG-CODE", + }, + }, + }, + }, + ], + }, + {"role": "user", "content": "what was that result"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + native_keys = [list(p.keys())[0] for p in parts if isinstance(p, dict)] + assert "executableCode" in native_keys, parts + assert "codeExecutionResult" in native_keys, parts + assert not any( + "functionCall" in p + and (p["functionCall"] or {}).get("name") == "code_execution" + for p in parts + ), parts + exec_part = next(p for p in parts if "executableCode" in p) + assert exec_part.get("thoughtSignature") == "SIG-CODE", exec_part + + +def test_image_generation_tool_call_replays_native_inline_data(monkeypatch): + """An assistant tool_call with toolName=image_generation and + extra_content.google.native_part.inlineData must replay the prior + image as a native Gemini inlineData part (not a generic + functionCall) so multi-turn image editing keeps the image + context.""" + pixel = base64.b64encode(b"PNG").decode() + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + messages = [ + {"role": "user", "content": "make a circle"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "img_a", + "type": "function", + "function": { + "name": "image_generation", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "inlineData": { + "mimeType": "image/png", + "data": pixel, + }, + "thoughtSignature": "SIG-IMG", + }, + }, + }, + }, + ], + }, + {"role": "user", "content": "now make it blue"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + inline_parts = [p for p in parts if "inlineData" in p] + assert inline_parts, parts + assert inline_parts[0]["inlineData"]["mimeType"] == "image/png" + assert inline_parts[0]["inlineData"]["data"] == pixel + assert inline_parts[0].get("thoughtSignature") == "SIG-IMG", inline_parts + assert not any( + "functionCall" in p + and (p["functionCall"] or {}).get("name") == "image_generation" + for p in parts + ), parts + + +def test_assistant_text_thought_signature_replays_on_outbound_text_part(monkeypatch): + """Assistant text with extra_content.google.thought_signature must + attach `thoughtSignature` to the LAST text part of the replayed + Gemini history. Gemini 3 strict function-calling rejects history + that drops returned signatures, so the frontend stows the latest + signed-text signature and the backend pins it on the next turn.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "hello"}, + ], + "extra_content": { + "google": {"thought_signature": "SIG-TEXT"}, + }, + }, + {"role": "user", "content": "again"}, + ], + ) + assistant_turn = captured["body"]["contents"][1] + assert assistant_turn["role"] == "model" + parts = assistant_turn["parts"] + text_parts = [p for p in parts if "text" in p] + assert text_parts, parts + assert text_parts[-1].get("thoughtSignature") == "SIG-TEXT", text_parts + + +def test_function_declarations_strip_openai_only_schema_keys(monkeypatch): + """OpenAI strict tools commonly include `additionalProperties`, + `$schema`, `$defs`, `strict`, etc. Gemini's Schema rejects those + with INVALID_ARGUMENT, so the translator must strip them while + keeping properties..type intact.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a value.", + "parameters": { + "type": "object", + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": False, + "strict": True, + "properties": { + "key": { + "type": "string", + "additionalProperties": False, + }, + }, + "required": ["key"], + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), + None, + ) + assert decls is not None, captured["body"] + params = decls[0]["parameters"] + assert "additionalProperties" not in params + assert "$schema" not in params + assert "strict" not in params + assert params["type"] == "object" + assert params["properties"]["key"]["type"] == "string" + assert "additionalProperties" not in params["properties"]["key"] + assert params["required"] == ["key"] + + +def test_function_declarations_inline_local_refs_into_gemini_schema(monkeypatch): + """Round 25: Pydantic-generated tool schemas hoist nested object + shapes into `$defs` and reference them with `{"$ref": "#/$defs/..."}`. + Gemini's OpenAPI subset has no $ref, so a naive allowlist sanitizer + drops the reference and reduces the nested property to `{}`, losing + its type, fields, and required keys. The sanitizer must resolve + local `#/...` pointers and inline the referenced schema.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "set_user", + "description": "Persist a user.", + "parameters": { + "type": "object", + "$defs": { + "Address": { + "type": "object", + "properties": { + "street": {"type": "string"}, + "zip": {"type": "string"}, + }, + "required": ["street", "zip"], + }, + }, + "properties": { + "name": {"type": "string"}, + "address": {"$ref": "#/$defs/Address"}, + }, + "required": ["name", "address"], + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), + None, + ) + assert decls is not None, captured["body"] + params = decls[0]["parameters"] + assert "$defs" not in params + address = params["properties"]["address"] + assert address.get("type") == "object", address + assert address.get("properties", {}).get("street", {}).get("type") == "string" + assert address.get("properties", {}).get("zip", {}).get("type") == "string" + assert address.get("required") == ["street", "zip"] + + +def test_function_declarations_inline_local_refs_in_anyof_and_items(monkeypatch): + """The recursive inliner must reach through `anyOf` branches and + `items` (array element schemas) as well, not just top-level + property refs.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "bulk_set", + "parameters": { + "type": "object", + "$defs": { + "Address": { + "type": "object", + "properties": {"zip": {"type": "string"}}, + "required": ["zip"], + }, + }, + "properties": { + "primary": { + "anyOf": [ + {"$ref": "#/$defs/Address"}, + {"type": "null"}, + ], + }, + "extras": { + "type": "array", + "items": {"$ref": "#/$defs/Address"}, + }, + }, + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), + None, + ) + assert decls is not None + params = decls[0]["parameters"] + primary = params["properties"]["primary"] + # anyOf with single non-null branch + null collapses to inline + + # nullable: true, and the inlined branch must contain the resolved + # Address shape. + assert primary.get("nullable") is True + assert primary.get("type") == "object" + assert primary.get("properties", {}).get("zip", {}).get("type") == "string" + extras = params["properties"]["extras"] + assert extras.get("type") == "array" + assert extras.get("items", {}).get("type") == "object" + assert ( + extras.get("items", {}).get("properties", {}).get("zip", {}).get("type") + == "string" + ) + + +def test_function_declarations_self_referential_schema_terminates(monkeypatch): + """Self-referential / cyclic JSON Schemas (a `Node` that contains + `children: [Node]`) must not infinite-loop. The inliner tracks the + set of refs in flight and short-circuits to `{}` on a cycle.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "set_tree", + "parameters": { + "type": "object", + "$defs": { + "Node": { + "type": "object", + "properties": { + "value": {"type": "string"}, + "children": { + "type": "array", + "items": {"$ref": "#/$defs/Node"}, + }, + }, + }, + }, + "properties": { + "root": {"$ref": "#/$defs/Node"}, + }, + }, + }, + } + ], + ) + tools_arr = captured["body"].get("tools") or [] + decls = next( + ( + t.get("functionDeclarations") + for t in tools_arr + if "functionDeclarations" in t + ), + None, + ) + assert decls is not None + root = decls[0]["parameters"]["properties"]["root"] + assert root.get("type") == "object" + assert root.get("properties", {}).get("value", {}).get("type") == "string" + + +def test_gemini_native_skips_orphan_function_response_for_dropped_builtin( + monkeypatch, +): + """Round 26: when the assistant-side synthetic web_search/web_fetch + tool_call is dropped from native Gemini history, the matching + role="tool" follow-up must also be dropped. Otherwise the outbound + body carries an orphan functionResponse with no preceding + functionCall, which 400s the Gemini turn.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_s", + "type": "function", + "function": { + "name": "web_search", + "arguments": ('{"_server_tool": true, "query": "x"}'), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_s", + "content": "[search result]", + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + captured = _capture_body(monkeypatch, messages = built) + contents = captured["body"].get("contents") or [] + for entry in contents: + for part in entry.get("parts", []): + fr = part.get("functionResponse") + if isinstance(fr, dict): + assert fr.get("name") != "web_search", contents + + +def test_gemini_native_skips_orphan_function_response_for_native_part_replay( + monkeypatch, +): + """Round 26: code_execution / image_generation tool_calls are + replayed as Gemini-native executableCode / codeExecutionResult / + inlineData parts. The matching role="tool" follow-up must NOT then + be emitted as a functionResponse named code_execution -- there is + no declared user function with that name, and Gemini's history + rules already attribute the result to the native parts above.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "plot something"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "parts": [ + { + "executableCode": { + "language": "PYTHON", + "code": "print(2)", + } + }, + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "2\n", + } + }, + ] + } + } + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ], + ) + contents = captured["body"].get("contents") or [] + saw_native = False + for entry in contents: + for part in entry.get("parts", []): + if "executableCode" in part or "codeExecutionResult" in part: + saw_native = True + fr = part.get("functionResponse") + if isinstance(fr, dict): + assert fr.get("name") != "code_execution", contents + assert saw_native, contents + + +def test_gemini_native_part_falls_back_to_args_google(monkeypatch): + """Round 27: a direct OpenAI-compat API caller (or imported third- + party thread) cannot use Studio's non-standard + `tool_calls[].extra_content` field, so the native_part payload + round-trips through `function.arguments` as + `{"google": {"native_part": {...}}}`. The synthetic-builtin + detector recognizes that location, but the replay branch was only + reading from `tc.extra_content.google.native_part`. Result: the + round-25 guard saw a synthetic builtin with no _native_part and + dropped the entire assistant turn, losing the prior code/image + context. The translator must fall back to args.google.native_part + and still emit the native executableCode / inlineData parts.""" + import json as _json + + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "draw a cat"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_img", + "type": "function", + "function": { + "name": "image_generation", + "arguments": _json.dumps( + { + "google": { + "native_part": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "AAAA", + } + } + ] + } + } + } + ), + }, + } + ], + }, + {"role": "user", "content": "now make it a dog"}, + ], + ) + contents = captured["body"].get("contents") or [] + saw_inline = False + for entry in contents: + for part in entry.get("parts", []): + if "inlineData" in part: + saw_inline = True + assert saw_inline, contents + + +def test_gemini_native_skips_synthetic_server_builtin_replay(monkeypatch): + """Round 25: Marked server-side builtin tool_calls (web_search / + web_fetch with `_server_tool` or `args.google.native_part`) must + not fall through to the generic Gemini `functionCall` replay path + when no replayable native part exists. Without this guard the + outbound body contains a fake `functionCall` whose name is not a + declared user function, and the Gemini turn 400s.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_s", + "type": "function", + "function": { + "name": "web_search", + "arguments": ('{"_server_tool": true, "query": "x"}'), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_s", + "content": "[search result]", + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + captured = _capture_body(monkeypatch, messages = built) + contents = captured["body"].get("contents") or [] + for entry in contents: + for part in entry.get("parts", []): + fc = part.get("functionCall") + if isinstance(fc, dict): + assert fc.get("name") != "web_search", contents + + +def test_chat_message_extra_content_round_trips_through_validation(): + """Round 9: ChatMessage was missing `extra_content`, so Pydantic + discarded the field during request validation and the text-part + signature replay path read nothing. The field must survive + model_validate and pass through _build_external_messages.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + req = ChatCompletionRequest.model_validate( + { + "model": "gemini-2.5-flash", + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "hello"}, + ], + "extra_content": { + "google": {"thought_signature": "SIG-TEXT"}, + }, + }, + {"role": "user", "content": "again"}, + ], + "max_tokens": 64, + "stream": True, + } + ) + assistant_msg = req.messages[1] + assert assistant_msg.extra_content == { + "google": {"thought_signature": "SIG-TEXT"}, + } + built = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + assistant_out = built[1] + assert assistant_out["extra_content"] == { + "google": {"thought_signature": "SIG-TEXT"}, + } + # Non-Gemini providers must NOT receive extra_content; Google's + # thought_signature field is unknown to OpenAI / Mistral / etc. + built_openai = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + ) + assert "extra_content" not in built_openai[1], built_openai[1] + # Custom non-Google Gemini bases (LiteLLM / OAI-compat gateways) + # also must not receive Gemini-only extra_content because the + # backend dispatches them through /chat/completions. + built_custom = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://litellm.example/v1", + ) + assert "extra_content" not in built_custom[1], built_custom[1] + + +def test_parallel_tool_results_group_into_one_user_block(monkeypatch): + """Round 14: Gemini docs show parallel functionResponses grouped + in a single subsequent user content with multiple + functionResponse parts. Consecutive OpenAI role="tool" messages + must merge into one Gemini user block, not split into separate + user turns.""" + captured = _capture_body( + monkeypatch, + messages = [ + {"role": "user", "content": "compute"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "add", "arguments": '{"x":1}'}, + }, + { + "id": "call_b", + "type": "function", + "function": {"name": "mul", "arguments": '{"x":2}'}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "add", + "content": "2", + }, + { + "role": "tool", + "tool_call_id": "call_b", + "name": "mul", + "content": "4", + }, + ], + ) + contents = captured["body"]["contents"] + # Initial user, model with two functionCalls, ONE user with two + # functionResponses. + tool_result_users = [ + c + for c in contents + if c.get("role") == "user" + and all( + isinstance(p, dict) and "functionResponse" in p + for p in (c.get("parts") or []) + ) + ] + assert len(tool_result_users) == 1, contents + fr_parts = tool_result_users[0]["parts"] + assert len(fr_parts) == 2, fr_parts + names = [p["functionResponse"]["name"] for p in fr_parts] + assert names == ["add", "mul"], names + + +def test_function_schema_nullable_type_array_flattens(monkeypatch): + """Round 14: OpenAI strict tools commonly use + `"type": ["string", "null"]` for optional fields. Gemini's + OpenAPI-style Schema rejects union types and expects + `"type": "string"` with `"nullable": true`. The sanitizer must + translate the union form.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "city": {"type": ["string", "null"]}, + "score": {"type": ["number", "null"]}, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + params = decls[0]["parameters"]["properties"] + assert params["city"]["type"] == "string" + assert params["city"]["nullable"] is True + assert params["score"]["type"] == "number" + assert params["score"]["nullable"] is True + + +def test_image_picker_model_with_search_off_pill_strips_text_tools(monkeypatch): + """Round 11: image-tier model id rejects text-only tools and + thinkingConfig at the model level regardless of whether the Images + pill is on. Selecting gemini-2.5-flash-image + enabled_tools= + ["web_search"] with no image_generation must NOT forward + googleSearch or thinkingConfig (Gemini 400s on text tools for + legacy image ids).""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["web_search"], + reasoning_effort = "high", + ) + body = captured["body"] + assert "tools" not in body, body.get("tools") + assert "thinkingConfig" not in body.get("generationConfig", {}), body[ + "generationConfig" + ] + + +def test_image_models_drop_function_declarations(monkeypatch): + """Image-mode requests cannot mix tools with responseModalities so + user-supplied function declarations must be dropped.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tools = [ + { + "type": "function", + "function": {"name": "noop", "parameters": {"type": "object"}}, + } + ], + ) + assert captured["body"].get("tools") is None + assert captured["body"]["generationConfig"]["responseModalities"] == [ + "TEXT", + "IMAGE", + ] + + +def test_safe_fetch_image_rejects_malformed_bracketed_url(): + """Round 17: bracketed IPv6 garbage like `https://[bad/x.png` makes + urlparse raise ValueError. The fetch helper must catch it and drop + the image rather than crashing the request mid-build.""" + res = _drive(ep_mod._safe_fetch_image_for_gemini("https://[bad/x.png", "image/png")) + assert res is None + + +def test_safe_fetch_image_pins_validated_ip_no_hostname_in_request( + monkeypatch, +): + """Round 17: the fetch helper must pin the validated IP into the + outgoing request URL (with a Host header carrying the original + hostname). A second hostname-style getaddrinfo after the validate + step would be a DNS-rebinding gap, so we assert the urllib opener + is called with an IP-rewritten URL.""" + import socket + + captured: dict = {"requests": []} + + # Public IP during validate; record every getaddrinfo call. + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + captured.setdefault("dns", []).append(host) + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("8.8.8.8", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + headers = {"content-type": "image/png", "content-length": "3"} + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + return b"PNG" + + class _StubOpener: + def open(self, req, timeout = None): + captured["requests"].append( + { + "url": req.full_url, + "host_header": req.get_header("Host"), + } + ) + return _StubResp() + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) + assert res is not None + assert res[0] == "image/png" + # The outgoing URL must use the pinned IP literal, not the hostname. + assert any("8.8.8.8" in r["url"] for r in captured["requests"]), captured + assert all( + "cdn.example.com" not in r["url"] for r in captured["requests"] + ), captured + # Host header still carries the original hostname for vhost/SNI. + assert captured["requests"][0]["host_header"] == "cdn.example.com" + + +def test_safe_fetch_image_redirect_to_private_host_rejected(monkeypatch): + """Round 17: each redirect hop must re-validate the new host. A + public hop that redirects to an internal address must be dropped.""" + import socket + import urllib.error + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + if host == "internal.bad": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("10.0.0.5", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubOpener: + def open(self, req, timeout = None): + # Simulate a 302 to a private host. + raise urllib.error.HTTPError( + req.full_url, + 302, + "Found", + {"Location": "https://internal.bad/secret.png"}, + None, + ) + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) + assert res is None + + +def test_files_api_substring_url_not_misclassified_as_filedata(monkeypatch): + """Round 17: a CDN URL whose path/query merely contains the Files + API substring must NOT be sent as `fileData.fileUri`; it must be + routed through the safe-fetch path. Previously the substring check + `"generativelanguage.googleapis.com/" in url.lower()` matched any + URL carrying that text anywhere.""" + captured_outbound: dict = {} + fetch_calls: list[str] = [] + + async def fake_fetch(url, fallback_mime, max_bytes = None): + fetch_calls.append(url) + return "image/png", base64.b64encode(b"DATA").decode("ascii") + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + captured_outbound["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + { + "type": "image_url", + "image_url": { + # Looks like a Files API URL in the path + # but the host is an attacker CDN. + "url": "https://evil.example/path/generativelanguage.googleapis.com/v1beta/files/abc.png", + }, + }, + { + "type": "image_url", + "image_url": { + # Looks YouTube-ish in the path. + "url": "https://cdn.example.com/youtube.com/cat.png", + }, + }, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + parts = captured_outbound["body"]["contents"][-1]["parts"] + assert not any("fileData" in p for p in parts), parts + inline_count = sum(1 for p in parts if "inlineData" in p) + assert inline_count == 2, parts + assert len(fetch_calls) == 2, fetch_calls + + +def test_function_schema_anyof_null_variant_flattens_to_nullable(monkeypatch): + """Round 17: OpenAI/Pydantic emit `anyOf: [{X}, {"type":"null"}]` + for Optional[X]. Gemini's OpenAPI subset rejects `"type":"null"` + inside anyOf. The sanitizer must collapse a singleton-plus-null + union back to the non-null branch with `nullable: true`.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "label": { + "anyOf": [ + {"type": "string"}, + {"type": "null"}, + ] + }, + "count": { + "anyOf": [ + {"type": "integer"}, + {"type": "null"}, + ] + }, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + params = decls[0]["parameters"]["properties"] + assert params["label"]["type"] == "string" + assert params["label"]["nullable"] is True + assert "anyOf" not in params["label"] + assert params["count"]["type"] == "integer" + assert params["count"]["nullable"] is True + + +def test_legacy_gemini3_pro_medium_coerced_to_high(monkeypatch): + """Round 17: legacy `gemini-3-pro*` (including `-preview`, shut down + 2026-03-09) only accepted low/high. 3.1+ Pro added medium. The + backend must coerce medium → high for the legacy model so stale UI + state does not 400 the request.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3-pro-preview", + reasoning_effort = "medium", + ) + assert captured["body"]["generationConfig"]["thinkingConfig"] == { + "thinkingLevel": "high", + } + + +def test_gemini_3_1_pro_medium_passes_through(monkeypatch): + """Round 17 regression: 3.1+ Pro accepts medium; coercion must NOT + apply when the model id is gemini-3.1-pro*.""" + captured = _capture_body( + monkeypatch, + model = "gemini-3.1-pro-preview", + reasoning_effort = "medium", + ) + assert captured["body"]["generationConfig"]["thinkingConfig"] == { + "thinkingLevel": "medium", + } + + +def test_tool_calls_extra_content_stripped_for_non_native_gemini(): + """Round 17: per-tool-call `extra_content` (Gemini thoughtSignature + carrier) must not leak through `_build_external_messages` to + non-native-Gemini providers; OpenAI / Anthropic / custom Gemini + OAI-compat gateways would 400 on the unknown key.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "extra_content": { + "google": {"thought_signature": "SIG"}, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + + # Non-native providers (openai, custom Gemini OAI-compat proxy) + # must have extra_content stripped from the tool_call entry. + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + assert len(result) == 1 + tc = result[0]["tool_calls"][0] + assert "extra_content" not in tc, (provider_type, tc) + + # Native Gemini still receives extra_content for the round-trip. + result_native = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + tc_native = result_native[0]["tool_calls"][0] + assert tc_native["extra_content"]["google"]["thought_signature"] == "SIG" + + +def test_user_function_named_with_server_tool_arg_not_dropped(monkeypatch): + """Round 17: the OpenAI Responses translator must NOT drop a user + function whose JSON arguments happen to contain `_server_tool: + true` UNLESS the function name is also one of the canonical + builtin names. Otherwise a user schema with an `_server_tool` field + becomes invisible to the model.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "user_function", + "arguments": json.dumps( + {"_server_tool": True, "q": "x"} + ), + }, + } + ], + }, + { + "role": "tool", + "content": "result", + "tool_call_id": "call_user", + "name": "user_function", + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + fn_outs = [i for i in items if i.get("type") == "function_call_output"] + # User function call must survive (matching call + output). + assert any(c.get("name") == "user_function" for c in fn_calls), items + assert len(fn_outs) == 1, items + + +def test_builtin_named_with_server_tool_marker_dropped(monkeypatch): + """Round 17 control: a builtin (web_search) tagged with + `_server_tool: true` continues to be filtered from outbound + history.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps( + {"_server_tool": True, "query": "x"} + ), + }, + } + ], + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + # Builtin server-side tool call must be filtered out. + assert all(c.get("name") != "web_search" for c in fn_calls), items + + +def test_gemini_tool_choice_none_disables_hosted_builtins(monkeypatch): + """Round 18: `tool_choice="none"` must drop hosted Google Search / + code execution from the outbound Gemini body, not just user + function declarations. Otherwise an API client that opted out of + tool use still triggers grounded search (privacy + billing).""" + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search", "code_execution"], + tool_choice = "none", + ) + assert captured["body"].get("tools") is None, captured["body"] + + +def test_gemini_tool_choice_none_disables_function_declarations(monkeypatch): + """Round 18: `tool_choice="none"` must drop user function + declarations as well as hosted builtins from the Gemini body.""" + captured = _capture_body( + monkeypatch, + tool_choice = "none", + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + ) + assert captured["body"].get("tools") is None, captured["body"] + + +def test_schema_anyof_multitype_with_null_keeps_anyof_and_nullable( + monkeypatch, +): + """Round 18: multi-branch unions with null (e.g. + `Union[str, int, None]`) must keep the slim anyOf without the null + branch and add `nullable: true`; Gemini rejects + `{"type":"null"}` inside anyOf.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "either": { + "anyOf": [ + {"type": "string"}, + {"type": "integer"}, + {"type": "null"}, + ] + }, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + either = decls[0]["parameters"]["properties"]["either"] + assert either.get("nullable") is True + inner = either.get("anyOf") + assert isinstance(inner, list) and len(inner) == 2, either + assert all( + not (isinstance(b, dict) and b.get("type") == "null") for b in inner + ), inner + + +def test_safe_fetch_image_redirect_malformed_url_no_crash(monkeypatch): + """Round 18: when the upstream 302 Location is a malformed + bracketed-IPv6 URL, the helper must return None instead of letting + a urlparse ValueError abort the chat stream.""" + import socket + import urllib.error + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubOpener: + def open(self, req, timeout = None): + raise urllib.error.HTTPError( + req.full_url, + 302, + "Found", + {"Location": "https://[bad/x.png"}, + None, + ) + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/x.png", "image/png" + ) + ) + assert res is None + + +def test_safe_fetch_image_malformed_port_no_crash(): + """Round 18: a URL with a non-numeric port (`https://h:bad/x.png`) + must not raise; urlparse's port property lazily ValueErrors.""" + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://example.com:bad/x.png", "image/png" + ) + ) + assert res is None + + +def test_safe_fetch_image_missing_content_type_uses_fallback(monkeypatch): + """Round 18: when the server returns image bytes but no + Content-Type header, the helper must use the caller-provided + fallback MIME (guessed from URL extension) instead of dropping the + image as `non-image content-type=`.""" + import socket + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("1.1.1.1", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + headers = {"content-length": "3"} + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + return b"PNG" + + class _StubOpener: + def open(self, req, timeout = None): + return _StubResp() + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/cat.png", "image/png" + ) + ) + assert res is not None + assert res[0] == "image/png" + + +def test_anthropic_translates_openai_tool_calls_into_tool_use_blocks(monkeypatch): + """Round 18: an assistant turn with OpenAI-style top-level + `tool_calls` must be translated into Anthropic native + `{type:"tool_use", id, name, input}` content blocks before being + forwarded. The OpenAI `role="tool"` follow-up must become a + `role:"user"` message with a `tool_result` content block.""" + captured: dict = {"messages": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["messages"] = body.get("messages") + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "look up X"}, + { + "role": "assistant", + "content": "let me check", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"q":"x"}', + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_a", + "name": "lookup", + }, + {"role": "user", "content": "summarise"}, + ], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + msgs = captured["messages"] or [] + # No top-level tool_calls should remain. + assert all("tool_calls" not in m for m in msgs), msgs + # The assistant turn must now have content blocks including a + # tool_use block. + asst = [m for m in msgs if m.get("role") == "assistant"] + assert asst and isinstance(asst[0]["content"], list), asst + tool_uses = [b for b in asst[0]["content"] if b.get("type") == "tool_use"] + assert len(tool_uses) == 1, asst[0] + assert tool_uses[0]["name"] == "lookup" + assert tool_uses[0]["input"] == {"q": "x"} + # The role="tool" message must become a user/tool_result message. + tool_results: list[dict] = [] + for m in msgs: + if m.get("role") == "user" and isinstance(m.get("content"), list): + tool_results.extend( + b for b in m["content"] if b.get("type") == "tool_result" + ) + assert any( + tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text" + for tr in tool_results + ), msgs + + +def test_unmarked_user_web_search_function_survives_serialization(): + """Round 18: a user-defined function literally named `web_search` + with NO `_server_tool` marker must survive `_build_external_messages` + when forwarded to a non-native provider; only marked synthetic + builtin cards may be dropped.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"query": "x"}', + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert len(result) == 1, result + tcs = result[0].get("tool_calls") or [] + assert len(tcs) == 1, result + assert tcs[0]["function"]["name"] == "web_search" + + +def test_marked_server_builtin_dropped_from_build_external_messages(): + """Round 18: when a Gemini-native turn carrying a marked + `image_generation` server-tool card is forwarded to OpenAI / a + custom Gemini OAI-compat proxy, the tool_call must be dropped, not + just have its extra_content stripped. Forwarding an orphan + `image_generation` tool_call would 400 the receiving API.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "kind": "image"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "image_generation", + "arguments": marked_args, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + # Non-native providers: marked builtin tool_call must be dropped + # AND if it was the only payload, the whole message disappears. + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + # Empty assistant turn with only synthetic tool_call dropped. + assert result == [] or all(not (m.get("tool_calls") or []) for m in result), ( + provider_type, + result, + ) + + # Native Gemini preserves it (round-trips via extra_content). + result_native = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "gemini", + base_url = "https://generativelanguage.googleapis.com/v1beta", + ) + assert len(result_native) == 1 + assert result_native[0]["tool_calls"][0]["function"]["name"] == "image_generation" + + +def test_openai_responses_tool_choice_none_drops_hosted_tools(monkeypatch): + """Round 18: `tool_choice="none"` must also drop hosted OpenAI + Responses builtins (web_search, code execution shell, image + generation), not just user function tools.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + enabled_tools = ["web_search", "code_execution", "image_generation"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("tools") in (None, []), body + + +def test_anthropic_tool_choice_none_drops_hosted_tools(monkeypatch): + """Round 19: tool_choice="none" must opt out of Anthropic hosted + builtins (web_search, web_fetch, code_execution) just like it does + for Gemini and OpenAI Responses.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "web_fetch", "code_execution"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("tools") in (None, []), body + + +def test_openrouter_tool_choice_none_drops_web_plugin(monkeypatch): + """Round 19: tool_choice="none" must drop the OpenRouter web + plugin so a request that opted out of tool use does not still + trigger hosted web search.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("plugins") in (None, []), body + + +def test_kimi_tool_choice_none_skips_web_search_helper(monkeypatch): + """Round 19: when tool_choice="none" plus enabled_tools= + ["web_search"] on Kimi, the dispatcher must NOT route into + `_stream_kimi_web_search`. Falling through to the generic OAI- + compat path is the expected behavior.""" + routed_to_helper = {"called": False} + + real_helper = ExternalProviderClient._stream_kimi_web_search + + async def fake_helper(self, *args, **kwargs): # noqa: ARG001 + routed_to_helper["called"] = True + if False: + yield "" # pragma: no cover + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + fake_helper, + ) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "kimi", + base_url = "https://api.moonshot.ai/v1", + api_key = "sk-kimi-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "kimi-k2.6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + pass + await client.close() + + _drive(run()) + assert routed_to_helper["called"] is False + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + real_helper, + ) + + +def test_user_code_execution_function_not_dropped(): + """Round 19: a user-declared function literally named + `code_execution` with normal `code` arguments must survive + `_build_external_messages` -- round 17's shape heuristic dropped + it, which broke function-calling round-trips.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_user", + "type": "function", + "function": { + "name": "code_execution", + "arguments": '{"code": "print(1)"}', + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert len(result) == 1, result + tcs = result[0].get("tool_calls") or [] + assert len(tcs) == 1, result + assert tcs[0]["function"]["name"] == "code_execution" + + +def test_native_part_code_execution_treated_as_server_side(): + """Round 19: a Gemini `code_execution` card persists its replay + payload at `args.google.native_part` (no `_server_tool` marker on + pre-PR cards). The backend filter must still drop it for non-native + providers because it is a synthetic card, not a real user function.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + args_with_native_part = json.dumps( + { + "google": { + "native_part": { + "executableCode": { + "language": "PYTHON", + "code": "print(1)", + } + } + } + } + ) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_x", + "type": "function", + "function": { + "name": "code_execution", + "arguments": args_with_native_part, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + assert result == [] or all(not (m.get("tool_calls") or []) for m in result), result + + +def test_remote_image_fetch_attempt_cap_includes_failures(monkeypatch): + """Round 19: the per-request image fetch count cap must count + ATTEMPTS, not just successes. Otherwise a request with 100 + failing/slow URLs runs 100 fetches each up to the 15s timeout.""" + fetch_calls: list[str] = [] + + async def fake_fetch(url, fallback_mime, max_bytes = None): + fetch_calls.append(url) + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _gemini_sse( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "ok"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + } + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + image_parts = [ + { + "type": "image_url", + "image_url": {"url": f"https://cdn.example.com/img{idx}.png"}, + } + for idx in range(20) + ] + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + *image_parts, + ], + } + ], + model = "gemini-2.5-flash", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + assert len(fetch_calls) <= 8, len(fetch_calls) + + +def test_orphan_function_call_output_dropped_when_call_skipped(monkeypatch): + """Round 19: when a marked server-side builtin `function_call` is + dropped from OpenAI Responses input items, the matching role=tool + follow-up must also be dropped to avoid an orphan + `function_call_output`.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "search please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps( + {"_server_tool": True, "query": "x"} + ), + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_b", + "name": "web_search", + }, + {"role": "user", "content": "continue"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + fn_calls = [i for i in items if i.get("type") == "function_call"] + fn_outs = [i for i in items if i.get("type") == "function_call_output"] + assert all(c.get("call_id") != "call_b" for c in fn_calls), items + assert all(o.get("call_id") != "call_b" for o in fn_outs), items + + +def test_schema_multitype_union_with_null_preserves_anyof(monkeypatch): + """Round 19: a JSON Schema `"type": ["string","integer","null"]` + must be sanitized to anyOf:[{string},{integer}] + nullable:true. + Flattening to just `{"type":"string"}` silently drops the integer + branch and changes the function contract.""" + captured = _capture_body( + monkeypatch, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "either": {"type": ["string", "integer", "null"]}, + }, + }, + }, + } + ], + ) + decls = next( + t["functionDeclarations"] + for t in captured["body"].get("tools") or [] + if "functionDeclarations" in t + ) + either = decls[0]["parameters"]["properties"]["either"] + assert either.get("nullable") is True + inner = either.get("anyOf") + assert isinstance(inner, list) and len(inner) == 2, either + types = sorted( + b.get("type") for b in inner if isinstance(b, dict) and b.get("type") + ) + assert types == ["integer", "string"], inner + + +def test_invalid_gemini_model_rejected_before_image_fetch(monkeypatch): + """Round 19: invalid Gemini model IDs are rejected at the top of + `_stream_gemini`, BEFORE any user-controlled remote image fetch + runs.""" + fetch_calls: list[str] = [] + + async def fake_fetch(url, fallback_mime, max_bytes = None): + fetch_calls.append(url) + return None + + monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = _make_gemini_client() + async for _ in client.stream_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + { + "type": "image_url", + "image_url": {"url": "https://cdn.example.com/x.png"}, + }, + ], + } + ], + model = "../cachedContents/leak", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + assert fetch_calls == [], fetch_calls + + +def test_empty_assistant_turn_skipped_after_synthetic_tool_calls_dropped(): + """Round 20: when `_filter_tool_calls` drops every synthetic + server-builtin tool_call on an empty-content assistant turn, the + whole message must be skipped. Forwarding + `{"role":"assistant","content":""}` is rejected by several + providers as an empty assistant turn.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "kind": "image"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "image_generation", + "arguments": marked_args, + }, + } + ], + } + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + for provider_type, base_url in [ + ("openai", None), + ("gemini", "https://litellm.example/v1"), + ]: + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = provider_type, + base_url = base_url, + ) + # The empty assistant turn (only a synthetic builtin) must + # NOT appear in the output at all. + assert result == [], (provider_type, result) + + +def test_role_tool_dropped_when_matching_synthetic_call_filtered(): + """Round 20: `_build_external_messages` drops the matching role= + tool follow-up when its tool_call was a synthetic builtin that + `_filter_tool_calls` removed. Otherwise the receiving provider + sees an orphan tool_result with no tool_call.""" + from models.inference import ChatCompletionRequest + from routes.inference import _build_external_messages + + marked_args = json.dumps({"_server_tool": True, "query": "x"}) + payload = { + "model": "gpt-5.5", + "messages": [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "web_search", + "arguments": marked_args, + }, + } + ], + }, + { + "role": "tool", + "content": "result_text", + "tool_call_id": "call_b", + "name": "web_search", + }, + {"role": "user", "content": "continue"}, + ], + "stream": True, + } + req = ChatCompletionRequest.model_validate(payload) + result = _build_external_messages( + req.messages, + supports_vision = True, + provider_type = "openai", + base_url = None, + ) + # Only the user "continue" message survives. + roles = [m.get("role") for m in result] + assert roles == ["user"], result + + +def test_openrouter_no_synthetic_web_search_event_on_tool_choice_none( + monkeypatch, +): + """Round 20: OpenRouter dispatcher must not emit synthetic + web_search tool_start / tool_end events when tool_choice="none"; + otherwise the chat UI shows a search card for a search that + never happened.""" + captured_events: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = "none", + ): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :].strip() + if not payload or payload == "[DONE]": + continue + try: + obj = json.loads(payload) + except Exception: + continue + # Backend emits synthetic tool events as a top-level + # `_toolEvent` on the SSE payload (not nested inside + # `delta`). Read both shapes so a future format change + # cannot mask this regression. + evt = obj.get("_toolEvent") + if isinstance(evt, dict): + captured_events.append(evt) + for ch in obj.get("choices") or []: + delta = ch.get("delta") or {} + nested = delta.get("_toolEvent") if isinstance(delta, dict) else None + if isinstance(nested, dict): + captured_events.append(nested) + await client.close() + + _drive(run()) + # No synthetic web_search tool_start / tool_end emitted. + assert all( + e.get("tool_name") != "web_search" for e in captured_events + ), captured_events + + +def test_anthropic_role_tool_list_content_translates_to_tool_result( + monkeypatch, +): + """Round 20: an OpenAI-shape role=tool message with list content + (`content=[{"type":"text","text":"result"}]`) must be translated + into Anthropic's native tool_result block, not forwarded as an + invalid role=tool message.""" + captured: dict = {"messages": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["messages"] = body.get("messages") + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "look up X"}, + { + "role": "assistant", + "content": "let me check", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"q":"x"}', + }, + } + ], + }, + { + "role": "tool", + "content": [{"type": "text", "text": "result_text"}], + "tool_call_id": "call_a", + "name": "lookup", + }, + {"role": "user", "content": "summarise"}, + ], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ): + pass + await client.close() + + _drive(run()) + + msgs = captured["messages"] or [] + assert all(m.get("role") != "tool" for m in msgs), msgs + tool_results: list[dict] = [] + for m in msgs: + if m.get("role") == "user" and isinstance(m.get("content"), list): + tool_results.extend( + b for b in m["content"] if b.get("type") == "tool_result" + ) + assert any( + tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text" + for tr in tool_results + ), msgs + + +def test_data_url_non_image_mime_dropped(monkeypatch): + """Round 20: a `data:text/html;base64,...` image_url must be + dropped from the outbound Gemini body, not forwarded as + `inlineData.mimeType="text/html"` which Gemini rejects.""" + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + { + "type": "image_url", + "image_url": { + "url": "data:text/html;base64,PGgxPmhpPC9oMT4=", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + assert not any("inlineData" in p for p in parts), parts + + +def test_youtube_filedata_uses_video_mime(monkeypatch): + """Round 20: YouTube `fileData.fileUri` must declare a video + mimeType, not `image/jpeg` guessed from the URL path.""" + captured = _capture_body( + monkeypatch, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarise"}, + { + "type": "image_url", + "image_url": { + "url": "https://www.youtube.com/watch?v=abc", + }, + }, + ], + } + ], + ) + parts = captured["body"]["contents"][-1]["parts"] + yt = next((p for p in parts if "fileData" in p), None) + assert yt is not None, parts + assert yt["fileData"]["mimeType"].startswith("video/"), yt + + +def test_openai_responses_assistant_text_serialized_before_function_call( + monkeypatch, +): + """Round 20: in OpenAI Responses history, the assistant's + visible text for a turn that ALSO emitted a function_call must + serialize BEFORE the function_call item, matching the prior + response.output sequence. Otherwise function_call_output (the + role=tool follow-up) appears to follow an unrelated assistant + message.""" + captured: dict = {"input_items": None} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content.decode("utf-8")) + captured["input_items"] = body.get("input") + return httpx.Response( + 200, + content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-test", + ) + async for _ in client.stream_chat_completion( + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "Let me check that.", + "tool_calls": [ + { + "id": "call_w", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{}", + }, + } + ], + }, + { + "role": "tool", + "content": "sunny", + "tool_call_id": "call_w", + "name": "get_weather", + }, + {"role": "user", "content": "thanks"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 1.0, + max_tokens = 16, + ): + pass + await client.close() + + _drive(run()) + + items = captured["input_items"] or [] + types = [i.get("type") or i.get("role") for i in items] + # Expected order: + # user ("weather?") + # assistant ("Let me check that.") + # function_call (get_weather) + # function_call_output (sunny) + # user ("thanks") + assert types == [ + "user", + "assistant", + "function_call", + "function_call_output", + "user", + ], items + + +def test_gemini_tool_choice_none_disables_image_generation(monkeypatch): + """Round 21: `tool_choice="none"` must also flip the implicit + image-generation hosted tool off on image-tier models. Otherwise + `responseModalities=["TEXT","IMAGE"]` still rides on the outbound + body and the provider can generate (and bill for) image output + despite the explicit OpenAI tool opt-out.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tool_choice = "none", + ) + body = captured["body"] + assert body["generationConfig"].get("responseModalities") == ["TEXT"], body + + +def test_gemini_forced_function_tool_choice_drops_hosted_builtins(monkeypatch): + """Round 21: forced-function `tool_choice` (e.g. + `{"type":"function","function":{"name":"lookup"}}`) must suppress + hosted Google Search / code execution. Gemini's toolConfig only + constrains function declarations, not hosted tools, so leaving + `googleSearch`/`codeExecution` in `tools[]` lets them fire despite + the caller pinning a specific user function.""" + captured = _capture_body( + monkeypatch, + enabled_tools = ["web_search", "code_execution"], + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + tool_choice = { + "type": "function", + "function": {"name": "lookup"}, + }, + ) + body = captured["body"] + tool_kinds = [list(t.keys())[0] for t in (body.get("tools") or [])] + assert "googleSearch" not in tool_kinds, body + assert "codeExecution" not in tool_kinds, body + # User function declaration still survives. + assert "functionDeclarations" in tool_kinds, body + + +def test_gemini_forced_function_tool_choice_drops_image_generation(monkeypatch): + """Round 21: forced-function `tool_choice` must also flip the + implicit image-generation hosted tool off on image-tier models.""" + captured = _capture_body( + monkeypatch, + model = "gemini-2.5-flash-image", + enabled_tools = ["image_generation"], + tool_choice = { + "type": "function", + "function": {"name": "lookup"}, + }, + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ], + ) + body = captured["body"] + assert body["generationConfig"].get("responseModalities") == ["TEXT"], body + + +def test_gemini_code_execution_native_part_list_replays_per_part_signatures( + monkeypatch, +): + """Round 21: merged code-execution history must replay per-part + `thoughtSignature`s, not fan one top-level signature across every + native subpart. Gemini 3 strict validators reject a signature + placed on the wrong part.""" + history = [ + {"role": "user", "content": "plot 1+1"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(1+1)", + }, + "thoughtSignature": "SIG-EXEC", + }, + { + "codeExecutionResult": { + "id": "res_a", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + }, + ], + }, + }, + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_a", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + # Locate the assistant turn replayed as native code-exec parts. + assistant_turn = next(c for c in contents if c["role"] == "model") + parts = assistant_turn["parts"] + exec_parts = [p for p in parts if "executableCode" in p] + result_parts = [p for p in parts if "codeExecutionResult" in p] + assert exec_parts and result_parts, parts + assert exec_parts[0].get("thoughtSignature") == "SIG-EXEC", exec_parts[0] + # codeExecutionResult had no signature -- must NOT inherit one. + assert "thoughtSignature" not in result_parts[0], result_parts[0] + + +def test_gemini_code_execution_legacy_merged_signature_only_on_executable( + monkeypatch, +): + """Round 21: backward compatibility for pre-round-21 persisted + history that stored merged `native_part` as a single object plus a + top-level `thoughtSignature`. The replay branch must attach that + signature only to `executableCode` (where Gemini 3 emits it), not + fan it across `codeExecutionResult` / `inlineData`.""" + history = [ + {"role": "user", "content": "plot 1+1"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_b", + "type": "function", + "function": { + "name": "code_execution", + "arguments": "{}", + }, + "extra_content": { + "google": { + "native_part": { + "executableCode": { + "id": "code_b", + "language": "PYTHON", + "code": "print(1+1)", + }, + "codeExecutionResult": { + "id": "res_b", + "outcome": "OUTCOME_OK", + "output": "2\n", + }, + "thoughtSignature": "LEGACY-SIG", + }, + }, + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_b", + "name": "code_execution", + "content": "2", + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + assistant_turn = next(c for c in contents if c["role"] == "model") + exec_parts = [p for p in assistant_turn["parts"] if "executableCode" in p] + result_parts = [p for p in assistant_turn["parts"] if "codeExecutionResult" in p] + assert exec_parts[0].get("thoughtSignature") == "LEGACY-SIG", exec_parts[0] + assert "thoughtSignature" not in result_parts[0], result_parts[0] + + +def test_gemini_role_tool_list_content_flattens_to_result_text(monkeypatch): + """Round 21: OpenAI-shape role=tool messages may carry list content + like `[{"type":"text","text":"result"}]`. Forwarding those parts + verbatim into `functionResponse.response.result` yields a list of + content-part objects instead of the actual tool output text.""" + history = [ + {"role": "user", "content": "look up"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "lookup", + "arguments": json.dumps({"q": "x"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "name": "lookup", + "content": [{"type": "text", "text": "answer-text"}], + }, + {"role": "user", "content": "next"}, + ] + captured = _capture_body(monkeypatch, messages = history) + contents = captured["body"]["contents"] + fn_response = None + for c in contents: + for p in c.get("parts") or []: + if isinstance(p, dict) and "functionResponse" in p: + fn_response = p["functionResponse"] + break + if fn_response: + break + assert fn_response is not None, contents + assert fn_response["response"] == {"result": "answer-text"}, fn_response + + +def test_safe_fetch_image_threads_per_request_byte_budget(monkeypatch): + """Round 21: the aggregate per-request byte cap must be passed into + `_safe_fetch_image_for_gemini` so an oversize URL is refused via + Content-Length (short-circuit) rather than fully downloaded then + discarded after the fact.""" + import socket + + captured: dict = {"reads": 0, "content_length_seen": None} + + original_getaddrinfo = socket.getaddrinfo + + def fake_getaddrinfo(host, *args, **kwargs): + if host == "cdn.example.com": + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + 0, + "", + ("8.8.8.8", 0), + ) + ] + return original_getaddrinfo(host, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + class _StubResp: + status = 200 + # Declared 5 MiB, but caller passes a 1 MiB remaining budget. + headers = { + "content-type": "image/png", + "content-length": str(5 * 1024 * 1024), + } + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, _n = None): + captured["reads"] += 1 + return b"\x00" * (5 * 1024 * 1024) + + class _StubOpener: + def open(self, req, timeout = None): + return _StubResp() + + monkeypatch.setattr( + "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener() + ) + + res = _drive( + ep_mod._safe_fetch_image_for_gemini( + "https://cdn.example.com/big.png", + "image/png", + max_bytes = 1 * 1024 * 1024, + ) + ) + assert res is None + # Refused via Content-Length pre-check, never read. + assert captured["reads"] == 0 + + +def test_openai_chat_delta_type_includes_tool_calls_and_extra_content(): + """Round 21: the frontend `OpenAIChatDelta` interface must expose + `tool_calls` and `extra_content` so TypeScript callers can consume + the Gemini-native stream fields without `any` casts. This test is + a static-string assertion against the .ts source; mirrors how other + frontend wire-contract tests are pinned from the backend suite.""" + import os + + here = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + types_path = os.path.join( + here, "frontend", "src", "features", "chat", "types", "api.ts" + ) + with open(types_path, "r", encoding = "utf-8") as f: + src = f.read() + assert "tool_calls?: OpenAIToolCallPart[]" in src, src[:200] + assert "extra_content?: Record" in src, src[:200] + assert "boolean | string | null" in src, src[:200] + + +def test_anthropic_forced_function_tool_choice_drops_hosted_tools(monkeypatch): + """Round 22: forced-function tool_choice must suppress Anthropic + hosted builtins the same way it does for Gemini. Pinning a user + function (`tool_choice={"type":"function","function":{"name":...}}`) + while also passing `enabled_tools=["web_search","web_fetch", + "code_execution"]` should not still fire those server-side.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com", + api_key = "sk-ant-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "claude-sonnet-4-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "web_fetch", "code_execution"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + # No hosted tools should be in the body — only the caller's user- + # function declarations (which this test doesn't pass any of). + tools = body.get("tools") or [] + hosted_tool_names = {"web_search", "web_fetch", "code_execution"} + for tool in tools: + assert tool.get("name") not in hosted_tool_names, body + + +def test_openrouter_forced_function_tool_choice_drops_web_plugin(monkeypatch): + """Round 22: forced-function tool_choice must drop the OpenRouter + web plugin too — caller pinned a user function, OpenRouter must not + still attach the hosted web-search plugin.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + assert body.get("plugins") in (None, []), body + + +def test_kimi_forced_function_tool_choice_skips_web_search_helper(monkeypatch): + """Round 22: forced-function tool_choice plus enabled_tools= + ["web_search"] on Kimi must NOT route into `_stream_kimi_web_search`. + Caller pinned a user function; hosted $web_search should be + suppressed for the same privacy/billing reason.""" + routed_to_helper = {"called": False} + + async def fake_helper(self, *args, **kwargs): # noqa: ARG001 + routed_to_helper["called"] = True + if False: + yield "" # pragma: no cover + + monkeypatch.setattr( + ExternalProviderClient, + "_stream_kimi_web_search", + fake_helper, + ) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = b"data: [DONE]\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "kimi", + base_url = "https://api.moonshot.ai/v1", + api_key = "sk-kimi-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "kimi-k2.6", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + assert not routed_to_helper["called"] + + +def test_openai_responses_forced_function_tool_choice_drops_hosted_tools(monkeypatch): + """Round 23: forced-function tool_choice on the OpenAI Responses + path must suppress hosted builtins (web_search, shell, + image_generation) the same way it does for Gemini / Anthropic / + OpenRouter / Kimi. User-defined function tools still flow through + so the pinned function can resolve.""" + captured: dict = {"body": None} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b"event: response.completed\ndata: {}\n\n", + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openai", + base_url = "https://api.openai.com/v1", + api_key = "sk-openai-test", + ) + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "gpt-5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search", "code_execution", "image_generation"], + tools = [ + { + "type": "function", + "function": { + "name": "lookup_record", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + pass + await client.close() + + _drive(run()) + body = captured["body"] or {} + tools = body.get("tools") or [] + hosted_types = {"web_search", "shell", "image_generation"} + hosted_seen = {t.get("type") for t in tools if isinstance(t, dict)} + assert not (hosted_seen & hosted_types), body + # The user function declaration must still be present so the pin + # has something to target. + user_function_seen = any( + isinstance(t, dict) and t.get("type") == "function" for t in tools + ) + assert user_function_seen, body + # And the forced-function tool_choice must be forwarded in Responses + # shape: `{type:"function", name:"..."}`. + tc = body.get("tool_choice") + assert isinstance(tc, dict) and tc.get("type") == "function", body + assert tc.get("name") == "lookup_record", body + + +def test_strip_provider_synthetic_tool_history_drops_text_only_extra_content(): + """Round 24: a plain text Gemini reply (no tool_calls) carrying + `extra_content.google.thought_signature` must still have that + metadata stripped before being forwarded to a local llama-server + backend. Without this, switching a Gemini thread mid-stream to a + local GGUF model leaks Gemini-only fields to llama-server.""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "Hello!", + "extra_content": {"google": {"thought_signature": "SIG_ABC"}}, + }, + {"role": "user", "content": "now in pirate voice"}, + ] + out = _strip_provider_synthetic_tool_history(messages) + # Same three turns, but the assistant's `extra_content` is gone. + assert [m["role"] for m in out] == ["user", "assistant", "user"] + assistant = out[1] + assert "extra_content" not in assistant, assistant + assert assistant["content"] == "Hello!" + + +def test_validate_and_resolve_host_blocks_shared_address_space(): + """Round 24 SSRF P1: 100.64.0.0/10 carrier-grade NAT addresses are + `is_private=False` AND `is_global=False` per Python's ipaddress + docs. The previous denylist (is_private/loopback/link_local/etc.) + missed them. Adding `not ip.is_global` as the primary gate covers + all non-public ranges, current and future.""" + import socket as _socket + from core.inference import tools as _tools + + orig_getaddrinfo = _socket.getaddrinfo + + def fake_getaddrinfo(hostname, port, *args, **kwargs): + if hostname == "shared.example": + return [ + ( + _socket.AF_INET, + _socket.SOCK_STREAM, + 0, + "", + ("100.64.0.1", port), + ), + ] + return orig_getaddrinfo(hostname, port, *args, **kwargs) + + _socket.getaddrinfo = fake_getaddrinfo + try: + ok, reason, _ip = _tools._validate_and_resolve_host("shared.example", 443) + finally: + _socket.getaddrinfo = orig_getaddrinfo + assert ok is False, (ok, reason) + assert "non-public" in reason.lower() or "100.64.0.1" in reason + + +def test_gemini_custom_oai_compat_base_skips_native_allowlist(): + """Round 24: a custom Gemini OAI-compatible base (LiteLLM/proxy) + must NOT have its model list filtered through the native Gemini + allowlist regex. A LiteLLM gateway returning + `["google/gemini-2.5-flash", "my-team/gemini", "gemini-2.5-flash"]` + should be passed through; the native filter would strip the + prefixed IDs even though the chat dispatch routes them via the + OpenAI-compatible client.""" + import asyncio as _asyncio + + from routes import providers as _providers + from routes.providers import ( + ProviderModelsRequest, + list_provider_models, + ) + + captured: dict = {"base": None} + + class _FakeClient: + def __init__(self, *, base_url, **kwargs): + captured["base"] = base_url + + async def list_models(self): + return [ + {"id": "google/gemini-2.5-flash"}, + {"id": "my-team/gemini"}, + {"id": "gemini-2.5-flash"}, + ] + + async def close(self): + return None + + orig = _providers.ExternalProviderClient + _providers.ExternalProviderClient = _FakeClient + try: + req = ProviderModelsRequest( + provider_type = "gemini", + base_url = "https://litellm.example/v1", + ) + result = _asyncio.run(list_provider_models(req, current_subject = "unsloth")) + finally: + _providers.ExternalProviderClient = orig + ids = {m.id for m in result} + # All three IDs survive — the native allowlist was bypassed. + assert "google/gemini-2.5-flash" in ids, ids + assert "my-team/gemini" in ids, ids + assert "gemini-2.5-flash" in ids, ids + + +def test_strip_provider_synthetic_tool_history_drops_synthetic_only(): + """Round 22: switching a thread from native Gemini (code_execution + / image_generation tool_cards in history) to a local GGUF backend + must strip the synthetic tool_calls + matching role=tool replies + before llama-server sees them. Real user-function tool_calls and + their matching tool replies must survive.""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "let me run it", + "tool_calls": [ + { + "id": "synth_ce_1", + "type": "function", + "function": { + "name": "code_execution", + "arguments": json.dumps( + { + "_server_tool": True, + "google": {"native_part": {"parts": []}}, + } + ), + }, + "extra_content": {"google": {"thought_signature": "abc"}}, + }, + { + "id": "real_lookup", + "type": "function", + "function": { + "name": "lookup_user", + "arguments": json.dumps({"id": 42}), + }, + }, + ], + "extra_content": {"google": {"thought_signature": "msglevel"}}, + }, + { + "role": "tool", + "tool_call_id": "synth_ce_1", + "content": "Gemini-only result text", + }, + { + "role": "tool", + "tool_call_id": "real_lookup", + "content": '{"name": "alice"}', + }, + ] + out = _strip_provider_synthetic_tool_history(messages) + assistant = next(m for m in out if m.get("role") == "assistant") + tcs = assistant["tool_calls"] + assert len(tcs) == 1, tcs + assert tcs[0]["id"] == "real_lookup" + assert "extra_content" not in tcs[0] + assert "extra_content" not in assistant + tool_msgs = [m for m in out if m.get("role") == "tool"] + assert len(tool_msgs) == 1 + assert tool_msgs[0]["tool_call_id"] == "real_lookup" + + +def test_strip_provider_synthetic_tool_history_drops_empty_assistant(): + """If every tool_call was synthetic and the assistant turn had no + content, the entire turn must be dropped (llama-server rejects + empty assistant messages with no tool_calls).""" + from routes.inference import _strip_provider_synthetic_tool_history + + messages = [ + {"role": "user", "content": "draw a sloth"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "synth_imggen", + "type": "function", + "function": { + "name": "image_generation", + "arguments": json.dumps( + { + "google": { + "native_part": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "Zm9v", + } + } + ] + } + } + } + ), + }, + } + ], + }, + {"role": "tool", "tool_call_id": "synth_imggen", "content": "(image)"}, + {"role": "user", "content": "now try in pirate voice"}, + ] + out = _strip_provider_synthetic_tool_history(messages) + roles = [m.get("role") for m in out] + # The synthetic assistant + its tool reply are both gone; only the + # two user turns survive. + assert roles == ["user", "user"], out + + +def test_openrouter_no_synthetic_web_search_event_on_forced_function_tool_choice( + monkeypatch, +): + """Round 22 sibling of the round-20 `tool_choice='none'` test: when + the caller forces a specific function via `tool_choice={"type": + "function", ...}` AND passes `enabled_tools=["web_search"]`, the + OpenRouter path must NOT synthesize a fake `web_search` tool card. + The plugin was not attached upstream so the UI must not see a + server-tool card.""" + captured_events: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = ( + b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n' + b"data: [DONE]\n\n" + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http(monkeypatch, handler) + + async def run(): + client = ExternalProviderClient( + provider_type = "openrouter", + base_url = "https://openrouter.ai/api/v1", + api_key = "sk-or-test", + ) + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "openai/gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 16, + enabled_tools = ["web_search"], + tool_choice = { + "type": "function", + "function": {"name": "lookup_record"}, + }, + ): + payload = line.strip().removeprefix("data: ") + if payload and payload != "[DONE]": + try: + captured_events.append(json.loads(payload)) + except Exception: + pass + await client.close() + + _drive(run()) + for evt in captured_events: + for choice in evt.get("choices") or []: + delta = choice.get("delta") or {} + extra = delta.get("extra_content") or {} + tool_event = extra.get("toolEvent") if isinstance(extra, dict) else None + if isinstance(tool_event, dict): + assert tool_event.get("tool_name") != "web_search", evt diff --git a/studio/backend/tests/test_openai_code_execution.py b/studio/backend/tests/test_openai_code_execution.py index 3d179371e3..0b1a65c69e 100644 --- a/studio/backend/tests/test_openai_code_execution.py +++ b/studio/backend/tests/test_openai_code_execution.py @@ -269,7 +269,15 @@ def test_shell_call_emits_tool_start_and_end(monkeypatch): assert len(ends) == 1 assert starts[0]["tool_name"] == "code_execution" assert starts[0]["tool_call_id"] == "scall_1" - assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la"} + # `_server_tool: True` is the synthetic-builtin marker the + # backend stamps onto every provider-side tool_start so the + # frontend serializer can distinguish hosted tools from + # user-declared functions on history replay. + assert starts[0]["arguments"] == { + "kind": "bash", + "command": "ls -la", + "_server_tool": True, + } assert ends[0]["tool_call_id"] == "scall_1" assert "total 24" in ends[0]["result"] diff --git a/studio/backend/tests/test_openai_image_generation.py b/studio/backend/tests/test_openai_image_generation.py index 6d415316d4..f5dee2561a 100644 --- a/studio/backend/tests/test_openai_image_generation.py +++ b/studio/backend/tests/test_openai_image_generation.py @@ -207,9 +207,12 @@ def test_image_generation_done_emits_tool_event_chunks(monkeypatch): ends = [e for e in image_events if e.get("type") == "tool_end"] assert len(starts) == 1, image_events assert len(ends) == 1, image_events + # `_server_tool: True` marks this as a provider-side synthetic + # tool card on the frontend's history serializer. assert starts[0]["arguments"] == { "kind": "image", "prompt": "A photorealistic cat sitting", + "_server_tool": True, "openai_image_generation_call_id": "img_abc", } assert ends[0]["image_b64"] == "AAAA" diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py index 22ccba7058..f177ed5ef3 100644 --- a/studio/backend/tests/test_openai_responses_translation.py +++ b/studio/backend/tests/test_openai_responses_translation.py @@ -215,6 +215,254 @@ def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch): assert payloads[-1] == "[DONE]" +def test_responses_function_call_output_translates_to_delta_tool_calls(monkeypatch): + """Round 12: caller-supplied function tools forwarded into /v1/responses + must have their `function_call` output items translated back into Chat + Completions delta.tool_calls, and the terminal chunk must emit + finish_reason="tool_calls" so the frontend's accumulator runs the + function instead of seeing finish_reason="stop".""" + + def handler(request: httpx.Request) -> httpx.Response: + events = [ + {"type": "response.created"}, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_abc", + "call_id": "call_xyz", + "name": "get_weather", + "arguments": '{"city":"SF"}', + }, + }, + {"type": "response.completed", "response": {}}, + ] + return httpx.Response( + 200, + content = _responses_sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + lines = await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "weather?"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } + ], + ) + ) + await client.close() + return lines + + lines = _drive(run()) + payloads = [ + json.loads(line[len("data:") :].strip()) + for line in lines + if line.startswith("data:") and line[len("data:") :].strip() != "[DONE]" + ] + tool_call_deltas = [ + p + for p in payloads + if isinstance(p, dict) + and p.get("choices") + and p["choices"][0].get("delta", {}).get("tool_calls") + ] + assert tool_call_deltas, payloads + tc = tool_call_deltas[0]["choices"][0]["delta"]["tool_calls"][0] + assert tc["id"] == "call_xyz" + assert tc["function"]["name"] == "get_weather" + assert tc["function"]["arguments"] == '{"city":"SF"}' + # Final chunk reports tool_calls instead of stop. + terminal = next( + p + for p in payloads + if isinstance(p, dict) + and p.get("choices") + and p["choices"][0].get("finish_reason") in ("stop", "tool_calls") + ) + assert terminal["choices"][0]["finish_reason"] == "tool_calls", payloads + + +def test_responses_parallel_function_calls_get_distinct_indices(monkeypatch): + """Round 13: parallel function_call items must land on distinct + delta.tool_calls[].index slots so index-keyed clients don't + collapse the second call into the first.""" + + def handler(request: httpx.Request) -> httpx.Response: + events = [ + {"type": "response.created"}, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_a", + "call_id": "call_a", + "name": "lookup_a", + "arguments": "{}", + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "function_call", + "id": "fc_b", + "call_id": "call_b", + "name": "lookup_b", + "arguments": "{}", + }, + }, + {"type": "response.completed", "response": {}}, + ] + return httpx.Response( + 200, + content = _responses_sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + lines = await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "x"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + tools = [ + { + "type": "function", + "function": { + "name": "lookup_a", + "parameters": {"type": "object"}, + }, + }, + { + "type": "function", + "function": { + "name": "lookup_b", + "parameters": {"type": "object"}, + }, + }, + ], + ) + ) + await client.close() + return lines + + lines = _drive(run()) + indices: list[int] = [] + for raw in lines: + if not raw.startswith("data:"): + continue + payload = raw[len("data:") :].strip() + if payload == "[DONE]": + continue + try: + obj = json.loads(payload) + except Exception: + continue + delta = (obj.get("choices") or [{}])[0].get("delta") or {} + for tc in delta.get("tool_calls") or []: + indices.append(tc.get("index")) + assert indices == [0, 1], indices + + +def test_responses_follow_up_tool_result_uses_function_call_output_items(monkeypatch): + """Round 13: a second turn after a Responses function call must + serialize the tool_calls history and tool result as Responses + `function_call` / `function_call_output` input items, not as + Chat Completions role="tool" content.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = _responses_sse( + [ + {"type": "response.created"}, + {"type": "response.completed", "response": {}}, + ] + ), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + await _collect( + client._stream_openai_responses( + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_xyz", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"SF"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_xyz", + "content": "sunny", + }, + {"role": "user", "content": "thanks"}, + ], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = None, + enable_thinking = None, + reasoning_effort = None, + ) + ) + await client.close() + + _drive(run()) + items = captured["body"]["input"] + types = [it.get("type") or it.get("role") for it in items] + assert "function_call" in types, items + assert "function_call_output" in types, items + fc = next(it for it in items if it.get("type") == "function_call") + assert fc["call_id"] == "call_xyz" + assert fc["name"] == "get_weather" + assert fc["arguments"] == '{"city":"SF"}' + fco = next(it for it in items if it.get("type") == "function_call_output") + assert fco["call_id"] == "call_xyz" + assert fco["output"] == "sunny" + + def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch): def handler(request: httpx.Request) -> httpx.Response: events = [ diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index a5bb4029d1..51cf000863 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -802,6 +802,7 @@ const ReasoningToggle: FC = () => { { isReasoningProvider: selectedExternalProvider?.isReasoningModel === true, + baseUrl: selectedExternalProvider?.baseUrl ?? null, }, ) : null; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index ce59429762..2ef14b5326 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -25,6 +25,7 @@ import { getExternalMinOutputTokens, getExternalReasoningCapabilities, getProviderCapabilities, + isGeminiCustomOpenAICompatBase, providerSupportsBuiltinCodeExecution, providerSupportsBuiltinImageGeneration, providerSupportsBuiltinWebFetch, @@ -105,6 +106,33 @@ type RunMessage = RunMessages[number]; /** Tracks which user messages were sent with an audio file (messageId → filename). */ export const sentAudioNames = new Map(); +// Synthetic provider-side tool names; backend stamps args._server_tool +// so user functions with the same name aren't dropped. Mirror of +// _SERVER_SIDE_BUILTIN_TOOL_NAMES on the backend. +const SERVER_SIDE_BUILTIN_TOOL_NAMES = new Set([ + "web_search", + "web_fetch", + "code_execution", + "image_generation", +]); + +/** + * Whether a persisted tool-call part is provider-side synthetic and + * should be stripped from outbound history. Match on the + * args._server_tool marker or a Gemini native_part payload — no shape + * heuristic, because user functions can legitimately share a name. + */ +function isServerSideBuiltinToolPart( + toolNameLower: string, + _argsObj: Record | null, + hasServerToolMarker: boolean, + hasNativePart: boolean, +): boolean { + if (!SERVER_SIDE_BUILTIN_TOOL_NAMES.has(toolNameLower)) return false; + if (hasServerToolMarker) return true; + return hasNativePart; +} + /** * Match error messages that indicate the request filled or would fill * the KV cache, so the UI can show a dedicated toast pointing at the @@ -508,21 +536,200 @@ function isAnthropicRefusalMessage(message: RunMessage): boolean { return metadata?.custom?.anthropicRefusal === true; } -function toOpenAIMessage(message: RunMessage): { - role: "system" | "user" | "assistant"; - content: OpenAIMessageContent; -} | null { +function collectAssistantToolCalls( + message: RunMessage, +): Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; +}> { + const out: Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + }> = []; + for (const part of message.content ?? []) { + if (part.type !== "tool-call") continue; + const tc = part as ToolCallMessagePart & { + argsText?: string; + extra_content?: unknown; + }; + const toolNameLower = (tc.toolName ?? "").toLowerCase(); + const argsObj = + tc.args && typeof tc.args === "object" + ? (tc.args as Record) + : null; + const argsGoogle = + argsObj && typeof argsObj.google === "object" && argsObj.google !== null + ? (argsObj.google as Record) + : null; + const hasNativePart = Boolean( + argsGoogle && + typeof argsGoogle.native_part === "object" && + argsGoogle.native_part !== null, + ); + const hasServerToolMarker = Boolean( + argsObj && (argsObj as Record)._server_tool === true, + ); + const isServerSideBuiltin = isServerSideBuiltinToolPart( + toolNameLower, + argsObj, + hasServerToolMarker, + hasNativePart, + ); + if (isServerSideBuiltin) { + // Gemini code_execution / image_generation still need to round- + // trip the native_part payload for native replay; drop the rest. + if (!hasNativePart) continue; + } + const argumentsStr = + typeof tc.argsText === "string" && tc.argsText.length > 0 + ? tc.argsText + : JSON.stringify(tc.args ?? {}); + const entry: { + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + } = { + id: tc.toolCallId, + type: "function" as const, + function: { + name: tc.toolName ?? "", + arguments: argumentsStr, + }, + }; + // Promote args.google to extra_content.google so the backend + // native_part replay branch can find it. The backend only inspects + // extra_content, not function.arguments. + if (tc.extra_content !== undefined) { + entry.extra_content = tc.extra_content; + } else if (argsGoogle) { + entry.extra_content = { google: argsGoogle }; + } + out.push(entry); + } + return out; +} + +function collectToolResultMessages( + message: RunMessage, +): Array<{ + role: "tool"; + content: string; + tool_call_id: string; + name?: string; +}> { + const out: Array<{ + role: "tool"; + content: string; + tool_call_id: string; + name?: string; + }> = []; + for (const part of message.content ?? []) { + if (part.type !== "tool-call") continue; + const tc = part as ToolCallMessagePart; + const result = (tc as { result?: unknown }).result; + // Skip provider-side builtins; see isServerSideBuiltinToolPart. + const argsObj = + tc.args && typeof tc.args === "object" + ? (tc.args as Record) + : null; + const argsGoogle = + argsObj && typeof argsObj.google === "object" && argsObj.google !== null + ? (argsObj.google as Record) + : null; + const toolNameLower = (tc.toolName ?? "").toLowerCase(); + const hasServerToolMarker = Boolean( + argsObj && argsObj._server_tool === true, + ); + const hasNativePart = Boolean( + argsGoogle && + typeof argsGoogle.native_part === "object" && + argsGoogle.native_part !== null, + ); + if ( + isServerSideBuiltinToolPart( + toolNameLower, + argsObj, + hasServerToolMarker, + hasNativePart, + ) + ) { + continue; + } + if (result === undefined || result === null) continue; + let content: string; + if (typeof result === "string") { + // Backend ChatMessage validator rejects role="tool" with empty + // content; serialise a sentinel JSON so legitimately empty tool + // outputs still round-trip the follow-up turn to the provider. + content = result.length > 0 ? result : JSON.stringify({ result: "" }); + } else { + try { + content = JSON.stringify(result); + } catch { + content = String(result); + } + } + out.push({ + role: "tool", + content, + tool_call_id: tc.toolCallId, + ...(tc.toolName ? { name: tc.toolName } : {}), + }); + } + return out; +} + +type SerializedMessage = { + role: "system" | "user" | "assistant" | "tool"; + content: OpenAIMessageContent | null; + tool_calls?: Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + }>; + tool_call_id?: string; + name?: string; + /** + * Gemini text-part thoughtSignature stashed during streaming on the + * last text MessagePart. Backend reads + * `extra_content.google.thought_signature` and attaches it to the + * matching Gemini text part on the outbound turn. + */ + extra_content?: unknown; +}; + +function collectAssistantTextThoughtSignature( + message: RunMessage, +): string | undefined { + if (!Array.isArray(message.content)) return undefined; + for (let i = message.content.length - 1; i >= 0; i -= 1) { + const part = message.content[i] as { type?: string } & Record< + string, + unknown + >; + if (part?.type !== "text") continue; + const sig = part._google_thought_signature; + if (typeof sig === "string" && sig) return sig; + } + return undefined; +} + +function toOpenAIMessages(message: RunMessage): SerializedMessage[] { if ( message.role !== "system" && message.role !== "user" && message.role !== "assistant" ) { - return null; + return []; } let textContent = collectTextParts(message).join("\n"); - // Strip inline audio base64 from prior assistant messages to avoid - // inflating token counts (e.g. audio-player responses with embedded WAV). if (message.role === "assistant") { textContent = textContent.replace( /data:audio\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/g, @@ -531,25 +738,70 @@ function toOpenAIMessage(message: RunMessage): { if (isAnthropicRefusalMessage(message)) { // Prune refused assistant turn from outbound history; the // rendered transcript still shows the user-visible notice. - return null; + return []; } } const imageParts = collectImageParts(message); - if (imageParts.length > 0) { - return { - role: message.role, - content: [ - ...(textContent ? [{ type: "text" as const, text: textContent }] : []), - ...imageParts, - ], - }; + const toolCalls = + message.role === "assistant" ? collectAssistantToolCalls(message) : []; + const toolResults = + message.role === "assistant" ? collectToolResultMessages(message) : []; + + const base: SerializedMessage = { + role: message.role, + content: + imageParts.length > 0 + ? [{ type: "text", text: textContent }, ...imageParts] + : textContent, + }; + if (toolCalls.length > 0) { + base.tool_calls = toolCalls; + // OpenAI requires content === null on assistant turns whose + // payload is entirely tool_calls (matches the wire shape Gemini + // expects for the next functionCall replay). + if (!textContent && imageParts.length === 0) { + base.content = null; + } + } + if (message.role === "assistant") { + const sig = collectAssistantTextThoughtSignature(message); + if (sig) { + base.extra_content = { google: { thought_signature: sig } }; + } } - if (!textContent) { + return toolResults.length > 0 ? [base, ...toolResults] : [base]; +} + +// Thin singular wrapper: returns only the first serialized message +// (without tool_calls or tool follow-ups) so the OpenAI image-edit +// replay path can map a thread to flat OpenAI chat messages without +// pulling in tool history. +function toOpenAIMessage(message: RunMessage): { + role: "system" | "user" | "assistant"; + content: OpenAIMessageContent; +} | null { + const serialized = toOpenAIMessages(message); + if (serialized.length === 0) return null; + const first = serialized[0]; + if ( + first.role !== "system" && + first.role !== "user" && + first.role !== "assistant" + ) { return null; } - return { role: message.role, content: textContent }; + if (first.content === null || first.content === undefined) { + return null; + } + if (typeof first.content === "string" && !first.content) { + return null; + } + return { + role: first.role, + content: first.content as OpenAIMessageContent, + }; } function extractImageBase64(input: string): string | undefined { @@ -1084,11 +1336,21 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { clearSelectedImageEditReference(); throw new Error("Connection not found."); } - // Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers. + // Local providers and custom Gemini bases allow an empty key. const externalProviderIsCustom = externalProvider ? isCustomProviderType(externalProvider.providerType) : false; - if (isExternalRequest && !externalApiKey && !externalProviderIsCustom) { + const externalProviderIsGeminiCustomBase = Boolean( + externalProvider && + externalProvider.providerType === "gemini" && + isGeminiCustomOpenAICompatBase(externalProvider.baseUrl), + ); + if ( + isExternalRequest && + !externalApiKey && + !externalProviderIsCustom && + !externalProviderIsGeminiCustomBase + ) { toast.error("Missing API key for selected connection.", { description: "Open Settings → Connections and set the API key again.", }); @@ -1096,15 +1358,38 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { throw new Error("Missing connection API key."); } + // Image-generation flag (OpenAI cloud + Responses-capable model). + // Computed first so Gemini image mode can suppress Search/Code. + const imageGenerationEnabledForThisTurn = Boolean( + externalProvider && + externalSelection && + imageToolsEnabled && + providerSupportsBuiltinImageGeneration( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ), + ); + // Per-model Search/Code allowances live in + // providerSupportsBuiltin*; this flag just signals image-mode. + const geminiImageModeForThisTurn = + externalProvider?.providerType === "gemini" && + imageGenerationEnabledForThisTurn; const webSearchEnabledForThisTurn = Boolean( externalProvider && + externalSelection && toolsEnabled && - providerSupportsBuiltinWebSearch(externalProvider.providerType), + providerSupportsBuiltinWebSearch( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ), ); const codeExecEnabledForThisTurn = Boolean( externalProvider && externalSelection && codeToolsEnabled && + !geminiImageModeForThisTurn && providerSupportsBuiltinCodeExecution( externalProvider.providerType, externalSelection.modelId, @@ -1124,20 +1409,6 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider && providerSupportsBuiltinWebFetch(externalProvider.providerType), ); - // OpenAI Responses-API image_generation server tool. Pill is - // gated on OpenAI cloud + a Responses-API model id; the backend - // additionally re-checks is_openai_cloud before appending - // {type:"image_generation"} to the request tools array. - const imageGenerationEnabledForThisTurn = Boolean( - externalProvider && - externalSelection && - imageToolsEnabled && - providerSupportsBuiltinImageGeneration( - externalProvider.providerType, - externalSelection.modelId, - externalProvider.baseUrl, - ), - ); if (selectedImageEditReference && !imageGenerationEnabledForThisTurn) { clearSelectedImageEditReference(); @@ -1148,11 +1419,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { throw new Error("Image generation edit unavailable."); } - // Two-pass build: a refused assistant turn also drops the user - // prompt that triggered it (leaving it in context re-triggers - // the classifier). Refusal flag rides assistant - // metadata.custom.anthropicRefusal, set out-of-band from the - // backend _toolEvent. + // Drop refused assistant turns + their triggering user prompt; + // otherwise context re-triggers the classifier. const survivingMessages: RunMessage[] = []; for (const message of messages) { if (isAnthropicRefusalMessage(message)) { @@ -1165,8 +1433,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { survivingMessages.push(message); } + // toOpenAIMessages emits assistant tool_calls + role="tool" + // follow-ups; the backend Gemini translator rebuilds the + // functionCall/functionResponse parts (with thoughtSignature). const outboundMessages = survivingMessages - .map(toOpenAIMessage) + .flatMap(toOpenAIMessages) .filter((message): message is NonNullable => Boolean(message), ); @@ -1189,7 +1460,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { break; } } - outboundMessages.splice(insertAt, 0, referenceMessage); + // OpenAIChatMessage is a structural superset of SerializedMessage + // for the role/content axis the outbound pipeline consumes; cast + // through unknown since referenceMessage carries no tool_calls + // (the image_edit reference is a plain assistant turn). + outboundMessages.splice( + insertAt, + 0, + referenceMessage as unknown as SerializedMessage, + ); } const safeSystemPrompt = @@ -1271,7 +1550,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { outboundMessages[0] = { ...firstMessage, content: [ - ...firstMessage.content, + ...(Array.isArray(firstMessage.content) + ? firstMessage.content + : []), { type: "text", text: `\n\n${disabledToolGuard}` }, ], }; @@ -1420,20 +1701,30 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { let cumulativeText = ""; let reasoningStartAt: number | null = null; let reasoningDuration = 0; - // Tracks whether we are currently inside a `` block opened by - // a `delta.reasoning_content` chunk. Kimi (kimi-k2.6, kimi-k2-thinking) - // and DeepSeek's reasoner stream their thinking as a separate - // `reasoning_content` field on the chat-completion delta — not as - // `content`, not as a structured part. We wrap those chunks with - // inline `...` so the existing parseAssistantContent - // lifts them into the reasoning panel the same way it does for - // local Harmony models. State has to live outside the SSE loop - // because the close tag fires when the next chunk carries content - // (or when the stream ends). + // True while wrapping a `delta.reasoning_content` stream in + // ... for parseAssistantContent. Lives outside + // the SSE loop because the close tag fires when content arrives. let reasoningContentOpen = false; - // Tool call content parts — accumulated and yielded cumulatively. - // result is set directly on the tool-call part when tool_end arrives. + // Tool call parts, cumulative; result lands on tool_end. const toolCallParts: ToolCallMessagePart[] = []; + // Latest Gemini text-part thoughtSignature; pinned onto the final + // text MessagePart so next-turn replay carries it. + let latestTextThoughtSignature: string | undefined; + const pinTextThoughtSignature = ( + parts: T[], + ): T[] => { + if (!latestTextThoughtSignature || parts.length === 0) return parts; + for (let i = parts.length - 1; i >= 0; i -= 1) { + if (parts[i].type === "text") { + parts[i] = { + ...parts[i], + _google_thought_signature: latestTextThoughtSignature, + } as T; + break; + } + } + return parts; + }; const orderAssistantContent = ( textParts: ReturnType, ) => { @@ -1525,6 +1816,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { { isReasoningProvider: externalProvider.isReasoningModel === true, + baseUrl: externalProvider.baseUrl ?? null, }, ) : { @@ -1557,13 +1849,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { forceRefreshPublicKey = false, ): Promise => { if (externalSelection && externalProvider) { - // OpenAI shell-tool container reuse: pull the per-thread - // container_id (if any) so subsequent turns in the same - // thread reference the existing container instead of - // auto-creating a fresh one. Empty string / undefined → - // backend falls back to container_auto. Anthropic uses - // the parallel `anthropicCodeExecContainerId` field below - // (sent as `container` on /v1/messages). + // Per-thread container reuse; empty/undefined falls back to + // container_auto. Anthropic uses anthropicCodeExecContainerId. let openaiCodeExecContainerId: string | null = null; let anthropicCodeExecContainerId: string | null = null; if (codeExecEnabledForThisTurn && resolvedThreadId) { @@ -1577,17 +1864,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { openaiCodeExecContainerId = null; anthropicCodeExecContainerId = null; } - // Pre-send container validation (OpenAI only). The list - // endpoint already filters status==="expired" server-side - // (studio/backend/routes/inference.py — list_openai_containers), - // so membership in this set means "OpenAI will accept it - // as container_reference". A stale id silently dropped here - // falls through to the inheritance + lazy-create logic - // below, so the user never sees "Container is expired" in - // the chat thread. On list-call failure we leave - // activeContainerIds null and skip validation — the - // backend's transparent retry path is the safety net for - // that case. + // Pre-send container validation (OpenAI). Stale ids drop + // silently and fall through to lazy-create. On list-call + // failure, skip and rely on the backend's retry path. let activeContainerIds: Set | null = null; if (externalProvider.providerType === "openai") { try { @@ -1610,15 +1889,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { openaiCodeExecContainerId = null; } } - // Cross-thread inheritance: when the active thread has - // no container yet, default to the one most recently - // used on *any* other thread (provider-scoped). - // Matches what the Code Execution settings section - // shows in the picker, and keeps the user from getting - // a fresh container on every new thread. The picker - // can still be set to "Auto-create per thread" - // explicitly to opt into a fresh container — but - // that's done via the dropdown, not silently. + // Cross-thread inheritance: reuse the most recently used + // container from any other thread; opt-out via the picker. if ( !openaiCodeExecContainerId && externalProvider.providerType === "openai" @@ -1651,13 +1923,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { /* fall through to lazy-create below */ } } - // Lazy pre-create when there's no inherited container. - // We always POST /v1/containers ourselves (rather than - // letting the backend send container_auto) so every - // container shows up in the picker with a friendly - // English-word name and the user's configured TTL. - // Falls back to container_auto only if the POST fails - // — keeps the chat moving in that case. + // Pre-create our own container (vs container_auto) so it + // shows up in the picker with a friendly name and the + // configured TTL. Falls back to container_auto on failure. if ( !openaiCodeExecContainerId && externalProvider.providerType === "openai" @@ -1716,22 +1984,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalSelection?.modelId, ), ), - // Only forward sampling knobs the provider actually accepts; the - // backend's external-provider proxy is param-permissive and would - // surface a 400 from providers that reject unknown fields (e.g. - // OpenAI rejects top_k, Anthropic/DeepSeek reject presence_penalty). + // Forward only sampling knobs the provider accepts. ...(externalCapabilities?.topK ? { top_k: params.topK } : {}), ...(externalCapabilities?.presencePenalty ? { presence_penalty: params.presencePenalty } : {}), - // Built-in tools: Search pill maps to provider-side - // web_search (currently OpenAI / Anthropic / OpenRouter / - // Kimi); Code pill maps to Anthropic's server-side - // code_execution_20250825 tool (Anthropic is the only - // external provider that ships one today). Backend - // translates enabled_tools into each provider's tool - // schema — for Anthropic that's the entries appended to - // body["tools"] inside _stream_anthropic. + // Compose the enabled_tools list from the active pills; + // backend maps each name to the provider's tool schema. ...(webSearchEnabledForThisTurn || webFetchEnabledForThisTurn || codeExecEnabledForThisTurn || @@ -1740,14 +1999,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { enable_tools: true, enabled_tools: [ ...(webSearchEnabledForThisTurn ? ["web_search"] : []), - // web_fetch has its own Fetch pill, independent - // of Search. Anthropic-only today. ...(webFetchEnabledForThisTurn ? ["web_fetch"] : []), ...(codeExecEnabledForThisTurn ? ["code_execution"] : []), - // OpenAI Responses-API only: `image_generation` - // returns inline image_generation_call output - // items; the backend's _stream_openai_responses - // path translates them to assistant tool events. ...(imageGenerationEnabledForThisTurn ? ["image_generation"] : []), @@ -1783,12 +2036,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider.enablePromptCaching ?? true, } : {}), - // Anthropic-only: pass the cache TTL the user picked in - // Configuration → Provider. Omitted = inherit the default - // 5-minute pool. The backend's `_stream_anthropic` only - // attaches `cache_control.ttl` when the value is one of - // "5m" / "1h" (see external_provider.py near line 1375), - // so unknown values are a no-op end-to-end. + // Anthropic prompt-cache TTL; unknown values no-op on backend. ...(supportsProviderPromptCacheTtl( externalProvider.providerType, ) && @@ -1896,11 +2144,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { chunk as unknown as { _toolEvent?: Record } )._toolEvent; if (toolEvent !== undefined) { - // OpenAI shell-tool container persistence — see - // ThreadRecord.openaiCodeExecContainerId. The backend - // emits these synthetic events on the OpenAI Responses - // SSE stream after capturing the container_id from a - // response, or detecting an expired-container error. + // Persist container_id onto the thread (OpenAI / Anthropic). if (toolEvent.type === "container_ready") { const newContainerId = toolEvent.container_id as | string @@ -1997,12 +2241,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { typeof imageB64 === "string" && imageB64 ) { - // OpenAI Responses image_generation_call: the - // backend stashes the base64 PNG/WebP/JPEG on - // separate `image_b64` / `image_mime` fields on - // the synthetic _toolEvent so the JSON result - // string stays small enough to log. Repackage as - // a structured result for the dedicated tool UI. + // Backend keeps base64 on separate image_b64 / + // image_mime fields so logs stay small; repackage. parsedResult = { image_b64: imageB64, image_mime: @@ -2029,27 +2269,104 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } else { parsedResult = rawResult; } + // Merge tool_end args first, then Gemini native_part. const nextArgs = toolEvent.arguments && typeof toolEvent.arguments === "object" ? (toolEvent.arguments as ToolCallMessagePart["args"]) : undefined; - const mergedArgs = nextArgs - ? { ...(toolCallParts[idx].args ?? {}), ...nextArgs } - : toolCallParts[idx].args; + const mergedArgs: ToolCallMessagePart["args"] = { + ...(toolCallParts[idx].args ?? {}), + ...(nextArgs ?? {}), + } as ToolCallMessagePart["args"]; + // Merge tool_end native_part into args.google so the + // outbound translator replays both start (executableCode) + // and end (result / inlineData) on the same turn. + // Concatenate parts so each keeps its own thoughtSignature. + const endGoogle = ( + toolEvent as { google?: { native_part?: unknown } } + ).google; + if ( + endGoogle && + typeof endGoogle === "object" && + endGoogle.native_part && + typeof endGoogle.native_part === "object" + ) { + const argsObj = mergedArgs as Record; + const existingGoogle = (argsObj.google ?? {}) as Record< + string, + unknown + >; + const existingNative = + (existingGoogle.native_part as Record< + string, + unknown + >) ?? {}; + const endNative = endGoogle.native_part as Record< + string, + unknown + >; + // Extract part entries from either parts:[...] or + // legacy single-object native_part. Legacy + // thoughtSignature always belongs on executableCode. + const collectParts = ( + native: Record, + ): Record[] => { + if (Array.isArray(native.parts)) { + return (native.parts as unknown[]).filter( + (entry): entry is Record => + Boolean(entry) && + typeof entry === "object" && + !Array.isArray(entry), + ); + } + const out: Record[] = []; + const legacySig = + typeof native.thoughtSignature === "string" + ? native.thoughtSignature + : typeof native.thought_signature === "string" + ? (native.thought_signature as string) + : null; + for (const key of [ + "executableCode", + "codeExecutionResult", + "inlineData", + ] as const) { + const sub = native[key]; + if (sub && typeof sub === "object") { + const entry: Record = { + [key]: sub, + }; + if (key === "executableCode" && legacySig) { + entry.thoughtSignature = legacySig; + } + out.push(entry); + } + } + return out; + }; + const mergedParts = [ + ...collectParts(existingNative), + ...collectParts(endNative), + ]; + argsObj.google = { + ...existingGoogle, + native_part: { parts: mergedParts }, + }; + } toolCallParts[idx] = { ...toolCallParts[idx], args: mergedArgs, - argsText: mergedArgs - ? JSON.stringify(mergedArgs) - : toolCallParts[idx].argsText, + argsText: JSON.stringify(mergedArgs ?? {}), result: parsedResult, }; } } - // Yield cumulative state so tool UI updates. Search/code tools stay - // before the text, while generated images sit after the answer. - const textParts = parseAssistantContent(cumulativeText); + // Cumulative yield. orderAssistantContent puts search/ + // code before text and generated images after. + const textParts = pinTextThoughtSignature( + parseAssistantContent(cumulativeText), + ); yield { content: orderAssistantContent(textParts), metadata: { @@ -2076,11 +2393,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } totalChunks += 1; - // OpenRouter's free router (openrouter/free) picks a different - // underlying free model per request and reports it in every - // chunk's top-level `model` field. Latch the first non-empty - // value that differs from the requested checkpoint so the - // header chip can render "openrouter/free:". + // Latch the chunk's `model` field so the openrouter/free + // chip can show the chosen underlying model. if ( isExternalRequest && externalProvider?.providerType === "openrouter" && @@ -2099,30 +2413,37 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } } const rawDelta = chunk.choices?.[0]?.delta?.content; - // Providers like Mistral's magistral return delta.content as an - // array of structured parts; normalize to text (with thinking - // parts re-wrapped as inline tags) so the rest of the - // accumulator stays string-based. + // Normalize structured delta.content (mistral magistral) to text. const delta = extractDeltaText(rawDelta); - // Kimi (kimi-k2.6, kimi-k2-thinking) and DeepSeek reasoner - // stream thinking via `delta.reasoning_content` as a plain - // string field — separate from `delta.content` which carries - // the answer. Wrap reasoning chunks inline as ... - // so parseAssistantContent treats them like any - // other reasoning. The close tag fires when the next chunk - // brings content, or when the stream ends. + // Latest Gemini text-part thoughtSignature for next-turn replay. + const deltaExtraContent = ( + chunk.choices?.[0]?.delta as + | { extra_content?: unknown } + | undefined + )?.extra_content; + if ( + deltaExtraContent && + typeof deltaExtraContent === "object" + ) { + const eGoogle = (deltaExtraContent as Record) + .google; + if (eGoogle && typeof eGoogle === "object") { + const sig = (eGoogle as Record) + .thought_signature; + if (typeof sig === "string" && sig) { + latestTextThoughtSignature = sig; + } + } + } + // Kimi / DeepSeek stream thinking via delta.reasoning_content. + // Wrap inline as ... for parseAssistantContent. const rawReasoning = ( chunk.choices?.[0]?.delta as | { reasoning_content?: unknown } | undefined )?.reasoning_content; - // OpenRouter uses a third reasoning shape: a structured - // `delta.reasoning_details` array of parts (each carrying - // `text`). The router emits this regardless of which - // underlying provider it picked, so we extract here and - // merge into the same ... wrap path used - // for Kimi / DeepSeek reasoning_content. See - // https://openrouter.ai/docs/guides/best-practices/reasoning-tokens + // OpenRouter ships reasoning as delta.reasoning_details[] + // regardless of underlying provider; merge into the same wrap path. const rawReasoningDetails = ( chunk.choices?.[0]?.delta as | { reasoning_details?: unknown } @@ -2140,6 +2461,138 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const reasoning = (typeof rawReasoning === "string" ? rawReasoning : "") + reasoningFromDetails; + // OpenAI delta.tool_calls: streams fragments by index; + // accumulate into one part. extra_content carries Gemini 3 + // thoughtSignature for next-turn replay. + const rawDeltaToolCalls = ( + chunk.choices?.[0]?.delta as + | { tool_calls?: unknown } + | undefined + )?.tool_calls; + if ( + Array.isArray(rawDeltaToolCalls) && + rawDeltaToolCalls.length > 0 + ) { + for (const tc of rawDeltaToolCalls) { + if (!tc || typeof tc !== "object") continue; + const call = tc as { + id?: string; + index?: number; + function?: { name?: string; arguments?: string }; + extra_content?: unknown; + }; + const idx = + typeof call.index === "number" ? call.index : undefined; + const stableId = call.id; + // Match an existing fragment by id first (canonical), + // then by index slot. Fall back to a freshly-minted + // tool_call_ id for streams that send neither. + let existing = stableId + ? toolCallParts.find((p) => p.toolCallId === stableId) + : undefined; + if (!existing && idx !== undefined) { + existing = toolCallParts.find( + (p) => + ( + p as ToolCallMessagePart & { _delta_index?: number } + )._delta_index === idx, + ); + } + const argsFragment = call.function?.arguments ?? ""; + if (existing) { + const prevName = existing.toolName ?? ""; + const nextName = call.function?.name ?? prevName; + const merged = + (existing.argsText ?? "") + argsFragment; + let parsedArgs: + ToolCallMessagePart["args"] = existing.args ?? {}; + if (merged) { + try { + parsedArgs = JSON.parse( + merged, + ) as ToolCallMessagePart["args"]; + } catch { + parsedArgs = { + _raw: merged, + } as ToolCallMessagePart["args"]; + } + } + const prevExtra = ( + existing as ToolCallMessagePart & { + extra_content?: unknown; + } + ).extra_content; + const updated: ToolCallMessagePart & { + _delta_index?: number; + extra_content?: unknown; + } = { + ...(existing as ToolCallMessagePart), + toolName: nextName, + argsText: merged, + args: parsedArgs, + ...(call.extra_content !== undefined + ? { extra_content: call.extra_content } + : prevExtra !== undefined + ? { extra_content: prevExtra } + : {}), + ...(idx !== undefined ? { _delta_index: idx } : {}), + }; + const replaceIdx = toolCallParts.indexOf(existing); + if (replaceIdx >= 0) { + toolCallParts[replaceIdx] = updated; + } + } else { + const callId = + stableId || + `tool_call_${idx ?? toolCallParts.length}`; + const argsText = argsFragment; + let parsedArgs: ToolCallMessagePart["args"] = {}; + if (argsText) { + try { + parsedArgs = JSON.parse( + argsText, + ) as ToolCallMessagePart["args"]; + } catch { + parsedArgs = { + _raw: argsText, + } as ToolCallMessagePart["args"]; + } + } + const fresh: ToolCallMessagePart & { + _delta_index?: number; + extra_content?: unknown; + } = { + type: "tool-call" as const, + toolCallId: callId, + toolName: call.function?.name ?? "", + argsText, + args: parsedArgs, + ...(call.extra_content !== undefined + ? { extra_content: call.extra_content } + : {}), + ...(idx !== undefined ? { _delta_index: idx } : {}), + }; + toolCallParts.push(fresh); + } + } + yield { + content: [ + ...toolCallParts, + ...pinTextThoughtSignature( + parseAssistantContent(cumulativeText), + ), + ], + metadata: { + timing: buildTiming( + streamStartTime, + totalChunks, + firstTokenTime, + ), + custom: { reasoningDuration }, + }, + }; + continue; + } if (!delta && !reasoning) { continue; } @@ -2165,21 +2618,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } cumulativeText += delta; } - // Mistral's magistral occasionally emits a trailing - // template-literal artifact (e.g. "${response}") at the end of - // an otherwise complete answer. It is never part of a real - // reply, so strip a trailing `${...}` token from external - // provider streams. The regex anchors to end-of-string and is - // idempotent — fragments mid-stream (e.g. "${re") leave the - // string untouched and only collapse once the closing brace - // arrives. Local-model output is left alone. + // Strip a trailing ${...} template-literal artifact from + // external streams (mistral magistral occasionally emits one). if (isExternalRequest) { cumulativeText = cumulativeText.replace( /\s*\$\{[^}]*\}\s*$/, "", ); } - const parts = parseAssistantContent(cumulativeText); + const parts = pinTextThoughtSignature( + parseAssistantContent(cumulativeText), + ); if ( parts.some((part) => part.type === "reasoning") && @@ -2295,7 +2744,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { yield { content: [ - ...orderAssistantContent(parseAssistantContent(cumulativeText)), + ...orderAssistantContent( + pinTextThoughtSignature(parseAssistantContent(cumulativeText)), + ), ...sourceParts, ...documentCitationParts, ], diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 1783a08281..c4a5f34273 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -728,7 +728,10 @@ export function ChatPage(): ReactElement { const reasoningCaps = getExternalReasoningCapabilities( provider?.providerType, selection.modelId, - { isReasoningProvider: provider?.isReasoningModel === true }, + { + isReasoningProvider: provider?.isReasoningModel === true, + baseUrl: provider?.baseUrl ?? null, + }, ); const state = useChatRuntimeStore.getState(); const preferredEffort = state.reasoningEffort; @@ -769,6 +772,8 @@ export function ChatPage(): ReactElement { : state.reasoningEffort; const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( provider?.providerType, + selection.modelId, + provider?.baseUrl, ); const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( provider?.providerType, @@ -968,6 +973,7 @@ export function ChatPage(): ReactElement { { isReasoningProvider: selectedProvider?.isReasoningModel === true, + baseUrl: selectedProvider?.baseUrl ?? null, }, ); const preferredEffort = store.reasoningEffort; @@ -1009,6 +1015,8 @@ export function ChatPage(): ReactElement { store.setCheckpoint(value, null); const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( selectedProvider?.providerType, + selectedExternal?.modelId, + selectedProvider?.baseUrl, ); const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( selectedProvider?.providerType, diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index dddf529068..b10253c280 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -37,6 +37,13 @@ export interface ExternalProviderConfig { updatedAt: number; } +// Gemini supports prompt caching, but the wire flow requires a +// separate POST to /v1beta/cachedContents to create the cache before +// the generateContent call can reference it; the boolean Studio +// currently emits on enable_prompt_caching is not enough on its own. +// Until that two-step orchestration ships we keep the picker off so +// the toggle does not silently no-op for Gemini users. See +// https://ai.google.dev/gemini-api/docs/caching. const PROMPT_CACHING_PROVIDER_TYPES = new Set(["openai", "anthropic"]); export function supportsProviderPromptCaching( diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index 5adc01ea2b..f550279826 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -189,7 +189,27 @@ function _inferProviderFromOpenrouterId( */ export function providerSupportsBuiltinWebSearch( providerType: string | null | undefined, + modelId?: string | null | undefined, + baseUrl?: string | null | undefined, ): boolean { + // Gemini ships grounded search via `tools: [{googleSearch: {}}]` on + // every chat-capable model. Most image-tier ids (`-image`, + // `nano-banana`) reject text-tool wiring because the + // responseModalities path is mutually exclusive with text tools, but + // Google explicitly documents Search grounding on the Gemini 3 image + // family (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + // nano-banana-pro). Allow Search on those; hide on older image ids. + // Custom Gemini OpenAI-compat proxies (non-Google bases) skip the + // native translator on the backend, so native tool envelopes never + // reach them -- hide the pill there. + if (providerType === "gemini") { + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + const normalized = modelId?.trim().toLowerCase() ?? ""; + if (normalized && isGeminiImageModel(normalized)) { + return geminiImageModelAllowsGoogleSearch(normalized); + } + return true; + } return ( providerType === "openai" || providerType === "anthropic" || @@ -319,6 +339,20 @@ export function providerSupportsBuiltinCodeExecution( normalized.startsWith(prefix), ); } + if (providerType === "gemini") { + // Gemini's `tools: [{codeExecution: {}}]` is supported on every + // chat-capable model. Image-tier ids (`-image`, `nano-banana`) + // reject text-tool wiring because the inline-image path is + // mutually exclusive with codeExecution. Custom Gemini + // OpenAI-compat proxies skip the native translator on the + // backend, so native codeExecution envelopes do not reach them. + // Wire-up lives in `_stream_gemini` on the backend; output comes + // back inline as executableCode/codeExecutionResult parts. See + // https://ai.google.dev/gemini-api/docs/code-execution. + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + if (isGeminiImageModel(normalized)) return false; + return normalized.startsWith("gemini-"); + } return false; } @@ -351,12 +385,75 @@ export function providerSupportsBuiltinImageGeneration( modelId: string | null | undefined, baseUrl?: string | null, ): boolean { - if (providerType !== "openai") return false; - if (!isOpenAICloudBaseUrl(baseUrl)) return false; const normalized = modelId?.trim().toLowerCase() ?? ""; if (!normalized) return false; - return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) => - normalized.startsWith(prefix), + if (providerType === "openai") { + if (!isOpenAICloudBaseUrl(baseUrl)) return false; + return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) => + normalized.startsWith(prefix), + ); + } + if (providerType === "gemini") { + // Gemini's Nano Banana image-output ids carry either `-image` (e.g. + // `gemini-2.5-flash-image`, `gemini-3.1-flash-image-preview`) or the + // `nano-banana` alias (`nano-banana-pro-preview`). The backend flips + // generationConfig.responseModalities to ["TEXT", "IMAGE"] when one + // is picked, and translates inlineData parts into the same image_b64 + // tool_end envelope the OpenAI path emits so the chat UI renders the + // picture inline. Custom Gemini OpenAI-compat proxies skip the + // native translator on the backend, so hide the image pill there. + // See https://ai.google.dev/gemini-api/docs/image-generation. + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + return normalized.includes("-image") || normalized.includes("nano-banana"); + } + return false; +} + +/** + * Whether `modelId` is a Gemini image-output id (Nano Banana family). + * Mirrors the backend's `is_image_picker_model` guard so the frontend + * hides text-only tool pills (web_search, code_execution) for these. + */ +function isGeminiImageModel(modelId: string): boolean { + const m = modelId.toLowerCase(); + return m.includes("-image") || m.includes("nano-banana"); +} + +/** + * Whether the saved Gemini connection points at a custom + * OpenAI-compatible gateway (any non-Google host). The backend + * `_is_openai_compatible` mirrors this to route those connections + * through `/chat/completions` instead of the native translator, so + * native Gemini tool envelopes (googleSearch, codeExecution, + * responseModalities) never reach them. Hide the corresponding + * Studio pills here so the request, builder, and UI agree. + */ +export function isGeminiCustomOpenAICompatBase( + baseUrl: string | null | undefined, +): boolean { + if (!baseUrl) return false; + try { + const host = new URL(baseUrl).hostname.toLowerCase(); + return host.length > 0 && host !== "generativelanguage.googleapis.com"; + } catch { + return false; + } +} + +/** + * Whether the given Gemini image model supports `tools: [{googleSearch: {}}]`. + * Google documents Search grounding on the Gemini 3 image family + * (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + * "Nano Banana Pro"); older image ids (gemini-2.5-flash-image) reject + * it with "Search as tool is not enabled for this model". + */ +function geminiImageModelAllowsGoogleSearch(modelId: string): boolean { + const m = modelId.toLowerCase(); + return ( + m.startsWith("gemini-3-pro-image") || + m.startsWith("gemini-3.1-flash-image") || + m.startsWith("nano-banana-pro") || + m.startsWith("nano-banana-2") ); } @@ -431,7 +528,20 @@ const PROVIDER_CAPABILITIES: Record = { presencePenalty: false, }, mistral: OPENAI_COMPAT_BASE, - gemini: OPENAI_COMPAT_BASE, + // Gemini's native generationConfig accepts temperature, topP, topK and + // presencePenalty (plus a separate frequencyPenalty we do not surface + // today). minP and repetitionPenalty are not part of the contract -- + // see https://ai.google.dev/api/rest/v1beta/GenerationConfig. Backend + // request shaping lives in _stream_gemini in + // studio/backend/core/inference/external_provider.py. + gemini: { + temperature: true, + topP: true, + topK: true, + minP: false, + repetitionPenalty: false, + presencePenalty: true, + }, // Kimi k2.5/k2.6 are reasoning-class — the API locks temperature and // top_p to fixed defaults and 400s on any other value: // "invalid temperature: only 1 is allowed for this model". @@ -643,6 +753,119 @@ function resolveKimiReasoningCapabilities(modelId: string): ExternalReasoningCap return withEnableThinkingStyle(); } +// Gemini's thinking ladder. +// - Gemini 3.x (3 / 3.1 / 3.5, Pro + Flash + Flash-Lite) and the +// gemini-pro-latest / gemini-flash-latest aliases use the new +// `thinkingConfig.thinkingLevel` string field (LOW/MEDIUM/HIGH/ +// MINIMAL). Pro tier rejects MINIMAL. +// - Gemini 2.5 Flash + 2.5 Pro stay on the integer +// `thinkingConfig.thinkingBudget` (0=off on Flash, -1=dynamic, +// N>0=cap; Pro rejects 0). +// - 2.5 Flash-Lite: no native thinking surfaced; leave it off. +// - Image-tier ids (`*-image*`, `nano-banana-pro-preview`): image +// generation path -- no reasoning controls. +const GEMINI3_PRO_PREFIXES = [ + "gemini-3.5-pro", + "gemini-3.1-pro", + "gemini-3-pro-preview", + "gemini-pro-latest", +]; +const GEMINI3_FLASH_PREFIXES = [ + "gemini-3.5-flash", + "gemini-3.1-flash", + "gemini-3-flash", + "gemini-flash-latest", + "gemini-flash-lite-latest", +]; +const GEMINI25_PRO_PREFIXES = [ + "gemini-2.5-pro", +]; +const GEMINI25_FLASH_PREFIXES = [ + "gemini-2.5-flash", +]; +const GEMINI_IMAGE_HINTS = [ + "-image", + "nano-banana", +]; +function resolveGeminiReasoningCapabilities( + modelId: string, +): ExternalReasoningCapabilities { + const m = modelId.toLowerCase(); + if (GEMINI_IMAGE_HINTS.some((h) => m.includes(h))) { + // Image generation; no thinking knob. + return withEnableThinkingStyle(); + } + // Gemini 2.5 Flash-Lite supports `thinkingBudget` with `0` = off and + // a positive range starting at 512 (the backend maps "minimal" to + // that floor at external_provider._stream_gemini). Check this branch + // BEFORE the broader `gemini-2.5-flash` prefix. + // https://ai.google.dev/gemini-api/docs/thinking + if (m.startsWith("gemini-2.5-flash-lite")) { + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: true, + reasoningEffortLevels: [ + "none", + "minimal", + "low", + "medium", + "high", + "max", + ] as const, + }); + } + if (GEMINI3_PRO_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 3.x Pro: thinkingLevel supports low/medium/high per + // https://ai.google.dev/gemini-api/docs/thinking and + // https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro. + // Cannot fully disable thinking; "minimal" is rejected on Pro. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: ["low", "medium", "high"] as const, + }); + } + if (GEMINI3_FLASH_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 3 Flash: thinkingLevel minimal/low/medium/high. Minimal + // is the closest to "off" Google offers on Gemini 3. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: [ + "minimal", + "low", + "medium", + "high", + ] as const, + }); + } + if (GEMINI25_PRO_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 2.5 Pro: thinkingBudget cannot be 0 (API rejects with + // "only works in thinking mode"); backend coerces to a small + // positive budget. The picker still hides the off switch. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: ["low", "medium", "high", "max"] as const, + }); + } + if (GEMINI25_FLASH_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 2.5 Flash: thinkingBudget supports 0 = off cleanly. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: true, + reasoningEffortLevels: [ + "none", + "low", + "medium", + "high", + "max", + ] as const, + }); + } + return withEnableThinkingStyle(); +} + function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoningCapabilities { if (modelId === "magistral-medium-latest") { return withReasoningEffortStyle({ @@ -665,6 +888,8 @@ function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoning export interface ExternalReasoningResolveOptions { /** vLLM connection flagged as a reasoning model in provider config. */ isReasoningProvider?: boolean; + /** Provider base URL; used to detect custom Gemini OAI-compat gateways. */ + baseUrl?: string | null; } // vLLM has no per-model reasoning signal on OpenAI-compat — pin via user toggle. @@ -740,6 +965,16 @@ export function getExternalReasoningCapabilities( } if (isKimiProvider) return resolveKimiReasoningCapabilities(modelForMatching); if (isMistralProvider) return resolveMistralReasoningCapabilities(modelForMatching); + if (normalizedProvider === "gemini") { + // Custom Gemini OAI-compat gateways (LiteLLM, proxies) route + // through /chat/completions which drops the Gemini-native + // thinkingConfig payload. Hide the native thinking ladder so the + // UI does not advertise a control the backend cannot honor. + if (isGeminiCustomOpenAICompatBase(options?.baseUrl)) { + return withEnableThinkingStyle(); + } + return resolveGeminiReasoningCapabilities(modelForMatching); + } if (!isOpenAIProvider && !isAnthropicProvider) { return withEnableThinkingStyle(); } diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index ba4a897f63..a811125a20 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -396,6 +396,7 @@ export function SharedComposer({ { isReasoningProvider: selectedExternalProvider?.isReasoningModel === true, + baseUrl: selectedExternalProvider?.baseUrl ?? null, }, ) : null; @@ -449,16 +450,36 @@ export function SharedComposer({ const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch( selectedExternalProvider?.providerType, ); - const searchDisabled = - !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); - const codeDisabled = - !modelLoaded || !(supportsTools || supportsBuiltinCodeExecution); - // Images pill is only ever lit on OpenAI cloud's Responses-API models. - // No local tool runtime fallback because the only image-generation - // server tool we wire today is OpenAI's; local models cannot dispatch - // it. Hidden entirely when the active model does not advertise it so - // the pill row stays compact for providers without the capability. + // Gemini rejects codeExecution alongside image modalities. Search is + // blocked on older Gemini image ids but allowed on Gemini 3 image + // models -- supportsBuiltinWebSearch already encodes the per-model + // allowance, so we only disable Code unconditionally in Gemini + // image mode. + const isExternalGemini = selectedExternalProvider?.providerType === "gemini"; const imageDisabled = !modelLoaded || !supportsBuiltinImageGeneration; + const imageModeDisablesCode = + isExternalGemini && imageToolsEnabled && !imageDisabled; + // Image-tier Gemini models always reject codeExecution and reject + // web_search on older ids (Gemini 3.x Pro/Flash allow it -- encoded + // in supportsBuiltinWebSearch). Don't let the local `supportsTools` + // runtime flag re-enable a pill the Gemini backend will silently + // drop. Detect "external provider is Gemini AND model is image-tier" + // and gate strictly on the provider builtin support. + const isGeminiImageTier = + isExternalGemini && supportsBuiltinImageGeneration; + const searchDisabled = + !modelLoaded || + (isGeminiImageTier + ? !supportsBuiltinWebSearch + : !(supportsTools || supportsBuiltinWebSearch)); + const codeDisabled = + !modelLoaded || + (isGeminiImageTier + ? true + : !(supportsTools || supportsBuiltinCodeExecution)) || + imageModeDisablesCode; + // Images pill is only ever lit on OpenAI cloud's Responses-API models + // and Gemini Nano Banana family. No local tool runtime fallback. const showImagePill = supportsBuiltinImageGeneration; // Fetch pill: Anthropic-only (web_fetch_20250910 / web_fetch_20260209). const webFetchDisabled = !modelLoaded || !supportsBuiltinWebFetch; diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 5238875b71..d313b43438 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -219,9 +219,34 @@ export type OpenAIMessageContentPart = export type OpenAIMessageContent = string | OpenAIMessageContentPart[]; +/** + * OpenAI Chat Completions tool_call shape. Assistant turns echo back + * function/tool calls as `tool_calls`; the matching tool result rides + * on a separate `role="tool"` message keyed by `tool_call_id`. + * `extra_content.google.thought_signature` is the Gemini-specific + * round-trip field the backend translator both emits (on `delta. + * tool_calls`) and consumes (when rebuilding the native functionCall + * part on the next turn). + */ +export interface OpenAIToolCallPart { + id?: string; + type?: "function"; + function?: { + name?: string; + arguments?: string; + }; + extra_content?: unknown; +} + export interface OpenAIChatMessage { - role: "system" | "user" | "assistant"; - content: OpenAIMessageContent; + role: "system" | "user" | "assistant" | "tool"; + content: OpenAIMessageContent | null; + /** Assistant tool-call deltas, when the turn invoked a function tool. */ + tool_calls?: OpenAIToolCallPart[]; + /** `role="tool"` only: id matching `assistant.tool_calls[].id`. */ + tool_call_id?: string; + /** `role="tool"` only: name of the function that produced the result. */ + name?: string; } export interface OpenAIChatCompletionsRequest { @@ -262,7 +287,14 @@ export interface OpenAIChatCompletionsRequest { external_model?: string; encrypted_api_key?: string; provider_base_url?: string | null; - enable_prompt_caching?: boolean | null; + /** + * Boolean toggle for OpenAI/Anthropic ephemeral cache_control. For + * Gemini the backend also accepts the cached-content resource name + * (`cachedContents/...`) as a string, which is forwarded as + * `generationConfig.cachedContent` on the native streamGenerateContent + * request. + */ + enable_prompt_caching?: boolean | string | null; /** * OpenAI shell-tool container id captured from the prior response in * this chat thread. When set and the Code pill is on, the backend @@ -292,7 +324,20 @@ export interface OpenAIChatCompletionsRequest { export interface OpenAIChatDelta { role?: string; - content?: string; + content?: string | null; + /** + * Streamed assistant tool calls. The Gemini and OpenAI Responses + * translators emit incremental `tool_calls` deltas (function name + + * arguments fragments) so the chat-adapter can render tool cards as + * they arrive. + */ + tool_calls?: OpenAIToolCallPart[]; + /** + * Provider-specific passthrough. Gemini ships `thoughtSignature`, + * citations, `native_part`, etc., here so the round-trip can replay + * them on follow-up turns without bleeding into other providers. + */ + extra_content?: Record; } export interface OpenAIChatChunkChoice {