fix(ui): expose api_base on Google AI Studio credential form - #24
Merged
Conversation
The Google_AI_Studio entry in provider_create_fields.json only exposed
api_key, so admins using a Gemini-compatible gateway (anything that
hosts the /v1beta/models/{model}:generateContent path on a custom host)
had no way to set a custom api_base via the UI. Forced workarounds:
either edit deployment-level litellm_params via raw API, or mislabel
the credential as "OpenAI" to borrow that form's api_base field.
The runtime gemini provider already handles custom api_base correctly
via vertex_llm_base._check_custom_proxy:415 — it builds
{api_base}/models/{model}:generateContent and uses x-goog-api-key. The
gap was UI-only.
Add an optional api_base field to the Google AI Studio credential form,
with a default value of the canonical Google AI Studio endpoint so
leaving the default behaves identically to leaving the field blank.
Place api_base before api_key in the field order to match Anthropic /
OpenAI / AI21 conventions (admin sees URL override before the key).
Add a regression test mirroring test_anthropic_provider_fields_support_byok
that asserts: api_base exists, optional, text type, default value
matches the canonical Google endpoint, and orders before api_key.
Test plan:
- python3 -m pytest tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py -v
(25 passed including new test_google_ai_studio_provider_fields_expose_api_base)
- black tests/.../test_public_endpoints.py (no changes after first run)
Adds the runtime half of PR #24's coverage: end-to-end validation that a deployment with `litellm_params.model = gemini/...` plus a custom `api_base` correctly routes through the gemini provider and surfaces the custom URL in `x-litellm-model-api-base`. Guards the path that the new UI api_base field on the Google_AI_Studio credential form will exercise once admins start using it. - e2e/tools/proxy: render a `gemini-custom-base` deployment when GEMINI_API_KEY is set in .env (gated so other cases don't fail with missing key). Reuses MODEL_GEMINI for the upstream model id with a gemini/gemini-3.1-pro-preview default. - e2e/cases/22_gemini_credential_custom_api_base.md: runbook with goal/preconditions/expected/failure-modes. - e2e/cases/data/22_gemini_credential_custom_api_base.sh: fixture asserts HTTP 200, x-litellm-model-api-base matches GEMINI_API_BASE from .env, x-litellm-model-group matches deployment name, and the response body shape. Exits 77 (SKIP) when GEMINI_API_KEY is unset. Test plan: - GEMINI_API_KEY + GEMINI_API_BASE set in e2e/.env pointing at a Gemini-compatible gateway: - e2e/tools/proxy restart - bash e2e/cases/data/22_gemini_credential_custom_api_base.sh → 3 PASS lines, exit 0 - Without GEMINI_API_KEY: deployment is omitted from rendered config; the case detects this via /v1/models and SKIPs (exit 77).
This was referenced May 28, 2026
songkuan-zheng
added a commit
that referenced
this pull request
Jun 4, 2026
* feat(proxy): add PublicReqMiddleware for X-Public-Req gated safeguards
Internal extension that lets a single LiteLLM proxy serve both external
users (through a public Nginx ingress) and internal services (direct
ClusterIP) from one process. The public ingress injects
`X-Public-Req: 1`; this middleware keys off that header to apply two
safeguards on external requests while leaving internal traffic
untouched.
When `X-Public-Req: 1` is present:
- Strip every inbound `x-litellm-*` request header. These let callers
override internal proxy behavior (`x-litellm-api-key`,
`x-litellm-num-retries`, `x-litellm-spend-logs-metadata`,
`x-litellm-mock-response`, etc.) and must never be honored from
untrusted sources.
- Silently drop forbidden query parameters on `/v1/models` and
`/models` (`include_metadata`, `fallback_type`,
`include_model_access_groups`, `only_model_access_groups`). These
expose router fallback chains and access-group naming. Other query
params are preserved byte-for-byte.
- Strip every outbound `x-litellm-*` response header (model deployment
IDs, cache-hit flags, cost/budget annotations, ratelimit-remaining
fields) before they leave the proxy.
When the header is absent (internal callers) the middleware is a
no-op: scope passes through, response is unchanged, internal observability
headers and the `x-litellm-spend-logs-metadata` write path keep working.
Implementation:
- Pure ASGI middleware (not `BaseHTTPMiddleware`) — does not buffer
response bodies. Modifies headers only on the single
`http.response.start` ASGI message; `http.response.body` chunks
pass through verbatim. Streaming TTFT is unaffected.
- Installed via a wrapper entrypoint, `litellm_extras.entrypoint`,
invoked in place of the upstream `litellm` console script. The
wrapper imports `litellm.proxy.proxy_server` first (constructing
the FastAPI `app`), calls `add_middleware`, then delegates to the
existing Click `run_server` CLI. Python's module cache means
uvicorn's `from .proxy_server import app` returns the same mutated
instance.
Coverage:
- 12 unit tests in `tests/test_litellm/test_public_req_middleware.py`
drive the middleware directly via fake ASGI scope/send, covering
internal passthrough, inbound + outbound strip, streaming chunk
pass-through (regression guard against accidental BaseHTTPMiddleware
switch), query-strip variations, marker case-insensitivity, and
non-HTTP scope routing.
- E2E case 18 (`e2e/cases/18_public_req_middleware.md` + fixture)
exercises the running proxy:
A1 SSE stream produces >=3 chunks
A2 streaming_phase (wall - ttfb) > 200ms (proves no buffering)
A3 public response: 0 x-litellm-* headers
A4 internal response: >=1 x-litellm-* header (control)
A5 public /v1/models?include_metadata=true returns 200 with the
same body as bare /v1/models (silent query strip)
A6 internal /v1/models?include_metadata=true returns 200 with a
different body (metadata expansion still works for internal)
A7 internal x-litellm-spend-logs-metadata reaches spend_logs
(control — proves the header would otherwise be honored)
A8 public x-litellm-spend-logs-metadata absent from spend_logs
(inbound strip verified at the DB layer)
All 8 assertions pass; the full e2e suite remains 18/18 PASS.
* fix(deploy): install PublicReqMiddleware at the Docker entrypoint, not docker-compose
The previous commit installed PublicReqMiddleware by overriding the
`entrypoint:` in `e2e/_config/docker-compose.yml`. That works for the
e2e harness, but for production it means every Kubernetes manifest /
Helm chart / docker run command has to remember to swap the entrypoint
too. A missed swap silently disables every X-Public-Req safeguard —
including the inbound header strip — and there is no runtime error to
flag the regression.
Move the wrapper invocation into `docker/prod_entrypoint.sh`, the
image's baked-in ENTRYPOINT. The upstream script previously dispatched
to `litellm "$@"`; it now dispatches to
`python -m litellm_extras.entrypoint "$@"` under both the plain and
the `USE_DDTRACE=true` paths. Every deployment of this image now
loads PublicReqMiddleware automatically with no extra configuration.
Drop the corresponding override in `e2e/_config/docker-compose.yml` so
e2e exercises the same entrypoint path that production does. If the
wrapper breaks in either, the other surface catches it.
Verified:
- `docker inspect litellm-e2e` shows entrypoint `docker/prod_entrypoint.sh`
(no docker-compose override), confirming the image's default path
is what runs.
- E2E case 18: 8/8 assertions PASS — middleware still strips inbound
+ outbound `x-litellm-*`, strips `/v1/models` forbidden query, and
the spend_logs marker check confirms the inbound header never
reached the proxy core.
- Full e2e suite: 18/18 PASS, no regressions.
Upstream-rebase note: `docker/prod_entrypoint.sh` is a small (14-line)
infrastructure file that rarely changes upstream. Future merge
conflicts on this file resolve by keeping our `python -m
litellm_extras.entrypoint` substitution in both branches of the
USE_DDTRACE conditional.
* fix(ui): expose api_base field on Google AI Studio credential form
The Google_AI_Studio entry in provider_create_fields.json only exposed
api_key, so admins using a Gemini-compatible gateway (anything that
hosts the /v1beta/models/{model}:generateContent path on a custom host)
had no way to set a custom api_base via the UI. Forced workarounds:
either edit deployment-level litellm_params via raw API, or mislabel
the credential as "OpenAI" to borrow that form's api_base field.
The runtime gemini provider already handles custom api_base correctly
via vertex_llm_base._check_custom_proxy:415 — it builds
{api_base}/models/{model}:generateContent and uses x-goog-api-key. The
gap was UI-only.
Add an optional api_base field to the Google AI Studio credential form,
with a default value of the canonical Google AI Studio endpoint so
leaving the default behaves identically to leaving the field blank.
Place api_base before api_key in the field order to match Anthropic /
OpenAI / AI21 conventions (admin sees URL override before the key).
Add a regression test mirroring test_anthropic_provider_fields_support_byok
that asserts: api_base exists, optional, text type, default value
matches the canonical Google endpoint, and orders before api_key.
Test plan:
- python3 -m pytest tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py -v
(25 passed including new test_google_ai_studio_provider_fields_expose_api_base)
- black tests/.../test_public_endpoints.py (no changes after first run)
* fix(ui): reset credential form state when switching providers
The credential modal's provider Select wired onChange to only update
`selectedProvider` + write the new value to `custom_llm_provider`. It
left every other field's Antd Form state untouched, so values seeded
by the previous provider's `default_value` carried over.
Most visibly: PR #24 added an `api_base` field to Google AI Studio with
a default of `https://generativelanguage.googleapis.com/v1beta`, but
because OpenAI is the modal's initial provider, OpenAI's `api_base`
default (`https://api.openai.com/v1`) populated Antd Form's state on
mount. Switching to Google AI Studio re-rendered ProviderSpecificFields
with the Google default, but Antd Form's controlled value still pointed
at the OpenAI URL — the user saw OpenAI's URL in the api_base field
when adding a Google AI Studio credential.
Fix: extract a `resetCredentialFormOnProviderChange` helper that
resetFields then restores the provider-agnostic fields (credential_name
+ custom_llm_provider) so the newly rendered ProviderSpecificFields can
apply its own defaults from a clean slate. Same call wires both modals
(Add + Edit) since they had the same bug.
Test plan:
- New unit test exercises the helper with 4 scenarios: (1) clears
provider-specific fields, (2) preserves credential_name, (3) writes
custom_llm_provider + invokes setSelectedProvider, (4) does NOT touch
credential_name when it was unset (avoids spurious "required"
validation on a brand-new modal).
- Existing AddCredentialModal + EditCredentialModal smoke tests still
pass (4 tests).
- vitest run src/components/model_add/credential_form_helpers.test.ts
AddCredentialModal.test.tsx EditCredentialModal.test.tsx -> 8 passed
Note: a direct end-to-end "render modal, click Provider Select, switch
to Google AI Studio, assert api_base" test would be more comprehensive
but Antd Select's portal/dropdown behavior is unreliable in jsdom —
attempted patterns (fireEvent.mouseDown, getByRole combobox,
document.body portal query) all timed out finding the dropdown
options. Testing the extracted helper directly is the more reliable
unit-level surrogate.
* fix(ui): logging callbacks table reads backend `type` for Mode badge (#35)
The `/get_callbacks` proxy endpoint returns each callback registration
as `{name, type, variables}` where `type` is `"success"` or `"failure"`.
The same callback name (e.g. `generic_api`) can appear twice — once
per event class — and the two entries fire on disjoint events, not
double-fire on a single event.
`LoggingCallbacksTable` ignored the `type` field and read
`record.mode`, which was always `undefined`, so every row fell back
to the "Success" badge. A `generic_api` callback registered for both
success and failure events showed up as two identical "Success"
rows, plus React emitted a duplicate-key warning because `rowKey` was
derived from `name` alone.
Changes
-------
* `types.ts`: `AlertingObject` gains an optional `type` field with a
comment explaining the success vs failure registration distinction.
* `LoggingCallbacksTable.tsx`:
- Mode column reads `record.type` first, falling back to
`record.mode` for newly-added (not-yet-server-acknowledged) rows.
- Composite rowKey ``${name}-${type ?? mode ?? 'success'}`` prevents
the duplicate-key warning when the same name appears twice.
- Removed leftover `console.log("availableCallbacks", ...)` debug
line that was firing on every render.
* `LoggingCallbacksTable.test.tsx`: adds a regression test that
renders two `generic_api` rows (one success, one failure) and
asserts both the "Success" and "Failure" badges are visible AND
that the shared display name appears twice.
Backend callback firing behaviour was already correct — this is a
display-only fix. Verified locally with the in-network mock provider
(separate PR): the mock's `/api/hooks/spend-log` receives exactly one
POST per success event and one per failure event, never both, even
when the UI was showing duplicate "Success" badges.
Test plan:
* `vitest run src/components/Settings/LoggingAndAlerts/LoggingCallbacks/`
→ 4 passed (3 existing + new regression).
* `prettier --check` clean.
* test(e2e): add case 22 — Gemini provider with custom api_base
Adds the runtime half of PR #24's coverage: end-to-end validation that a
deployment with `litellm_params.model = gemini/...` plus a custom
`api_base` correctly routes through the gemini provider and surfaces
the custom URL in `x-litellm-model-api-base`. Guards the path that the
new UI api_base field on the Google_AI_Studio credential form will
exercise once admins start using it.
- e2e/tools/proxy: render a `gemini-custom-base` deployment when
GEMINI_API_KEY is set in .env (gated so other cases don't fail with
missing key). Reuses MODEL_GEMINI for the upstream model id with a
gemini/gemini-3.1-pro-preview default.
- e2e/cases/22_gemini_credential_custom_api_base.md: runbook with
goal/preconditions/expected/failure-modes.
- e2e/cases/data/22_gemini_credential_custom_api_base.sh: fixture
asserts HTTP 200, x-litellm-model-api-base matches GEMINI_API_BASE
from .env, x-litellm-model-group matches deployment name, and the
response body shape. Exits 77 (SKIP) when GEMINI_API_KEY is unset.
Test plan:
- GEMINI_API_KEY + GEMINI_API_BASE set in e2e/.env pointing at a
Gemini-compatible gateway:
- e2e/tools/proxy restart
- bash e2e/cases/data/22_gemini_credential_custom_api_base.sh
→ 3 PASS lines, exit 0
- Without GEMINI_API_KEY: deployment is omitted from rendered config;
the case detects this via /v1/models and SKIPs (exit 77).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The Google AI Studio credential form in the dashboard UI only exposed
api_key, so admins using a Gemini-compatible gateway (any host serving/v1beta/models/{model}:generateContentwithx-goog-api-keyauth — e.g. self-hosted Raven Router, anispark router, etc.) had no way to set a customapi_basevia the UI.Forced workarounds today:
litellm_params.api_basevia raw admin APIapi_basefield, leavingModel.litellm_params.modelwith thegemini/prefix so the runtime still routes through the gemini providerBoth work because the runtime gemini provider already handles custom
api_basecorrectly vialitellm/llms/vertex_ai/vertex_llm_base.py:415— it builds{api_base}/models/{model}:generateContentand authenticates withx-goog-api-key. The gap was UI-only.Change
litellm/proxy/public_endpoints/provider_create_fields.json— add an optionalapi_basefield to theGoogle_AI_Studioentry, withdefault_value = "https://generativelanguage.googleapis.com/v1beta"(the canonical Google AI Studio endpoint LiteLLM talks to when api_base is unset, so leaving the default behaves identically to leaving the field blank).Field order is
api_basethenapi_key, matching Anthropic / OpenAI / AI21 conventions (admin sees the URL override before the key field).Test plan
python3 -m pytest tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py -v— 25 passed including newtest_google_ai_studio_provider_fields_expose_api_baseblack tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py— clean after one normalization pass/v1beta/models/{model}:generateContent: a Credential withprovider=Google AI Studio+api_base=<gw>/v1beta+api_key=<key>paired withModel.litellm_params.model = gemini/gemini-3.1-pro-previewreturns 200 from the gateway (gemini provider path,x-goog-api-keyheader)Note
This is a candidate upstream PR — the same gap exists on
BerriAI/litellmmain. Happy to mirror to upstream after this lands.