From 414a032865160d77e2e69afbfdad50b01dfdc4ba Mon Sep 17 00:00:00 2001
From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com>
Date: Fri, 22 May 2026 08:54:47 +0000
Subject: [PATCH 1/6] feat(proxy): add PublicReqMiddleware for X-Public-Req
 gated safeguards
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 e2e/_config/docker-compose.yml                |   7 +
 e2e/cases/18_public_req_middleware.md         | 120 ++++++
 e2e/cases/README.md                           |   1 +
 e2e/cases/data/18_public_req_middleware.sh    | 232 ++++++++++++
 e2e/tools/run-all-cases                       |  20 +
 litellm_extras/__init__.py                    |   7 +
 litellm_extras/entrypoint.py                  |  51 +++
 litellm_extras/public_req_middleware.py       | 134 +++++++
 .../test_public_req_middleware.py             | 358 ++++++++++++++++++
 9 files changed, 930 insertions(+)
 create mode 100644 e2e/cases/18_public_req_middleware.md
 create mode 100755 e2e/cases/data/18_public_req_middleware.sh
 create mode 100644 litellm_extras/__init__.py
 create mode 100644 litellm_extras/entrypoint.py
 create mode 100644 litellm_extras/public_req_middleware.py
 create mode 100644 tests/test_litellm/test_public_req_middleware.py

diff --git a/e2e/_config/docker-compose.yml b/e2e/_config/docker-compose.yml
index 7df49baa6f0..1626c7cd526 100644
--- a/e2e/_config/docker-compose.yml
+++ b/e2e/_config/docker-compose.yml
@@ -70,6 +70,13 @@ services:
       # Rendered at proxy-start time by `e2e/tools/proxy` from .env values.
       # See ../tools/proxy `render_config()`. The file is gitignored.
       - ./.litellm.rendered.yaml:/app/config.yaml:ro
+    # Invoke the internal entrypoint wrapper so PublicReqMiddleware is
+    # installed before the LiteLLM CLI starts uvicorn. Without this, the
+    # X-Public-Req gating in litellm_extras/ never runs.
+    entrypoint:
+      - python
+      - -m
+      - litellm_extras.entrypoint
     command:
       - --config=/app/config.yaml
       - --port=4000
diff --git a/e2e/cases/18_public_req_middleware.md b/e2e/cases/18_public_req_middleware.md
new file mode 100644
index 00000000000..cc6fd529b76
--- /dev/null
+++ b/e2e/cases/18_public_req_middleware.md
@@ -0,0 +1,120 @@
+# Case 18 — `X-Public-Req` middleware: streaming + header / query gating
+
+## Goal
+
+Verify the `PublicReqMiddleware` (installed via
+`litellm_extras.entrypoint`) correctly differentiates external vs internal
+requests **without buffering streaming bodies**.
+
+Two-pronged check:
+
+1. **Streaming integrity.** SSE chunks arrive incrementally; TTFB is
+   small relative to total wall time. This is the regression guard
+   against accidentally switching to `BaseHTTPMiddleware` or otherwise
+   awaiting on body messages — both would collapse TTFB onto end-of-stream.
+2. **Differential behavior.** `X-Public-Req: 1` strips `x-litellm-*`
+   from response headers and blocks `/v1/models?include_metadata=true`.
+   The same requests without the header pass through unchanged.
+
+Functional logic (inbound header strip, case-insensitive matching,
+forbidden-query parsing, non-HTTP scope handling) is covered by
+`tests/test_litellm/test_public_req_middleware.py`. This e2e adds the
+wire-level confirmation.
+
+## Origin
+
+We expose the LiteLLM proxy to external users behind a public Nginx
+ingress. The ingress injects `X-Public-Req: 1` on every forwarded
+request; internal services reach the proxy on a separate Service that
+does not pass through the ingress. The middleware sits in the proxy's
+ASGI stack and applies safeguards conditional on that header.
+
+## Preconditions
+
+- `e2e/tools/proxy status` reports `ready`
+- The proxy image was built **after** `litellm_extras/` was added —
+  run `e2e/tools/proxy rebuild` if unsure
+- `ANTHROPIC_API_KEY` set (one ~50-word streaming completion;
+  cost ≈ $0.002)
+
+## Steps
+
+```bash
+bash e2e/cases/data/18_public_req_middleware.sh
+echo "exit=$?"
+```
+
+The fixture executes six assertions. Each prints `PASS:` or `FAIL:`:
+
+### Streaming integrity (paid)
+
+- **A1** SSE response yields ≥ 3 `data:` chunks
+- **A2** Streaming phase (wall − TTFB) > 200 ms — chunks spread over
+  time rather than dumped at end-of-stream. Provider TTFT variance can
+  push the ratio above 0.5 on short completions, so the absolute
+  duration is more robust than a ratio threshold; a buffering
+  middleware would collapse this to ≤ 5 ms regardless of provider.
+
+### Header gating (paid)
+
+- **A3** Public mode: zero `x-litellm-*` headers in response
+- **A4** Internal mode (no `X-Public-Req`): at least one
+  `x-litellm-*` header in response (control)
+
+### Query gating on `/v1/models` (free)
+
+- **A5** Public mode: `GET /v1/models?include_metadata=true` returns
+  HTTP 200 with the **same** body as `GET /v1/models` (forbidden query
+  silently stripped before reaching the proxy).
+- **A6** Internal mode: same request returns HTTP 200 with a body that
+  **differs** from the bare-models response (metadata expansion still
+  works for internal callers).
+
+### Inbound `x-litellm-*` strip (free — uses `mock_response`)
+
+Two chat completions are made with `mock_response: "pong"` (no upstream
+provider call). Each carries `X-Litellm-Spend-Logs-Metadata` with a
+unique marker. LiteLLM persists honored header values to
+`metadata.spend_logs_metadata` in the spend_logs row, so the marker's
+presence/absence in Postgres is a definitive signal of whether the
+header reached the proxy core.
+
+- **A7** Internal mode (control): marker **must** appear in spend_logs
+  within 15 s (proves the header would otherwise be honored)
+- **A8** Public mode (`X-Public-Req: 1`): marker **must not** appear in
+  any spend_logs row (proves the middleware stripped the header before
+  LiteLLM saw it)
+
+## Expected — GREEN
+
+```
+A1 PASS: 7 SSE chunks received
+A2 PASS: ttfb=8294ms wall=10171ms streaming_phase=1877ms
+A3 PASS: 0 x-litellm-* headers in public response
+A4 PASS: 7 x-litellm-* headers in internal response
+A5 PASS: /v1/models?include_metadata=true returned 200 with query stripped (public)
+A6 PASS: /v1/models?include_metadata=true returned 200 with metadata expanded (internal)
+A7 PASS: internal mode: x-litellm-spend-logs-metadata reached spend_logs (1 row)
+A8 PASS: public mode: inbound x-litellm-spend-logs-metadata stripped (marker absent from spend_logs)
+```
+
+## Failure modes
+
+| Symptom | Likely cause |
+|---|---|
+| A1 fails with 0-1 chunks | Middleware is buffering — switched to `BaseHTTPMiddleware`, or `send_wrapper` is awaiting on body |
+| A2 `streaming_phase ≤ 5ms` | Middleware buffered the body and dumped it on close — same root cause as A1 |
+| A3 fails (still see `x-litellm-*` in public response) | `send_wrapper` not wired, or middleware not installed; check `proxy logs` for `PublicReqMiddleware` |
+| A4 fails (no `x-litellm-*` in internal) | Middleware is stripping for *all* requests; check `_is_public` returns False without the header |
+| A5 fails — public bodies differ between `?include_metadata=true` and bare `/v1/models` | Query strip did not run — `/v1/models` not in `MODELS_PATHS`, middleware not installed, or `parse_qsl`/`urlencode` lost the rewrite |
+| A5 fails — public returns 4xx | Middleware reverted to the old reject-with-400 behavior; revert the strip refactor |
+| A6 fails — internal bodies identical | Middleware running for internal calls — `_is_public` defaulting to True; strip is happening when it shouldn't |
+| A7 fails — internal marker missing from spend_logs | Async spend logger lag or DB schema drift — not a middleware bug. Raise the 15 s poll if reproducible |
+| A8 fails — public marker present in spend_logs | Inbound `x-litellm-*` strip is NOT running. Verify the middleware is installed and that `LITELLM_HEADER_PREFIX` matching is case-insensitive |
+
+## Cross-reference
+
+- `litellm_extras/public_req_middleware.py` — middleware under test
+- `litellm_extras/entrypoint.py` — wrapper that installs the middleware
+- `e2e/_config/docker-compose.yml` — `command:` invokes the wrapper
+- `tests/test_litellm/test_public_req_middleware.py` — functional unit tests
diff --git a/e2e/cases/README.md b/e2e/cases/README.md
index b4a0bb8e6d1..930fba4a316 100644
--- a/e2e/cases/README.md
+++ b/e2e/cases/README.md
@@ -23,6 +23,7 @@ Humans can execute them too — every step is a concrete shell command.
 | 11 | `11_error_information_message_populated.md` | (none — invalid key) | `spend_logs.metadata.error_information.error_message` non-empty on failure | ✓ |
 | 12 | `12_custom_pricing_must_honor_cache_tokens.md` | (none — direct calc) | `custom_cost_per_token` short-circuit must include cache pricing (Bug #2 root cause) | — |
 | 16 | `16_budget_reset_no_prisma_error.md` | (none — proxy only) | `ResetBudgetJob.reset_budget_windows` background tick must not raise `prisma.errors.MissingRequiredValueError` on `Json?` null-filter. Regression for BerriAI/litellm#26346 | ✓ |
+| 18 | `18_public_req_middleware.md` | Anthropic | `litellm_extras.PublicReqMiddleware` keeps streaming responses incremental, strips `x-litellm-*` under `X-Public-Req: 1`, and rejects sensitive `/v1/models` query params | — |
 
 ## How to invoke
 
diff --git a/e2e/cases/data/18_public_req_middleware.sh b/e2e/cases/data/18_public_req_middleware.sh
new file mode 100755
index 00000000000..a015e326b51
--- /dev/null
+++ b/e2e/cases/data/18_public_req_middleware.sh
@@ -0,0 +1,232 @@
+#!/usr/bin/env bash
+# Case 18 fixture — see e2e/cases/18_public_req_middleware.md
+#
+# Asserts that the PublicReqMiddleware installed via
+# litellm_extras.entrypoint:
+#   - keeps streaming responses incremental (A1, A2)
+#   - strips x-litellm-* response headers under X-Public-Req: 1 (A3)
+#   - leaves x-litellm-* response headers intact otherwise (A4)
+#   - rejects sensitive query params on /v1/models in public mode (A5)
+#   - accepts them in internal mode (A6)
+#
+# Exits 0 on PASS, 77 on SKIP (missing API key), anything else on FAIL.
+
+set -u
+
+PROXY="${PROXY_URL:-http://localhost:4011}"
+DB_CONTAINER="${DB_CONTAINER:-litellm-e2e-db}"
+DB_USER="${DB_USER:-litellm}"
+DB_NAME="${DB_NAME:-litellm}"
+KEY="${MASTER_KEY:-sk-e2e-test}"
+MODEL="${MODEL_E2E_NAME:-claude-sonnet-cache}"
+FAILED=0
+
+pass() { echo "PASS: $*"; }
+fail() { echo "FAIL: $*"; FAILED=1; }
+skip() { echo "SKIP: $*"; exit 77; }
+
+# ---- precondition: Anthropic key wired? ----------------------------------
+# The proxy reads ANTHROPIC_API_KEY at startup; if we don't have it the
+# streaming assertions cannot run. (/v1/models assertions would still work
+# but the fixture is treated as a unit.)
+if ! curl -sSf -o /dev/null -m 3 "$PROXY/health/readiness"; then
+    skip "proxy not ready at $PROXY"
+fi
+
+# ---- shared streaming request --------------------------------------------
+# A single Anthropic streaming call serves A1+A2+A3.
+PUB_BODY=$(mktemp); PUB_HDRS=$(mktemp); PUB_TIMING=$(mktemp)
+trap 'rm -f "$PUB_BODY" "$PUB_HDRS" "$PUB_TIMING" "$INT_HDRS" 2>/dev/null' EXIT
+
+curl -sN -D "$PUB_HDRS" -o "$PUB_BODY" \
+     -w "%{time_starttransfer} %{time_total} %{http_code}\n" \
+     -H "Authorization: Bearer $KEY" \
+     -H "X-Public-Req: 1" \
+     -H "Content-Type: application/json" \
+     -d "{
+        \"model\":\"$MODEL\",
+        \"messages\":[{\"role\":\"user\",\"content\":\"Write a 50-word essay about clouds.\"}],
+        \"max_tokens\":200,
+        \"stream\":true
+     }" "$PROXY/v1/chat/completions" > "$PUB_TIMING"
+
+read TTFB_S WALL_S STATUS < "$PUB_TIMING"
+if [ "$STATUS" != "200" ]; then
+    # Upstream auth failure → treat as SKIP not FAIL (no API key configured).
+    if [ "$STATUS" = "401" ] || [ "$STATUS" = "403" ]; then
+        skip "upstream returned $STATUS — ANTHROPIC_API_KEY likely missing"
+    fi
+    fail "streaming POST returned HTTP $STATUS"
+    echo "--- response body (truncated) ---"
+    head -c 400 "$PUB_BODY"
+    echo
+    exit 1
+fi
+
+TTFB_MS=$(awk "BEGIN{printf \"%d\", $TTFB_S*1000}")
+WALL_MS=$(awk "BEGIN{printf \"%d\", $WALL_S*1000}")
+
+# ---- A1: stream produced >= 3 SSE chunks ---------------------------------
+CHUNKS=$(grep -c '^data:' "$PUB_BODY" || true)
+if [ "$CHUNKS" -ge 3 ]; then
+    pass "$CHUNKS SSE chunks received"
+else
+    fail "only $CHUNKS SSE chunks (need >=3) — middleware may be buffering"
+fi
+
+# ---- A2: streaming phase visible (wall - ttfb > 200ms) -------------------
+#
+# A buffering middleware would emit ALL chunks at end-of-stream, collapsing
+# (wall - ttfb) toward 0. Provider TTFT variance can push the *ratio* above
+# 0.5 on short completions (slow first token + fast last tokens), so we use
+# the absolute streaming-phase duration instead of a ratio. Even ~200 ms of
+# streaming phase across multiple SSE chunks is unambiguous evidence that
+# the middleware did not buffer.
+STREAM_PHASE_MS=$(awk "BEGIN{printf \"%d\", ($WALL_S - $TTFB_S) * 1000}")
+if [ "$STREAM_PHASE_MS" -gt 200 ]; then
+    pass "ttfb=${TTFB_MS}ms wall=${WALL_MS}ms streaming_phase=${STREAM_PHASE_MS}ms"
+else
+    fail "ttfb=${TTFB_MS}ms wall=${WALL_MS}ms streaming_phase=${STREAM_PHASE_MS}ms (need >200ms; buffering suspected)"
+fi
+
+# ---- A3: public response has zero x-litellm-* headers -------------------
+PUB_LITELLM_COUNT=$(grep -ic '^x-litellm-' "$PUB_HDRS" || true)
+if [ "$PUB_LITELLM_COUNT" -eq 0 ]; then
+    pass "0 x-litellm-* headers in public response"
+else
+    fail "$PUB_LITELLM_COUNT x-litellm-* headers leaked (expected 0)"
+    grep -i '^x-litellm-' "$PUB_HDRS" | sed 's/^/    /'
+fi
+
+# ---- A4: internal response keeps x-litellm-* (control) ------------------
+INT_HDRS=$(mktemp)
+curl -sN -D "$INT_HDRS" -o /dev/null \
+     -H "Authorization: Bearer $KEY" \
+     -H "Content-Type: application/json" \
+     -d "{
+        \"model\":\"$MODEL\",
+        \"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],
+        \"max_tokens\":5,
+        \"stream\":true
+     }" "$PROXY/v1/chat/completions"
+
+INT_LITELLM_COUNT=$(grep -ic '^x-litellm-' "$INT_HDRS" || true)
+if [ "$INT_LITELLM_COUNT" -ge 1 ]; then
+    pass "$INT_LITELLM_COUNT x-litellm-* headers in internal response"
+else
+    fail "no x-litellm-* in internal response — middleware over-strips"
+fi
+
+# ---- A5: public /v1/models — forbidden query silently stripped ---------
+# Strategy: compare bodies for `?include_metadata=true` in public mode vs
+# internal mode. In public mode the middleware drops the parameter before
+# the proxy sees it, so the response must match the internal `no metadata`
+# baseline — neither expanded with fallback chains nor a 4xx error.
+PUB_NO_META=$(mktemp); PUB_WITH_META=$(mktemp); INT_NO_META=$(mktemp); INT_WITH_META=$(mktemp)
+
+S=$(curl -s -o "$PUB_WITH_META" -w "%{http_code}" \
+    -H "Authorization: Bearer $KEY" \
+    -H "X-Public-Req: 1" \
+    "$PROXY/v1/models?include_metadata=true")
+S_PUB_NOMETA=$(curl -s -o "$PUB_NO_META" -w "%{http_code}" \
+    -H "Authorization: Bearer $KEY" \
+    -H "X-Public-Req: 1" \
+    "$PROXY/v1/models")
+S_INT_META=$(curl -s -o "$INT_WITH_META" -w "%{http_code}" \
+    -H "Authorization: Bearer $KEY" \
+    "$PROXY/v1/models?include_metadata=true")
+S_INT_NOMETA=$(curl -s -o "$INT_NO_META" -w "%{http_code}" \
+    -H "Authorization: Bearer $KEY" \
+    "$PROXY/v1/models")
+
+if [ "$S" = "200" ] && [ "$S_PUB_NOMETA" = "200" ] \
+   && diff -q "$PUB_WITH_META" "$PUB_NO_META" >/dev/null; then
+    pass "/v1/models?include_metadata=true returned 200 with query stripped (public)"
+else
+    fail "/v1/models?include_metadata=true public status=$S vs no-meta=$S_PUB_NOMETA; bodies differ → strip failed"
+    diff "$PUB_WITH_META" "$PUB_NO_META" | head -5
+fi
+
+# ---- A6: internal /v1/models — metadata still expanded ------------------
+# Internal mode (no X-Public-Req) must NOT strip. The metadata-enriched
+# response must differ from the bare-models response.
+if [ "$S_INT_META" = "200" ] && [ "$S_INT_NOMETA" = "200" ] \
+   && ! diff -q "$INT_WITH_META" "$INT_NO_META" >/dev/null; then
+    pass "/v1/models?include_metadata=true returned 200 with metadata expanded (internal)"
+else
+    fail "internal status: with-meta=$S_INT_META no-meta=$S_INT_NOMETA; bodies identical → middleware over-strips"
+fi
+
+rm -f "$PUB_NO_META" "$PUB_WITH_META" "$INT_NO_META" "$INT_WITH_META"
+
+# ---- A7/A8: inbound x-litellm-* strip vs preserve ------------------------
+# Send two mock_response chat completions (free, no provider call) each
+# carrying x-litellm-spend-logs-metadata with a unique marker. The marker
+# only reaches LiteLLM's pre-call hooks if the header was honored — i.e.,
+# the proxy's spend_logs row will contain it under
+# metadata.spend_logs_metadata.case18_marker.
+#
+# Expected:
+#   A7 (public  + X-Public-Req: 1): row exists, marker ABSENT  → strip ran
+#   A8 (internal, no X-Public-Req): row exists, marker PRESENT → control
+PUB_MARKER="case18-pub-$(date +%s%N)"
+INT_MARKER="case18-int-$(date +%s%N)"
+
+# A7: public call — header must be stripped by the middleware
+curl -sS -o /dev/null \
+    -H "Authorization: Bearer $KEY" \
+    -H "X-Public-Req: 1" \
+    -H "X-Litellm-Spend-Logs-Metadata: {\"case18_marker\":\"$PUB_MARKER\"}" \
+    -H "Content-Type: application/json" \
+    -d "{
+        \"model\":\"$MODEL\",
+        \"messages\":[{\"role\":\"user\",\"content\":\"ping\"}],
+        \"mock_response\":\"pong\"
+    }" "$PROXY/v1/chat/completions"
+
+# A8: internal call (control) — header must be honored
+curl -sS -o /dev/null \
+    -H "Authorization: Bearer $KEY" \
+    -H "X-Litellm-Spend-Logs-Metadata: {\"case18_marker\":\"$INT_MARKER\"}" \
+    -H "Content-Type: application/json" \
+    -d "{
+        \"model\":\"$MODEL\",
+        \"messages\":[{\"role\":\"user\",\"content\":\"ping\"}],
+        \"mock_response\":\"pong\"
+    }" "$PROXY/v1/chat/completions"
+
+# Async spend logger flush — poll up to 15s for the internal-marker row to
+# appear, then make the public-marker assertion. If the internal row never
+# appears the logger is backed up and the absence of the public marker is
+# not yet proof of strip; in that case we surface a warning.
+INT_FOUND=0
+for _ in $(seq 1 15); do
+    sleep 1
+    INT_CNT=$(docker exec "$DB_CONTAINER" psql -U "$DB_USER" -d "$DB_NAME" -tA -c "
+SELECT COUNT(*) FROM \"LiteLLM_SpendLogs\"
+WHERE metadata::text LIKE '%$INT_MARKER%';
+" 2>/dev/null | tr -d ' ')
+    if [ "${INT_CNT:-0}" -ge 1 ]; then
+        INT_FOUND=1
+        break
+    fi
+done
+
+if [ "$INT_FOUND" -ne 1 ]; then
+    fail "A8 control: internal marker '$INT_MARKER' never reached spend_logs within 15s — async logger backed up?"
+else
+    pass "internal mode: x-litellm-spend-logs-metadata reached spend_logs ($INT_CNT row)"
+fi
+
+PUB_CNT=$(docker exec "$DB_CONTAINER" psql -U "$DB_USER" -d "$DB_NAME" -tA -c "
+SELECT COUNT(*) FROM \"LiteLLM_SpendLogs\"
+WHERE metadata::text LIKE '%$PUB_MARKER%';
+" 2>/dev/null | tr -d ' ')
+
+if [ "${PUB_CNT:-0}" -eq 0 ]; then
+    pass "public mode: inbound x-litellm-spend-logs-metadata stripped (marker absent from spend_logs)"
+else
+    fail "public mode: inbound x-litellm-spend-logs-metadata LEAKED — $PUB_CNT spend_logs row(s) carry marker '$PUB_MARKER'"
+fi
+
+exit $FAILED
diff --git a/e2e/tools/run-all-cases b/e2e/tools/run-all-cases
index f8161c1619f..6ec1366e50e 100755
--- a/e2e/tools/run-all-cases
+++ b/e2e/tools/run-all-cases
@@ -332,6 +332,21 @@ case_16() {
 }
 
 
+# ------------------------------------------------------------------ 18
+case_18() {
+    echo "[18] X-Public-Req middleware streaming + header/query gating..."
+    local out=/tmp/e2e_case18.out
+    bash e2e/cases/data/18_public_req_middleware.sh > "$out" 2>&1
+    local rc=$?
+    if [ "$rc" -eq 0 ]; then
+        ok "18 public-req-middleware: $(grep -c '^PASS' "$out") assertions PASS"
+    elif [ "$rc" -eq 77 ]; then
+        skip "18 public-req-middleware" "$(grep -m1 SKIP "$out")"
+    else
+        fail "18 public-req-middleware" "$(grep -m1 '^FAIL' "$out" || tail -2 "$out")"
+    fi
+}
+
 # Pre-flight: proxy must be ready
 if ! curl -sSL --max-time 3 -o /dev/null -w '%{http_code}' \
         "$PROXY/health/readiness" 2>/dev/null | grep -q '^2'; then
@@ -369,6 +384,11 @@ else
 fi
 case_12
 case_16
+if [ "$SKIP_PAID" -eq 1 ]; then
+    skip "18 public-req-middleware" "skipped (--skip-paid)"
+else
+    case_18
+fi
 
 echo
 echo "============ SUMMARY ============"
diff --git a/litellm_extras/__init__.py b/litellm_extras/__init__.py
new file mode 100644
index 00000000000..3edddae54de
--- /dev/null
+++ b/litellm_extras/__init__.py
@@ -0,0 +1,7 @@
+"""Internal LiteLLM extensions kept out of the upstream tree.
+
+This package hosts deployment-specific middleware and entrypoints layered on
+top of the vendored LiteLLM source. Modules here MUST NOT be imported by any
+file under ``litellm/`` itself — that would couple upstream code to internal
+add-ons and complicate future rebases.
+"""
diff --git a/litellm_extras/entrypoint.py b/litellm_extras/entrypoint.py
new file mode 100644
index 00000000000..60c3ecead50
--- /dev/null
+++ b/litellm_extras/entrypoint.py
@@ -0,0 +1,51 @@
+"""Wrapper entrypoint that installs internal middleware before the LiteLLM CLI runs.
+
+Invoke this module in place of the plain ``litellm`` console script::
+
+    python -m litellm_extras.entrypoint --config=/app/config.yaml --port=4000
+
+It works by importing :mod:`litellm.proxy.proxy_server` first (which has the
+side effect of constructing the FastAPI ``app`` instance), attaching
+:class:`~litellm_extras.public_req_middleware.PublicReqMiddleware`, and only
+then delegating to the existing Click-based ``run_server`` CLI. Because
+Python caches imported modules, the subsequent ``from .proxy_server import
+app`` inside ``run_server`` returns the same ``app`` object with our
+middleware already inserted.
+
+``FastAPI.add_middleware`` prepends to ``app.user_middleware``, so calling it
+after the proxy's own ``add_middleware`` lines puts ours on the OUTSIDE —
+which is what we want so the header strip runs before any LiteLLM
+auth/logging layer.
+"""
+
+import sys
+
+from litellm.proxy import proxy_server  # noqa: E402  (load order is intentional)
+from litellm.proxy.proxy_cli import run_server  # noqa: E402
+
+from litellm_extras.public_req_middleware import PublicReqMiddleware
+
+
+def install_middleware() -> None:
+    """Attach internal middleware to the proxy's FastAPI app.
+
+    Idempotent: if PublicReqMiddleware is already present we skip the
+    insert. This makes the entrypoint safe to invoke from tests that may
+    have imported ``litellm.proxy.proxy_server`` earlier.
+    """
+    existing = {
+        m.cls.__name__ for m in proxy_server.app.user_middleware if hasattr(m, "cls")
+    }
+    if PublicReqMiddleware.__name__ in existing:
+        return
+    proxy_server.app.add_middleware(PublicReqMiddleware)
+
+
+def main() -> int:
+    install_middleware()
+    run_server.main(args=sys.argv[1:], standalone_mode=False)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/litellm_extras/public_req_middleware.py b/litellm_extras/public_req_middleware.py
new file mode 100644
index 00000000000..d636f22c1ad
--- /dev/null
+++ b/litellm_extras/public_req_middleware.py
@@ -0,0 +1,134 @@
+"""ASGI middleware that gates public-vs-internal request behavior.
+
+The public Nginx ingress injects ``X-Public-Req: 1`` on every request it
+forwards. Internal callers reach the proxy directly (e.g. via the in-cluster
+Service ``litellm-internal``) and never carry that header. This middleware
+keys off the header to apply two safeguards for public requests only:
+
+1.  Strip every ``x-litellm-*`` request header — these headers (notably
+    ``x-litellm-api-key``, ``x-litellm-mock-response``,
+    ``x-litellm-num-retries``) let callers override internal proxy behavior
+    and must never be honored from untrusted sources.
+
+2.  Silently strip sensitive query parameters from ``/v1/models`` and
+    ``/models`` requests (``include_metadata``, ``fallback_type``,
+    ``include_model_access_groups``, ``only_model_access_groups``). These
+    parameters expose router fallback chains and access-group naming,
+    which are deployment-internal details. The request still reaches the
+    inner app and returns a normal 200 with the redacted view, so naive
+    public clients that always set these params do not break.
+
+For public responses the middleware additionally strips every
+``x-litellm-*`` response header (model deployment IDs, cache-hit flags,
+cost/budget annotations) before they leave the proxy.
+
+Internal requests (no header / header != ``"1"``) are passed through
+untouched so internal services can keep using the override headers and
+observability fields.
+
+The middleware is pure-ASGI (not Starlette ``BaseHTTPMiddleware``). It does
+not buffer response bodies, so streaming endpoints (chat completions with
+``stream=true``, ``/v1/messages``, ``/v1/realtime``) keep their original
+time-to-first-byte profile.
+"""
+
+from typing import Iterable, List, Optional, Tuple
+from urllib.parse import parse_qsl, urlencode
+
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+LITELLM_HEADER_PREFIX = b"x-litellm-"
+PUBLIC_REQ_HEADER = b"x-public-req"
+PUBLIC_REQ_VALUE = b"1"
+
+MODELS_PATHS = frozenset({"/v1/models", "/models"})
+MODELS_FORBIDDEN_QUERY_KEYS = frozenset(
+    {
+        "include_metadata",
+        "fallback_type",
+        "include_model_access_groups",
+        "only_model_access_groups",
+    }
+)
+
+
+class PublicReqMiddleware:
+    """Apply public-request safeguards keyed off ``X-Public-Req``.
+
+    Install order matters: this middleware should be the OUTERMOST layer so
+    its header strip runs before any LiteLLM auth/logging middleware reads
+    the request. ``FastAPI.add_middleware`` inserts at index 0, so calling
+    it after the proxy's own ``add_middleware`` lines puts this on the
+    outside automatically.
+    """
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        headers: List[Tuple[bytes, bytes]] = list(scope.get("headers") or [])
+        if not self._is_public(headers):
+            await self.app(scope, receive, send)
+            return
+
+        scope = dict(scope)
+
+        path = scope.get("path", "")
+        if path in MODELS_PATHS:
+            sanitized = self._strip_forbidden_query(scope.get("query_string", b""))
+            if sanitized is not None:
+                scope["query_string"] = sanitized
+
+        scope["headers"] = [
+            (name, value)
+            for name, value in headers
+            if not name.lower().startswith(LITELLM_HEADER_PREFIX)
+            and name.lower() != PUBLIC_REQ_HEADER
+        ]
+
+        async def send_wrapper(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                message = {
+                    **message,
+                    "headers": [
+                        (name, value)
+                        for name, value in message.get("headers", [])
+                        if not name.lower().startswith(LITELLM_HEADER_PREFIX)
+                    ],
+                }
+            await send(message)
+
+        await self.app(scope, receive, send_wrapper)
+
+    @staticmethod
+    def _is_public(headers: Iterable[Tuple[bytes, bytes]]) -> bool:
+        for name, value in headers:
+            if name.lower() == PUBLIC_REQ_HEADER:
+                return value.strip() == PUBLIC_REQ_VALUE
+        return False
+
+    @staticmethod
+    def _strip_forbidden_query(query_string: bytes) -> Optional[bytes]:
+        """Remove forbidden keys from a percent-encoded query string.
+
+        Returns the rewritten query string (possibly empty) when at least
+        one forbidden key was present, or ``None`` if the query string was
+        already safe. Returning ``None`` lets callers skip the
+        ``scope["query_string"]`` write and keep the original bytes
+        untouched — preserving exact ordering and any odd encoding the
+        client may have sent.
+        """
+        if not query_string:
+            return None
+        pairs = parse_qsl(
+            query_string.decode("latin-1"),
+            keep_blank_values=True,
+        )
+        kept = [(k, v) for k, v in pairs if k not in MODELS_FORBIDDEN_QUERY_KEYS]
+        if len(kept) == len(pairs):
+            return None
+        return urlencode(kept, doseq=True).encode("latin-1")
diff --git a/tests/test_litellm/test_public_req_middleware.py b/tests/test_litellm/test_public_req_middleware.py
new file mode 100644
index 00000000000..cd87580b167
--- /dev/null
+++ b/tests/test_litellm/test_public_req_middleware.py
@@ -0,0 +1,358 @@
+"""Unit tests for litellm_extras.public_req_middleware.PublicReqMiddleware."""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import pytest
+
+from litellm_extras.public_req_middleware import PublicReqMiddleware
+
+
+def _build_scope(
+    method: str = "GET",
+    path: str = "/v1/chat/completions",
+    query_string: bytes = b"",
+    headers: Optional[List[Tuple[bytes, bytes]]] = None,
+) -> Dict[str, Any]:
+    return {
+        "type": "http",
+        "asgi": {"version": "3.0"},
+        "http_version": "1.1",
+        "method": method,
+        "path": path,
+        "raw_path": path.encode(),
+        "query_string": query_string,
+        "headers": list(headers or []),
+        "client": ("127.0.0.1", 0),
+        "server": ("testserver", 80),
+        "scheme": "http",
+    }
+
+
+class _RecordingApp:
+    """ASGI app that records what scope it received and yields predetermined chunks."""
+
+    def __init__(
+        self,
+        response_headers: Optional[List[Tuple[bytes, bytes]]] = None,
+        body_chunks: Optional[List[bytes]] = None,
+        status: int = 200,
+    ) -> None:
+        self.response_headers = response_headers or []
+        self.body_chunks = body_chunks or [b'{"ok":true}']
+        self.status = status
+        self.received_scope: Optional[Dict[str, Any]] = None
+
+    async def __call__(self, scope: Dict[str, Any], receive: Any, send: Any) -> None:
+        self.received_scope = scope
+        await send(
+            {
+                "type": "http.response.start",
+                "status": self.status,
+                "headers": list(self.response_headers),
+            }
+        )
+        for i, chunk in enumerate(self.body_chunks):
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": chunk,
+                    "more_body": i < len(self.body_chunks) - 1,
+                }
+            )
+
+
+async def _drive(
+    middleware: PublicReqMiddleware,
+    scope: Dict[str, Any],
+) -> Tuple[List[Dict[str, Any]], _RecordingApp]:
+    """Run the middleware against an empty-receive client and capture sends."""
+    sent: List[Dict[str, Any]] = []
+
+    async def receive() -> Dict[str, Any]:
+        return {"type": "http.disconnect"}
+
+    async def send(message: Dict[str, Any]) -> None:
+        sent.append(message)
+
+    await middleware(scope, receive, send)
+    return sent, middleware.app  # type: ignore[return-value]
+
+
+def _start_msg(sent: List[Dict[str, Any]]) -> Dict[str, Any]:
+    starts = [m for m in sent if m["type"] == "http.response.start"]
+    assert starts, f"no response.start in {sent}"
+    return starts[0]
+
+
+def _body_msgs(sent: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    return [m for m in sent if m["type"] == "http.response.body"]
+
+
+@pytest.mark.asyncio
+async def test_internal_request_passes_through_unchanged() -> None:
+    """No X-Public-Req → request scope untouched, response headers preserved."""
+    app = _RecordingApp(
+        response_headers=[
+            (b"content-type", b"application/json"),
+            (b"x-litellm-model-id", b"gpt-4o-mini-internal"),
+            (b"x-litellm-cache-hit", b"true"),
+        ]
+    )
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(
+        headers=[
+            (b"authorization", b"Bearer sk-internal"),
+            (b"x-litellm-mock-response", b"trace:xyz"),
+            (b"x-litellm-tags", b"team:alpha"),
+        ]
+    )
+
+    sent, _ = await _drive(mw, scope)
+
+    received_hdrs = dict(app.received_scope["headers"])
+    assert received_hdrs.get(b"x-litellm-mock-response") == b"trace:xyz"
+    assert received_hdrs.get(b"x-litellm-tags") == b"team:alpha"
+
+    resp_hdrs = dict(_start_msg(sent)["headers"])
+    assert resp_hdrs.get(b"x-litellm-model-id") == b"gpt-4o-mini-internal"
+    assert resp_hdrs.get(b"x-litellm-cache-hit") == b"true"
+
+
+@pytest.mark.asyncio
+async def test_public_request_strips_inbound_litellm_headers() -> None:
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(
+        headers=[
+            (b"authorization", b"Bearer sk-user"),
+            (b"x-public-req", b"1"),
+            (b"x-litellm-mock-response", b"hack"),
+            (b"x-litellm-num-retries", b"50"),
+            (b"x-litellm-tags", b"team:victim"),
+            (b"content-type", b"application/json"),
+        ]
+    )
+
+    await _drive(mw, scope)
+
+    received = dict(app.received_scope["headers"])
+    assert b"x-litellm-mock-response" not in received
+    assert b"x-litellm-num-retries" not in received
+    assert b"x-litellm-tags" not in received
+    # X-Public-Req itself is stripped too (don't surface the marker to LiteLLM).
+    assert b"x-public-req" not in received
+    # Non-LiteLLM headers must survive.
+    assert received[b"authorization"] == b"Bearer sk-user"
+    assert received[b"content-type"] == b"application/json"
+
+
+@pytest.mark.asyncio
+async def test_public_request_strips_outbound_litellm_headers() -> None:
+    app = _RecordingApp(
+        response_headers=[
+            (b"content-type", b"application/json"),
+            (b"x-litellm-model-id", b"gpt-4o-mini-leak"),
+            (b"x-litellm-cache-hit", b"true"),
+            (b"x-litellm-response-cost", b"0.0001"),
+            (b"x-other", b"keep"),
+        ]
+    )
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(headers=[(b"x-public-req", b"1")])
+
+    sent, _ = await _drive(mw, scope)
+
+    resp = dict(_start_msg(sent)["headers"])
+    assert b"x-litellm-model-id" not in resp
+    assert b"x-litellm-cache-hit" not in resp
+    assert b"x-litellm-response-cost" not in resp
+    assert resp[b"x-other"] == b"keep"
+
+
+@pytest.mark.asyncio
+async def test_public_request_streaming_body_chunks_pass_through_unbuffered() -> None:
+    """Body chunks must reach send() one-by-one without coalescing.
+
+    Regression guard for the BaseHTTPMiddleware trap: a misimplemented
+    middleware would buffer the whole body before forwarding. We assert that
+    each chunk produced by the inner app surfaces as a separate
+    http.response.body message in original order.
+    """
+    chunks = [
+        b'data: {"a":1}\n\n',
+        b'data: {"b":2}\n\n',
+        b'data: {"c":3}\n\n',
+        b"data: [DONE]\n\n",
+    ]
+    app = _RecordingApp(
+        response_headers=[(b"content-type", b"text/event-stream")],
+        body_chunks=chunks,
+    )
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(headers=[(b"x-public-req", b"1")])
+
+    sent, _ = await _drive(mw, scope)
+
+    body_messages = _body_msgs(sent)
+    assert len(body_messages) == len(chunks)
+    assert [m["body"] for m in body_messages] == chunks
+    # All but the last chunk must signal more_body=True.
+    assert [m.get("more_body", False) for m in body_messages] == [
+        True,
+        True,
+        True,
+        False,
+    ]
+
+
+@pytest.mark.asyncio
+async def test_public_request_strips_forbidden_models_query() -> None:
+    """Public mode silently drops forbidden keys and proxies the rest."""
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(
+        method="GET",
+        path="/v1/models",
+        query_string=b"include_metadata=true&fallback_type=general&team_id=keep",
+        headers=[(b"x-public-req", b"1")],
+    )
+
+    sent, _ = await _drive(mw, scope)
+
+    assert _start_msg(sent)["status"] == 200, "request must reach inner app"
+    assert app.received_scope is not None
+    forwarded_qs = app.received_scope["query_string"]
+    assert b"include_metadata" not in forwarded_qs
+    assert b"fallback_type" not in forwarded_qs
+    assert b"team_id=keep" in forwarded_qs
+
+
+@pytest.mark.asyncio
+async def test_public_request_strip_leaves_query_when_only_forbidden_keys() -> None:
+    """If every query key was forbidden, the forwarded query is empty."""
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(
+        method="GET",
+        path="/v1/models",
+        query_string=b"include_metadata=true&only_model_access_groups=1",
+        headers=[(b"x-public-req", b"1")],
+    )
+
+    sent, _ = await _drive(mw, scope)
+
+    assert _start_msg(sent)["status"] == 200
+    assert app.received_scope is not None
+    assert app.received_scope["query_string"] == b""
+
+
+@pytest.mark.asyncio
+async def test_public_request_models_clean_query_unchanged() -> None:
+    """Safe query strings must be passed through byte-for-byte."""
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(
+        method="GET",
+        path="/v1/models",
+        query_string=b"team_id=alpha&scope=expand",
+        headers=[(b"x-public-req", b"1")],
+    )
+
+    sent, _ = await _drive(mw, scope)
+
+    assert _start_msg(sent)["status"] == 200
+    assert app.received_scope["query_string"] == b"team_id=alpha&scope=expand"
+
+
+@pytest.mark.asyncio
+async def test_public_request_allows_models_without_forbidden_query() -> None:
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(
+        method="GET",
+        path="/v1/models",
+        query_string=b"",
+        headers=[(b"x-public-req", b"1")],
+    )
+
+    sent, _ = await _drive(mw, scope)
+
+    assert _start_msg(sent)["status"] == 200
+    assert app.received_scope is not None
+
+
+@pytest.mark.asyncio
+async def test_internal_request_keeps_forbidden_query_on_models() -> None:
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(
+        method="GET",
+        path="/v1/models",
+        query_string=b"include_metadata=true",
+        headers=[],
+    )
+
+    sent, _ = await _drive(mw, scope)
+
+    assert _start_msg(sent)["status"] == 200
+
+
+@pytest.mark.asyncio
+async def test_marker_value_other_than_1_treated_as_internal() -> None:
+    """Defensive: only the literal "1" enables public mode."""
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    scope = _build_scope(
+        headers=[
+            (b"x-public-req", b"true"),
+            (b"x-litellm-tags", b"team:keep"),
+        ]
+    )
+
+    await _drive(mw, scope)
+
+    received = dict(app.received_scope["headers"])
+    assert received.get(b"x-litellm-tags") == b"team:keep"
+
+
+@pytest.mark.asyncio
+async def test_marker_is_case_insensitive_for_header_name() -> None:
+    """ASGI lowercases header names by spec, but defensive parsing matters."""
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    # ASGI gives us lowercase, but we double-check the comparison anyway.
+    scope = _build_scope(
+        headers=[
+            (b"X-Public-Req", b"1"),  # mixed case
+            (b"X-Litellm-Tags", b"team:victim"),
+        ]
+    )
+
+    await _drive(mw, scope)
+
+    received = dict(app.received_scope["headers"])
+    # Mixed-case x-litellm-tags should still be stripped.
+    assert b"X-Litellm-Tags" not in received
+    assert b"x-litellm-tags" not in received
+
+
+@pytest.mark.asyncio
+async def test_non_http_scope_is_passed_through() -> None:
+    """Websocket and lifespan must reach the inner app untouched."""
+    app = _RecordingApp()
+    mw = PublicReqMiddleware(app)
+    scope = {"type": "websocket", "path": "/v1/realtime", "headers": []}
+
+    async def receive() -> Dict[str, Any]:
+        return {"type": "websocket.connect"}
+
+    async def send(message: Dict[str, Any]) -> None:
+        return None
+
+    await mw(scope, receive, send)
+
+    assert app.received_scope is scope
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

From 4293d940f6e005cd8c984583d325cd14279171f1 Mon Sep 17 00:00:00 2001
From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com>
Date: Fri, 22 May 2026 09:12:38 +0000
Subject: [PATCH 2/6] fix(deploy): install PublicReqMiddleware at the Docker
 entrypoint, not docker-compose
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 docker/prod_entrypoint.sh      | 11 +++++++++--
 e2e/_config/docker-compose.yml | 12 +++++-------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/docker/prod_entrypoint.sh b/docker/prod_entrypoint.sh
index bd78bf6687b..941e2de3dca 100644
--- a/docker/prod_entrypoint.sh
+++ b/docker/prod_entrypoint.sh
@@ -1,8 +1,15 @@
 #!/bin/sh
 
+# Internal-fork modification: invoke `python -m litellm_extras.entrypoint`
+# instead of the upstream `litellm` console script. The wrapper installs
+# PublicReqMiddleware on the FastAPI app before delegating to the original
+# Click CLI, which is what enforces the X-Public-Req gating in production.
+# Routing through the wrapper at the entrypoint level (rather than via
+# docker-compose overrides) means every deployment of this image picks
+# the middleware up automatically — there is no per-deploy step to forget.
 if [ "$USE_DDTRACE" = "true" ]; then
     export DD_TRACE_OPENAI_ENABLED="False"
-    exec ddtrace-run litellm "$@"
+    exec ddtrace-run python -m litellm_extras.entrypoint "$@"
 else
-    exec litellm "$@"
+    exec python -m litellm_extras.entrypoint "$@"
 fi
diff --git a/e2e/_config/docker-compose.yml b/e2e/_config/docker-compose.yml
index 1626c7cd526..3477fd99f62 100644
--- a/e2e/_config/docker-compose.yml
+++ b/e2e/_config/docker-compose.yml
@@ -70,13 +70,11 @@ services:
       # Rendered at proxy-start time by `e2e/tools/proxy` from .env values.
       # See ../tools/proxy `render_config()`. The file is gitignored.
       - ./.litellm.rendered.yaml:/app/config.yaml:ro
-    # Invoke the internal entrypoint wrapper so PublicReqMiddleware is
-    # installed before the LiteLLM CLI starts uvicorn. Without this, the
-    # X-Public-Req gating in litellm_extras/ never runs.
-    entrypoint:
-      - python
-      - -m
-      - litellm_extras.entrypoint
+    # NOTE: no `entrypoint:` override here. The fork's
+    # docker/prod_entrypoint.sh already routes through
+    # `python -m litellm_extras.entrypoint`, so PublicReqMiddleware loads
+    # automatically. This keeps e2e behavior aligned with production —
+    # if the wrapper breaks in either path, both surfaces catch it.
     command:
       - --config=/app/config.yaml
       - --port=4000

From 767909d4d401662e1ba4b7385cbb506a4474114f Mon Sep 17 00:00:00 2001
From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com>
Date: Thu, 28 May 2026 09:59:35 +0000
Subject: [PATCH 3/6] fix(ui): expose api_base field on Google AI Studio
 credential form
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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)
---
 .../provider_create_fields.json               | 10 ++++
 .../public_endpoints/test_public_endpoints.py | 49 +++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json
index 163c9648de7..86c489e1823 100644
--- a/litellm/proxy/public_endpoints/provider_create_fields.json
+++ b/litellm/proxy/public_endpoints/provider_create_fields.json
@@ -1243,6 +1243,16 @@
     "provider_display_name": "Google AI Studio",
     "litellm_provider": "gemini",
     "credential_fields": [
+      {
+        "key": "api_base",
+        "label": "API Base",
+        "placeholder": "https://generativelanguage.googleapis.com/v1beta",
+        "tooltip": "Override only when using a Gemini-compatible gateway (e.g. /v1beta path on a self-hosted proxy). LiteLLM appends '/models/{model}:generateContent' so include '/v1beta' but not the trailing slash.",
+        "required": false,
+        "field_type": "text",
+        "options": null,
+        "default_value": "https://generativelanguage.googleapis.com/v1beta"
+      },
       {
         "key": "api_key",
         "label": "API Key",
diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
index f82da59899b..caabc7b5fdd 100644
--- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
+++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
@@ -201,6 +201,55 @@ def test_anthropic_provider_fields_support_byok():
     ), "api_base must appear before api_key in credential_fields (matches AI21 and ANTHROPIC_TEXT convention)."
 
 
+def test_google_ai_studio_provider_fields_expose_api_base():
+    """The Google AI Studio (gemini) credential form must let admins set a custom
+    api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted
+    proxy at /v1beta) without env var access.
+
+    The runtime gemini provider already supports custom api_base via
+    `vertex_llm_base._check_custom_proxy`; the UI just needs to expose the field.
+    """
+    app_instance = FastAPI()
+    app_instance.include_router(router)
+    test_client = TestClient(app_instance)
+
+    response = test_client.get("/public/providers/fields")
+    assert response.status_code == 200
+    providers = response.json()
+
+    google_ai = next(
+        (p for p in providers if p["provider"] == "Google_AI_Studio"), None
+    )
+    assert google_ai is not None, "Google_AI_Studio provider entry not found"
+    assert google_ai["litellm_provider"] == "gemini"
+
+    fields_by_key = {f["key"]: f for f in google_ai["credential_fields"]}
+    assert "api_key" in fields_by_key
+    assert "api_base" in fields_by_key, (
+        "Google_AI_Studio provider form must expose api_base so admins can "
+        "point at a Gemini-compatible gateway without env var access."
+    )
+
+    api_base_field = fields_by_key["api_base"]
+    assert api_base_field["required"] is False
+    assert api_base_field["field_type"] == "text"
+    # Default value should match the canonical Google AI Studio endpoint that
+    # LiteLLM's gemini provider talks to when api_base is unset, so leaving the
+    # default in the form behaves identically to leaving it blank.
+    assert (
+        api_base_field["default_value"]
+        == "https://generativelanguage.googleapis.com/v1beta"
+    )
+
+    # UI forms render fields in credential_fields order; api_base should come
+    # first so an admin sees the URL override before the key field (matches
+    # OpenAI and Anthropic conventions).
+    field_order = [f["key"] for f in google_ai["credential_fields"]]
+    assert field_order.index("api_base") < field_order.index(
+        "api_key"
+    ), "api_base must appear before api_key in credential_fields."
+
+
 def test_public_model_hub_with_healthy_model():
     """Test that health information is populated for a healthy model"""
     app = FastAPI()

From d6e6b0f4524653dd1c60a844273808698cdd49b4 Mon Sep 17 00:00:00 2001
From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com>
Date: Thu, 28 May 2026 11:27:01 +0000
Subject: [PATCH 4/6] fix(ui): reset credential form state when switching
 providers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 .../model_add/AddCredentialModal.tsx          |  4 +-
 .../model_add/EditCredentialModal.tsx         |  4 +-
 .../model_add/credential_form_helpers.test.ts | 86 +++++++++++++++++++
 .../model_add/credential_form_helpers.ts      | 33 +++++++
 4 files changed, 123 insertions(+), 4 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/model_add/credential_form_helpers.test.ts
 create mode 100644 ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts

diff --git a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx
index 694a98201c6..2324cbdd0e9 100644
--- a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx
+++ b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx
@@ -4,6 +4,7 @@ import type { UploadProps } from "antd/es/upload";
 import React, { useState } from "react";
 import ProviderSpecificFields from "../add_model/provider_specific_fields";
 import { Providers, providerLogoMap } from "../provider_info_helpers";
+import { resetCredentialFormOnProviderChange } from "./credential_form_helpers";
 const { Link } = Typography;
 
 interface AddCredentialsModalProps {
@@ -59,8 +60,7 @@ const AddCredentialsModal: React.FC<AddCredentialsModalProps> = ({ open, onCance
           <AntdSelect
             showSearch
             onChange={(value) => {
-              setSelectedProvider(value as Providers);
-              form.setFieldValue("custom_llm_provider", value);
+              resetCredentialFormOnProviderChange(form, value as Providers, setSelectedProvider);
             }}
           >
             {Object.entries(Providers).map(([providerEnum, providerDisplayName]) => (
diff --git a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx
index b206ed6c91d..f504ba7a78a 100644
--- a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx
+++ b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx
@@ -5,6 +5,7 @@ import { useEffect, useState } from "react";
 import ProviderSpecificFields from "../add_model/provider_specific_fields";
 import { CredentialItem } from "../networking";
 import { Providers, providerLogoMap } from "../provider_info_helpers";
+import { resetCredentialFormOnProviderChange } from "./credential_form_helpers";
 const { Link } = Typography;
 
 interface EditCredentialsModalProps {
@@ -92,8 +93,7 @@ export default function EditCredentialsModal({
           <AntdSelect
             showSearch
             onChange={(value) => {
-              setSelectedProvider(value as Providers);
-              form.setFieldValue("custom_llm_provider", value);
+              resetCredentialFormOnProviderChange(form, value as Providers, setSelectedProvider);
             }}
           >
             {Object.entries(Providers).map(([providerEnum, providerDisplayName]) => (
diff --git a/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.test.ts b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.test.ts
new file mode 100644
index 00000000000..9452cdb30d0
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.test.ts
@@ -0,0 +1,86 @@
+import type { FormInstance } from "antd";
+import { describe, expect, it, vi } from "vitest";
+import { Providers } from "../provider_info_helpers";
+import { resetCredentialFormOnProviderChange } from "./credential_form_helpers";
+
+/**
+ * Build a minimal FormInstance stub that records calls. We don't depend
+ * on the full Antd API surface — only the three methods the helper uses.
+ */
+function makeFormStub(initialFields: Record<string, unknown> = {}) {
+  const fields: Record<string, unknown> = { ...initialFields };
+  const stub = {
+    getFieldValue: vi.fn((key: string) => fields[key]),
+    setFieldValue: vi.fn((key: string, value: unknown) => {
+      fields[key] = value;
+    }),
+    resetFields: vi.fn(() => {
+      Object.keys(fields).forEach((k) => delete fields[k]);
+    }),
+  };
+  return { stub: stub as unknown as FormInstance, fields, calls: stub };
+}
+
+describe("resetCredentialFormOnProviderChange", () => {
+  it("clears all fields when switching providers", () => {
+    // Simulate the OpenAI->Google AI Studio leak: api_base picked up
+    // OpenAI's default value and the user typed a custom URL.
+    const { stub, fields, calls } = makeFormStub({
+      credential_name: "my-prod-key",
+      custom_llm_provider: "OpenAI",
+      api_base: "https://api.openai.com/v1",
+      api_key: "sk-stale-openai-key",
+      organization: "org-leak",
+    });
+    const setSelectedProvider = vi.fn();
+
+    resetCredentialFormOnProviderChange(stub, Providers.Google_AI_Studio, setSelectedProvider);
+
+    expect(calls.resetFields).toHaveBeenCalledTimes(1);
+    // Provider-specific fields must be gone so the next render starts
+    // from the new provider's default_value, not OpenAI's leftover.
+    expect(fields.api_base).toBeUndefined();
+    expect(fields.api_key).toBeUndefined();
+    expect(fields.organization).toBeUndefined();
+  });
+
+  it("preserves credential_name across the switch", () => {
+    // credential_name is user-supplied metadata, not provider-specific.
+    // The admin shouldn't have to retype it just because they re-picked
+    // the provider.
+    const { stub, fields } = makeFormStub({
+      credential_name: "my-prod-key",
+      custom_llm_provider: "OpenAI",
+      api_base: "https://api.openai.com/v1",
+    });
+
+    resetCredentialFormOnProviderChange(stub, Providers.Google_AI_Studio, vi.fn());
+
+    expect(fields.credential_name).toBe("my-prod-key");
+  });
+
+  it("updates custom_llm_provider and selectedProvider state to the new value", () => {
+    const { stub, fields } = makeFormStub({ credential_name: "x" });
+    const setSelectedProvider = vi.fn();
+
+    resetCredentialFormOnProviderChange(stub, Providers.Google_AI_Studio, setSelectedProvider);
+
+    expect(fields.custom_llm_provider).toBe(Providers.Google_AI_Studio);
+    expect(setSelectedProvider).toHaveBeenCalledExactlyOnceWith(Providers.Google_AI_Studio);
+  });
+
+  it("does not call setFieldValue('credential_name', undefined) when the name was unset", () => {
+    // Edge case: brand-new modal with no name typed yet. We shouldn't
+    // explicitly write `undefined` back into the form (Antd treats that
+    // as a touched empty field, triggering the "required" validation
+    // prematurely).
+    const { stub, calls } = makeFormStub({});
+
+    resetCredentialFormOnProviderChange(stub, Providers.Anthropic, vi.fn());
+
+    const credentialNameCalls = calls.setFieldValue.mock.calls.filter(
+      ([key]) => key === "credential_name",
+    );
+    expect(credentialNameCalls).toHaveLength(0);
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts
new file mode 100644
index 00000000000..5fb06e8e921
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts
@@ -0,0 +1,33 @@
+import type { FormInstance } from "antd";
+import { Providers } from "../provider_info_helpers";
+
+/**
+ * Reset the credential form when the user switches providers.
+ *
+ * Why: provider-specific fields (api_base, api_key, organization, ...)
+ * share a single Antd Form state across providers. Without this reset,
+ * the previous provider's values stick around — most visibly, OpenAI's
+ * default `api_base` (https://api.openai.com/v1) carries over when the
+ * user switches to Google AI Studio, overriding that provider's own
+ * default_value.
+ *
+ * Strategy: blow away the whole form, then restore the provider-agnostic
+ * fields (credential name + the new provider id) so the newly rendered
+ * `ProviderSpecificFields` can apply its own defaults from a clean slate.
+ *
+ * The credential name is preserved because it's a user-supplied label
+ * that shouldn't reset just because the admin re-selected a provider.
+ */
+export function resetCredentialFormOnProviderChange(
+  form: FormInstance,
+  newProvider: Providers,
+  setSelectedProvider: (p: Providers) => void,
+): void {
+  const preservedName = form.getFieldValue("credential_name");
+  form.resetFields();
+  if (preservedName !== undefined) {
+    form.setFieldValue("credential_name", preservedName);
+  }
+  setSelectedProvider(newProvider);
+  form.setFieldValue("custom_llm_provider", newProvider);
+}

From 895b8fd983262aa0e2517f3df76438c67955a159 Mon Sep 17 00:00:00 2001
From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com>
Date: Tue, 2 Jun 2026 12:28:00 +0800
Subject: [PATCH 5/6] fix(ui): logging callbacks table reads backend `type` for
 Mode badge (#35)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 .../LoggingCallbacksTable.test.tsx            | 35 +++++++++++++++++++
 .../LoggingCallbacksTable.tsx                 | 11 ++++--
 .../LoggingCallbacks/types.ts                 |  7 ++++
 3 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx
index a65a22edc85..0533b98b762 100644
--- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx
+++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.test.tsx
@@ -55,4 +55,39 @@ describe("LoggingCallbacksTable", () => {
     );
     expect(getByText("custom_callback_x")).toBeInTheDocument();
   });
+
+  // Regression: `/get_callbacks` returns the same `name` twice when a
+  // callback is registered for both success and failure (e.g. `generic_api`
+  // → POST to spend-log on both 200 and 4xx/5xx). The UI used to ignore
+  // the `type` field and render every row as "Success", masking the
+  // failure registration. Reading `record.type` fixes the badge AND
+  // composing the rowKey with type avoids React's duplicate-key warning.
+  it("renders distinct Success and Failure badges for same-name dual registration", () => {
+    const baseVars = {
+      SLACK_WEBHOOK_URL: null,
+      LANGFUSE_PUBLIC_KEY: null,
+      LANGFUSE_SECRET_KEY: null,
+      LANGFUSE_HOST: null,
+      OPENMETER_API_KEY: null,
+    };
+    const { getAllByText, getByText } = render(
+      <LoggingCallbacksTable
+        callbacks={[
+          { name: "generic_api", type: "success", variables: baseVars },
+          { name: "generic_api", type: "failure", variables: baseVars },
+        ]}
+        availableCallbacks={{
+          generic_api: {
+            litellm_callback_name: "generic_api",
+            litellm_callback_params: [],
+            ui_callback_name: "Custom Callback API",
+          },
+        }}
+      />,
+    );
+    // Both rows show the same display name, but distinct mode badges.
+    expect(getAllByText("Custom Callback API")).toHaveLength(2);
+    expect(getByText("Success")).toBeInTheDocument();
+    expect(getByText("Failure")).toBeInTheDocument();
+  });
 });
diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.tsx b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.tsx
index 8f332d0317a..70ec6599ca2 100644
--- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.tsx
+++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/LoggingCallbacksTable.tsx
@@ -48,7 +48,6 @@ export const LoggingCallbacksTable: React.FC<LoggingCallbacksProps> = ({
       key: "name",
       render: (_: string, record: CallbackRow) => {
         const id = record.name;
-        console.log("availableCallbacks", availableCallbacks);
         const displayName = availableCallbacks[id]?.ui_callback_name || id;
         return <div className="font-medium text-gray-800">{displayName}</div>;
       },
@@ -57,7 +56,10 @@ export const LoggingCallbacksTable: React.FC<LoggingCallbacksProps> = ({
       title: <span className="font-medium text-gray-700">Mode</span>,
       key: "mode",
       render: (_: unknown, record: CallbackRow) => {
-        const mode = record.mode || "success";
+        // Backend sends `type` (success | failure); legacy in-memory rows
+        // from add-callback flow set `mode`. Read both so newly-added rows
+        // and server-fetched rows both render correctly.
+        const mode = record.type || record.mode || "success";
         const label = CALLBACK_MODES.find((m) => m.value === mode)?.label || mode;
         const badgeClass =
           mode === "success"
@@ -109,7 +111,10 @@ export const LoggingCallbacksTable: React.FC<LoggingCallbacksProps> = ({
             <Table
               columns={columns}
               dataSource={callbacks as CallbackRow[]}
-              rowKey={(record) => record.name}
+              // `generic_api` can appear as both a success and a failure
+              // callback simultaneously — keying by `name` alone produced
+              // duplicate React keys. Compose with type to keep keys unique.
+              rowKey={(record) => `${record.name}-${record.type || record.mode || "success"}`}
               pagination={false}
               rowClassName={() => "hover:bg-gray-50"}
             />
diff --git a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/types.ts b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/types.ts
index 2fc180e49f3..5d265f95484 100644
--- a/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/types.ts
+++ b/ui/litellm-dashboard/src/components/Settings/LoggingAndAlerts/LoggingCallbacks/types.ts
@@ -1,5 +1,12 @@
 export interface AlertingObject {
   name: string;
+  // Backend distinguishes success vs failure callback registrations
+  // (`/get_callbacks` returns `type: "success" | "failure"`). Same callback
+  // (e.g. `generic_api`) can appear twice — once per event class — and
+  // those entries fire on disjoint events, not double-fire on one event.
+  // UI must read this to render the correct badge; missing it caused
+  // every row to render as "Success".
+  type?: "success" | "failure" | "success_and_failure";
   variables: AlertingVariables;
 }
 

From ede8b23eddfffbe68ad2c8d3f606388a70df8eb8 Mon Sep 17 00:00:00 2001
From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com>
Date: Thu, 28 May 2026 10:08:18 +0000
Subject: [PATCH 6/6] =?UTF-8?q?test(e2e):=20add=20case=2022=20=E2=80=94=20?=
 =?UTF-8?q?Gemini=20provider=20with=20custom=20api=5Fbase?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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).
---
 .../22_gemini_credential_custom_api_base.md   | 100 ++++++++++++++++++
 .../22_gemini_credential_custom_api_base.sh   |  95 +++++++++++++++++
 e2e/tools/proxy                               |  26 ++++-
 3 files changed, 220 insertions(+), 1 deletion(-)
 create mode 100644 e2e/cases/22_gemini_credential_custom_api_base.md
 create mode 100755 e2e/cases/data/22_gemini_credential_custom_api_base.sh

diff --git a/e2e/cases/22_gemini_credential_custom_api_base.md b/e2e/cases/22_gemini_credential_custom_api_base.md
new file mode 100644
index 00000000000..9c7560b43f0
--- /dev/null
+++ b/e2e/cases/22_gemini_credential_custom_api_base.md
@@ -0,0 +1,100 @@
+# Case 22 — Gemini provider with custom `api_base`
+
+## Goal
+
+End-to-end validation that the UI change in PR #24 (exposing `api_base`
+on the Google AI Studio credential form) reaches the gemini provider
+runtime correctly. Specifically:
+
+1. A deployment with `litellm_params.model = gemini/...` plus a custom
+   `api_base` pointing at a Gemini-compatible gateway resolves to the
+   gemini provider code path (no ADC, no Vertex auth).
+2. The outbound request hits `{api_base}/models/{model}:generateContent`
+   with an `x-goog-api-key` header — matching what a self-hosted gateway
+   like Raven Router or anispark's ai-router exposes.
+3. The response surface (status, `x-litellm-model-api-base`,
+   `x-litellm-model-group`) reflects the custom api_base, not Google's
+   default `generativelanguage.googleapis.com`.
+
+This case guards the runtime half of PR #24. The UI half is covered by
+`tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py::test_google_ai_studio_provider_fields_expose_api_base`.
+
+## Background
+
+Before this fix, the only LiteLLM UI options for "URL + API key Gemini"
+were:
+
+- **OpenAI credential** mislabeled (admin confusion, wrong logo/group in
+  the dashboard).
+- Raw admin API PATCH of `litellm_params.api_base` (no UI path).
+
+The runtime gemini provider already supported custom `api_base` via
+`litellm/llms/vertex_ai/vertex_llm_base.py:415` —
+`_check_custom_proxy` builds `{api_base}/models/{model}:{endpoint}` and
+attaches `x-goog-api-key: {gemini_api_key}` (line 484). The
+`_ensure_access_token_async` ADC path is skipped because
+`custom_llm_provider == "gemini"` (line 714). The UI just needed to
+expose the `api_base` field — see `provider_create_fields.json` →
+`Google_AI_Studio`.
+
+## Preconditions
+
+- `e2e/.env` has `GEMINI_API_KEY` and `GEMINI_API_BASE` set, pointing at
+  a Gemini-compatible gateway. Examples:
+  - `https://generativelanguage.googleapis.com/v1beta` (canonical Google
+    AI Studio; works with an `AIza...` key)
+  - `https://ai-router-hk.anispark.ai/v1beta` (anispark / Raven Router
+    with their own sk-style key)
+  - any self-hosted proxy that serves
+    `POST /v1beta/models/{model}:generateContent` with
+    `x-goog-api-key` auth.
+- Optionally `MODEL_GEMINI` to override the upstream model id (default
+  `gemini/gemini-3.1-pro-preview`).
+- `e2e/tools/proxy status` reports `ready` against an image built AFTER
+  this PR's branch (`e2e/tools/proxy build` if unsure — the rendered
+  config gains a `gemini-custom-base` deployment when `GEMINI_API_KEY`
+  is set).
+- Without `GEMINI_API_KEY` the case exits 77 (SKIP) — the proxy doesn't
+  render the deployment, so there's nothing to test.
+
+## Steps
+
+```bash
+bash e2e/cases/data/22_gemini_credential_custom_api_base.sh
+echo "exit=$?"
+```
+
+The fixture issues a single chat completion call against the
+`gemini-custom-base` deployment and inspects the response headers:
+
+- **C1** `POST /v1/chat/completions` model=`gemini-custom-base` →
+  expects HTTP 200, body has `choices`, header
+  `x-litellm-model-api-base` equals the value of `GEMINI_API_BASE` from
+  `e2e/.env`, header `x-litellm-model-group == gemini-custom-base`.
+
+## Expected — GREEN
+
+```
+PASS C1: HTTP 200, x-litellm-model-api-base=https://<gateway>/v1beta
+PASS C1: x-litellm-model-group=gemini-custom-base
+PASS C1: body has 'choices' array
+```
+
+## Failure modes
+
+| Symptom | Likely cause |
+|---|---|
+| `SKIP: GEMINI_API_KEY not set` | Add `GEMINI_API_KEY` and `GEMINI_API_BASE` to `e2e/.env`, then `e2e/tools/proxy restart`. |
+| `FAIL C1: HTTP 500 ... DefaultCredentialsError` | The deployment routed to `vertex_ai`/`vertex_ai_beta` instead of `gemini`. Verify `MODEL_GEMINI` does not have `vertex_ai/` prefix and the rendered config shows `model: gemini/...` for `gemini-custom-base`. |
+| `FAIL C1: HTTP 200 but body is HTML (<!doctype html>)` | `api_base` is the gateway's UI root, not the API path. Add `/v1beta` (or the equivalent for your gateway) to `GEMINI_API_BASE`. |
+| `FAIL C1: x-litellm-model-api-base != GEMINI_API_BASE` | LiteLLM is ignoring the custom api_base. Confirm the deployment block in the rendered config has `api_base: os.environ/GEMINI_API_BASE` and not the default Google endpoint. |
+| `FAIL C1: 401 / 403` from upstream | `GEMINI_API_KEY` is wrong for this gateway, or the gateway expects a different auth header (some non-Google gateways want `Authorization: Bearer` instead of `x-goog-api-key`). Confirm with a direct curl to `{api_base}/models/{model}:generateContent`. |
+
+## Cross-reference
+
+- PR #24 (`fix(ui): expose api_base on Google AI Studio credential form`)
+- `litellm/llms/vertex_ai/vertex_llm_base.py:415` —
+  `_check_custom_proxy` URL construction for gemini provider
+- `litellm/llms/vertex_ai/vertex_llm_base.py:714` — gemini provider
+  skipping ADC
+- `tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py::test_google_ai_studio_provider_fields_expose_api_base` — UI field metadata assertion
diff --git a/e2e/cases/data/22_gemini_credential_custom_api_base.sh b/e2e/cases/data/22_gemini_credential_custom_api_base.sh
new file mode 100755
index 00000000000..795121a2b0a
--- /dev/null
+++ b/e2e/cases/data/22_gemini_credential_custom_api_base.sh
@@ -0,0 +1,95 @@
+#!/usr/bin/env bash
+# Case 22 fixture — see e2e/cases/22_gemini_credential_custom_api_base.md
+#
+# Asserts that a deployment with `litellm_params.model = gemini/...`
+# plus a custom `api_base` routes through the gemini provider and the
+# x-litellm-model-api-base header surfaces the custom URL (proving the
+# UI api_base field added to the Google_AI_Studio credential form
+# actually reaches the runtime).
+#
+# Exits 0 on PASS, 77 on SKIP, anything else on FAIL.
+
+set -u
+
+PROXY="${PROXY_URL:-http://localhost:4011}"
+KEY="${MASTER_KEY:-sk-e2e-test}"
+MODEL="gemini-custom-base"
+FAILED=0
+
+pass() { echo "PASS: $*"; }
+fail() { echo "FAIL: $*"; FAILED=1; }
+skip() { echo "SKIP: $*"; exit 77; }
+
+# Pull GEMINI_API_BASE from .env so we can assert the response header
+# matches. We don't need the key here — the proxy holds it.
+ENV_FILE="$(dirname "$0")/../../.env"
+if [ -f "$ENV_FILE" ]; then
+    # Source the value without exporting other vars unintentionally.
+    EXPECTED_API_BASE=$(
+        grep -E '^GEMINI_API_BASE=' "$ENV_FILE" | tail -1 | sed -E 's/^GEMINI_API_BASE=//; s/^"//; s/"$//; s/^'"'"'//; s/'"'"'$//'
+    )
+else
+    EXPECTED_API_BASE=""
+fi
+
+if ! curl -sSf -o /dev/null -m 3 "$PROXY/health/readiness"; then
+    skip "proxy not ready at $PROXY"
+fi
+
+# If the proxy didn't render the deployment, GEMINI_API_KEY was unset.
+if ! curl -sSf -m 5 -H "Authorization: Bearer $KEY" "$PROXY/v1/models" \
+        2>/dev/null | grep -q '"id"[[:space:]]*:[[:space:]]*"'"$MODEL"'"'; then
+    skip "deployment '$MODEL' missing from /v1/models — set GEMINI_API_KEY (and optionally GEMINI_API_BASE / MODEL_GEMINI) in e2e/.env, then re-run e2e/tools/proxy restart"
+fi
+
+if [ -z "$EXPECTED_API_BASE" ]; then
+    skip "GEMINI_API_BASE not set in e2e/.env — case 22 needs a custom api_base to assert against"
+fi
+
+# ---- C1: chat/completions with the custom-api_base gemini deployment ----
+read -r -d '' BODY <<JSON || true
+{"model":"$MODEL","max_tokens":8,"messages":[{"role":"user","content":"hi"}]}
+JSON
+
+C1_BODY=$(mktemp); C1_HDRS=$(mktemp)
+C1_STATUS=$(curl -sS -D "$C1_HDRS" -o "$C1_BODY" -w "%{http_code}" \
+    -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
+    -d "$BODY" "$PROXY/v1/chat/completions")
+
+if [ "$C1_STATUS" != "200" ]; then
+    fail "HTTP $C1_STATUS (expected 200)"
+    echo "--- body (truncated) ---"
+    head -c 600 "$C1_BODY"; echo
+    rm -f "$C1_BODY" "$C1_HDRS"
+    exit $FAILED
+fi
+
+# Extract headers we care about. tr -d '\r' because curl preserves CRLF.
+ACTUAL_API_BASE=$(grep -i '^x-litellm-model-api-base:' "$C1_HDRS" \
+    | tail -1 | tr -d '\r' | sed -E 's/^[^:]+:[[:space:]]*//')
+ACTUAL_GROUP=$(grep -i '^x-litellm-model-group:' "$C1_HDRS" \
+    | tail -1 | tr -d '\r' | sed -E 's/^[^:]+:[[:space:]]*//')
+
+if [ "$ACTUAL_API_BASE" = "$EXPECTED_API_BASE" ]; then
+    pass "HTTP 200, x-litellm-model-api-base=$ACTUAL_API_BASE"
+else
+    fail "x-litellm-model-api-base mismatch: expected='$EXPECTED_API_BASE', got='$ACTUAL_API_BASE'"
+fi
+
+if [ "$ACTUAL_GROUP" = "$MODEL" ]; then
+    pass "x-litellm-model-group=$ACTUAL_GROUP"
+else
+    fail "x-litellm-model-group mismatch: expected='$MODEL', got='$ACTUAL_GROUP'"
+fi
+
+# Body shape sanity — gemini provider returns OpenAI-compatible JSON via LiteLLM.
+if python3 -c "import json,sys; d=json.load(open('$C1_BODY')); assert isinstance(d.get('choices'), list) and len(d['choices'])>=1" 2>/dev/null; then
+    pass "body has 'choices' array"
+else
+    fail "body missing or malformed 'choices' array"
+    echo "--- body (truncated) ---"
+    head -c 600 "$C1_BODY"; echo
+fi
+
+rm -f "$C1_BODY" "$C1_HDRS"
+exit $FAILED
diff --git a/e2e/tools/proxy b/e2e/tools/proxy
index 6f0cf64c296..78691f08c41 100755
--- a/e2e/tools/proxy
+++ b/e2e/tools/proxy
@@ -109,6 +109,12 @@ def render_config() -> Path:
     )
     anthropic_base = os.environ.get("ANTHROPIC_API_BASE", "").strip()
     openai_base = os.environ.get("OPENAI_API_BASE", "").strip()
+    gemini_base = os.environ.get("GEMINI_API_BASE", "").strip()
+    gemini_key_set = bool(os.environ.get("GEMINI_API_KEY", "").strip())
+    gemini_model = _ensure_provider_prefix(
+        _env_or("MODEL_GEMINI", "gemini/gemini-3.1-pro-preview"),
+        implied_provider="gemini",
+    )
 
     def _emit_model(
         model_name: str, model: str, key_env: str, base_env: str, has_base: bool
@@ -139,6 +145,19 @@ def render_config() -> Path:
             "OPENAI_API_KEY", "OPENAI_API_BASE", bool(openai_base),
         ),
     ]
+    # Dedicated deployment for case 22: gemini/ provider with a custom
+    # api_base, validating that the UI api_base field added to the
+    # Google_AI_Studio credential form actually flows through to a
+    # Gemini-compatible gateway (e.g. /v1beta path).
+    # Only emitted when GEMINI_API_KEY is configured so cases that don't
+    # need it (most of the suite) don't fail with "missing key".
+    if gemini_key_set:
+        blocks.append(
+            _emit_model(
+                "gemini-custom-base", gemini_model,
+                "GEMINI_API_KEY", "GEMINI_API_BASE", bool(gemini_base),
+            )
+        )
     body = (
         "# AUTOGENERATED by e2e/tools/proxy from e2e/.env at proxy-start time.\n"
         "# Do not edit by hand — edit e2e/.env and re-run `e2e/tools/proxy start`.\n"
@@ -147,7 +166,12 @@ def render_config() -> Path:
         f"#   claude-sonnet-cache -> {sonnet_model}\n"
         f"#   claude-haiku-cache  -> {haiku_model}\n"
         f"#   gpt-4o-mini-cache   -> {openai_model}\n"
-        "model_list:\n"
+        + (
+            f"#   gemini-custom-base  -> {gemini_model} (case 22; GEMINI_API_BASE={gemini_base or '<unset>'})\n"
+            if gemini_key_set
+            else ""
+        )
+        + "model_list:\n"
         + "\n\n".join(blocks)
         + "\n\n"
         "litellm_settings:\n"