feat(proxy): add PublicReqMiddleware for X-Public-Req gated safeguards

docker/prod_entrypoint.sh

@@ -6,9 +6,16 @@ if [ "$SEPARATE_HEALTH_APP" = "1" ]; then
     exec supervisord -c /etc/supervisord.conf
 fi
 
+# 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

e2e/_config/docker-compose.yml

@@ -70,6 +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
+    # 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

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

e2e/cases/README.md

@@ -27,6 +27,7 @@ Humans can execute them too — every step is a concrete shell command.
 | 15 | `15_v1_models_user_filter.md` | (none — proxy only) | `GET /v1/models` honors `LiteLLM_UserTable.models` (Personal Models). Regression for BerriAI/litellm#26420 | ✓ |
 | 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 | ✓ |
 | 17 | `17_model_info_user_filter.md` | (none — proxy only) | `GET /v1/model/info` (Path B) and `GET /v2/model/info` (every flag combo) honor `LiteLLM_UserTable.models`. Extends PR #10 fix from `/v1/models` to the two info endpoints | ✓ |
+| 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
 

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