diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index d766897ab3c..f815601a29c 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -93,6 +93,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac Otherwise proof must be fully e2e, no mocks, with real provider calls where applicable. Pytest alone is insufficient. The following format applies to live proof Show only the latest proof. Before uses the merge base and After the tested tip. Rerun affected QA after behavior changes. Verified neutral diffs may carry proof with the tested SHA, current tip, equivalence reason and focused checks, following AGENTS.md and /qa Use matching Before/After case headings with numbered commands and observed output. Put shared setup above them. For one case, omit case headings + A reviewer must be able to rerun every case from the body alone, so the shared setup gives the proxy start command, the config.yaml, env var names (never values) and the client code or curl used, all exactly as run. Wrap any block longer than about 20 lines in a collapsed
with a one-line summary ### Before () diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index 3e2bca8779d..d769b6ab8af 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -53,7 +53,7 @@ jobs: if: >- github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository - timeout-minutes: 30 + timeout-minutes: 60 permissions: contents: read steps: diff --git a/.github/workflows/interactions-spec-canary.yml b/.github/workflows/interactions-spec-canary.yml index a2e4646dce5..89ebdd87a7c 100644 --- a/.github/workflows/interactions-spec-canary.yml +++ b/.github/workflows/interactions-spec-canary.yml @@ -49,9 +49,9 @@ jobs: id: compliance run: | uv run --no-sync pytest tests/unit/interactions/test_openapi_compliance.py \ - --interactions-spec-url=https://ai.google.dev/static/api/interactions.openapi.json \ + --interactions-spec-url=https://aistudio.google.com/api-reference/interactions.openapi.json \ --tb=short -vv -p no:rerunfailures -p no:pytest-retry \ - --allow-hosts=127.0.0.1,::1,localhost,ai.google.dev + --allow-hosts=127.0.0.1,::1,localhost,ai.google.dev,aistudio.google.com tracking-issue: name: tracking-issue diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index cbeea71ad0f..fd1d2973951 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -21,7 +21,7 @@ jobs: pull-requests: read actions: read runs-on: ubuntu-latest - timeout-minutes: 20 + timeout-minutes: 30 services: redis: image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index f1456b1fde8..8ac157c94f5 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -45,7 +45,7 @@ jobs: id: rust-bridge-key if: steps.changes.outputs.decision != 'skip' env: - RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml') }} + RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml', 'packaging/litellm-proxy/pyproject.toml') }} run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT" - name: Set up Python @@ -98,6 +98,18 @@ jobs: mkdir -p rust-bridge-dist/litellm/rust_bridge cp litellm/rust_bridge/_native.abi3.so rust-bridge-dist/litellm/rust_bridge/_native.abi3.so + - name: Build the gateway wheel once for test shards + if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' + timeout-minutes: 15 + working-directory: packaging/litellm-proxy + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" + run: uv run --no-project --with maturin==1.15.0 maturin build --profile dev --strip --out ../../rust-bridge-dist + - name: Upload Rust bridge artifact if: steps.changes.outputs.decision != 'skip' timeout-minutes: 10 @@ -1024,7 +1036,7 @@ jobs: if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))" @@ -1486,7 +1498,7 @@ jobs: if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so fi diff --git a/AGENTS.md b/AGENTS.md index 9aa9ab32528..01127303da3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -45,7 +45,7 @@ Run QA proxies with `python litellm/proxy/proxy_cli.py --config litellm/proxy/de Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at `, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030) -For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce +For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. The body alone must let a reviewer rerun it, with the proxy command, config, env var names and client code exactly as run, and long blocks collapsed in
. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis diff --git a/Dockerfile b/Dockerfile index 4cef63c8d71..94557bcf68c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -76,6 +76,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ COPY pyproject.toml uv.lock ./ COPY enterprise/pyproject.toml enterprise/ COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/ +COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/ # Install third-party dependencies (cached unless pyproject.toml/uv.lock change) RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ @@ -101,7 +102,8 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh # Install project and workspace packages (fast - deps already cached) -RUN uv sync --frozen --no-default-groups --no-editable \ +RUN uv build --package litellm-proxy --wheel --out-dir /wheels +RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ --group admin-mcp \ @@ -111,6 +113,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra bedrock-realtime \ --python python3.13 +RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl + RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ prisma generate --schema=./schema.prisma diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 20f810686a3..2ebffc50b5e 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -74,6 +74,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ COPY pyproject.toml uv.lock ./ COPY enterprise/pyproject.toml enterprise/ COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/ +COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/ # Install third-party dependencies (cached unless pyproject.toml/uv.lock change) RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ @@ -99,7 +100,8 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh # Install project and workspace packages (fast - deps already cached) -RUN uv sync --frozen --no-default-groups --no-editable \ +RUN uv build --package litellm-proxy --wheel --out-dir /wheels +RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ --group admin-mcp \ @@ -109,6 +111,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra bedrock-realtime \ --python python3.13 +RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl + RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ prisma generate --schema=./schema.prisma diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 658bed831ba..1e1c7d73091 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -41,10 +41,14 @@ COPY pyproject.toml README.md rust-toolchain.toml ./ COPY .cargo ./.cargo COPY litellm-rust ./litellm-rust COPY model_prices_and_context_window.json ./ +COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/ RUN --mount=type=cache,target=/root/.cache/uv \ mkdir -p litellm/rust_bridge /wheels && \ touch litellm/__init__.py litellm/rust_bridge/__init__.py && \ + uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \ + python -c "import maturin; maturin.build_wheel('/wheels')" && \ + cd packaging/litellm-proxy && \ uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \ python -c "import maturin; maturin.build_wheel('/wheels')" @@ -85,6 +89,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,source=uv.lock,target=uv.lock \ --mount=type=bind,source=enterprise/pyproject.toml,target=enterprise/pyproject.toml \ --mount=type=bind,source=litellm-proxy-extras/pyproject.toml,target=litellm-proxy-extras/pyproject.toml \ + --mount=type=bind,source=packaging/litellm-proxy/pyproject.toml,target=packaging/litellm-proxy/pyproject.toml \ uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ @@ -98,7 +103,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ COPY . . RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --frozen --no-install-project --no-default-groups --no-editable \ + uv sync --frozen --no-install-project --no-install-package litellm-proxy --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ --group admin-mcp \ @@ -108,7 +113,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --python python3.13 RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \ - uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl + uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl /wheels/litellm_proxy-*.whl # PYTHONPATH=/app makes the source tree shadow the installed package, so the # compiled Rust extension must live next to the source or it is never imported. diff --git a/litellm-rust/.agents/skills/rust-gateway/SKILL.md b/litellm-rust/.agents/skills/rust-gateway/SKILL.md index 4fee0b54a68..923c9bf5ab6 100644 --- a/litellm-rust/.agents/skills/rust-gateway/SKILL.md +++ b/litellm-rust/.agents/skills/rust-gateway/SKILL.md @@ -7,7 +7,7 @@ metadata: # Rust gateway -`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching +`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. `gateway-forward` is the reverse proxy that relays every request to the Python proxy and has no routes of its own. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching - [surfaces](references/surfaces.md) - Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 23ebfc7d3cf..107b55e89cb 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -547,7 +547,7 @@ dependencies = [ "http 1.4.2", "http-body 0.4.6", "hyper 0.14.32", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.24.2", "hyper-rustls 0.27.9", "hyper-util", @@ -736,7 +736,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "itoa", "matchit", @@ -1000,7 +1000,7 @@ dependencies = [ "home", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-named-pipe", "hyper-rustls 0.27.9", "hyper-util", @@ -2255,7 +2255,7 @@ dependencies = [ "chrono", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.27.9", "hyper-util", "ring", @@ -2391,7 +2391,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "lazy_static", "opentelemetry", "opentelemetry-semantic-conventions", @@ -2833,9 +2833,9 @@ dependencies = [ [[package]] name = "hyper" -version = "1.10.1" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +checksum = "2c3e324da4c95177d6291d4c8730197c0d1822f8a9766814a4a44fa5ab797c9c" dependencies = [ "atomic-waker", "bytes", @@ -2860,7 +2860,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fab3637d6b04a8037af8a266fdf6cf92ea957e8c53981a2bf6136572531025bf" dependencies = [ "hex", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -2889,7 +2889,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ "http 1.4.2", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "rustls 0.23.42", "rustls-native-certs", @@ -2904,7 +2904,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" dependencies = [ - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -2923,7 +2923,7 @@ dependencies = [ "futures-util", "http 1.4.2", "http-body 1.1.0", - "hyper 1.10.1", + "hyper 1.12.0", "ipnet", "libc", "percent-encoding", @@ -2942,7 +2942,7 @@ checksum = "986c5ce3b994526b3cd75578e62554abd09f0899d6206de48b3e96ab34ccc8c7" dependencies = [ "hex", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -3769,6 +3769,7 @@ dependencies = [ "litellm-config", "litellm-error", "litellm-gateway-auth", + "litellm-gateway-forward", "litellm-gateway-inference", "litellm-gateway-mcp", "litellm-gateway-ui", @@ -3815,6 +3816,29 @@ dependencies = [ "veil", ] +[[package]] +name = "litellm-gateway-forward" +version = "0.1.0" +dependencies = [ + "axum", + "bytes", + "envy", + "futures-util", + "http-body-util", + "httparse", + "hyper 1.12.0", + "hyper-util", + "proptest", + "rstest", + "serde", + "thiserror 2.0.19", + "tokio", + "tokio-tungstenite", + "tokio-util", + "tower-service", + "tracing", +] + [[package]] name = "litellm-gateway-inference" version = "0.1.0" @@ -6031,7 +6055,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.27.9", "hyper-util", "js-sys", @@ -7512,6 +7536,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -7591,7 +7616,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-timeout", "hyper-util", "percent-encoding", @@ -8167,9 +8192,9 @@ dependencies = [ [[package]] name = "want" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +checksum = "ec4cdd0dd910afe868b7ef477227d8d538b46b3075031afee8a9f2acb0a2ed0b" dependencies = [ "try-lock", ] @@ -8551,7 +8576,7 @@ dependencies = [ "futures", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "log", "once_cell", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 71914a0bffe..7121df06bc0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -26,6 +26,7 @@ litellm-inference-ocr = { path = "crates/inference-ocr" } litellm-inference-testing = { path = "crates/inference-testing" } litellm-gateway-mcp = { path = "crates/gateway-mcp" } litellm-gateway = { path = "crates/gateway" } +litellm-gateway-forward = { path = "crates/gateway-forward" } litellm-gateway-inference = { path = "crates/gateway-inference" } litellm-gateway-auth = { path = "crates/gateway-auth" } litellm-gateway-management = { path = "crates/gateway-management" } diff --git a/litellm-rust/crates/gateway-forward/AGENTS.md b/litellm-rust/crates/gateway-forward/AGENTS.md new file mode 100644 index 00000000000..2135980266f --- /dev/null +++ b/litellm-rust/crates/gateway-forward/AGENTS.md @@ -0,0 +1,37 @@ +# Forwarding gateway + +`Forward` is a reverse proxy in front of the Python proxy on loopback. Python must see each request as if the client had connected to it directly, because every Python route (inference, MCP, pass-through, UI, WebSocket) is served behind it without Rust knowing which one it is + +- Serve HTTP/1.1 only, as uvicorn does. The supervisor never puts Rust in front of a Python server that speaks HTTP/2 +- Never branch on a route or path. A behavior one endpoint needs is a protocol behavior every endpoint gets +- Forward the method, request target, headers, and body byte for byte within the [URL contract](FORWARDING.md#url-fragments-and-literal-hashes). An unescaped `#` starts a fragment and is omitted with everything after it, while `%23` is preserved. Only drop hop-by-hop headers (RFC 9110 section 7.6.1) and set `x-litellm-rust-peer`, which replaces any client copy +- Stream both bodies. Never buffer, cap, or time out a body, since uploads and provider streams have no upper bound +- A client disconnect closes the upstream connection at every stage, because that is how Python learns to stop a provider stream +- An upstream that dies mid-response fails the client body. Never end it cleanly, or SSE clients treat a truncated stream as complete +- Never invent a Host. Python builds absolute URLs from it, so a missing one must never become the private upstream address +- A request Python could not have received unchanged (HTTP/1.1 without Host, a transfer coding other than chunked, a target without a path) gets the status uvicorn would have given, from Rust +- Pooled upstream connections are retired well before Python's keep-alive timeout, so a request never lands on a connection Python is closing +- Idle client connections close after Python's keep-alive timeout, as uvicorn closed them, so a client that vanishes cannot hold a socket forever. Only the wait for the next request is timed, never an exchange in progress +- WebSocket is the one protocol switch Python serves. Upgrades are relayed as raw bytes after the 101 and drained on shutdown + +## Layout + +Every forwarding decision has one home. `rewrite.rs` owns what Python and the client receive, as pure functions over requests and responses. `Rejection` in `error.rs` is the only place that picks a status for a request Rust answers itself. `settings.rs` parses the environment into typed values, so an `Upstream` is a loopback origin by construction. `server.rs` accepts client connections and owns their timeouts, since `axum::serve` exposes none, and `connections.rs` tracks every connection and relay for shutdown. `forward.rs` and `websocket.rs` only move requests, bodies and upgraded connections and never touch a header themselves. `connector.rs` keeps hyper reading Python's answer after Python stops reading the request + +## Environment + +`Settings` reads only these variables, and only when `LITELLM_FORWARD_UPSTREAM` is set. The Python supervisor in `litellm/proxy/rust_forward/` sets every one of them from the values Python actually runs with, so all of them are required and a missing or invalid one fails startup naming it. The crate never reads Python's own settings (`HOST`, `PORT`, `KEEPALIVE_TIMEOUT`, `GRACEFUL_SHUTDOWN_TIMEOUT`) and keeps no copy of Python's defaults, since Python can get those values from CLI flags the environment never shows + +| Variable | Meaning | +|---|---| +| `LITELLM_FORWARD_UPSTREAM` | Python's loopback origin, e.g. `http://127.0.0.1:53211` | +| `LITELLM_FORWARD_HOST` | public bind host | +| `LITELLM_FORWARD_PORT` | public bind port | +| `LITELLM_FORWARD_UPSTREAM_KEEPALIVE` | seconds uvicorn keeps an idle connection open. Pooled connections are retired at half of it, and idle client connections are closed after it | +| `LITELLM_FORWARD_SHUTDOWN_DRAIN` | seconds to drain on shutdown, already including Python's graceful timeout and teardown headroom | + +## Tests + +Test through real sockets with `tests/support`, never `oneshot`. `Gateway` runs the forwarder on a listener like production, `recording` is a well-behaved upstream that captures what Python would receive, and `raw_upstream` and `Raw` drive bytes by hand for framing, timing and disconnects a real server would not produce. Waits measured in minutes use paused tokio time instead of real sleeps + +The forwarder sees HTTP/1.1 messages, not routes, so the tests are split into one file per traffic shape uvicorn can emit: `exchange` (plain request and response), `framing` (requests whose framing a parser could read two ways), `slow_head` (minutes before response headers), `stream` (incremental response bodies), `idle_stream` (long-lived streams that go quiet), `upload`, `download`, `websocket`, `client_identity`, `lifecycle` (disconnects, pooling, idle timeouts, shutdown) and `roundtrip` (the forwarding invariant for generated messages). A new route never gets its own test. Add its behavior to the shape it already has, and when it needs a wire behavior none of the files covers, add a new shape file for it diff --git a/litellm-rust/crates/gateway-forward/Cargo.toml b/litellm-rust/crates/gateway-forward/Cargo.toml new file mode 100644 index 00000000000..2fef19dca34 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "litellm-gateway-forward" +version = "0.1.0" +edition.workspace = true +rust-version.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +axum.workspace = true +envy = "0.4.2" +hyper = { version = "1.12.0", features = ["http1", "server"] } +hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } +serde.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["io-util"] } +tokio-util = { version = "0.7", features = ["rt"] } +tower-service = "0.3.3" +tracing.workspace = true + +[dev-dependencies] +bytes.workspace = true +futures-util.workspace = true +http-body-util = "0.1" +httparse = "1" +hyper = { version = "1.12.0", features = ["client"] } +proptest.workspace = true +rstest.workspace = true +tokio = { workspace = true, features = ["io-util", "sync", "test-util"] } +tokio-tungstenite.workspace = true diff --git a/litellm-rust/crates/gateway-forward/FORWARDING.md b/litellm-rust/crates/gateway-forward/FORWARDING.md new file mode 100644 index 00000000000..532f341aba5 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/FORWARDING.md @@ -0,0 +1,31 @@ +# Python forwarding compatibility + +The opt-in forwarding gateway serves HTTP/1.1 and WebSockets in front of Python. Python still owns every route and policy. Its protocol contract and environment settings are in [AGENTS.md](AGENTS.md) + +## URL fragments and literal hashes + +The gateway follows Hyper's `http::Uri` parsing. If a raw client sends an unescaped `#` in the request target, only the path and query before the first `#` reach Python. For example, `/v1/files?purpose=batch#retry&limit=1` becomes `/v1/files?purpose=batch`, and `/gemini/v1beta/files/a#b?key=k` becomes `/gemini/v1beta/files/a` + +This is accepted behavior because `#` starts a URL fragment, which is separated before dereferencing ([RFC 3986 section 3.5](https://www.rfc-editor.org/rfc/rfc3986.html#section-3.5)). HTTP origin-form contains only the path and optional query, not a fragment ([RFC 9112 section 3.2.1](https://www.rfc-editor.org/rfc/rfc9112.html#section-3.2.1)). Browsers and URL-based clients omit the fragment before sending the request, so UI hash routing is unaffected + +A literal hash in path or query data must be encoded as `%23`, which the gateway forwards unchanged. Direct Python with h11 also accepts unescaped hashes as data, but that permissive raw-target behavior is outside this URL contract. Active socket tests cover both [fragment omission](tests/exchange.rs) and [encoded-hash preservation](tests/exchange.rs) + +## Known differences from uvicorn with h11 + +These cases remain executable as ignored socket tests. They are not passing parity checks + +| Input | Direct Python | Through Rust | Regression test | +|---|---|---|---| +| Header after a removed hop-by-hop field | Keeps the remaining header order | `HeaderMap::remove` can move a later header | [dropping_a_hop_by_hop_header_keeps_the_others_in_order](tests/exchange.rs) | +| Repeated header names interleaved with other names | Keeps the interleaving | Groups values by name, keeping their order within that name | [interleaved_repeated_headers_reach_python_interleaved](tests/exchange.rs) | +| Obsolete folded header, such as `x-a: b\r\n c` | Unfolds to `b c` | Rejects with 400 | [a_folded_header_reaches_python_unfolded](tests/exchange.rs) | +| Identical Content-Length list, such as `7, 7` | Accepts length 7 | Rejects with 400 | [a_content_length_list_of_one_value_reaches_python](tests/framing.rs) | +| HTTP/1.0 request with chunked transfer coding | Decodes the body | Rejects with 400 | [an_http10_chunked_upload_reaches_python](tests/framing.rs) | + +Clients can avoid the parser rejections by sending unfolded headers, one Content-Length value, and HTTP/1.1 for chunked uploads. Integrations must not depend on header ordering across different field names + +Run `cargo test -p litellm-gateway-forward --test exchange --test framing -- --ignored` from this directory to reproduce the differences + +## Connection header handling + +The gateway strips standard hop-by-hop fields and fields named by Connection, except routing, representation, credential, forwarding, and WebSocket handshake headers. A client cannot nominate those end-to-end fields for removal. The private `x-litellm-rust-peer` always replaces any client-supplied copy diff --git a/litellm-rust/crates/gateway-forward/src/connections.rs b/litellm-rust/crates/gateway-forward/src/connections.rs new file mode 100644 index 00000000000..10a4867fa51 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/connections.rs @@ -0,0 +1,58 @@ +use tokio_util::{ + sync::{CancellationToken, DropGuard}, + task::TaskTracker, +}; + +/// Every client connection and WebSocket relay the gateway runs, so shutdown can wait for them +/// and then close whatever outlives the drain. +#[derive(Clone, Default)] +pub(crate) struct Connections { + tasks: TaskTracker, + draining: CancellationToken, + closed: CancellationToken, +} + +impl Connections { + /// Runs `work` until it finishes or the gateway closes every connection. + pub fn spawn(&self, work: impl Future + Send + 'static) { + let closed = self.closed.clone(); + self.tasks.spawn(async move { + tokio::select! { + () = closed.cancelled() => {} + () = work => {} + } + }); + } + + pub fn begin_shutdown(&self) { + self.draining.cancel(); + } + + pub async fn draining(&self) { + self.draining.cancelled().await; + } + + pub fn is_draining(&self) -> bool { + self.draining.is_cancelled() + } + + pub async fn until_closed(&self, future: impl Future) -> Option { + tokio::select! { + () = self.closed.cancelled() => None, + output = future => Some(output), + } + } + + pub fn close(&self) { + self.closed.cancel(); + } + + pub fn close_on_drop(&self) -> DropGuard { + self.closed.clone().drop_guard() + } + + pub async fn wait(&self) { + self.tasks.close(); + self.tasks.wait().await; + } +} diff --git a/litellm-rust/crates/gateway-forward/src/connector.rs b/litellm-rust/crates/gateway-forward/src/connector.rs new file mode 100644 index 00000000000..5eb998db0c0 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/connector.rs @@ -0,0 +1,122 @@ +use std::{ + io::{self, IoSlice}, + pin::Pin, + task::{Context, Poll, ready}, +}; + +use axum::http::Uri; +use hyper::rt::{Read, ReadBufCursor, Write}; +use hyper_util::{ + client::legacy::connect::{Connected, Connection, HttpConnector}, + rt::TokioIo, +}; +use tokio::net::TcpStream; + +/// Connects to Python like `HttpConnector`, but keeps hyper reading Python's answer after Python +/// stops reading the request. uvicorn rejects an upload (413, 401) as soon as it reads the head +/// and closes with the body unread, so writing the rest fails, and hyper would fail the request +/// and drop the answer it already received. Like curl and httpcore, the forwarder discards the +/// rest of the request instead and reads on. A connection that never answered still fails, on the +/// read, and no request is ever sent twice. +#[derive(Clone)] +pub(crate) struct Connector(pub HttpConnector); + +impl tower_service::Service for Connector { + type Response = AnswerAfterRefusal; + type Error = >::Error; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll> { + self.0.poll_ready(context) + } + + fn call(&mut self, uri: Uri) -> Self::Future { + let connecting = self.0.call(uri); + Box::pin(async move { + Ok(AnswerAfterRefusal { + stream: connecting.await?, + refused: false, + }) + }) + } +} + +pub(crate) struct AnswerAfterRefusal { + stream: TokioIo, + refused: bool, +} + +impl AnswerAfterRefusal { + fn written(&mut self, written: io::Result, length: usize) -> io::Result { + match written { + Err(error) + if matches!( + error.kind(), + io::ErrorKind::BrokenPipe | io::ErrorKind::ConnectionReset + ) => + { + self.refused = true; + Ok(length) + } + written => written, + } + } +} + +impl Read for AnswerAfterRefusal { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: ReadBufCursor<'_>, + ) -> Poll> { + Pin::new(&mut self.stream).poll_read(context, buffer) + } +} + +impl Write for AnswerAfterRefusal { + fn poll_write( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + if self.refused { + return Poll::Ready(Ok(buffer.len())); + } + let written = ready!(Pin::new(&mut self.stream).poll_write(context, buffer)); + Poll::Ready(self.written(written, buffer.len())) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffers: &[IoSlice<'_>], + ) -> Poll> { + let length = buffers.iter().map(|buffer| buffer.len()).sum(); + if self.refused { + return Poll::Ready(Ok(length)); + } + let written = ready!(Pin::new(&mut self.stream).poll_write_vectored(context, buffers)); + Poll::Ready(self.written(written, length)) + } + + fn is_write_vectored(&self) -> bool { + self.stream.is_write_vectored() + } + + fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if self.refused { + return Poll::Ready(Ok(())); + } + Pin::new(&mut self.stream).poll_flush(context) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_shutdown(context) + } +} + +impl Connection for AnswerAfterRefusal { + fn connected(&self) -> Connected { + self.stream.connected() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/error.rs b/litellm-rust/crates/gateway-forward/src/error.rs new file mode 100644 index 00000000000..7dbdde68a85 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/error.rs @@ -0,0 +1,91 @@ +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, +}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid forwarding upstream URI")] + UpstreamUri(#[source] axum::http::uri::InvalidUri), + #[error( + "LITELLM_FORWARD_UPSTREAM must be a loopback IPv4 http origin without a path, query or credentials" + )] + Upstream, + #[error("{0} must be set in forward mode")] + MissingSetting(String), + #[error("invalid forwarding setting: {0}")] + InvalidSetting(String), + #[error("forwarding gateway I/O failed: {0}")] + Io(#[from] std::io::Error), +} + +/// Why a request got an answer from Rust instead of from Python. +#[derive(Debug, thiserror::Error)] +pub(crate) enum Rejection { + #[error("HTTP/1.1 request without a Host header")] + MissingHost, + #[error("request target has no path to forward")] + NoPath, + #[error("request uses a transfer coding other than chunked")] + UnsupportedTransferCoding, + #[error("client request body failed: {0}")] + ClientBody(#[source] hyper_util::client::legacy::Error), + #[error("python proxy unreachable: {0}")] + Unreachable(#[source] hyper_util::client::legacy::Error), + #[error("python proxy switched protocols without being asked to")] + UnrequestedUpgrade, + #[error("python proxy websocket upgrade failed: {0}")] + UpgradeFailed(#[source] hyper::Error), + #[error("forwarding gateway is shutting down")] + ShuttingDown, + #[error("could not build the request for the python proxy")] + Internal, +} + +impl Rejection { + /// A client body that ends early or breaks its framing fails the upstream request too, but + /// the fault is the client's, not Python's. + pub fn from_upstream(error: hyper_util::client::legacy::Error) -> Self { + let client_body = std::error::Error::source(&error) + .and_then(|source| source.downcast_ref::()) + .is_some_and(hyper::Error::is_user); + match client_body { + true => Self::ClientBody(error), + false => Self::Unreachable(error), + } + } + + fn status(&self) -> StatusCode { + match self { + Self::MissingHost + | Self::NoPath + | Self::UnsupportedTransferCoding + | Self::ClientBody(_) => StatusCode::BAD_REQUEST, + Self::Unreachable(_) | Self::UnrequestedUpgrade | Self::UpgradeFailed(_) => { + StatusCode::BAD_GATEWAY + } + Self::ShuttingDown => StatusCode::SERVICE_UNAVAILABLE, + Self::Internal => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn is_gateway_fault(&self) -> bool { + matches!( + self, + Self::Unreachable(_) + | Self::UnrequestedUpgrade + | Self::UpgradeFailed(_) + | Self::Internal + ) + } +} + +impl IntoResponse for Rejection { + fn into_response(self) -> Response { + match self.is_gateway_fault() { + true => tracing::warn!(error = %self, "request not forwarded"), + false => tracing::debug!(error = %self, "request not forwarded"), + } + self.status().into_response() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/forward.rs b/litellm-rust/crates/gateway-forward/src/forward.rs new file mode 100644 index 00000000000..69fb63ada7d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/forward.rs @@ -0,0 +1,169 @@ +use std::{ + pin::Pin, + task::{Context, Poll}, + time::Duration, +}; + +use axum::{ + body::Body, + extract::Request, + http::{StatusCode, Uri}, + response::{IntoResponse, Response}, +}; +use hyper::body::{Frame, Incoming, SizeHint}; +use hyper_util::{ + client::legacy::{Client, connect::HttpConnector}, + rt::{TokioExecutor, TokioTimer}, +}; +use tokio::net::TcpListener; + +use crate::{ + Rejection, Upstream, connections::Connections, connector::Connector, rewrite, server, + server::ClientSocket, websocket, +}; + +// h11 counts no headers and only caps a head at 16 KiB, so it accepts as many headers as fit in +// that many shortest lines (`a:\r\n`). Python copies every provider header into its response +// too, so hyper's default of 100 would turn ordinary traffic into 431s and 502s. +pub(crate) const MAX_HEADERS: usize = 16 * 1024 / b"a:\r\n".len(); + +#[derive(Clone)] +pub struct Forward { + upstream: Uri, + client: Client, + connections: Connections, + keepalive: Duration, +} + +impl Forward { + /// `upstream_keepalive` is how long the upstream keeps an idle connection open. Pooled + /// connections are dropped well before that, so a request never lands on one it is closing, + /// and idle client connections are closed after it, as uvicorn did. + pub fn new(upstream: &Upstream, upstream_keepalive: Duration) -> Self { + let mut connector = HttpConnector::new(); + connector.set_nodelay(true); + let client = Client::builder(TokioExecutor::new()) + .http1_max_headers(MAX_HEADERS) + .pool_timer(TokioTimer::new()) + .pool_idle_timeout(upstream_keepalive / 2) + .set_host(false) + .build(Connector(connector)); + Self { + upstream: upstream.uri().clone(), + client, + connections: Connections::default(), + keepalive: upstream_keepalive, + } + } + + /// Serves until `shutdown` resolves, then lets open requests and WebSockets finish for up to + /// `drain` before closing whatever is left. Python shuts down alongside and ends long streams + /// itself, so `drain` only backstops a Python that never does. + pub async fn serve( + self, + listener: TcpListener, + shutdown: impl Future, + drain: Duration, + ) { + server::serve(self, listener, shutdown, drain).await; + } + + pub(crate) fn connections(&self) -> &Connections { + &self.connections + } + + pub(crate) fn keepalive(&self) -> Duration { + self.keepalive + } + + pub(crate) async fn handle(&self, client: ClientSocket, request: Request) -> Response { + self.forward(client, request) + .await + .unwrap_or_else(IntoResponse::into_response) + } + + async fn forward( + &self, + client: ClientSocket, + mut request: Request, + ) -> Result { + if let Some(downstream) = websocket::take_handshake(&self.connections, &mut request)? { + return self.forward_websocket(client, request, downstream).await; + } + let request = rewrite::upstream_request(request, &self.upstream, client, false)?; + let response = self + .client + .request(request) + .await + .map_err(Rejection::from_upstream)?; + http_response(response) + } + + async fn forward_websocket( + &self, + client: ClientSocket, + request: Request, + downstream: hyper::upgrade::OnUpgrade, + ) -> Result { + let request = rewrite::upstream_request(request, &self.upstream, client, true)?; + let response = self + .connections + .until_closed(self.client.request(request)) + .await + .ok_or(Rejection::ShuttingDown)? + .map_err(Rejection::from_upstream)?; + if response.status() == StatusCode::SWITCHING_PROTOCOLS { + return websocket::switch(&self.connections, downstream, response).await; + } + http_response(response) + } +} + +fn http_response(response: Response) -> Result { + if response.status() == StatusCode::SWITCHING_PROTOCOLS { + return Err(Rejection::UnrequestedUpgrade); + } + Ok( + rewrite::downstream_response(response) + .map(|body| Body::new(Relayed { body, failed: None })), + ) +} + +/// Python's response body on its way to the client. +struct Relayed { + body: Incoming, + failed: Option, +} + +impl hyper::body::Body for Relayed { + type Data = ::Data; + type Error = hyper::Error; + + // hyper aborts a client connection whose response body fails without flushing the head it + // buffered, so a head Python sent right before dying would never arrive and SDKs would retry + // a request Python already ran. Holding the error for one poll lets hyper flush the head. + fn poll_frame( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + if let Some(error) = self.failed.take() { + return Poll::Ready(Some(Err(error))); + } + match Pin::new(&mut self.body).poll_frame(context) { + Poll::Ready(Some(Err(error))) => { + self.failed = Some(error); + context.waker().wake_by_ref(); + Poll::Pending + } + frame => frame, + } + } + + fn is_end_stream(&self) -> bool { + self.failed.is_none() && self.body.is_end_stream() + } + + fn size_hint(&self) -> SizeHint { + self.body.size_hint() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/lib.rs b/litellm-rust/crates/gateway-forward/src/lib.rs new file mode 100644 index 00000000000..4ac17b50544 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/lib.rs @@ -0,0 +1,29 @@ +mod connections; +mod connector; +mod error; +mod forward; +mod rewrite; +mod server; +mod settings; +mod websocket; + +pub use error::Error; +pub(crate) use error::Rejection; +pub use forward::Forward; +pub use rewrite::PEER_HEADER; +pub use settings::{Settings, Upstream}; + +/// Binds the public listener and forwards to Python until `shutdown` resolves, then drains for +/// `settings.shutdown_drain`. +pub async fn run( + settings: Settings, + shutdown: impl Future + Send + 'static, +) -> Result<(), Error> { + let forward = Forward::new(&settings.upstream, settings.upstream_keepalive); + let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; + tracing::info!(address = %listener.local_addr()?, "forwarding gateway listening"); + forward + .serve(listener, shutdown, settings.shutdown_drain) + .await; + Ok(()) +} diff --git a/litellm-rust/crates/gateway-forward/src/rewrite.rs b/litellm-rust/crates/gateway-forward/src/rewrite.rs new file mode 100644 index 00000000000..fa22bf19037 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/rewrite.rs @@ -0,0 +1,161 @@ +use axum::{ + body::HttpBody, + extract::Request, + http::{HeaderMap, HeaderValue, Response, Uri, Version, header}, +}; + +use crate::{Rejection, server::ClientSocket}; + +pub const PEER_HEADER: &str = "x-litellm-rust-peer"; + +/// The request Python receives: the client's method, target, headers and body unchanged, sent to +/// `upstream` over HTTP/1.1 without the client connection's own headers, plus the client's TCP +/// peer in `PEER_HEADER`, which replaces any copy the client sent. +pub(crate) fn upstream_request( + mut request: Request, + upstream: &Uri, + client: ClientSocket, + websocket: bool, +) -> Result { + if !only_chunked(request.headers()) { + return Err(Rejection::UnsupportedTransferCoding); + } + let path_and_query = request + .uri() + .path_and_query() + .cloned() + .ok_or(Rejection::NoPath)?; + if let Some(host) = public_host(&request, client)? { + request.headers_mut().insert(header::HOST, host); + } + let mut uri = upstream.clone().into_parts(); + uri.path_and_query = Some(path_and_query); + *request.uri_mut() = Uri::from_parts(uri).map_err(|_| Rejection::Internal)?; + *request.version_mut() = Version::HTTP_11; + let chunked = request.body().size_hint().exact().is_none(); + let headers = request.headers_mut(); + strip_hop_by_hop(headers); + if chunked { + // hyper sends no body for a GET with an unknown length unless told it is chunked. + headers.insert( + header::TRANSFER_ENCODING, + HeaderValue::from_static("chunked"), + ); + } + if websocket { + restore_websocket_handshake(headers); + } + headers.insert( + PEER_HEADER, + HeaderValue::from_str(&client.peer.to_string()).map_err(|_| Rejection::Internal)?, + ); + Ok(request) +} + +// Python builds absolute URLs from Host and, without a usable one, from the address it listens +// on, which behind Rust is the private loopback port. A client that sent no usable Host gets the +// public address it connected to, which is what Python used before. h11 rejects HTTP/1.1 without +// Host. +fn public_host(request: &Request, client: ClientSocket) -> Result, Rejection> { + let host = request.headers().get(header::HOST); + if host.is_some_and(|host| !host.is_empty()) { + return Ok(None); + } + if host.is_none() && request.version() == Version::HTTP_11 { + return Err(Rejection::MissingHost); + } + HeaderValue::from_str(&client.local.to_string()) + .map(Some) + .map_err(|_| Rejection::Internal) +} + +// hyper only decodes chunked framing, so anything else would reach Python still coded but framed +// as a plain body. uvicorn rejects it with the 400 it sends for every h11 parse error. +fn only_chunked(headers: &HeaderMap) -> bool { + headers + .get_all(header::TRANSFER_ENCODING) + .iter() + .flat_map(|value| value.as_bytes().split(|byte| *byte == b',')) + .all(|coding| coding.trim_ascii().eq_ignore_ascii_case(b"chunked")) +} + +/// The response the client receives: Python's, without the upstream connection's own headers. +pub(crate) fn downstream_response(mut response: Response) -> Response { + strip_hop_by_hop(response.headers_mut()); + response +} + +/// Python's 101 as the client receives it, still naming the protocol both sides switch to. +pub(crate) fn switched_response(mut response: Response) -> Response { + strip_hop_by_hop(response.headers_mut()); + restore_websocket_handshake(response.headers_mut()); + response +} + +pub(crate) fn is_websocket_handshake(headers: &HeaderMap) -> bool { + headers + .get(header::UPGRADE) + .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket")) + && headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .any(|token| token.trim().eq_ignore_ascii_case("upgrade")) +} + +const HOP_BY_HOP: [&str; 9] = [ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + "proxy-connection", +]; + +// Per-connection headers (RFC 9110 section 7.6.1) describe the client's connection to Rust, not +// Rust's connection to Python, so a proxy must drop them instead of forwarding them. +fn strip_hop_by_hop(headers: &mut HeaderMap) { + let named: Vec<_> = headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .filter_map(|name| header::HeaderName::from_bytes(name.trim().as_bytes()).ok()) + .filter(|name| !must_forward(name)) + .collect(); + for name in named { + headers.remove(name); + } + for name in HOP_BY_HOP { + headers.remove(name); + } +} + +fn must_forward(name: &header::HeaderName) -> bool { + matches!( + name.as_str(), + "host" + | "authorization" + | "cookie" + | "set-cookie" + | "www-authenticate" + | "authentication-info" + | "api-key" + | "x-api-key" + | "litellm-api-key" + | "forwarded" + | PEER_HEADER + ) || name.as_str().starts_with("content-") + || name.as_str().starts_with("x-forwarded-") + || name.as_str().starts_with("sec-websocket-") +} + +// Stripping drops the handshake's Connection and Upgrade, so each hop gets them back. +fn restore_websocket_handshake(headers: &mut HeaderMap) { + headers.insert(header::CONNECTION, HeaderValue::from_static("upgrade")); + headers.insert(header::UPGRADE, HeaderValue::from_static("websocket")); +} diff --git a/litellm-rust/crates/gateway-forward/src/server.rs b/litellm-rust/crates/gateway-forward/src/server.rs new file mode 100644 index 00000000000..0388642446a --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/server.rs @@ -0,0 +1,116 @@ +use std::{convert::Infallible, io, net::SocketAddr, pin::pin, time::Duration}; + +use axum::{body::Body, extract::Request}; +use hyper::{body::Incoming, server::conn::http1, service::service_fn}; +use hyper_util::rt::{TokioIo, TokioTimer}; +use tokio::net::TcpListener; + +use crate::{Forward, forward::MAX_HEADERS}; + +/// The client's TCP connection to the public listener. +#[derive(Clone, Copy, Debug)] +pub(crate) struct ClientSocket { + pub peer: SocketAddr, + pub local: SocketAddr, +} + +/// Serves each public connection with hyper directly, because `axum::serve` cannot configure +/// the connection timeouts uvicorn applied to the same clients. +pub(crate) async fn serve( + forward: Forward, + listener: TcpListener, + shutdown: impl Future, + drain: Duration, +) { + let connections = forward.connections().clone(); + let _close_on_drop = connections.close_on_drop(); + let keepalive = forward.keepalive(); + let builder = builder(keepalive); + let mut shutdown = pin!(shutdown); + loop { + let (stream, peer) = tokio::select! { + () = &mut shutdown => break, + accepted = listener.accept() => match accepted { + Ok(accepted) => accepted, + Err(error) => { + back_off(error).await; + continue; + } + }, + }; + let local = match stream.local_addr() { + Ok(local) => local, + Err(error) => { + tracing::debug!(%error, "client connection closed before it was served"); + continue; + } + }; + if let Err(error) = stream.set_nodelay(true) { + tracing::warn!(%error, "could not disable Nagle on a client connection"); + } + let socket = ClientSocket { peer, local }; + let forward = forward.clone(); + let service = service_fn(move |request: Request| { + let forward = forward.clone(); + async move { Ok::<_, Infallible>(forward.handle(socket, request.map(Body::new)).await) } + }); + let connection = builder + .serve_connection(TokioIo::new(stream), service) + .with_upgrades(); + let draining = connections.clone(); + connections.spawn(async move { + let mut connection = pin!(connection); + let ended = tokio::select! { + ended = connection.as_mut() => ended, + () = draining.draining() => { + connection.as_mut().graceful_shutdown(); + connection.await + } + }; + if let Err(error) = ended { + tracing::debug!(%error, "client connection ended with an error"); + } + }); + } + drop(listener); + connections.begin_shutdown(); + if tokio::time::timeout(drain, connections.wait()) + .await + .is_err() + { + tracing::warn!( + ?drain, + "closing connections still open after the shutdown drain" + ); + } + connections.close(); + connections.wait().await; +} + +// uvicorn closes a client connection that sits idle for `--keepalive_timeout` between requests. +// Without the same limit, a client that vanishes without closing (NAT expiry, a crashed host) +// would hold its socket for the life of the process until the gateway runs out of descriptors. +// Operators already tune that timeout against their load balancer, so the public side keeps +// Python's value instead of introducing its own. hyper only runs the header timer while a +// connection waits for its next request, so it never cuts a long stream, upload or WebSocket. +fn builder(keepalive: Duration) -> http1::Builder { + let mut builder = http1::Builder::new(); + builder + .max_headers(MAX_HEADERS) + .timer(TokioTimer::new()) + .header_read_timeout(keepalive); + builder +} + +async fn back_off(error: io::Error) { + if matches!( + error.kind(), + io::ErrorKind::ConnectionRefused + | io::ErrorKind::ConnectionAborted + | io::ErrorKind::ConnectionReset + ) { + return; + } + tracing::error!(%error, "could not accept a client connection"); + tokio::time::sleep(Duration::from_secs(1)).await; +} diff --git a/litellm-rust/crates/gateway-forward/src/settings.rs b/litellm-rust/crates/gateway-forward/src/settings.rs new file mode 100644 index 00000000000..de01a96c519 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/settings.rs @@ -0,0 +1,113 @@ +use std::{ffi::OsString, net::SocketAddr, str::FromStr, time::Duration}; + +use axum::http::Uri; +use serde::Deserialize; + +use crate::Error; + +const UPSTREAM: &str = "LITELLM_FORWARD_UPSTREAM"; + +/// The Python proxy's private origin. Only a loopback IPv4 address parses, so nothing can point +/// the forwarder at another host. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Upstream(Uri); + +impl Upstream { + pub(crate) fn uri(&self) -> &Uri { + &self.0 + } +} + +impl FromStr for Upstream { + type Err = Error; + + fn from_str(value: &str) -> Result { + let uri: Uri = value.parse().map_err(Error::UpstreamUri)?; + let address = uri + .authority() + .filter(|_| uri.scheme_str() == Some("http")) + .filter(|authority| !authority.as_str().contains('@')) + .filter(|_| uri.path() == "/" && uri.query().is_none()) + .and_then(|authority| authority.as_str().parse::().ok()); + match address { + Some(SocketAddr::V4(address)) if address.ip().is_loopback() => Ok(Self(uri)), + _ => Err(Error::Upstream), + } + } +} + +#[derive(Deserialize)] +struct Environment { + litellm_forward_upstream: String, + litellm_forward_host: String, + litellm_forward_port: u16, + litellm_forward_upstream_keepalive: f64, + litellm_forward_shutdown_drain: f64, +} + +/// Everything forward mode runs with. The Python supervisor resolves each value from the settings +/// Python actually runs with, so none of them is guessed or defaulted here. +/// Setting `LITELLM_FORWARD_UPSTREAM` requires every other `LITELLM_FORWARD_*` setting below. +#[derive(Debug, Clone, PartialEq)] +pub struct Settings { + pub upstream: Upstream, + pub host: String, + pub port: u16, + pub upstream_keepalive: Duration, + pub shutdown_drain: Duration, +} + +impl Settings { + /// `None` when `LITELLM_FORWARD_UPSTREAM` is unset, which means the gateway is not forwarding. + pub fn from_env() -> Result, Error> { + Self::from_vars(std::env::vars_os()) + } + + pub fn from_vars( + variables: impl IntoIterator, + ) -> Result, Error> { + let variables: Vec<(String, String)> = variables + .into_iter() + .filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?))) + .collect(); + if !variables.iter().any(|(key, _)| key == UPSTREAM) { + return Ok(None); + } + let environment: Environment = envy::from_iter(variables).map_err(|error| match error { + envy::Error::MissingValue(field) => Error::MissingSetting(field.to_uppercase()), + envy::Error::Custom(message) => Error::InvalidSetting(message), + })?; + Ok(Some(Self { + upstream: environment.litellm_forward_upstream.parse()?, + host: environment.litellm_forward_host, + port: environment.litellm_forward_port, + upstream_keepalive: positive_seconds( + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE", + environment.litellm_forward_upstream_keepalive, + )?, + shutdown_drain: seconds( + "LITELLM_FORWARD_SHUTDOWN_DRAIN", + environment.litellm_forward_shutdown_drain, + )?, + })) + } +} + +fn seconds(variable: &str, value: f64) -> Result { + Duration::try_from_secs_f64(value).map_err(|_| { + Error::InvalidSetting(format!( + "{variable} must be a non-negative number of seconds, got {value}" + )) + }) +} + +/// The keep-alive also bounds the wait for a client's first request, so zero would close every +/// client connection before it could send one. +fn positive_seconds(variable: &str, value: f64) -> Result { + match seconds(variable, value)? { + Duration::ZERO => Err(Error::InvalidSetting(format!( + "{variable} must be a positive number of seconds, got {value}" + ))), + duration => Ok(duration), + } +} diff --git a/litellm-rust/crates/gateway-forward/src/websocket.rs b/litellm-rust/crates/gateway-forward/src/websocket.rs new file mode 100644 index 00000000000..84d89c75607 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/websocket.rs @@ -0,0 +1,86 @@ +use axum::{ + body::Body, + extract::Request, + http::{Method, Version}, + response::Response, +}; +use hyper::{ + body::Incoming, + upgrade::{OnUpgrade, Upgraded}, +}; +use hyper_util::rt::TokioIo; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +use crate::{Rejection, connections::Connections, rewrite}; + +pub(crate) fn take_handshake( + connections: &Connections, + request: &mut Request, +) -> Result, Rejection> { + if request.method() != Method::GET + || request.version() != Version::HTTP_11 + || !rewrite::is_websocket_handshake(request.headers()) + { + return Ok(None); + } + if connections.is_draining() { + return Err(Rejection::ShuttingDown); + } + Ok(Some(hyper::upgrade::on(request))) +} + +pub(crate) async fn switch( + connections: &Connections, + downstream: OnUpgrade, + mut response: Response, +) -> Result { + if !rewrite::is_websocket_handshake(response.headers()) { + return Err(Rejection::UnrequestedUpgrade); + } + let upstream = connections + .until_closed(hyper::upgrade::on(&mut response)) + .await + .ok_or(Rejection::ShuttingDown)? + .map_err(Rejection::UpgradeFailed)?; + connections.spawn(async move { + match downstream.await { + Ok(downstream) => relay(downstream, upstream).await, + Err(_) => tracing::debug!("websocket upgrade did not complete"), + } + }); + Ok(rewrite::switched_response(response).map(|_| Body::empty())) +} + +// uvicorn never half-closes a WebSocket, so Python's EOF means the socket is gone and the relay +// ends. Only a client that is gone ends it from the other side: the client's EOF just half-closes +// Python's side, and a write Python refuses means it closed with its last frames still in flight +// to the client. +async fn relay(downstream: Upgraded, upstream: Upgraded) { + let (mut client_read, mut client_write) = tokio::io::split(TokioIo::new(downstream)); + let (mut python_read, mut python_write) = tokio::io::split(TokioIo::new(upstream)); + let to_python = async { + let mut buffer = vec![0; 8192]; + loop { + let read = match client_read.read(&mut buffer).await { + Ok(0) => break, + Ok(read) => read, + Err(_) => return, + }; + if python_write.write_all(&buffer[..read]).await.is_err() { + return std::future::pending().await; + } + } + python_write.shutdown().await.ok(); + std::future::pending().await + }; + let to_client = async { + tokio::io::copy(&mut python_read, &mut client_write) + .await + .ok(); + client_write.shutdown().await.ok(); + }; + tokio::select! { + () = to_python => tracing::debug!("websocket client disconnected"), + () = to_client => {} + } +} diff --git a/litellm-rust/crates/gateway-forward/tests/client_identity.rs b/litellm-rust/crates/gateway-forward/tests/client_identity.rs new file mode 100644 index 00000000000..c484c89dc1d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/client_identity.rs @@ -0,0 +1,223 @@ +mod support; + +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use axum::{body::Body, http::StatusCode, response::Response}; +use litellm_gateway_forward::{Forward, PEER_HEADER}; +use rstest::rstest; +use support::{DRAIN, Gateway, Raw, UVICORN_KEEPALIVE, raw_upstream, recording, within}; +use tokio::net::{TcpListener, TcpStream}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +#[rstest] +#[case::ordinary("")] +#[case::connection_nominated( + "connection: X-Forwarded-For, X-Forwarded-Proto, Forwarded, X-Litellm-Rust-Peer\r\n" +)] +#[tokio::test] +async fn python_sees_forwarding_headers_untouched_and_the_tcp_peer_rust_saw( + #[case] connection: &str, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "GET /sso/callback HTTP/1.1\r\n\ + host: public.example\r\n\ + x-forwarded-for: 192.0.2.1\r\n\ + x-forwarded-for: 192.0.2.2\r\n\ + x-forwarded-proto: https\r\n\ + forwarded: for=192.0.2.1;proto=https\r\n\ + x-litellm-rust-peer: 10.1.2.3:1\r\n{connection}\r\n", + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + + let request = received.recv().await.unwrap(); + let forwarded_for: Vec<_> = request.headers.get_all("x-forwarded-for").iter().collect(); + assert_eq!(forwarded_for, ["192.0.2.1", "192.0.2.2"]); + assert_eq!(request.headers["x-forwarded-proto"], "https"); + assert_eq!(request.headers["forwarded"], "for=192.0.2.1;proto=https"); + let peers: Vec<_> = request.headers.get_all(PEER_HEADER).iter().collect(); + assert_eq!(peers, [raw.local_address().to_string().as_str()]); +} + +/// Python builds redirects and OAuth metadata from Host, and without one from the address it +/// listens on, which behind Rust is its private loopback port. So a request without a usable Host +/// carries the public address the client connected to, except HTTP/1.1, which h11 rejects. +#[rstest] +#[case::http11_without_host("HTTP/1.1", "", false)] +#[case::http10_without_host("HTTP/1.0", "", true)] +#[case::http11_with_an_empty_host("HTTP/1.1", "host:\r\n", true)] +#[tokio::test] +async fn a_request_without_a_usable_host_never_names_the_private_upstream( + #[case] version: &str, + #[case] host: &str, + #[case] forwarded: bool, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!("GET /sso/key/generate {version}\r\n{host}\r\n")) + .await; + + let status = within(raw.response_head()).await.unwrap().status; + let hosts = received.try_recv().ok().map(|request| { + request + .headers + .get_all("host") + .iter() + .map(|value| value.to_str().unwrap().to_owned()) + .collect::>() + }); + let expected = if forwarded { + (StatusCode::OK, Some(vec![gateway.address.to_string()])) + } else { + (StatusCode::BAD_REQUEST, None) + }; + assert_eq!((status, hosts), expected); +} + +#[rstest] +#[case::uppercase_with_port("PUBLIC.Example:4000")] +#[case::ipv6_literal("[::1]:4000")] +#[case::ipv4_with_port("192.0.2.10:8443")] +#[tokio::test] +async fn python_sees_the_host_header_byte_for_byte(#[case] host: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!("GET /ui HTTP/1.1\r\nHost: {host}\r\n\r\n")) + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let hosts: Vec<_> = request.headers.get_all("host").iter().collect(); + assert_eq!(hosts, [host]); +} + +/// Python's IP allowlists and rate limits read the peer exactly as uvicorn would have seen it on +/// the public socket, so a dual-stack listener reports IPv4 clients in their mapped form. +#[rstest] +#[case::ipv6("[::1]:0", "[::1]", "[::1]")] +#[case::ipv4_on_a_dual_stack_listener("[::]:0", "127.0.0.1", "[::ffff:127.0.0.1]")] +#[tokio::test] +async fn ipv6_listeners_report_the_peer_bracketed_with_its_port( + #[case] bind: &str, + #[case] connect: &str, + #[case] peer_ip: &str, +) { + let (upstream, mut received) = recording(ok).await; + let forward = Forward::new( + &format!("http://{upstream}").parse().unwrap(), + UVICORN_KEEPALIVE, + ); + let public = TcpListener::bind(bind).await.unwrap(); + let port = public.local_addr().unwrap().port(); + tokio::spawn(forward.serve(public, std::future::pending(), DRAIN)); + let mut raw = Raw::new( + TcpStream::connect(format!("{connect}:{port}")) + .await + .unwrap(), + ); + + raw.send("GET / HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!( + received.recv().await.unwrap().headers[PEER_HEADER], + format!("{peer_ip}:{}", raw.local_address().port()).as_str() + ); +} + +/// Without the peer header Python would take Rust's loopback hop as the client, which allowlists +/// often trust, so no spelling, repetition or Connection token may remove or replace it. +#[rstest] +#[case::named_as_hop_by_hop( + "connection: x-litellm-rust-peer\r\nx-litellm-rust-peer: 127.0.0.1:1\r\n" +)] +#[case::mixed_case_copies( + "X-LiteLLM-Rust-Peer: 127.0.0.1:1\r\nx-litellm-rust-peer: 192.0.2.7:2\r\n" +)] +#[case::comma_joined("x-litellm-rust-peer: 127.0.0.1:1, 192.0.2.7:2\r\n")] +#[tokio::test] +async fn a_client_can_neither_strip_nor_spoof_the_peer(#[case] spoof: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!( + "GET /key/info HTTP/1.1\r\nhost: public.example\r\n{spoof}\r\n" + )) + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let peers: Vec<_> = request.headers.get_all(PEER_HEADER).iter().collect(); + assert_eq!(peers, [raw.local_address().to_string().as_str()]); +} + +/// Pooled upstream connections carry requests from many clients in turn, so the peer belongs to +/// each request and never to the connection it travels on. +#[rstest] +#[tokio::test] +async fn each_request_carries_its_own_peer_over_a_shared_upstream_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + while let Some(head) = python.request_head().await { + let peer = head.header(PEER_HEADER).unwrap().to_vec(); + python + .send( + [ + format!("HTTP/1.1 200 OK\r\ncontent-length: {}\r\n\r\n", peer.len()) + .as_bytes(), + &peer, + ] + .concat(), + ) + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut clients = [ + gateway.connect().await, + gateway.connect().await, + gateway.connect().await, + ]; + + for index in [0, 1, 2, 2, 1, 0] { + let raw = &mut clients[index]; + raw.send("GET /v1/models HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + let peer = raw.body(head.content_length().unwrap()).await; + assert_eq!(peer, raw.local_address().to_string().into_bytes()); + } + assert_eq!(connections.load(Ordering::SeqCst), 1); +} diff --git a/litellm-rust/crates/gateway-forward/tests/download.rs b/litellm-rust/crates/gateway-forward/tests/download.rs new file mode 100644 index 00000000000..9707d2f53b7 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/download.rs @@ -0,0 +1,192 @@ +mod support; + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{body::Body, extract::Request, http::StatusCode, response::Response}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, recording, within}; +use tokio::sync::{Notify, mpsc}; + +#[rstest] +#[tokio::test] +async fn streams_a_body_without_a_content_length_as_python_produces_it() { + let first: &[u8] = &[0xff, 0xfb, 0x90, 0x64, 0x00, 0x0d, 0x0a]; + let rest = vec![0xa5; 64 * 1024]; + let finish = Arc::new(Notify::new()); + let finished = finish.clone(); + let tail = rest.clone(); + let upstream = raw_upstream(move |mut python| { + let finished = finished.clone(); + let tail = tail.clone(); + async move { + python.request_head().await.unwrap(); + python + .send([&b"HTTP/1.1 200 OK\r\ncontent-type: audio/mpeg\r\ntransfer-encoding: chunked\r\n\r\n7\r\n"[..], first, b"\r\n"].concat()) + .await; + finished.notified().await; + python + .send([format!("{:x}\r\n", tail.len()).as_bytes(), &tail, b"\r\n0\r\n\r\n"].concat()) + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut response = within( + client().request( + Request::post(gateway.url("/v1/audio/speech")) + .body(Body::empty()) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert!(!response.headers().contains_key("content-length")); + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), first); + finish.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, rest); +} + +#[rstest] +#[tokio::test] +async fn range_requests_and_partial_content_pass_unchanged() { + let (upstream, mut received) = recording(|| -> Response { + Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header("content-range", "bytes 2-5/10") + .header("accept-ranges", "bytes") + .body(Body::from("2345")) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/videos/video_1/content HTTP/1.1\r\nhost: x\r\nrange: bytes=2-5\r\nif-range: \"v1\"\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::PARTIAL_CONTENT); + assert_eq!(head.header("content-range"), Some(&b"bytes 2-5/10"[..])); + assert_eq!(raw.body(head.content_length().unwrap()).await, b"2345"); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["range"], "bytes=2-5"); + assert_eq!(request.headers["if-range"], "\"v1\""); +} + +#[rstest] +#[tokio::test] +async fn a_large_download_is_paced_by_the_client_and_keeps_its_content_length() { + const LENGTH: usize = 128 * 1024 * 1024; + let written = Arc::new(AtomicUsize::new(0)); + let writing = written.clone(); + let upstream = raw_upstream(move |mut python| { + let writing = writing.clone(); + async move { + python.request_head().await.unwrap(); + python + .send(format!("HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: {LENGTH}\r\n\r\n")) + .await; + let chunk = [0xa5; 64 * 1024]; + while writing.load(Ordering::SeqCst) < LENGTH { + python.send(chunk).await; + writing.fetch_add(chunk.len(), Ordering::SeqCst); + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/files/file-1/content HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.content_length(), Some(LENGTH)); + assert_eq!(head.header("transfer-encoding"), None); + tokio::time::sleep(Duration::from_secs(1)).await; + assert!( + written.load(Ordering::SeqCst) < LENGTH / 4, + "the forwarder buffered a download the client was not reading" + ); + let mut remaining = LENGTH; + while remaining > 0 { + remaining -= within(raw.body(remaining.min(1 << 20))).await.len(); + } +} + +/// The UI revalidates its static assets, and a 304 must carry exactly the validators Python sent, +/// with no length invented for a body it never has. +#[rstest] +#[tokio::test] +async fn a_conditional_request_gets_python_304_with_its_validators_only() { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + event + .send(head.header("if-none-match").map(<[u8]>::to_vec)) + .unwrap(); + python + .send("HTTP/1.1 304 Not Modified\r\netag: \"5f1-abc\"\r\ncache-control: no-cache\r\nvary: Accept-Encoding\r\ndate: Tue, 03 Nov 2026 10:00:00 GMT\r\n\r\n") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + raw.send("GET /ui/_next/static/chunks/main.js HTTP/1.1\r\nhost: x\r\nif-none-match: \"5f1-abc\"\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::NOT_MODIFIED); + let mut names: Vec<_> = head.headers.iter().map(|(name, _)| name.as_str()).collect(); + names.sort_unstable(); + assert_eq!(names, ["cache-control", "date", "etag", "vary"]); + assert_eq!(head.header("etag"), Some(&b"\"5f1-abc\""[..])); + assert_eq!( + head.header("date"), + Some(&b"Tue, 03 Nov 2026 10:00:00 GMT"[..]) + ); + assert_eq!( + within(events.recv()).await.unwrap().as_deref(), + Some(&b"\"5f1-abc\""[..]) + ); + } +} + +/// A client knows a download was cut short only because fewer bytes than the Content-Length +/// arrived before the connection closed, so the length must survive and the connection must end. +#[rstest] +#[tokio::test] +async fn a_download_python_cuts_short_closes_the_client_connection_before_its_length() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send([&b"HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: 1000\r\n\r\n"[..], &[0xa5; 500]].concat()) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/files/file-1/content HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.content_length(), Some(1000)); + assert_eq!(within(raw.until_closed()).await, [0xa5; 500]); +} diff --git a/litellm-rust/crates/gateway-forward/tests/exchange.rs b/litellm-rust/crates/gateway-forward/tests/exchange.rs new file mode 100644 index 00000000000..8fed7b18453 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/exchange.rs @@ -0,0 +1,855 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::{HeaderValue, StatusCode}, + response::Response, +}; +use litellm_gateway_forward::PEER_HEADER; +use rstest::rstest; +use support::{Gateway, Head, client, raw_upstream, recording, serve, within}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + sync::mpsc, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +/// An upstream that reports each bodiless request head exactly as it arrived on the wire. +async fn heads() -> (std::net::SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = raw_upstream(move |mut python| { + let sender = sender.clone(); + async move { + while let Some(head) = python.request_head().await { + sender.send(head).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + (address, receiver) +} + +#[rstest] +#[tokio::test] +async fn forwards_the_request_as_the_client_sent_it_minus_hop_headers() { + let (upstream, mut received) = recording(|| { + Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header("connection", "x-upstream-hop") + .header("x-upstream-hop", "remove") + .header("x-response", "keep") + .body(Body::from("python error")) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + "PATCH /v1/messages?x=a%2Fb&x=c HTTP/1.1\r\n\ + host: public.example\r\n\ + authorization: Bearer test\r\n\ + connection: x-remove, keep-alive\r\n\ + x-remove: remove\r\n\ + proxy-authorization: remove\r\n\ + te: trailers\r\n\ + content-length: 7\r\n\r\npayload", + ) + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::UNAUTHORIZED); + assert_eq!(head.header("x-response"), Some(&b"keep"[..])); + assert_eq!(head.header("x-upstream-hop"), None); + let body = raw.body(head.content_length().unwrap()).await; + assert_eq!(body, b"python error"); + + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "PATCH"); + assert_eq!(request.uri, "/v1/messages?x=a%2Fb&x=c"); + assert_eq!(request.headers["host"], "public.example"); + assert_eq!(request.headers["authorization"], "Bearer test"); + assert!(!request.headers.contains_key("x-remove")); + assert!(!request.headers.contains_key("proxy-authorization")); + assert!(!request.headers.contains_key("te")); + assert!(!request.headers.contains_key("connection")); + assert_eq!(request.body, "payload"); +} + +#[rstest] +#[case::host("HoSt", "public.example")] +#[case::length("Content-Length", "7")] +#[case::authorization("AUTHORIZATION", "Bearer test")] +#[case::cookie("Cookie", "session=test")] +#[case::api_key("X-Api-Key", "test")] +#[case::azure_key("Api-Key", "test")] +#[case::litellm_key("Litellm-Api-Key", "test")] +#[case::forwarded("Forwarded", "for=192.0.2.1;proto=https")] +#[case::forwarded_for("X-Forwarded-For", "192.0.2.1")] +#[case::forwarded_proto("X-Forwarded-Proto", "https")] +#[case::forwarded_host("X-Forwarded-Host", "public.example")] +#[case::content_type("Content-Type", "application/octet-stream")] +#[case::content_encoding("Content-Encoding", "identity")] +#[tokio::test] +async fn connection_cannot_remove_end_to_end_headers(#[case] name: &str, #[case] value: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let host = if name.eq_ignore_ascii_case("host") { + "" + } else { + "host: public.example\r\n" + }; + let length = if name.eq_ignore_ascii_case("content-length") { + "" + } else { + "content-length: 7\r\n" + }; + raw.send(format!( + "POST / HTTP/1.1\r\n{host}{length}{name}: {value}\r\nconnection: x-remove\r\nconnection: {name}, keep-alive\r\nx-remove: remove\r\n\r\npayload" + )).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = within(received.recv()).await.unwrap(); + assert_eq!(request.headers[name], value); + assert_eq!(request.body, "payload"); + assert!(!request.headers.contains_key("connection")); + assert!(!request.headers.contains_key("x-remove")); +} + +#[rstest] +#[tokio::test] +async fn connection_cannot_remove_response_credentials_or_representation_metadata() { + let (upstream, _received) = recording(|| { + Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header("connection", "Content-Length, Set-Cookie, WWW-Authenticate, Authentication-Info, Content-Type, Content-Encoding, x-remove") + .header("content-length", "2") + .header("set-cookie", "session=; Max-Age=0") + .header("www-authenticate", "Bearer") + .header("authentication-info", "nextnonce=test") + .header("content-type", "text/plain") + .header("content-encoding", "identity") + .header("x-remove", "remove") + .body(Body::from("ok")) + .unwrap() + }).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::UNAUTHORIZED); + assert_eq!(head.content_length(), Some(2)); + assert_eq!(head.header("set-cookie"), Some(&b"session=; Max-Age=0"[..])); + assert_eq!(head.header("www-authenticate"), Some(&b"Bearer"[..])); + assert_eq!( + head.header("authentication-info"), + Some(&b"nextnonce=test"[..]) + ); + assert_eq!(head.header("content-type"), Some(&b"text/plain"[..])); + assert_eq!(head.header("content-encoding"), Some(&b"identity"[..])); + assert_eq!(head.header("x-remove"), None); + assert_eq!(raw.body(2).await, b"ok"); +} + +#[rstest] +#[case::get("GET")] +#[case::post("POST")] +#[case::put("PUT")] +#[case::delete("DELETE")] +#[case::options("OPTIONS")] +#[case::extension("PURGE")] +#[case::webdav("PROPFIND")] +#[case::lowercase("get")] +#[case::punctuated("M-SEARCH")] +#[tokio::test] +async fn forwards_every_method(#[case] method: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let response = client() + .request( + Request::builder() + .method(method) + .uri(gateway.url("/mcp")) + .body(Body::from("body")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, method); + assert_eq!(request.body, "body"); +} + +#[rstest] +#[case::double_slash("//v1//models")] +#[case::encoded_slash("/v1/files/a%2Fb/content")] +#[case::encoded_hash_in_query("/v1/files?purpose=batch%23retry&limit=1")] +#[case::encoded_hash_in_path("/v1/files/report%231.pdf/content")] +#[case::encoded_hash_before_query("/gemini/v1beta/files/a%23b?key=k")] +#[case::encoded_query("/v1/models?q=%E2%9C%93&empty=&flag")] +#[case::utf8_bytes_escaped("/mcp/%F0%9F%A6%80")] +#[case::lowercase_encoded_slash("/v1/files/a%2fb/content")] +#[case::dot_segments("/openai/./v1/../v1/models")] +#[case::encoded_dot_segments("/openai/%2e%2e/v1/%2E/models")] +#[case::matrix_params("/langfuse/api;v=1/public;x=y/traces")] +#[case::plus_and_encoded_space("/gemini/v1beta/models?q=a+b&r=a%20b&s=%2B")] +#[case::empty_query("/v1/models?")] +#[case::question_marks_in_query("/v1/models??a=b?c")] +#[case::colons_in_segments( + "/v1/projects/p/locations/l/publishers/google/models/m:streamGenerateContent?alt=sse" +)] +#[case::sub_delims("/~u/@:!$&'()*+,=/x")] +#[case::raw_bytes_browsers_leave_unescaped("/a{b}|c^d\\e\"f?g={}|^`")] +#[case::invalid_percent_escapes("/a%zz%00b?%")] +#[tokio::test] +async fn forwards_the_request_target_byte_for_byte(#[case] target: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().uri, target); +} + +#[rstest] +#[tokio::test] +async fn forwards_a_query_string_as_long_as_uvicorn_accepts() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let target = format!("/vertex_ai/v1/models?key=%2F{}&end", "a%20b".repeat(3000)); + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().uri, target.as_str()); +} + +#[rstest] +#[case::in_query("/v1/files?purpose=batch#retry&limit=1", "/v1/files?purpose=batch")] +#[case::in_path("/v1/files/report#1.pdf/content", "/v1/files/report")] +#[case::before_query("/gemini/v1beta/files/a#b?key=k", "/gemini/v1beta/files/a")] +#[tokio::test] +async fn forwards_only_the_path_and_query_before_a_fragment( + #[case] target: &str, + #[case] expected: &str, +) { + let (upstream, mut received) = heads().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(within(received.recv()).await.unwrap().path, expected); +} + +#[rstest] +#[tokio::test] +async fn absolute_form_targets_reach_python_as_origin_form() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + "GET http://elsewhere.example/v1/models?x=1 HTTP/1.1\r\nhost: elsewhere.example\r\n\r\n", + ) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.uri, "/v1/models?x=1"); +} + +#[rstest] +#[tokio::test] +async fn relays_a_response_with_more_headers_than_hyper_allows_by_default() { + let upstream = serve(axum::Router::new().fallback(|| async { + let headers = (0..150).map(|index| (format!("llm_provider-x-{index}"), "v")); + let mut response = ok(); + for (name, value) in headers { + response.headers_mut().insert( + axum::http::HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_static(value), + ); + } + response + })) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send("GET /v1/chat/completions HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.header("llm_provider-x-149"), Some(&b"v"[..])); +} + +/// uvicorn's h11 parser has no header count limit, only a 16 KiB cap on the whole head. +#[rstest] +#[tokio::test] +async fn forwards_a_request_with_as_many_headers_as_uvicorn_accepts() { + let (sender, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |python| { + let sender = sender.clone(); + async move { + let mut stream = python.into_stream(); + let mut head = Vec::new(); + while !head.ends_with(b"\r\n\r\n") { + head.push(stream.read_u8().await.unwrap()); + } + sender.send(head).unwrap(); + stream + .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await + .unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let headers: String = (0..1000).map(|index| format!("x-{index}: v\r\n")).collect(); + raw.send(format!( + "GET /v1/chat/completions HTTP/1.1\r\nhost: x\r\n{headers}\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let head = String::from_utf8(within(received.recv()).await.unwrap()).unwrap(); + assert!(head.contains(&headers)); +} + +async fn header_fields_python_receives(fields: &str) -> Vec<(String, Vec)> { + let (upstream, mut received) = heads().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "GET /anthropic/v1/messages HTTP/1.1\r\nhost: x\r\n{fields}\r\n" + )) + .await; + within(raw.response_head()).await.unwrap(); + within(received.recv()) + .await + .unwrap() + .headers + .into_iter() + .filter(|(name, _)| name.starts_with("x-") && name != PEER_HEADER) + .collect() +} + +fn fields(expected: &[(&str, &str)]) -> Vec<(String, Vec)> { + expected + .iter() + .map(|(name, value)| ((*name).to_owned(), value.as_bytes().to_vec())) + .collect() +} + +/// Python reads `scope["headers"]` as the list the client sent, in that order. +#[rstest] +#[ignore = "divergence: HeaderMap::remove can move the headers after a dropped hop-by-hop one"] +#[tokio::test] +async fn dropping_a_hop_by_hop_header_keeps_the_others_in_order() { + let received = header_fields_python_receives( + "x-a: 1\r\nconnection: keep-alive\r\nx-b: 2\r\nx-c: 3\r\nx-d: 4\r\n", + ) + .await; + assert_eq!( + received, + fields(&[("x-a", "1"), ("x-b", "2"), ("x-c", "3"), ("x-d", "4")]) + ); +} + +/// uvicorn speaks only HTTP/1.1, so an HTTP/2 prior-knowledge client never reaches Python through +/// Rust either. hyper closes the connection where h11 answers 400, which no HTTP/2 client reads. +#[rstest] +#[tokio::test] +async fn an_http2_client_never_reaches_python() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00") + .await; + + assert!(within(raw.until_closed()).await.is_empty()); + assert!(received.try_recv().is_err()); +} + +#[rstest] +#[ignore = "divergence: HeaderMap groups every value of a name at its first position"] +#[tokio::test] +async fn interleaved_repeated_headers_reach_python_interleaved() { + let received = header_fields_python_receives("x-a: 1\r\nx-b: 2\r\nx-a: 3\r\n").await; + assert_eq!( + received, + fields(&[("x-a", "1"), ("x-b", "2"), ("x-a", "3")]) + ); +} + +/// h11 trims optional whitespace around a value and keeps everything inside it, including empty +/// values and names using every token character. +#[rstest] +#[tokio::test] +async fn python_sees_header_values_as_h11_would_parse_them() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nx-empty:\r\nx-blank: \r\nx-inner: a \t b \t\r\nX-Token!#$%&'*+.^_`|~: v\r\nx-empty: second\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let values = |name: &str| -> Vec<&[u8]> { + request + .headers + .get_all(name) + .iter() + .map(HeaderValue::as_bytes) + .collect() + }; + assert_eq!(values("x-empty"), [&b""[..], b"second"]); + assert_eq!(values("x-blank"), [&b""[..]]); + assert_eq!(values("x-inner"), [&b"a \t b"[..]]); + assert_eq!(values("x-token!#$%&'*+.^_`|~"), [&b"v"[..]]); +} + +/// RFC 9112 section 5.2 lets a server either reject obs-fold or unfold it, and h11 unfolds. +#[rstest] +#[ignore = "divergence: hyper rejects obs-fold with a 400 where uvicorn unfolds it to one space"] +#[tokio::test] +async fn a_folded_header_reaches_python_unfolded() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nx-a: b\r\n c\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().headers["x-a"], "b c"); +} + +#[rstest] +#[tokio::test] +async fn keeps_repeated_response_headers_and_opaque_header_bytes() { + let (upstream, mut received) = recording(|| { + Response::builder() + .header("set-cookie", "a=1; Path=/") + .header("set-cookie", "b=2; HttpOnly") + .header("x-opaque", HeaderValue::from_bytes(b"caf\xe9").unwrap()) + .body(Body::empty()) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(b"GET / HTTP/1.1\r\nhost: x\r\nx-opaque: caf\xe9\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + let cookies: Vec<_> = head + .headers + .iter() + .filter(|(name, _)| name == "set-cookie") + .map(|(_, value)| value.as_slice()) + .collect(); + assert_eq!(cookies, [&b"a=1; Path=/"[..], b"b=2; HttpOnly"]); + assert_eq!(head.header("x-opaque"), Some(&b"caf\xe9"[..])); + assert_eq!( + received.recv().await.unwrap().headers["x-opaque"].as_bytes(), + b"caf\xe9" + ); +} + +/// An HTTP/1.0 client cannot decode chunked framing, so a streamed answer reaches it delimited by +/// the connection closing, as uvicorn sends it. +#[rstest] +#[tokio::test] +async fn an_http10_client_gets_a_streamed_answer_unframed_and_python_gets_its_body() { + let (upstream, mut received) = recording(|| { + Response::new(Body::from_stream(futures_util::stream::iter( + ["str", "eamed"].map(|piece| Ok::<_, Infallible>(Bytes::from_static(piece.as_bytes()))), + ))) + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/completions HTTP/1.0\r\nhost: x\r\ncontent-length: 7\r\n\r\npayload") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.header("transfer-encoding"), None); + assert_eq!(within(raw.until_closed()).await, b"streamed"); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "POST"); + assert_eq!(request.body, "payload"); +} + +#[rstest] +#[tokio::test] +async fn answers_pipelined_requests_in_order() { + let upstream = serve( + Router::new().fallback(|request: Request| async move { request.uri().path().to_owned() }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /first HTTP/1.1\r\nhost: x\r\n\r\nHEAD /second HTTP/1.1\r\nhost: x\r\n\r\nGET /third HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let first = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(first.content_length().unwrap()).await, b"/first"); + let second = within(raw.response_head()).await.unwrap(); + assert_eq!(second.content_length(), Some(7)); + let third = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(third.content_length().unwrap()).await, b"/third"); +} + +#[rstest] +#[case::no_content(StatusCode::NO_CONTENT)] +#[case::not_modified(StatusCode::NOT_MODIFIED)] +#[tokio::test] +async fn bodiless_statuses_keep_the_connection_reusable(#[case] status: StatusCode) { + let upstream = serve(Router::new().fallback(move |request: Request| async move { + match request.uri().path() { + "/next" => (StatusCode::OK, "next"), + _ => (status, ""), + } + })) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /empty HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + assert_eq!(within(raw.response_head()).await.unwrap().status, status); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +#[rstest] +#[tokio::test] +async fn informational_responses_keep_the_connection_reusable() { + let upstream = raw_upstream(|mut python| async move { + while let Some(head) = python.request_head().await { + match head.path.as_str() { + "/hints" => { + python + .send("HTTP/1.1 103 Early Hints\r\nlink: ; rel=preload\r\n\r\nHTTP/1.1 200 OK\r\ncontent-length: 5\r\n\r\nfinal") + .await + } + _ => { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await + } + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /hints HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let mut hinted = within(raw.response_head()).await.unwrap(); + while hinted.status.is_informational() { + hinted = within(raw.response_head()).await.unwrap(); + } + assert_eq!(hinted.status, StatusCode::OK); + assert_eq!(raw.body(hinted.content_length().unwrap()).await, b"final"); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +#[rstest] +#[case::found( + StatusCode::FOUND, + "https://idp.example/authorize?state=a%2Fb&redirect_uri=https%3A%2F%2Fpublic.example%2Fsso%2Fcallback" +)] +#[case::see_other(StatusCode::SEE_OTHER, "/ui/?login=success")] +#[case::temporary(StatusCode::TEMPORARY_REDIRECT, "./mcp/oauth/authorize?client_id=x")] +#[tokio::test] +async fn redirects_reach_the_client_with_location_and_cookies_unchanged( + #[case] status: StatusCode, + #[case] location: &'static str, +) { + let (upstream, _received) = recording(move || { + Response::builder() + .status(status) + .header("location", location) + .header( + "set-cookie", + "token=a%2Fb; Path=/; HttpOnly; Secure; SameSite=Lax", + ) + .header("set-cookie", "state=; Max-Age=0") + .header( + "set-cookie", + "session=x; Expires=Wed, 04 Nov 2026 10:00:00 GMT; Path=/", + ) + .body(Body::empty()) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /sso/key/generate HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, status); + assert_eq!(head.header("location"), Some(location.as_bytes())); + let cookies: Vec<_> = head + .headers + .iter() + .filter(|(name, _)| name == "set-cookie") + .map(|(_, value)| value.as_slice()) + .collect(); + assert_eq!( + cookies, + [ + &b"token=a%2Fb; Path=/; HttpOnly; Secure; SameSite=Lax"[..], + b"state=; Max-Age=0", + b"session=x; Expires=Wed, 04 Nov 2026 10:00:00 GMT; Path=/", + ] + ); +} + +#[rstest] +#[tokio::test] +async fn non_websocket_upgrades_reach_python_as_plain_requests() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nconnection: upgrade, http2-settings\r\nupgrade: h2c\r\nhttp2-settings: AAMAAABkAAQAoAAAAAIAAAAA\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert!(!request.headers.contains_key("upgrade")); + assert!(!request.headers.contains_key("connection")); + assert!(!request.headers.contains_key("http2-settings")); +} + +#[rstest] +#[tokio::test] +async fn forwards_an_asterisk_form_options_request() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("OPTIONS * HTTP/1.1\r\nhost: x\r\n\r\n").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "OPTIONS"); + assert_eq!(request.uri, "*"); +} + +#[rstest] +#[tokio::test] +async fn an_authority_form_target_is_a_400_python_never_sees() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("CONNECT api.example:443 HTTP/1.1\r\nhost: api.example:443\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + assert!(received.try_recv().is_err()); +} + +#[rstest] +#[tokio::test] +async fn a_client_connection_close_ends_only_the_client_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + while python.request_head().await.is_some() { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..2 { + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nconnection: close\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("connection"), Some(&b"close"[..])); + assert_eq!(raw.body(2).await, b"ok"); + within(raw.closed()).await; + } + assert_eq!(connections.load(Ordering::SeqCst), 1); +} + +#[rstest] +#[tokio::test] +async fn a_python_connection_close_ends_only_the_upstream_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 500 Internal Server Error\r\ncontent-length: 5\r\nconnection: close\r\n\r\nerror") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + for _ in 0..2 { + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(head.header("connection"), None); + assert_eq!(raw.body(5).await, b"error"); + } + assert_eq!(connections.load(Ordering::SeqCst), 2); +} + +/// uvicorn answers HEAD on a streaming route with chunked framing and no body, so the client +/// must not be told the body is empty, and must not wait for one before its next request. +#[rstest] +#[tokio::test] +async fn head_answered_with_chunked_framing_claims_no_length_and_keeps_the_connection() { + let upstream = raw_upstream(|mut python| async move { + while let Some(head) = python.request_head().await { + match head.path.as_str() { + "/next" => { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await + } + _ => { + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n") + .await + } + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("HEAD /mcp HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.content_length(), None); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +const GZIP: &[u8] = &[ + 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x4b, 0x49, 0x2c, 0x49, 0xb4, 0x52, + 0x30, 0xe4, 0xe2, 0x02, 0x00, 0xce, 0xcd, 0x78, 0x54, 0x09, 0x00, 0x00, 0x00, +]; + +#[rstest] +#[case::fixed_length(false)] +#[case::streamed(true)] +#[tokio::test] +async fn content_codings_cross_unchanged_both_ways(#[case] streamed: bool) { + let (upstream, mut received) = recording(move || { + let body = match streamed { + true => Body::from_stream(futures_util::stream::iter( + GZIP.chunks(5) + .map(|piece| Ok::<_, Infallible>(Bytes::from_static(piece))), + )), + false => Body::from(GZIP), + }; + Response::builder() + .header("content-encoding", "gzip") + .header("vary", "Accept-Encoding") + .body(body) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + [ + format!( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-encoding: gzip\r\naccept-encoding: gzip, br;q=0.5\r\ncontent-length: {}\r\n\r\n", + GZIP.len() + ) + .as_bytes(), + GZIP, + ] + .concat(), + ) + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("content-encoding"), Some(&b"gzip"[..])); + assert_eq!(head.header("vary"), Some(&b"Accept-Encoding"[..])); + let body = match head.content_length() { + Some(length) => raw.body(length).await, + None => within(raw.chunked_body()).await, + }; + assert_eq!(body, GZIP); + assert_eq!(head.content_length().is_none(), streamed); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["content-encoding"], "gzip"); + assert_eq!(request.headers["accept-encoding"], "gzip, br;q=0.5"); + assert_eq!(request.body, GZIP); +} + +#[rstest] +#[tokio::test] +async fn python_sees_no_accept_encoding_the_client_did_not_send() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /metrics/ HTTP/1.1\r\nhost: x\r\n\r\n").await; + within(raw.response_head()).await.unwrap(); + assert!( + !received + .recv() + .await + .unwrap() + .headers + .contains_key("accept-encoding") + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/framing.rs b/litellm-rust/crates/gateway-forward/tests/framing.rs new file mode 100644 index 00000000000..32a2c3f2b8d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/framing.rs @@ -0,0 +1,259 @@ +mod support; + +use std::{collections::BTreeSet, net::SocketAddr}; + +use axum::http::StatusCode; +use rstest::rstest; +use support::{Gateway, Head, Raw, raw_upstream, within}; +use tokio::sync::mpsc; + +struct Captured { + head: Head, + body: Vec, +} + +/// A keep-alive upstream that decodes framing like h11 and reports every request it reads. +async fn capturing() -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = raw_upstream(move |mut python: Raw| { + let sender = sender.clone(); + async move { + while let Some(head) = python.request_head().await { + let body = match head.header("transfer-encoding") { + Some(_) => python.chunked_body().await, + None => python.body(content_length(&head)).await, + }; + sender.send(Captured { head, body }).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + (address, receiver) +} + +/// The length h11 reads from every Content-Length field, which may repeat one value. +fn content_length(head: &Head) -> usize { + let lengths: BTreeSet = head + .headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .flat_map(|(_, value)| value.split(|byte| *byte == b',')) + .map(|length| { + std::str::from_utf8(length.trim_ascii()) + .unwrap() + .parse() + .unwrap() + }) + .collect(); + assert!( + lengths.len() < 2, + "python got conflicting lengths {lengths:?}" + ); + lengths.first().copied().unwrap_or(0) +} + +/// RFC 9112 section 6.3: Transfer-Encoding overrides Content-Length. Forwarding both would let +/// Python split the stream at a different place than Rust did. Section 6.1 then closes the +/// connection after the response, so a request smuggled behind it never runs. +#[rstest] +#[tokio::test] +async fn a_request_framed_two_ways_reaches_python_framed_by_transfer_encoding_and_closes() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 3\r\ntransfer-encoding: chunked\r\n\r\n5\r\nhello\r\n0\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + within(raw.until_closed()).await; + + let upload = received.recv().await.unwrap(); + assert_eq!(upload.head.header("content-length"), None); + assert_eq!( + upload.head.header("transfer-encoding"), + Some(&b"chunked"[..]) + ); + assert_eq!(upload.body, b"hello"); + assert!( + received.try_recv().is_err(), + "python ran the smuggled request" + ); +} + +#[rstest] +#[case::conflicting_content_lengths("content-length: 4\r\ncontent-length: 40\r\n")] +#[case::content_length_list("content-length: 4, 40\r\n")] +#[case::signed_content_length("content-length: +4\r\n")] +#[case::chunked_not_last("transfer-encoding: chunked, identity\r\ncontent-length: 4\r\n")] +#[case::unknown_coding("transfer-encoding: xchunked\r\ncontent-length: 4\r\n")] +#[case::space_before_colon("transfer-encoding : chunked\r\ncontent-length: 4\r\n")] +#[case::folded_transfer_encoding("transfer-encoding:\r\n chunked\r\ncontent-length: 4\r\n")] +#[case::bare_cr("x-a: b\rx-b: c\r\ncontent-length: 4\r\n")] +#[case::nul_byte("x-a: b\x00c\r\ncontent-length: 4\r\n")] +#[tokio::test] +async fn ambiguous_framing_never_reaches_python(#[case] headers: &str) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\nbodyGET /smuggled HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + within(raw.closed()).await; + assert!(received.try_recv().is_err(), "python received a request"); +} + +#[rstest] +#[case::bare_lf("4\nbody\n0\n\n")] +#[case::invalid_size("zz\r\nbody\r\n0\r\n\r\n")] +#[case::size_overflow("10000000000000004\r\nbody\r\n0\r\n\r\n")] +#[tokio::test] +async fn a_malformed_request_body_is_a_400_not_a_502(#[case] chunks: &str) { + let upstream = raw_upstream(|mut python| async move { python.closed().await }).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{chunks}" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); +} + +/// hyper decodes only the chunked layer, so any other coding would reach Python still coded but +/// framed as a plain body. uvicorn turns every h11 parse error into a 400, ignoring the 501 h11 +/// suggests for codings. +#[rstest] +#[case::listed_before_chunked("transfer-encoding: gzip, chunked\r\n")] +#[case::split_across_fields("transfer-encoding: gzip\r\ntransfer-encoding: chunked\r\n")] +#[tokio::test] +async fn a_transfer_coding_other_than_chunked_is_the_400_uvicorn_sends_python_never_sees( + #[case] headers: &str, +) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\n4\r\nbody\r\n0\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + assert!( + received.try_recv().is_err(), + "python received a still-coded body" + ); +} + +async fn assert_python_reads_the_body_and_the_next_request(headers: &str) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\npayloadGET /next HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + for _ in 0..2 { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + } + assert_eq!(received.recv().await.unwrap().body, b"payload"); + assert_eq!(received.recv().await.unwrap().head.path, "/next"); +} + +#[rstest] +#[case::repeated_fields("content-length: 7\r\ncontent-length: 7\r\n")] +#[case::leading_zeros("content-length: 007\r\n")] +#[tokio::test] +async fn a_content_length_h11_reads_one_way_reaches_python_framed_the_same(#[case] headers: &str) { + assert_python_reads_the_body_and_the_next_request(headers).await; +} + +/// h11 accepts a Content-Length list that repeats one value, as RFC 9110 section 8.6 allows. +#[rstest] +#[ignore = "divergence: hyper's server parses each Content-Length field as one number and 400s a list"] +#[tokio::test] +async fn a_content_length_list_of_one_value_reaches_python() { + assert_python_reads_the_body_and_the_next_request("content-length: 7, 7\r\n").await; +} + +#[rstest] +#[case::leading_zeros("0004\r\nbody\r\n0\r\n\r\n", "body")] +#[case::uppercase_hex("A\r\nbody\r\nbody\r\n0\r\n\r\n", "body\r\nbody")] +#[case::whitespace_after_size("4 \r\nbody\r\n0\r\n\r\n", "body")] +#[case::extension_and_trailer("4;a=\"b;c\"\r\nbody\r\n0\r\nx-sum: 1\r\n\r\n", "body")] +#[tokio::test] +async fn chunk_framing_h11_accepts_reaches_python_decoded( + #[case] chunks: &str, + #[case] body: &str, +) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{chunks}GET /next HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + for _ in 0..2 { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + } + assert_eq!(received.recv().await.unwrap().body, body.as_bytes()); + assert_eq!(received.recv().await.unwrap().head.path, "/next"); +} + +/// h11 reads a head whose lines end in a bare LF, so Python gets the same request through Rust. +#[rstest] +#[tokio::test] +async fn a_head_with_bare_lf_line_endings_reaches_python() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.1\nhost: x\nx-a: b\ncontent-length: 4\n\nbody") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let upload = received.recv().await.unwrap(); + assert_eq!(upload.head.header("x-a"), Some(&b"b"[..])); + assert_eq!(upload.body, b"body"); +} + +/// h11 decodes a chunked body from an HTTP/1.0 client, and RFC 9112 section 6.1 only asks the +/// server to close the connection afterwards. +#[rstest] +#[ignore = "divergence: hyper rejects Transfer-Encoding on HTTP/1.0 with a 400 that uvicorn never sends"] +#[tokio::test] +async fn an_http10_chunked_upload_reaches_python() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.0\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n4\r\nbody\r\n0\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, b"body"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/idle_stream.rs b/litellm-rust/crates/gateway-forward/tests/idle_stream.rs new file mode 100644 index 00000000000..a07c0466481 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/idle_stream.rs @@ -0,0 +1,102 @@ +mod support; + +use std::time::Duration; + +use axum::{ + body::Body, + extract::Request, + http::StatusCode, + response::{IntoResponse, Response}, +}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, recording, within}; +use tokio::{sync::mpsc, time::Instant}; + +const SSE_HEAD: &str = "HTTP/1.1 200 OK\r\n\ + date: Tue, 03 Nov 2026 00:00:00 GMT\r\n\ + server: uvicorn\r\n\ + cache-control: no-cache, no-transform\r\n\ + connection: keep-alive\r\n\ + content-type: text/event-stream\r\n\ + mcp-session-id: session-1\r\n\ + x-accel-buffering: no\r\n\ + transfer-encoding: chunked\r\n\r\n"; + +const PING: &str = ": ping - 2026-11-03 00:20:00.000000+00:00\r\n\r\n"; + +const MESSAGE: &str = "event: message\r\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\r\n\r\n"; + +fn chunk(data: &str) -> String { + format!("{:x}\r\n{data}\r\n", data.len()) +} + +/// An MCP client holds a GET stream open for server-initiated messages, which may not come for +/// many minutes, so neither hop may close a stream for being idle. +#[rstest] +#[tokio::test(start_paused = true)] +async fn an_idle_mcp_stream_stays_open_in_both_directions() { + let idle = Duration::from_secs(20 * 60); + let (session, mut sessions) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let session = session.clone(); + async move { + let head = python.request_head().await.unwrap(); + session + .send(head.header("mcp-session-id").map(<[u8]>::to_vec)) + .unwrap(); + python.send(SSE_HEAD).await; + tokio::time::sleep(idle).await; + python.send(chunk(PING)).await; + python.send(chunk(MESSAGE)).await; + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let started = Instant::now(); + + let response = client() + .request( + Request::get(gateway.url("/github/mcp")) + .header("accept", "text/event-stream") + .header("mcp-session-id", "session-1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.headers()["mcp-session-id"], "session-1"); + assert_eq!(response.headers()["x-accel-buffering"], "no"); + assert!(!response.headers().contains_key("connection")); + assert_eq!( + sessions.recv().await.unwrap().as_deref(), + Some(&b"session-1"[..]) + ); + let body = response.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, [PING, MESSAGE].concat()); + assert!(started.elapsed() >= idle); +} + +#[rstest] +#[tokio::test] +async fn ends_a_session_with_a_bodiless_delete() { + let (upstream, mut received) = recording(|| -> Response { + ([("mcp-session-id", "session-1")], StatusCode::NO_CONTENT).into_response() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("DELETE /mcp HTTP/1.1\r\nhost: x\r\nmcp-session-id: session-1\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::NO_CONTENT); + assert_eq!(head.header("mcp-session-id"), Some(&b"session-1"[..])); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "DELETE"); + assert_eq!(request.headers["mcp-session-id"], "session-1"); + assert!(!request.headers.contains_key("transfer-encoding")); + assert!(request.body.is_empty()); +} diff --git a/litellm-rust/crates/gateway-forward/tests/lifecycle.rs b/litellm-rust/crates/gateway-forward/tests/lifecycle.rs new file mode 100644 index 00000000000..85d10a37be4 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/lifecycle.rs @@ -0,0 +1,473 @@ +mod support; + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use axum::{Router, body::Body, extract::Request, http::StatusCode, response::Response}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{ + DRAIN, Gateway, Raw, client, raw_upstream, recording, serve, unreachable_upstream, within, +}; +use tokio::{ + net::{TcpListener, TcpStream}, + sync::{Notify, mpsc}, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +/// Python stops a provider stream when it sees the client go away, so a client disconnect must +/// close Rust's connection to Python at every stage of the exchange. +#[rstest] +#[case::waiting_for_headers( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n", + false +)] +#[case::mid_stream( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n", + true +)] +#[case::mid_upload( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 1000000\r\n\r\npartial", + false +)] +#[case::waiting_for_websocket_handshake( + "GET /v1/realtime HTTP/1.1\r\nhost: x\r\nconnection: Upgrade\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: dGhlIHNhbXBsZSBub25jZQ==\r\n\r\n", + false +)] +#[tokio::test] +async fn client_disconnect_closes_the_upstream_connection( + #[case] request: &'static str, + #[case] headers_sent: bool, +) { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + python.request_head().await.unwrap(); + if headers_sent { + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n6\r\ndata: \r\n") + .await; + } + event.send("received").unwrap(); + python.closed().await; + event.send("closed").unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(request).await; + assert_eq!(within(events.recv()).await, Some("received")); + if headers_sent { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + } + drop(raw); + assert_eq!(within(events.recv()).await, Some("closed")); +} + +/// Python treats EOF on its socket as the client leaving, so Rust must not half-close the +/// upstream connection once it has sent the request body. +#[rstest] +#[tokio::test] +async fn python_sees_no_eof_while_its_response_is_pending() { + let (verdict, mut verdicts) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let verdict = verdict.clone(); + async move { + let head = python.request_head().await.unwrap(); + python.body(head.content_length().unwrap()).await; + let eof = tokio::time::timeout(Duration::from_millis(200), python.closed()) + .await + .is_ok(); + verdict.send(eof).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + + let response = within( + client().request( + Request::post(gateway.url("/mcp")) + .body(Body::from("{}")) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(verdicts.recv().await, Some(false)); +} + +/// uvicorn closes a connection that sat idle past its keep-alive timeout. A request Rust sends on +/// that connection as uvicorn closes it is lost, so Rust must retire idle connections first. +#[rstest] +#[case::reused_while_fresh(Duration::from_millis(0))] +#[case::past_the_pool_timeout(Duration::from_millis(250))] +#[case::past_python_keepalive(Duration::from_millis(450))] +#[tokio::test] +async fn never_reuses_a_connection_python_is_closing_for_idleness(#[case] idle: Duration) { + let keepalive = Duration::from_millis(400); + let upstream = raw_upstream(move |mut python: Raw| async move { + let mut idle_since = None::; + while python.request_head().await.is_some() { + if idle_since.is_some_and(|since| since.elapsed() >= keepalive) { + return; + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + idle_since = Some(Instant::now()); + } + }) + .await; + let gateway = Gateway::with_keepalive(upstream, keepalive).await; + let client = client(); + let first = within(client.request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + first.into_body().collect().await.unwrap(); + tokio::time::sleep(idle).await; + let second = + within(client.request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(second.status(), StatusCode::OK); +} + +#[rstest] +#[case::health("/health/liveliness")] +#[case::messages("/v1/messages")] +#[tokio::test] +async fn unreachable_python_is_a_502(#[case] path: &str) { + let gateway = Gateway::start(unreachable_upstream()).await; + let response = + within(client().request(Request::get(gateway.url(path)).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); +} + +#[rstest] +#[case::bad_status_line(b"HTTP/1.1 OK\r\ncontent-length: 0\r\n\r\n".to_vec())] +#[case::header_without_colon(b"HTTP/1.1 200 OK\r\nbroken header\r\ncontent-length: 0\r\n\r\n".to_vec())] +#[case::not_http(b"SSH-2.0-OpenSSH_9.6\r\n".to_vec())] +#[case::conflicting_lengths( + b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\ncontent-length: 3\r\n\r\nok".to_vec() +)] +#[case::oversized_head( + [&b"HTTP/1.1 200 OK\r\nx-big: "[..], &vec![b'a'; 1024 * 1024], b"\r\n\r\n"].concat() +)] +#[case::closed_without_reply(Vec::new())] +#[tokio::test] +async fn an_upstream_that_breaks_http_is_a_502_and_the_client_connection_survives( + #[case] reply: Vec, +) { + let upstream = raw_upstream(move |mut python: Raw| { + let reply = reply.clone(); + async move { + python.request_head().await.unwrap(); + if reply.is_empty() { + return; + } + python.send(reply).await; + std::future::pending::<()>().await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::BAD_GATEWAY); + assert_eq!(head.content_length(), Some(0)); + } +} + +#[rstest] +#[tokio::test] +async fn python_restarting_is_a_502_until_it_listens_again() { + let address = unreachable_upstream(); + let gateway = Gateway::start(address).await; + let get = || client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()); + assert_eq!( + within(get()).await.unwrap().status(), + StatusCode::BAD_GATEWAY + ); + + let listener = TcpListener::bind(address).await.unwrap(); + let router = Router::new().fallback(|| async { "up" }); + tokio::spawn(async move { axum::serve(listener, router).await }); + + assert_eq!(within(get()).await.unwrap().status(), StatusCode::OK); +} + +/// A request Python already read may have run, so it is answered with a 502 and never replayed. +#[rstest] +#[tokio::test] +async fn a_request_python_read_before_closing_is_a_502_and_never_replayed() { + let (event, mut seen) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + python.body(head.content_length().unwrap()).await; + event.send(head.path.clone()).unwrap(); + if head.path == "/v1/chat/completions" { + return; + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /warm HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + let warm = within(raw.response_head()).await.unwrap(); + raw.body(warm.content_length().unwrap()).await; + + raw.send("POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 2\r\n\r\n{}") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::BAD_GATEWAY); + let paths: Vec<_> = std::iter::from_fn(|| seen.try_recv().ok()).collect(); + assert_eq!(paths, ["/warm", "/v1/chat/completions"]); +} + +#[rstest] +#[tokio::test] +async fn shutdown_lets_an_in_flight_stream_finish() { + let release = Arc::new(Notify::new()); + let released = release.clone(); + let upstream = raw_upstream(move |mut python: Raw| { + let released = released.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n") + .await; + released.notified().await; + python.send("4\r\nlast\r\n0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(!server.is_finished()); + release.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, "firstlast"); + within(server).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn shutdown_stops_waiting_for_a_stream_still_open_after_the_drain() { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n") + .await; + std::future::pending::<()>().await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + let started = std::time::Instant::now(); + + within(gateway.shutdown()).await.unwrap(); + + assert!(started.elapsed() >= DRAIN); + drop(response); +} + +#[rstest] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_keep_alive_connections_each_get_their_own_responses() { + let upstream = serve(Router::new().fallback(|request: Request| async move { + let index: usize = request.uri().query().unwrap().parse().unwrap(); + Response::builder() + .header("etag", format!("\"{index}\"")) + .body(Body::from(index.to_string().repeat(1 + index * 1000))) + .unwrap() + })) + .await; + let gateway = Gateway::start(upstream).await; + let connections = 6; + + let tasks: Vec<_> = (0..connections) + .map(|first| { + let address = gateway.address; + tokio::spawn(async move { + let mut raw = Raw::new(TcpStream::connect(address).await.unwrap()); + for index in (first..50).step_by(connections) { + raw.send(format!( + "GET /ui/chunk.js?{index} HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("etag"), Some(format!("\"{index}\"").as_bytes())); + let body = raw.body(head.content_length().unwrap()).await; + assert_eq!( + body, + index.to_string().repeat(1 + index * 1000).into_bytes() + ); + } + }) + }) + .collect(); + + for task in tasks { + within(task).await.unwrap(); + } +} + +// Python's `--keepalive_timeout` governed idle client connections before Rust took the public +// port, and operators tune it against their load balancer's idle timeout, so it still does. +// Without it, a client that vanishes without closing its connection holds a socket forever and +// the gateway eventually runs out of file descriptors. +const IDLE: Duration = Duration::from_millis(400); + +#[rstest] +#[tokio::test] +async fn closes_a_client_connection_idle_past_python_keepalive() { + let (upstream, _received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + tokio::time::sleep(IDLE / 2).await; + raw.send("GET /health/liveliness HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + raw.body(head.content_length().unwrap()).await; + } + let idle_since = Instant::now(); + within(raw.closed()).await; + + assert!(idle_since.elapsed() >= IDLE * 3 / 4); +} + +/// uvicorn would hold a connection that never sends a request, which lets a client park sockets +/// for free, so the same limit applies before the first request too. +#[rstest] +#[tokio::test] +async fn closes_a_client_connection_that_never_sends_a_request() { + let (upstream, _received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + within(raw.closed()).await; +} + +/// The limit only covers the wait for the next request. A response that goes quiet for longer +/// belongs to Python, which decides how long it may take. +#[rstest] +#[tokio::test] +async fn keeps_a_connection_open_while_a_response_goes_quiet() { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + tokio::time::sleep(IDLE * 3).await; + python.send("4\r\ndone\r\n0\r\n\r\n").await; + }) + .await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + raw.send("GET /v1/chat/completions HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + within(raw.response_head()).await.unwrap(); + assert_eq!(within(raw.chunked_body()).await, b"done"); +} + +/// Likewise for an upload that pauses longer than the limit before it finishes. +#[rstest] +#[tokio::test] +async fn keeps_a_connection_open_while_an_upload_goes_quiet() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + raw.send("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 8\r\n\r\nhalf") + .await; + tokio::time::sleep(IDLE * 3).await; + raw.send("done").await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "halfdone"); +} + +/// uvicorn answers a request it rejects (413, 401) as soon as it reads the head and then closes, +/// leaving the upload unread. The client must get that answer, not a 502 or a reset. +#[rstest] +#[case::sized("content-length: 16777216\r\n", vec![b'x'; 256 * 1024])] +#[case::chunked( + "transfer-encoding: chunked\r\n", + [&b"40000\r\n"[..], &[b'x'; 256 * 1024], b"\r\n"].concat() +)] +#[tokio::test] +async fn python_rejecting_an_upload_unread_and_closing_still_answers_the_client( + #[case] framing: &'static str, + #[case] partial_upload: Vec, +) { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 413 Payload Too Large\r\ncontent-type: application/json\r\ncontent-length: 26\r\nconnection: close\r\n\r\n{\"error\":\"file too large\"}") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{framing}\r\n" + )) + .await; + raw.send(partial_upload).await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!( + raw.body(head.content_length().unwrap()).await, + b"{\"error\":\"file too large\"}" + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/roundtrip.rs b/litellm-rust/crates/gateway-forward/tests/roundtrip.rs new file mode 100644 index 00000000000..c220c7fe8e9 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/roundtrip.rs @@ -0,0 +1,360 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{Arc, Mutex}, +}; + +use axum::{ + body::{Body, Bytes}, + http::{HeaderMap, HeaderName, HeaderValue, Method, Request, StatusCode}, + response::Response, +}; +use http_body_util::BodyExt; +use litellm_gateway_forward::PEER_HEADER; +use proptest::{prelude::*, test_runner::TestRunner}; +use support::{Gateway, client, recording, within}; + +#[derive(Clone, Debug)] +struct Message { + headers: Vec<(String, Vec>)>, + body: Vec, + chunked: bool, +} + +#[derive(Clone, Debug)] +struct Exchange { + method: Method, + target: String, + request: Message, + status: StatusCode, + response: Message, +} + +/// Any value h11 reads back unchanged: visible ASCII, obs-text and interior whitespace. +fn header_value() -> impl Strategy> { + prop::collection::vec( + prop_oneof![8 => 0x21u8..=0x7e, 2 => 0x80u8..=0xff, 1 => Just(b' '), 1 => Just(b'\t')], + 0..24, + ) + .prop_map(|value| value.trim_ascii().to_vec()) +} + +fn method() -> impl Strategy { + prop_oneof![ + Just(Method::GET), + Just(Method::POST), + Just(Method::PUT), + Just(Method::PATCH), + Just(Method::DELETE), + Just(Method::OPTIONS), + "[A-Za-z0-9!#$%&'*+.^_`|~-]{1,16}" + .prop_filter("HEAD and CONNECT frame differently", |method| { + !["HEAD", "CONNECT"].contains(&method.as_str()) + }) + .prop_map(|method| Method::from_bytes(method.as_bytes()).unwrap()), + ] +} + +fn message() -> impl Strategy { + ( + prop::collection::btree_map( + "x-[a-z0-9!#$%&'*+.^_`|~-]{1,12}", + prop::collection::vec(header_value(), 1..4), + 0..12, + ), + prop::collection::vec(any::(), 0..32 * 1024), + any::(), + ) + .prop_map(|(headers, body, chunked)| Message { + headers: headers.into_iter().collect(), + body, + chunked, + }) +} + +fn target() -> impl Strategy { + let segment = prop_oneof![ + "[A-Za-z0-9._~!$&'()*+,;=:@-]{0,8}", + "%[0-9A-Fa-f]{2}", + Just(".".to_owned()), + Just("..".to_owned()), + Just(String::new()), + ]; + ( + prop::collection::vec(segment, 0..6), + prop::option::of("[A-Za-z0-9=&._~%/?:@!$'()*+,;-]{0,40}"), + ) + .prop_map(|(segments, query)| { + let path = format!("/{}", segments.join("/")); + match query { + Some(query) => format!("{path}?{query}"), + None => path, + } + }) +} + +fn exchange() -> impl Strategy { + let status = (200u16..600) + .prop_filter("bodiless status", |code| ![204, 205, 304].contains(code)) + .prop_map(|code| StatusCode::from_u16(code).unwrap()); + (method(), target(), message(), status, message()).prop_map( + |(method, target, request, status, response)| Exchange { + method, + target, + request, + status, + response, + }, + ) +} + +fn body(message: &Message) -> Body { + if !message.chunked { + return Body::from(message.body.clone()); + } + let chunks: Vec> = message + .body + .chunks(1000) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect(); + Body::from_stream(futures_util::stream::iter(chunks)) +} + +fn apply(headers: &mut HeaderMap, message: &Message) { + if message.chunked { + headers.insert("transfer-encoding", HeaderValue::from_static("chunked")); + } + for (name, values) in &message.headers { + for value in values { + headers.append( + HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_bytes(value).unwrap(), + ); + } + } +} + +fn assert_carried(headers: &HeaderMap, message: &Message, framing_and_hop: &[&str]) { + for (name, values) in &message.headers { + let carried: Vec<&[u8]> = headers + .get_all(name) + .iter() + .map(HeaderValue::as_bytes) + .collect(); + let sent: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect(); + assert_eq!(carried, sent, "header {name}"); + } + for name in headers.keys() { + assert!( + framing_and_hop.contains(&name.as_str()) + || message + .headers + .iter() + .any(|(sent, _)| sent == name.as_str()), + "unexpected header {name}" + ); + } +} + +/// Any request a client can send and any response Python can give cross the forwarder unchanged, +/// apart from the framing and peer headers the forwarder owns. +#[test] +fn every_exchange_crosses_the_forwarder_unchanged() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let current = Arc::new(Mutex::new(None::)); + let responding = current.clone(); + let (gateway, received) = runtime.block_on(async { + let (upstream, received) = recording(move || { + let exchange = responding.lock().unwrap().clone().unwrap(); + let mut response = Response::new(body(&exchange.response)); + *response.status_mut() = exchange.status; + apply(response.headers_mut(), &exchange.response); + response + }) + .await; + ( + Gateway::start(upstream).await, + tokio::sync::Mutex::new(received), + ) + }); + let client = client(); + TestRunner::new(ProptestConfig::with_cases(256)) + .run(&exchange(), |exchange| { + *current.lock().unwrap() = Some(exchange.clone()); + runtime.block_on(async { + while received.lock().await.try_recv().is_ok() {} + let mut request = Request::builder() + .method(exchange.method.clone()) + .uri(gateway.url(&exchange.target)) + .body(body(&exchange.request)) + .unwrap(); + apply(request.headers_mut(), &exchange.request); + let response = support::within(client.request(request)).await.unwrap(); + + let python = received.lock().await.recv().await.unwrap(); + assert_eq!(python.method, exchange.method); + assert_eq!(python.uri, exchange.target.as_str()); + assert_eq!(python.body, exchange.request.body); + assert_carried( + &python.headers, + &exchange.request, + &["host", "content-length", "transfer-encoding", PEER_HEADER], + ); + assert_eq!(python.headers.get_all(PEER_HEADER).iter().count(), 1); + + assert_eq!(response.status(), exchange.status); + let headers = response.headers().clone(); + assert_carried( + &headers, + &exchange.response, + &["date", "content-length", "transfer-encoding"], + ); + let body = support::within(response.into_body().collect()) + .await + .unwrap(); + assert_eq!(body.to_bytes(), exchange.response.body); + }); + Ok(()) + }) + .unwrap(); +} + +#[derive(Clone, Debug)] +enum Framing { + Unframed, + Length, + Chunked(usize), +} + +#[derive(Clone, Debug)] +struct WireRequest { + method: Method, + target: String, + headers: Vec<(String, Vec)>, + body: Vec, + framing: Framing, +} + +fn wire_request() -> impl Strategy { + let framing = prop_oneof![ + Just(Framing::Unframed), + Just(Framing::Length), + (1usize..2048).prop_map(Framing::Chunked), + ]; + ( + method(), + target(), + prop::collection::vec( + ("[xX]-[A-Za-z0-9!#$%&'*+.^_`|~-]{1,12}", header_value()), + 0..24, + ), + prop::collection::vec(any::(), 0..8 * 1024), + framing, + ) + .prop_map(|(method, target, headers, body, framing)| WireRequest { + method, + target, + headers, + body: match framing { + Framing::Unframed => Vec::new(), + _ => body, + }, + framing, + }) +} + +fn encode(request: &WireRequest) -> Vec { + let framing = match request.framing { + Framing::Unframed => String::new(), + Framing::Length => format!("Content-Length: {}\r\n", request.body.len()), + Framing::Chunked(_) => "Transfer-Encoding: chunked\r\n".to_owned(), + }; + let fields = request + .headers + .iter() + .flat_map(|(name, value)| [name.as_bytes(), b": ", value, b"\r\n"].concat()); + let body = match request.framing { + Framing::Chunked(size) => request + .body + .chunks(size) + .flat_map(|chunk| { + [format!("{:x}\r\n", chunk.len()).as_bytes(), chunk, b"\r\n"].concat() + }) + .chain(*b"0\r\n\r\n") + .collect(), + _ => request.body.clone(), + }; + [ + format!( + "{} {} HTTP/1.1\r\nHost: x\r\n", + request.method, request.target + ) + .into_bytes(), + fields.collect(), + framing.into_bytes(), + b"\r\n".to_vec(), + body, + ] + .concat() +} + +fn by_name<'a>(fields: impl Iterator) -> Vec<(String, &'a [u8])> { + let mut fields: Vec<_> = fields + .map(|(name, value)| (name.to_ascii_lowercase(), value)) + .collect(); + fields.sort_by(|(left, _), (right, _)| left.cmp(right)); + fields +} + +/// Pipelined requests in any framing, written in arbitrary pieces, reach Python one by one with +/// the bytes h11 would have read from the client directly. +#[test] +fn pipelined_wire_requests_reach_python_as_sent() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let (gateway, received) = runtime.block_on(async { + let (upstream, received) = recording(|| Response::new(Body::from("ok"))).await; + ( + Gateway::start(upstream).await, + tokio::sync::Mutex::new(received), + ) + }); + let requests = (prop::collection::vec(wire_request(), 1..5), 1usize..4096); + TestRunner::new(ProptestConfig::with_cases(256)) + .run(&requests, |(requests, piece)| { + runtime.block_on(async { + let mut received = received.lock().await; + while received.try_recv().is_ok() {} + let mut raw = gateway.connect().await; + let wire: Vec = requests.iter().flat_map(encode).collect(); + for bytes in wire.chunks(piece) { + raw.send(bytes).await; + } + for request in &requests { + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK, "{request:?}"); + assert_eq!(raw.body(2).await, b"ok"); + let python = within(received.recv()).await.unwrap(); + assert_eq!(python.method, request.method); + assert_eq!(python.uri, request.target.as_str()); + assert_eq!(python.body, request.body); + assert_eq!(python.headers["host"], "x"); + let forwarded = python.headers.iter().filter(|(name, _)| { + !["host", "content-length", "transfer-encoding", PEER_HEADER] + .contains(&name.as_str()) + }); + assert_eq!( + by_name(forwarded.map(|(name, value)| (name.as_str(), value.as_bytes()))), + by_name( + request + .headers + .iter() + .map(|(name, value)| (name.as_str(), value.as_slice())) + ) + ); + } + }); + Ok(()) + }) + .unwrap(); +} diff --git a/litellm-rust/crates/gateway-forward/tests/settings.rs b/litellm-rust/crates/gateway-forward/tests/settings.rs new file mode 100644 index 00000000000..ce0d95e723f --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/settings.rs @@ -0,0 +1,110 @@ +use std::{ffi::OsString, time::Duration}; + +use litellm_gateway_forward::{Error, Settings}; +use rstest::rstest; + +const COMPLETE: [(&str, &str); 5] = [ + ("LITELLM_FORWARD_UPSTREAM", "http://127.0.0.1:53211"), + ("LITELLM_FORWARD_HOST", "0.0.0.0"), + ("LITELLM_FORWARD_PORT", "4000"), + ("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "2"), + ("LITELLM_FORWARD_SHUTDOWN_DRAIN", "40"), +]; + +fn parse(variables: &[(&str, &str)]) -> Result, Error> { + Settings::from_vars( + variables + .iter() + .map(|(key, value)| (OsString::from(key), OsString::from(value))), + ) +} + +fn without(variable: &str) -> Vec<(&'static str, &'static str)> { + COMPLETE + .into_iter() + .filter(|(key, _)| *key != variable) + .collect() +} + +fn with<'a>(variable: &'a str, value: &'a str) -> Vec<(&'a str, &'a str)> { + without(variable) + .into_iter() + .chain([(variable, value)]) + .collect() +} + +#[rstest] +fn reads_every_variable() { + assert_eq!( + parse(&COMPLETE).unwrap(), + Some(Settings { + upstream: "http://127.0.0.1:53211".parse().unwrap(), + host: "0.0.0.0".into(), + port: 4000, + upstream_keepalive: Duration::from_secs(2), + shutdown_drain: Duration::from_secs(40), + }) + ); +} + +#[rstest] +fn without_an_upstream_the_gateway_is_not_forwarding() { + assert_eq!(parse(&without("LITELLM_FORWARD_UPSTREAM")).unwrap(), None); +} + +#[rstest] +#[case::host("LITELLM_FORWARD_HOST")] +#[case::port("LITELLM_FORWARD_PORT")] +#[case::upstream_keepalive("LITELLM_FORWARD_UPSTREAM_KEEPALIVE")] +#[case::shutdown_drain("LITELLM_FORWARD_SHUTDOWN_DRAIN")] +fn a_missing_variable_fails_naming_it(#[case] variable: &str) { + let error = parse(&without(variable)).unwrap_err(); + + assert!(error.to_string().contains(variable), "{error}"); +} + +#[rstest] +#[case::port_overflow("LITELLM_FORWARD_PORT", "65536")] +#[case::port_text("LITELLM_FORWARD_PORT", "http")] +#[case::keepalive_text("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "soon")] +#[case::keepalive_negative("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "-1")] +#[case::keepalive_zero("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "0")] +#[case::drain_text("LITELLM_FORWARD_SHUTDOWN_DRAIN", "soon")] +#[case::drain_negative("LITELLM_FORWARD_SHUTDOWN_DRAIN", "-1")] +#[case::drain_infinite("LITELLM_FORWARD_SHUTDOWN_DRAIN", "inf")] +fn an_invalid_variable_fails_naming_it(#[case] variable: &str, #[case] value: &str) { + let error = parse(&with(variable, value)).unwrap_err(); + + assert!(error.to_string().contains(variable), "{error}"); +} + +#[rstest] +fn fractional_drain_seconds_are_used_as_given() { + let settings = parse(&with("LITELLM_FORWARD_SHUTDOWN_DRAIN", "12.5")) + .unwrap() + .unwrap(); + + assert_eq!(settings.shutdown_drain, Duration::from_millis(12_500)); +} + +#[rstest] +#[case::keepalive("KEEPALIVE_TIMEOUT")] +#[case::graceful_shutdown("GRACEFUL_SHUTDOWN_TIMEOUT")] +#[case::host("HOST")] +#[case::port("PORT")] +fn python_settings_are_never_read(#[case] variable: &str) { + let mut variables = COMPLETE.to_vec(); + variables.push((variable, "9")); + + assert_eq!(parse(&variables).unwrap(), parse(&COMPLETE).unwrap()); +} + +#[rstest] +fn a_non_loopback_upstream_fails_startup() { + let error = parse(&with("LITELLM_FORWARD_UPSTREAM", "http://192.0.2.1:4000")).unwrap_err(); + + assert!( + error.to_string().contains("LITELLM_FORWARD_UPSTREAM"), + "{error}" + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/slow_head.rs b/litellm-rust/crates/gateway-forward/tests/slow_head.rs new file mode 100644 index 00000000000..764733aa524 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/slow_head.rs @@ -0,0 +1,41 @@ +mod support; + +use std::time::Duration; + +use axum::{body::Body, extract::Request, http::StatusCode}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream}; +use tokio::time::Instant; + +/// Reasoning models with `stream=false`, image and video generation, and model health checks send +/// no response headers until the provider finishes, which can take many minutes. +#[rstest] +#[tokio::test(start_paused = true)] +async fn waits_for_response_headers_as_long_as_python_takes() { + let thinking = Duration::from_secs(30 * 60); + let upstream = raw_upstream(move |mut python| async move { + python.request_head().await.unwrap(); + tokio::time::sleep(thinking).await; + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 5\r\n\r\nimage") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let started = Instant::now(); + + let response = client() + .request( + Request::post(gateway.url("/v1/images/generations")) + .body(Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert!(started.elapsed() >= thinking); + let body = response.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, "image"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/stream.rs b/litellm-rust/crates/gateway-forward/tests/stream.rs new file mode 100644 index 00000000000..eaa12f20b77 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/stream.rs @@ -0,0 +1,313 @@ +mod support; + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{body::Body, extract::Request, http::StatusCode}; +use http_body_util::BodyExt; +use hyper::body::Incoming; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, within}; +use tokio::{ + io::AsyncWriteExt, + sync::{Notify, mpsc}, +}; + +fn chunk(bytes: &[u8]) -> Vec { + [format!("{:x}\r\n", bytes.len()).as_bytes(), bytes, b"\r\n"].concat() +} + +const EVENTS: [&str; 3] = ["data: 1\n\n", "data: 2\n\n", "data: [DONE]\n\n"]; + +#[rstest] +#[tokio::test] +async fn each_event_reaches_the_client_before_python_produces_the_next( + #[values(true, false)] chunked: bool, +) { + let produce = Arc::new(Notify::new()); + let next = produce.clone(); + let upstream = raw_upstream(move |mut python| { + let next = next.clone(); + async move { + python.request_head().await.unwrap(); + let framing = match chunked { + true => "transfer-encoding: chunked".to_owned(), + false => format!("content-length: {}", EVENTS.concat().len()), + }; + python + .send(format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n{framing}\r\n\r\n" + )) + .await; + for event in EVENTS { + match chunked { + true => python.send(chunk(event.as_bytes())).await, + false => python.send(event).await, + } + next.notified().await; + } + if chunked { + python.send("0\r\n\r\n").await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut response = within( + client().request( + Request::post(gateway.url("/v1/chat/completions")) + .body(Body::from(r#"{"stream":true}"#)) + .unwrap(), + ), + ) + .await + .unwrap(); + + for event in EVENTS { + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), event); + produce.notify_one(); + } + let rest = within(response.into_body().collect()).await.unwrap(); + assert!(rest.to_bytes().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn relays_a_binary_framed_stream_unchanged() { + let message: Vec = [ + &[ + 0x00, 0x00, 0x00, 0x3d, 0x00, 0x00, 0x00, 0x20, 0x8a, 0x0d, 0x0a, 0xff, + ][..], + b"\x0b:event-type\x07\x00\x05chunk\x0d:content-type\x07\x00\x10application/json", + b"0\r\n\r\n{\"bytes\":\"eyJ0eXBlIjoiY29udGVudCJ9\"}", + &[0xc3, 0x28, 0x00, 0x0d], + ] + .concat(); + let sent = message.clone(); + let upstream = raw_upstream(move |mut python| { + let sent = sent.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: application/vnd.amazon.eventstream\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + for piece in sent.chunks(7) { + python.send(chunk(piece)).await; + } + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = within( + client().request( + Request::post(gateway.url("/bedrock/model/m/converse-stream")) + .body(Body::empty()) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert_eq!( + response.headers()["content-type"], + "application/vnd.amazon.eventstream" + ); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, message); +} + +#[rstest] +#[tokio::test] +async fn relays_a_body_delimited_by_the_upstream_closing() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\nconnection: close\r\n\r\nuntil eof") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()); + let body = within(response).await.unwrap().into_body().collect(); + assert_eq!(within(body).await.unwrap().to_bytes(), "until eof"); +} + +async fn drain(mut body: Incoming) -> (Vec, Result<(), hyper::Error>) { + let mut delivered = Vec::new(); + loop { + match within(body.frame()).await { + None => return (delivered, Ok(())), + Some(Err(error)) => return (delivered, Err(error)), + Some(Ok(frame)) => delivered.extend_from_slice(&frame.into_data().unwrap()), + } + } +} + +#[rstest] +#[case::after_the_head("transfer-encoding: chunked", "", "")] +#[case::between_events("transfer-encoding: chunked", "6\r\ndata: \r\n", "data: ")] +#[case::mid_event("transfer-encoding: chunked", "a\r\ndata:", "data:")] +#[case::short_of_its_content_length("content-length: 100", "data: 1\n\n", "data: 1\n\n")] +#[tokio::test] +async fn an_upstream_dying_mid_stream_fails_the_client_body_instead_of_ending_it( + #[case] framing: &'static str, + #[case] sent: &'static str, + #[case] delivered: &'static str, +) { + let upstream = raw_upstream(move |mut python| async move { + python.request_head().await.unwrap(); + python + .send(format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n{framing}\r\n\r\n{sent}" + )) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let (received, ended) = drain(response.into_body()).await; + assert_eq!(received, delivered.as_bytes()); + assert!(ended.is_err()); +} + +#[rstest] +#[tokio::test] +async fn response_headers_reach_the_client_before_python_sends_any_body() { + let release = Arc::new(Notify::new()); + let released = release.clone(); + let upstream = raw_upstream(move |mut python| { + let released = released.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream; charset=utf-8\r\nx-accel-buffering: no\r\nx-litellm-call-id: call-1\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + released.notified().await; + python.send(chunk(b"data: 1\n\n")).await; + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + + let response = + within(client().request(Request::post(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + + assert_eq!(response.headers()["x-litellm-call-id"], "call-1"); + assert_eq!(response.headers()["x-accel-buffering"], "no"); + release.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, "data: 1\n\n"); +} + +/// A connection closed without any response looks retryable to client SDKs, which would re-run a +/// request Python already ran. Python dying right after its head races the relay, hence the loop. +#[rstest] +#[tokio::test] +async fn a_head_python_sent_right_before_dying_still_reaches_the_client() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..20 { + let mut raw = gateway.connect().await; + raw.send("POST / HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + let head = within(raw.response_head()).await; + assert_eq!(head.map(|head| head.status), Some(StatusCode::OK)); + } +} + +#[rstest] +#[tokio::test] +async fn a_client_that_stops_reading_stalls_python_instead_of_being_buffered_for() { + let written = Arc::new(AtomicUsize::new(0)); + let counted = written.clone(); + let (failed, mut write_failed) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let counted = counted.clone(); + let failed = failed.clone(); + async move { + python.request_head().await.unwrap(); + let mut python = python.into_stream(); + python + .write_all(b"HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n") + .await + .unwrap(); + let piece = chunk(&[0x5a; 64 * 1024]); + while python.write_all(&piece).await.is_ok() { + counted.fetch_add(piece.len(), Ordering::SeqCst); + } + failed.send(()).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST / HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + within(raw.response_head()).await.unwrap(); + + tokio::time::sleep(Duration::from_secs(1)).await; + let stalled = written.load(Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(250)).await; + + assert_eq!(written.load(Ordering::SeqCst), stalled); + assert!(stalled < 64 * 1024 * 1024, "{stalled} bytes in flight"); + drop(raw); + within(write_failed.recv()).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn an_http10_client_receives_a_streamed_body_delimited_by_close() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send([ + &b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n"[..], + &chunk(b"data: 1\n"), + &chunk(b"data: 2\n"), + b"0\r\n\r\n", + ] + .concat()) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/stream HTTP/1.0\r\nhost: x\r\n\r\n").await; + + let response = String::from_utf8(within(raw.until_closed()).await).unwrap(); + let (head, body) = response.split_once("\r\n\r\n").unwrap(); + assert!(head.starts_with("HTTP/1.0 200 OK"), "{head}"); + assert!( + !head.to_ascii_lowercase().contains("transfer-encoding"), + "{head}" + ); + assert_eq!(body, "data: 1\ndata: 2\n"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/support/mod.rs b/litellm-rust/crates/gateway-forward/tests/support/mod.rs new file mode 100644 index 00000000000..d39532a456e --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/support/mod.rs @@ -0,0 +1,311 @@ +#![allow( + dead_code, + reason = "each test binary uses a different subset of the harness" +)] + +use std::{future::Future, net::SocketAddr, time::Duration}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::{HeaderMap, Method, StatusCode, Uri}, +}; +use http_body_util::BodyExt; +use hyper_util::{ + client::legacy::{Client, connect::HttpConnector}, + rt::TokioExecutor, +}; +use litellm_gateway_forward::Forward; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::{mpsc, oneshot}, + task::JoinHandle, +}; + +pub const WAIT: Duration = Duration::from_secs(5); +pub const UVICORN_KEEPALIVE: Duration = Duration::from_secs(5); +pub const DRAIN: Duration = Duration::from_secs(1); + +pub async fn within(future: impl Future) -> T { + tokio::time::timeout(WAIT, future).await.expect("timed out") +} + +/// A running forwarder in front of `upstream`, listening on a real socket like production. +pub struct Gateway { + pub address: SocketAddr, + stop: oneshot::Sender<()>, + server: JoinHandle<()>, +} + +impl Gateway { + pub async fn start(upstream: SocketAddr) -> Self { + Self::with_keepalive(upstream, UVICORN_KEEPALIVE).await + } + + pub async fn with_keepalive(upstream: SocketAddr, keepalive: Duration) -> Self { + let forward = Forward::new(&format!("http://{upstream}").parse().unwrap(), keepalive); + let public = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = public.local_addr().unwrap(); + let (stop, stopped) = oneshot::channel::<()>(); + let server = tokio::spawn(forward.serve( + public, + async { + stopped.await.ok(); + }, + DRAIN, + )); + Self { + address, + stop, + server, + } + } + + pub fn url(&self, path_and_query: &str) -> Uri { + format!("http://{}{path_and_query}", self.address) + .parse() + .unwrap() + } + + pub async fn connect(&self) -> Raw { + Raw::new(TcpStream::connect(self.address).await.unwrap()) + } + + /// Starts graceful shutdown and returns the server task, which ends once in-flight work drains + /// or `DRAIN` passes. + pub fn shutdown(self) -> JoinHandle<()> { + self.stop.send(()).unwrap(); + self.server + } +} + +pub fn client() -> Client { + Client::builder(TokioExecutor::new()).build_http() +} + +/// What the upstream received, after the forwarder rewrote it. +pub struct Received { + pub method: Method, + pub uri: Uri, + pub headers: HeaderMap, + pub body: Bytes, +} + +/// An axum upstream that records every request it fully reads and answers with `respond`. +pub async fn recording( + respond: impl Fn() -> axum::response::Response + Clone + Send + Sync + 'static, +) -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = serve(Router::new().fallback(move |request: Request| { + let sender = sender.clone(); + let respond = respond.clone(); + async move { + let (parts, body) = request.into_parts(); + let body = body.collect().await.unwrap().to_bytes(); + sender + .send(Received { + method: parts.method, + uri: parts.uri, + headers: parts.headers, + body, + }) + .unwrap(); + respond() + } + })) + .await; + (address, receiver) +} + +pub async fn serve(router: Router) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + address +} + +/// An upstream that hands each accepted connection to `handle` as raw bytes, for behavior a +/// well-behaved HTTP server would never produce. +pub async fn raw_upstream(handle: F) -> SocketAddr +where + F: Fn(Raw) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, +{ + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + tokio::spawn(handle(Raw::new(stream))); + } + }); + address +} + +pub fn unreachable_upstream() -> SocketAddr { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap() +} + +pub struct Head { + pub status: StatusCode, + pub headers: Vec<(String, Vec)>, + pub path: String, +} + +impl Head { + pub fn header(&self, name: &str) -> Option<&[u8]> { + self.headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_slice()) + } + + pub fn content_length(&self) -> Option { + self.header("content-length") + .map(|value| std::str::from_utf8(value).unwrap().parse().unwrap()) + } +} + +/// One TCP connection driven byte by byte, so tests control framing, timing and disconnects. +/// Bytes read past a message head stay buffered for the next read. +pub struct Raw { + stream: TcpStream, + buffer: Vec, +} + +impl Raw { + pub fn new(stream: TcpStream) -> Self { + Self { + stream, + buffer: Vec::new(), + } + } + + pub fn into_stream(self) -> TcpStream { + assert!(self.buffer.is_empty(), "buffered bytes would be lost"); + self.stream + } + + pub fn local_address(&self) -> SocketAddr { + self.stream.local_addr().unwrap() + } + + pub async fn send(&mut self, bytes: impl AsRef<[u8]>) { + self.stream.write_all(bytes.as_ref()).await.unwrap(); + } + + /// Reads one request head, or `None` when the peer closed the connection first. + pub async fn request_head(&mut self) -> Option { + self.head(|bytes| { + let mut headers = [httparse::EMPTY_HEADER; 256]; + let mut request = httparse::Request::new(&mut headers); + let httparse::Status::Complete(length) = request.parse(bytes).unwrap() else { + return None; + }; + Some(( + length, + Head { + status: StatusCode::OK, + headers: owned(request.headers), + path: request.path.unwrap().to_owned(), + }, + )) + }) + .await + } + + pub async fn response_head(&mut self) -> Option { + self.head(|bytes| { + let mut headers = [httparse::EMPTY_HEADER; 256]; + let mut response = httparse::Response::new(&mut headers); + let httparse::Status::Complete(length) = response.parse(bytes).unwrap() else { + return None; + }; + Some(( + length, + Head { + status: StatusCode::from_u16(response.code.unwrap()).unwrap(), + headers: owned(response.headers), + path: String::new(), + }, + )) + }) + .await + } + + async fn head(&mut self, parse: impl Fn(&[u8]) -> Option<(usize, Head)>) -> Option { + loop { + if let Some((length, head)) = parse(&self.buffer) { + self.buffer.drain(..length); + return Some(head); + } + if self.fill().await == 0 { + return None; + } + } + } + + /// Reads exactly `length` body bytes, using what is already buffered first. + pub async fn body(&mut self, length: usize) -> Vec { + while self.buffer.len() < length { + assert_ne!(self.fill().await, 0, "closed mid-body"); + } + self.buffer.drain(..length).collect() + } + + /// Reads a whole chunked body and returns its decoded bytes. + pub async fn chunked_body(&mut self) -> Vec { + let mut decoded = Vec::new(); + loop { + let size_line = self.line().await; + let size = + usize::from_str_radix(size_line.split(';').next().unwrap().trim(), 16).unwrap(); + if size == 0 { + while !self.line().await.is_empty() {} + return decoded; + } + decoded.extend(self.body(size).await); + assert_eq!(self.body(2).await, b"\r\n"); + } + } + + async fn line(&mut self) -> String { + loop { + if let Some(end) = self.buffer.windows(2).position(|window| window == b"\r\n") { + let line = String::from_utf8(self.buffer[..end].to_vec()).unwrap(); + self.buffer.drain(..end + 2); + return line; + } + assert_ne!(self.fill().await, 0, "closed mid-line"); + } + } + + async fn fill(&mut self) -> usize { + let mut chunk = [0; 8192]; + let read = self.stream.read(&mut chunk).await.unwrap_or(0); + self.buffer.extend_from_slice(&chunk[..read]); + read + } + + /// Everything the peer sends until it closes, including bytes already buffered. + pub async fn until_closed(&mut self) -> Vec { + while self.fill().await != 0 {} + std::mem::take(&mut self.buffer) + } + + /// Resolves once the peer closes the connection, discarding anything it still sends. + pub async fn closed(&mut self) { + let mut chunk = [0; 8192]; + while self.stream.read(&mut chunk).await.unwrap_or(0) != 0 {} + } +} + +fn owned(headers: &[httparse::Header<'_>]) -> Vec<(String, Vec)> { + headers + .iter() + .map(|header| (header.name.to_owned(), header.value.to_vec())) + .collect() +} diff --git a/litellm-rust/crates/gateway-forward/tests/upload.rs b/litellm-rust/crates/gateway-forward/tests/upload.rs new file mode 100644 index 00000000000..3dc865c9c86 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/upload.rs @@ -0,0 +1,339 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::StatusCode, + response::Response, +}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, Raw, client, raw_upstream, recording, serve, within}; +use tokio::{ + io::AsyncWriteExt, + sync::{Notify, mpsc}, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +#[rstest] +#[tokio::test] +async fn streams_upload_and_response_incrementally_without_a_body_limit() { + let upstream = serve( + Router::new().fallback(|request: Request| async { Response::new(request.into_body()) }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let (sender, receiver) = mpsc::channel::>(1); + let stream = futures_util::stream::unfold(receiver, |mut receiver| async { + receiver.recv().await.map(|chunk| (chunk, receiver)) + }); + sender.send(Ok(Bytes::from_static(b"first"))).await.unwrap(); + let mut response = within( + client().request( + Request::post(gateway.url("/v1/files")) + .body(Body::from_stream(stream)) + .unwrap(), + ), + ) + .await + .unwrap(); + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), "first"); + let large = Bytes::from(vec![b'x'; 8 * 1024 * 1024]); + sender.send(Ok(large.clone())).await.unwrap(); + drop(sender); + let rest = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(rest, large); +} + +/// hyper sends the 100 as soon as the request is forwarded, where uvicorn waits for the app to +/// read the body, so an upload Python rejects on headers alone is still sent in full. +#[rstest] +#[tokio::test] +async fn answers_expect_continue_so_large_uploads_start() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/audio/transcriptions HTTP/1.1\r\nhost: x\r\nexpect: 100-continue\r\ncontent-length: 4\r\n\r\n") + .await; + let interim = within(raw.response_head()).await.unwrap(); + assert_eq!(interim.status, StatusCode::CONTINUE); + raw.send("data").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "data"); +} + +/// curl sends `expect: 100-continue` for any body over 1 MiB and waits for the 100 before sending +/// the multipart body, which must reach Python intact even where it looks like framing. +#[rstest] +#[tokio::test] +async fn expect_continue_delivers_a_large_multipart_upload_byte_for_byte() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let boundary = "------------------------d74496d66958873e"; + let audio: Vec = (0..3 * 1024 * 1024) + .map(|index: u32| (index.wrapping_mul(2_654_435_761) >> 24) as u8) + .chain(*b"\r\n0\r\n\r\nPOST /smuggled HTTP/1.1\r\n\r\n--") + .collect(); + let body = [ + format!("--{boundary}\r\ncontent-disposition: form-data; name=\"model\"\r\n\r\nwhisper-1\r\n--{boundary}\r\ncontent-disposition: form-data; name=\"file\"; filename=\"a b.wav\"\r\ncontent-type: audio/wav\r\n\r\n").into_bytes(), + audio, + format!("\r\n--{boundary}--\r\n").into_bytes(), + ] + .concat(); + let content_type = format!("multipart/form-data; boundary={boundary}"); + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/audio/transcriptions HTTP/1.1\r\nhost: x\r\ncontent-type: {content_type}\r\nexpect: 100-continue\r\ncontent-length: {}\r\n\r\n", + body.len() + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::CONTINUE + ); + raw.send(&body).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["content-type"], content_type.as_str()); + assert_eq!(request.body, body); +} + +#[rstest] +#[case::get("GET")] +#[case::delete("DELETE")] +#[case::options("OPTIONS")] +#[case::post("POST")] +#[tokio::test] +async fn forwards_a_chunked_body_on_any_method(#[case] method: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "{method} /v1/search HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n4\r\nbody\r\n0\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "body"); +} + +#[rstest] +#[case::no_framing("", "", &[], "")] +#[case::empty_content_length("content-length: 0\r\n", "", &[("content-length", "0")], "")] +#[case::empty_chunked( + "transfer-encoding: chunked\r\n", + "0\r\n\r\n", + &[("transfer-encoding", "chunked")], + "" +)] +#[case::sized( + "content-length: 7\r\n", + "payload", + &[("content-length", "7")], + "payload" +)] +#[case::chunk_extensions_and_trailers( + "transfer-encoding: chunked\r\n", + "4;name=\"v\"\r\nbody\r\n0\r\nx-checksum: 1\r\n\r\n", + &[("transfer-encoding", "chunked")], + "body" +)] +#[tokio::test] +async fn a_body_reaches_python_with_the_framing_the_client_chose( + #[values("POST", "GET")] method: &str, + #[case] framing: &str, + #[case] wire_body: &str, + #[case] python_framing: &[(&str, &str)], + #[case] python_body: &str, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "{method} /v1/files HTTP/1.1\r\nhost: x\r\n{framing}\r\n{wire_body}" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + + let request = received.recv().await.unwrap(); + let forwarded: Vec<(&str, &str)> = ["content-length", "transfer-encoding"] + .into_iter() + .flat_map(|name| { + request + .headers + .get_all(name) + .iter() + .map(move |value| (name, value.to_str().unwrap())) + }) + .collect(); + assert_eq!(request.method, method); + assert_eq!(forwarded, python_framing); + assert_eq!(request.body, python_body); +} + +#[rstest] +#[tokio::test] +async fn a_large_upload_is_paced_by_python_reading_it() { + const LENGTH: usize = 128 * 1024 * 1024; + let start_reading = Arc::new(Notify::new()); + let reading = start_reading.clone(); + let upstream = raw_upstream(move |mut python| { + let reading = reading.clone(); + async move { + let head = python.request_head().await.unwrap(); + reading.notified().await; + let mut remaining = head.content_length().unwrap(); + while remaining > 0 { + remaining -= python.body(remaining.min(1 << 20)).await.len(); + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (read, mut write) = gateway.connect().await.into_stream().into_split(); + let sent = Arc::new(AtomicUsize::new(0)); + let sending = sent.clone(); + let upload = tokio::spawn(async move { + write + .write_all( + format!("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: {LENGTH}\r\n\r\n") + .as_bytes(), + ) + .await + .unwrap(); + let chunk = [0x5a; 64 * 1024]; + while sending.load(Ordering::SeqCst) < LENGTH { + write.write_all(&chunk).await.unwrap(); + sending.fetch_add(chunk.len(), Ordering::SeqCst); + } + write + }); + + tokio::time::sleep(Duration::from_secs(1)).await; + assert!( + sent.load(Ordering::SeqCst) < LENGTH / 4, + "the forwarder buffered an upload Python was not reading" + ); + start_reading.notify_one(); + let write = tokio::time::timeout(Duration::from_secs(60), upload) + .await + .unwrap() + .unwrap(); + let mut raw = Raw::new(read.reunite(write).unwrap()); + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); +} + +#[rstest] +#[case::between_chunks("5\r\nhello\r\n")] +#[case::mid_chunk("100\r\nhello")] +#[tokio::test] +async fn a_client_abandoning_a_chunked_upload_never_ends_python_body_cleanly( + #[case] sent: &'static str, +) { + let (forwarded, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let forwarded = forwarded.clone(); + async move { + python.request_head().await.unwrap(); + forwarded.send(python.until_closed().await).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{sent}" + )) + .await; + tokio::time::sleep(Duration::from_millis(100)).await; + drop(raw); + + let body = within(received.recv()).await.unwrap(); + assert!(body.starts_with(b"5\r\nhello\r\n")); + assert!( + !body.ends_with(b"0\r\n\r\n"), + "python saw a truncated upload end cleanly" + ); +} + +#[rstest] +#[tokio::test] +async fn an_early_answer_during_an_upload_keeps_both_connections_in_sync() { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + let length = head.content_length(); + event.send(head.path.into_bytes()).unwrap(); + let Some(length) = length else { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await; + continue; + }; + python + .send("HTTP/1.1 413 Payload Too Large\r\ncontent-length: 2\r\n\r\n{}") + .await; + event.send(python.body(length).await).unwrap(); + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let upload = b"GET /smuggled HTTP/1.1\r\nhost: x\r\n\r\n".repeat(1024); + let (first_half, second_half) = upload.split_at(upload.len() / 2); + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: {}\r\n\r\n", + upload.len() + )) + .await; + raw.send(first_half).await; + + let rejected = within(raw.response_head()).await.unwrap(); + assert_eq!(rejected.status, StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!(raw.body(2).await, b"{}"); + raw.send(second_half).await; + raw.send("GET /next HTTP/1.1\r\nhost: x\r\n\r\n").await; + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); + + assert_eq!(within(events.recv()).await.unwrap(), b"/v1/files"); + assert_eq!(within(events.recv()).await.unwrap(), upload); + assert_eq!(within(events.recv()).await.unwrap(), b"/next"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/upstream.rs b/litellm-rust/crates/gateway-forward/tests/upstream.rs new file mode 100644 index 00000000000..25e4768a54e --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/upstream.rs @@ -0,0 +1,21 @@ +use litellm_gateway_forward::Upstream; +use rstest::rstest; + +#[rstest] +#[case::https("https://127.0.0.1:4000")] +#[case::public("http://192.0.2.1:4000")] +#[case::hostname("http://localhost:4000")] +#[case::ipv6("http://[::1]:4000")] +#[case::path("http://127.0.0.1:4000/api")] +#[case::query("http://127.0.0.1:4000/?x=y")] +#[case::credentials("http://user@127.0.0.1:4000")] +fn rejects_anything_but_a_loopback_http_origin(#[case] upstream: &str) { + assert!(upstream.parse::().is_err()); +} + +#[rstest] +#[case::standard("http://127.0.0.1:4000")] +#[case::loopback_range("http://127.0.0.2:4000/")] +fn accepts_a_loopback_http_origin(#[case] upstream: &str) { + assert!(upstream.parse::().is_ok()); +} diff --git a/litellm-rust/crates/gateway-forward/tests/websocket.rs b/litellm-rust/crates/gateway-forward/tests/websocket.rs new file mode 100644 index 00000000000..ff64ddf38d3 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/websocket.rs @@ -0,0 +1,719 @@ +mod support; + +use std::{ + net::SocketAddr, + sync::Arc, + time::{Duration, Instant}, +}; + +use axum::{ + Router, + http::{HeaderValue, StatusCode}, + response::IntoResponse, +}; +use futures_util::{SinkExt, StreamExt, future::join_all}; +use litellm_gateway_forward::PEER_HEADER; +use rstest::rstest; +use support::{DRAIN, Gateway, Head, Raw, raw_upstream, serve, within}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, + sync::{Barrier, mpsc}, +}; +use tokio_tungstenite::{ + WebSocketStream, accept_hdr_async, client_async, + tungstenite::{ + Error, Message, + client::IntoClientRequest, + handshake::{ + derive_accept_key, + server::{ErrorResponse, Request as Handshake, Response as Accepted}, + }, + protocol::{CloseFrame, frame::coding::CloseCode}, + }, +}; + +const KEY: &str = "dGhlIHNhbXBsZSBub25jZQ=="; + +async fn connect( + address: SocketAddr, + connection: &'static str, +) -> Result, Box> { + let mut request = format!("ws://{address}/v1/realtime?model=m") + .into_client_request() + .unwrap(); + let headers = request.headers_mut(); + headers.insert("authorization", HeaderValue::from_static("Bearer test")); + headers.insert( + "sec-websocket-protocol", + HeaderValue::from_static("realtime"), + ); + headers.insert("connection", HeaderValue::from_static(connection)); + let stream = TcpStream::connect(address).await.unwrap(); + let (socket, response) = client_async(request, stream).await.map_err(Box::new)?; + assert_eq!(response.headers()["sec-websocket-protocol"], "realtime"); + Ok(socket) +} + +/// Greets, echoes every message, and closes with 4000 when told to "bye". +#[expect( + clippy::result_large_err, + reason = "tungstenite's handshake callback returns ErrorResponse" +)] +async fn echo() -> SocketAddr { + raw_upstream(|python| async move { + let callback = |request: &Handshake, mut response: Accepted| { + assert_eq!(request.uri(), "/v1/realtime?model=m"); + assert_eq!(request.headers()["authorization"], "Bearer test"); + assert!(request.headers().contains_key(PEER_HEADER)); + response.headers_mut().insert( + "sec-websocket-protocol", + HeaderValue::from_static("realtime"), + ); + Ok::<_, ErrorResponse>(response) + }; + let mut socket = accept_hdr_async(python.into_stream(), callback) + .await + .unwrap(); + socket.send(Message::text("greeting")).await.unwrap(); + while let Some(Ok(message)) = socket.next().await { + if message == Message::text("bye") { + socket + .close(Some(CloseFrame { + code: CloseCode::Library(4000), + reason: "done".into(), + })) + .await + .unwrap(); + continue; + } + if message.is_close() { + break; + } + socket.send(message).await.unwrap(); + } + }) + .await +} + +fn handshake(target: &str, headers: &str) -> String { + format!( + "GET {target} HTTP/1.1\r\nhost: public.example\r\nconnection: Upgrade\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: {KEY}\r\n{headers}\r\n" + ) +} + +/// The 101 uvicorn writes for `handshake`, with `headers` added. +fn switching(handshake: &Head, headers: &str) -> String { + let accept = derive_accept_key(handshake.header("sec-websocket-key").unwrap()); + format!( + "HTTP/1.1 101 Switching Protocols\r\nupgrade: websocket\r\nconnection: Upgrade\r\nsec-websocket-accept: {accept}\r\n{headers}\r\n" + ) +} + +fn values(head: &Head, name: &str) -> Vec { + head.headers + .iter() + .filter(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| String::from_utf8(value.clone()).unwrap()) + .collect() +} + +/// Accepts with `headers` in the 101, hands the handshake head to the test, then echoes bytes. +async fn byte_echo(headers: &'static str) -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, handshakes) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let sender = sender.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, headers)).await; + sender.send(handshake).unwrap(); + echo_bytes(python).await; + } + }) + .await; + (upstream, handshakes) +} + +async fn echo_bytes(python: Raw) { + let mut stream = python.into_stream(); + let (mut read, mut write) = stream.split(); + tokio::io::copy(&mut read, &mut write).await.ok(); +} + +async fn assert_echoes(socket: &mut TcpStream, message: &[u8]) { + socket.write_all(message).await.unwrap(); + let mut echoed = vec![0; message.len()]; + within(socket.read_exact(&mut echoed)).await.unwrap(); + assert_eq!(echoed, message); +} + +async fn open(gateway: &Gateway, handshake: String) -> (TcpStream, Head) { + let mut raw = gateway.connect().await; + raw.send(handshake).await; + let switched = within(raw.response_head()).await.unwrap(); + assert_eq!(switched.status, StatusCode::SWITCHING_PROTOCOLS); + (raw.into_stream(), switched) +} + +#[rstest] +#[tokio::test] +async fn relays_messages_both_ways( + #[values( + "Upgrade", + "keep-alive, Upgrade", + "upgrade", + "Upgrade, Sec-WebSocket-Key, Sec-WebSocket-Version" + )] + connection: &'static str, + #[values( + Message::text("hello"), + Message::binary(vec![0, 1, 2]), + Message::binary(vec![7; 4 * 1024 * 1024]) + )] + message: Message, +) { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, connection)).await.unwrap(); + assert_eq!( + socket.next().await.unwrap().unwrap(), + Message::text("greeting") + ); + socket.send(message.clone()).await.unwrap(); + assert_eq!(within(socket.next()).await.unwrap().unwrap(), message); + socket.close(None).await.unwrap(); + while socket.next().await.is_some() {} + drop(socket); + within(gateway.shutdown()).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn relays_a_close_python_initiates() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + socket.send(Message::text("bye")).await.unwrap(); + let Message::Close(Some(frame)) = within(socket.next()).await.unwrap().unwrap() else { + panic!("expected python's close frame"); + }; + assert_eq!(frame.code, CloseCode::Library(4000)); + assert_eq!(frame.reason, "done"); +} + +#[rstest] +#[tokio::test] +async fn passes_an_upstream_rejection_through() { + let upstream = serve( + Router::new() + .fallback(|| async { (StatusCode::FORBIDDEN, "origin not allowed").into_response() }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let error = within(connect(gateway.address, "Upgrade")) + .await + .unwrap_err(); + let Error::Http(response) = *error else { + panic!("expected the upstream rejection, got {error}"); + }; + assert_eq!(response.status(), StatusCode::FORBIDDEN); + assert_eq!(response.body().as_deref(), Some(&b"origin not allowed"[..])); +} + +/// The idle timeout only covers a connection waiting for its next HTTP request, and a realtime +/// session may sit silent between turns. +#[rstest] +#[tokio::test] +async fn keeps_a_quiet_socket_open_past_the_keepalive() { + let gateway = Gateway::with_keepalive(echo().await, Duration::from_millis(400)).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + tokio::time::sleep(Duration::from_millis(1200)).await; + socket.send(Message::text("still here")).await.unwrap(); + assert_eq!( + within(socket.next()).await.unwrap().unwrap(), + Message::text("still here") + ); +} + +/// Browsers carry the API key in a subprotocol, Python accepts with the first one it was offered, +/// and uvicorn negotiates permessage-deflate, so every handshake header must cross unchanged. +#[rstest] +#[tokio::test] +async fn handshake_headers_cross_unchanged_both_ways() { + let (upstream, mut handshakes) = byte_echo( + "sec-websocket-protocol: realtime\r\nsec-websocket-extensions: permessage-deflate; server_max_window_bits=12; client_max_window_bits=12\r\nset-cookie: a=1\r\nset-cookie: b=2\r\nx-litellm-call-id: abc\r\n", + ) + .await; + let gateway = Gateway::start(upstream).await; + let target = "/v1/realtime?model=gpt-realtime&api-key=sk%2Fa%2Bb%3D&dup=1&dup=2"; + let (_socket, switched) = open( + &gateway, + handshake( + target, + "sec-websocket-protocol: realtime, openai-insecure-api-key.sk-websocket-test, openai-beta.realtime-v1\r\nsec-websocket-protocol: second-line\r\nsec-websocket-extensions: permessage-deflate; client_max_window_bits\r\ncookie: token=abc\r\n", + ), + ) + .await; + let handshake = within(handshakes.recv()).await.unwrap(); + + assert_eq!(handshake.path, target); + assert_eq!( + values(&handshake, "sec-websocket-protocol"), + [ + "realtime, openai-insecure-api-key.sk-websocket-test, openai-beta.realtime-v1", + "second-line" + ] + ); + assert_eq!( + values(&handshake, "sec-websocket-extensions"), + ["permessage-deflate; client_max_window_bits"] + ); + assert_eq!(values(&handshake, "cookie"), ["token=abc"]); + assert_eq!(values(&switched, "sec-websocket-protocol"), ["realtime"]); + assert_eq!( + values(&switched, "sec-websocket-extensions"), + ["permessage-deflate; server_max_window_bits=12; client_max_window_bits=12"] + ); + assert_eq!(values(&switched, "set-cookie"), ["a=1", "b=2"]); + assert_eq!(values(&switched, "x-litellm-call-id"), ["abc"]); +} + +#[rstest] +#[tokio::test] +async fn invents_no_subprotocol_or_extension() { + let (upstream, mut handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (_socket, switched) = open(&gateway, handshake("/v1/realtime", "")).await; + let handshake = within(handshakes.recv()).await.unwrap(); + for name in ["sec-websocket-protocol", "sec-websocket-extensions"] { + assert!(handshake.header(name).is_none()); + assert!(switched.header(name).is_none()); + } +} + +#[rstest] +#[case::capitalized_upgrade("Upgrade: WebSocket\r\nConnection: keep-alive, Upgrade\r\n")] +#[case::uppercase("UPGRADE: WEBSOCKET\r\nCONNECTION: UPGRADE\r\n")] +#[case::split_connection("Upgrade: websocket\r\nConnection: keep-alive\r\nConnection: Upgrade\r\n")] +#[tokio::test] +async fn accepts_any_spelling_of_the_upgrade_headers(#[case] upgrade: &str) { + let (upstream, _handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open( + &gateway, + format!( + "GET /v1/realtime HTTP/1.1\r\nhost: x\r\n{upgrade}sec-websocket-version: 13\r\nsec-websocket-key: {KEY}\r\n\r\n" + ), + ) + .await; + socket.write_all(b"ping").await.unwrap(); + let mut echoed = [0; 4]; + within(socket.read_exact(&mut echoed)).await.unwrap(); + assert_eq!(&echoed, b"ping"); +} + +/// Compressed frames set RSV1 and every client frame is masked, so the socket is relayed as bytes, +/// never re-framed. +#[rstest] +#[tokio::test] +async fn relays_socket_bytes_exactly() { + let (upstream, _handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let sent: Vec = (0..3_000_017u32) + .map(|index| index.wrapping_mul(2_654_435_761).to_be_bytes()[1]) + .collect(); + let (mut read, mut write) = socket.into_split(); + let writing = sent.clone(); + let writer = tokio::spawn(async move { + for chunk in writing.chunks(7919) { + write.write_all(chunk).await.unwrap(); + } + write + }); + let mut echoed = vec![0; sent.len()]; + within(read.read_exact(&mut echoed)).await.unwrap(); + assert!(echoed == sent); + drop(writer.await.unwrap()); +} + +/// uvicorn writes the 101, an error event and a close frame back to back and then closes, so they +/// share a segment with the FIN, and a client may send its first frame with the handshake. +#[rstest] +#[tokio::test] +async fn keeps_bytes_that_arrive_with_the_handshake() { + let (early, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let early = early.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + let mut segment = switching(&handshake, "").into_bytes(); + segment.extend(b"\x81\x05error\x88\x08\x03\xf0denied"); + python.send(segment).await; + early.send(within(python.body(6)).await).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(handshake("/v1/realtime", "") + "early!").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::SWITCHING_PROTOCOLS + ); + assert_eq!(within(received.recv()).await.unwrap(), b"early!"); + assert_eq!( + within(raw.until_closed()).await, + b"\x81\x05error\x88\x08\x03\xf0denied" + ); +} + +/// Python frees a realtime session's budget and parallel slot when its socket closes, so a client +/// that disappears without a close frame must still close Python's socket. +#[rstest] +#[case::fin(false)] +#[case::reset(true)] +#[tokio::test] +async fn a_vanished_client_closes_the_python_socket(#[case] reset: bool) { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + event.send("open").unwrap(); + python.closed().await; + event.send("closed").unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_eq!(within(events.recv()).await, Some("open")); + if reset { + #[expect(deprecated, reason = "a zero linger is how a test sends a TCP reset")] + socket.set_linger(Some(Duration::ZERO)).unwrap(); + } + drop(socket); + assert_eq!(within(events.recv()).await, Some("closed")); +} + +#[rstest] +#[case::fin(false)] +#[case::reset(true)] +#[tokio::test] +async fn a_vanished_python_ends_the_client_socket(#[case] reset: bool) { + let upstream = raw_upstream(move |mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + let socket = python.into_stream(); + if reset { + #[expect(deprecated, reason = "a zero linger is how a test sends a TCP reset")] + socket.set_linger(Some(Duration::ZERO)).unwrap(); + } + drop(socket); + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let mut rest = Vec::new(); + assert!(matches!( + within(socket.read_to_end(&mut rest)).await, + Ok(0) | Err(_) + )); +} + +#[rstest] +#[tokio::test] +async fn shutdown_closes_open_sockets_after_drain() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + let started = Instant::now(); + within(gateway.shutdown()).await.unwrap(); + assert!(started.elapsed() >= DRAIN); + assert!(!matches!(socket.next().await, Some(Ok(Message::Text(_))))); +} + +#[rstest] +#[tokio::test] +async fn shutdown_relays_the_close_python_sends_while_draining() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + socket.send(Message::text("bye")).await.unwrap(); + let Message::Close(Some(frame)) = within(socket.next()).await.unwrap().unwrap() else { + panic!("expected python's close frame"); + }; + assert_eq!(frame.code, CloseCode::Library(4000)); + while socket.next().await.is_some() {} + drop(socket); + let started = Instant::now(); + within(server).await.unwrap(); + assert!(started.elapsed() < DRAIN / 2); +} + +#[rstest] +#[tokio::test] +async fn refuses_new_sockets_while_draining() { + let gateway = Gateway::start(echo().await).await; + let mut open = within(connect(gateway.address, "Upgrade")).await.unwrap(); + open.next().await.unwrap().unwrap(); + let address = gateway.address; + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!( + TcpStream::connect(address).await.is_err() || connect(address, "Upgrade").await.is_err() + ); + open.send(Message::text("still open")).await.unwrap(); + assert_eq!( + within(open.next()).await.unwrap().unwrap(), + Message::text("still open") + ); + drop(open); + within(server).await.unwrap(); +} + +/// uvicorn never half-closes: once Python closes the socket nothing more can cross it, so the +/// relay must end instead of waiting for the client to close its side too. +#[rstest] +#[tokio::test] +async fn relay_ends_when_python_closes_while_the_client_holds_its_socket() { + let upstream = raw_upstream(move |mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_eq!(within(socket.read(&mut [0; 1])).await.unwrap(), 0); + let started = Instant::now(); + within(gateway.shutdown()).await.unwrap(); + assert!(started.elapsed() < DRAIN / 2); + drop(socket); +} + +/// uvicorn answers a handshake Python closes before accepting with a 403, and one Python denies +/// through `websocket.http.response` with any status, headers and body, then drops the socket. The +/// client gets that response whole and keeps a connection that still reaches Python. +#[rstest] +#[case::closed_before_accept("403 Forbidden", None, "")] +#[case::denial_response( + "401 Unauthorized", + Some("Bearer realm=\"litellm\""), + "{\"error\":\"invalid api key\"}" +)] +#[tokio::test] +async fn a_denied_handshake_reaches_the_client_whole( + #[case] status: &'static str, + #[case] authenticate: Option<&'static str>, + #[case] body: &'static str, +) { + let upstream = raw_upstream(move |mut python: Raw| async move { + python.request_head().await.unwrap(); + let authenticate = authenticate + .map(|value| format!("www-authenticate: {value}\r\n")) + .unwrap_or_default(); + python + .send(format!( + "HTTP/1.1 {status}\r\ndate: Tue, 03 Nov 2026 08:00:00 GMT\r\nconnection: close\r\ncontent-length: {}\r\ncontent-type: text/plain; charset=utf-8\r\n{authenticate}\r\n{body}", + body.len() + )) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + for _ in 0..2 { + raw.send(handshake("/v1/realtime", "")).await; + let denied = within(raw.response_head()).await.unwrap(); + assert_eq!(denied.status.as_str(), &status[..3]); + assert_eq!( + values(&denied, "www-authenticate"), + Vec::from_iter(authenticate.map(String::from)) + ); + assert!(denied.header("sec-websocket-accept").is_none()); + let length = denied.content_length().unwrap(); + assert_eq!(within(raw.body(length)).await, body.as_bytes()); + } +} + +/// Only Python's EOF ends a relay, so a client that half-closes after its close frame still gets +/// everything Python sends before closing its own side. +#[rstest] +#[tokio::test] +async fn a_client_half_close_still_receives_what_python_sends_after_it() { + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + let received = python.until_closed().await; + python.send(received).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let close = b"\x88\x82\x01\x02\x03\x04\x02\xea"; + socket.write_all(close).await.unwrap(); + socket.shutdown().await.unwrap(); + let mut replied = Vec::new(); + within(socket.read_to_end(&mut replied)).await.unwrap(); + assert_eq!(replied, close); +} + +/// Python authenticates every handshake before accepting it, so sessions opened together must all +/// reach Python before any of them is accepted, and each must stay on its own socket. +#[rstest] +#[tokio::test] +async fn opens_many_sockets_at_once_without_crossing_them() { + const SOCKETS: usize = 64; + let everyone = Arc::new(Barrier::new(SOCKETS)); + let upstream = raw_upstream(move |mut python: Raw| { + let everyone = everyone.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + everyone.wait().await; + python.send(switching(&handshake, "")).await; + echo_bytes(python).await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let gateway = &gateway; + within(join_all((0..SOCKETS).map(|index| async move { + let (mut socket, _) = open(gateway, handshake("/v1/realtime", "")).await; + assert_echoes(&mut socket, format!("socket {index}").as_bytes()).await; + }))) + .await; +} + +/// Python resolves the key, budget and model before it accepts, and nothing bounds how long that +/// takes, so the idle timeout must not cut a handshake waiting on its 101. +#[rstest] +#[tokio::test] +async fn waits_for_a_101_slower_than_the_keepalive() { + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + tokio::time::sleep(Duration::from_millis(1200)).await; + python.send(switching(&handshake, "")).await; + echo_bytes(python).await; + }) + .await; + let gateway = Gateway::with_keepalive(upstream, Duration::from_millis(400)).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_echoes(&mut socket, b"late").await; +} + +/// A 101 the client did not ask for, or to a protocol other than the one it asked for, would hand +/// the client a socket it cannot speak, so it is a 502 and Python's socket is closed. +#[rstest] +#[case::plain_request("GET /v1/models HTTP/1.1\r\nhost: x\r\n\r\n".to_owned(), "websocket")] +#[case::other_protocol(handshake("/v1/realtime", ""), "h2c")] +#[tokio::test] +async fn an_unrequested_switch_is_a_502_that_closes_python( + #[case] request: String, + #[case] protocol: &'static str, +) { + let (closed, mut closes) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let closed = closed.clone(); + async move { + python.request_head().await.unwrap(); + python + .send(format!( + "HTTP/1.1 101 Switching Protocols\r\nconnection: upgrade\r\nupgrade: {protocol}\r\n\r\n" + )) + .await; + python.closed().await; + closed.send(()).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(request).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_GATEWAY + ); + within(closes.recv()).await.unwrap(); +} + +/// A realtime session streams audio for as long as the client reads it, so a client that stops +/// reading must stall Python's writes instead of piling them up in the gateway. +#[rstest] +#[tokio::test] +async fn a_client_that_stops_reading_stalls_python() { + let (sent, mut totals) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let sent = sent.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + python.body(2).await; + let mut stream = python.into_stream(); + let chunk = vec![0x82; 64 * 1024]; + let mut total = 0; + while let Ok(Ok(())) = + tokio::time::timeout(Duration::from_millis(300), stream.write_all(&chunk)).await + { + total += chunk.len(); + } + sent.send(total).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + socket.write_all(b"go").await.unwrap(); + let total = within(totals.recv()).await.unwrap(); + assert!( + total < 64 << 20, + "python wrote {total} bytes into a stalled socket" + ); +} + +/// Python closes a session that is still streaming audio by sending an error event and a close +/// frame and dropping its socket with the client's next frames unread, which resets it. Those last +/// bytes reached the client when it talked to uvicorn directly, so they must reach it through the +/// relay too, even though writing the client's frames to Python now fails. +#[rstest] +#[tokio::test] +async fn python_closing_mid_stream_still_delivers_its_last_frames() { + const LAST: &[u8] = b"\x81\x05error\x88\x02\x03\xf0"; + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + python.body(1).await; + tokio::time::sleep(Duration::from_millis(20)).await; + python.send(LAST).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..20 { + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let (mut read, mut write) = socket.split(); + let streaming = async { + while write.write_all(b"\x82\x80\0\0\0\0").await.is_ok() { + tokio::time::sleep(Duration::from_micros(200)).await; + } + std::future::pending().await + }; + let mut received = Vec::new(); + let receiving = async { + let mut chunk = [0; 4096]; + while let Ok(read @ 1..) = read.read(&mut chunk).await { + received.extend_from_slice(&chunk[..read]); + } + }; + within(async { + tokio::select! { + () = streaming => {} + () = receiving => {} + } + }) + .await; + assert_eq!(received, LAST); + } +} diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 6eeaa43e013..d0cedbb3b36 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -12,6 +12,7 @@ envy = "0.4.2" http-body-util = "0.1" litellm-inference.workspace = true litellm-gateway-inference.workspace = true +litellm-gateway-forward.workspace = true litellm-gateway-mcp.workspace = true tokio-util = "0.7" thiserror.workspace = true diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 47c1fe83f50..ff528723095 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -10,7 +10,7 @@ use serde_json::json; mod settings; -use settings::{Settings, UiSettings}; +use settings::{LogSettings, Settings, UiSettings}; struct StderrSink { level: Level, @@ -41,6 +41,15 @@ impl Sink for StderrSink { #[tokio::main] async fn main() -> Result<(), Box> { + if let Some(forward) = litellm_gateway_forward::Settings::from_env()? { + let log: LogSettings = settings::from_iter(std::env::vars_os())?; + Logger::new(StderrSink { + level: log.log_level(), + }) + .install_global()?; + litellm_gateway_forward::run(forward, shutdown_signal()).await?; + return Ok(()); + } let settings: Settings = settings::from_iter(std::env::vars_os())?; let level = settings.log_level(); Logger::new(StderrSink { level }).install_global()?; @@ -126,7 +135,20 @@ async fn shutdown_signal() { } } } - #[cfg(not(unix))] + // Windows support exists for local CLI runs only, the proxy is never deployed on Windows. The supervisor + // stops the gateway with CTRL_BREAK_EVENT, which has to drain the same way SIGTERM does. + #[cfg(windows)] + { + match tokio::signal::windows::ctrl_break() { + Ok(mut ctrl_break) => { + tokio::select! { _ = tokio::signal::ctrl_c() => (), _ = ctrl_break.recv() => () } + } + Err(_) => { + let _ = tokio::signal::ctrl_c().await; + } + } + } + #[cfg(not(any(unix, windows)))] { let _ = tokio::signal::ctrl_c().await; } diff --git a/litellm-rust/crates/gateway/src/settings.rs b/litellm-rust/crates/gateway/src/settings.rs index 445071504ed..c2c8bc30934 100644 --- a/litellm-rust/crates/gateway/src/settings.rs +++ b/litellm-rust/crates/gateway/src/settings.rs @@ -4,6 +4,23 @@ use litellm_secret_value::SecretValue; use litellm_tracing::Level; use serde::{Deserialize, de::DeserializeOwned}; +#[derive(Deserialize)] +pub(super) struct LogSettings { + rust_log: Option, +} + +impl LogSettings { + pub fn log_level(&self) -> Level { + log_level(self.rust_log.as_deref()) + } +} + +fn log_level(rust_log: Option<&str>) -> Level { + rust_log + .and_then(|value| value.parse().ok()) + .unwrap_or(Level::INFO) +} + #[derive(Deserialize)] #[serde(default)] pub(super) struct Settings { @@ -26,10 +43,7 @@ impl Default for Settings { impl Settings { pub fn log_level(&self) -> Level { - self.rust_log - .as_deref() - .and_then(|value| value.parse().ok()) - .unwrap_or(Level::INFO) + log_level(self.rust_log.as_deref()) } } diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs index c41f7e74d02..c60030d7ea6 100644 --- a/litellm-rust/crates/router-core/src/call.rs +++ b/litellm-rust/crates/router-core/src/call.rs @@ -4,13 +4,7 @@ //! `should_retry_this_error` and `_time_to_sleep_before_retry` in `litellm/router.py`, and //! `run_async_fallback` in `router_utils/fallback_event_handlers.py`. -use std::{ - borrow::Cow, - future::Future, - pin::Pin, - sync::Arc, - time::{Duration, Instant}, -}; +use std::{borrow::Cow, future::Future, pin::Pin, sync::Arc, time::Duration}; use serde_json::{Map, Value}; @@ -1247,7 +1241,7 @@ impl Core { .unrouted_attempt(snapshot, op, request, target, (retry, max_retries), at) .await; } - let pick_started = Instant::now(); + let pick_started = self.health.now_time(); let (deployment, candidates) = match self .picked(snapshot, op, request, target, inputs, retry, at) .await @@ -1264,7 +1258,11 @@ impl Core { self.host.observe(&RouterEvent::Selected { model_group: candidates.model_group.clone(), deployment_id: deployment.id.clone(), - pick: pick_started.elapsed(), + pick: self + .health + .now_time() + .duration_since(pick_started) + .unwrap_or_default(), }); if !self.host.supports(op, deployment.provider()) { return Err(Error::Config(ConfigError::Unsupported(Unsupported::Op { @@ -1290,7 +1288,7 @@ impl Core { }, None, snapshot.settings.dynamic.default_ref(), - jiff::Timestamp::now(), + self.health.timestamp(), ), Ok(None) ); diff --git a/litellm-rust/crates/router-core/src/health.rs b/litellm-rust/crates/router-core/src/health.rs index c8055eafe31..eaa61dac7ba 100644 --- a/litellm-rust/crates/router-core/src/health.rs +++ b/litellm-rust/crates/router-core/src/health.rs @@ -247,6 +247,10 @@ impl Health { self.clock.now() } + pub(crate) fn timestamp(&self) -> jiff::Timestamp { + jiff::Timestamp::try_from(self.clock.now()).unwrap_or(jiff::Timestamp::UNIX_EPOCH) + } + fn now(&self) -> f64 { unix_seconds(self.clock.now()) } diff --git a/litellm-rust/crates/router-core/src/lookup.rs b/litellm-rust/crates/router-core/src/lookup.rs index 39c235d6f7b..8fef0ba222f 100644 --- a/litellm-rust/crates/router-core/src/lookup.rs +++ b/litellm-rust/crates/router-core/src/lookup.rs @@ -433,6 +433,10 @@ impl<'a> Lookup<'a> { .collect() } + pub fn has_team_wildcards(&self, team_id: &str) -> bool { + self.host.team_patterns.iter().any(|(t, _)| t == team_id) + } + /// The deployments `team_id`'s own wildcard pattern routes `name` to /// (`_core_team_pattern_routers`, keyed by `team_public_model_name`). pub fn team_wildcard(&self, name: &str, team_id: &str) -> Vec> { diff --git a/litellm-rust/crates/router-core/src/resolve.rs b/litellm-rust/crates/router-core/src/resolve.rs index 9075fc36c53..d897716b1e5 100644 --- a/litellm-rust/crates/router-core/src/resolve.rs +++ b/litellm-rust/crates/router-core/src/resolve.rs @@ -223,14 +223,20 @@ pub(crate) fn wildcard_name<'a>( registry: &'a Registry, name: &'a str, specific_deployment: bool, + team_id: Option<&str>, ) -> Option<&'a str> { let model = registry.alias_of(name).unwrap_or(name); - (registry.has_wildcards() + let lookup = crate::lookup::Lookup::new(registry); + let team_patterns = team_id.is_some_and(|team| lookup.has_team_wildcards(team)); + let team_misses = || team_id.is_none_or(|team| lookup.team_wildcard(model, team).is_empty()); + let global_misses = || registry.has_wildcards() && registry.wildcard_for(model).is_none(); + ((team_patterns || registry.has_wildcards()) && !specific_deployment && !registry.has_name(model) && (registry.has_name(name) || registry.deployment(name).is_none()) && registry.routing_group_named(model).is_none() - && registry.wildcard_for(model).is_none()) + && team_misses() + && (team_patterns || global_misses())) .then_some(model) } @@ -339,15 +345,21 @@ pub fn resolve_request( let kept = reserved(model, deployments, caller, now)?; return resolved(ResolveKind::Group, model, kept); } - // `_core_wildcard_route_deployments`: the team's own pattern, else the proxy-wide one. - let team_owned = team_id - .map(|team| crate::lookup::Lookup::new(registry).team_wildcard(model, team)) - .filter(|deployments| !deployments.is_empty()); - if let Some(deployments) = team_owned { + // `_core_wildcard_route_deployments`: the team's own pattern, else the proxy-wide one, + // each tried as spelled and then under the name's provider (`get_pattern`). + let team_owned = team_id.and_then(|team| { + let lookup = crate::lookup::Lookup::new(registry); + [Some(model), caller.wildcard_request] + .into_iter() + .flatten() + .map(|request| (request, lookup.team_wildcard(request, team))) + .find(|(_, deployments)| !deployments.is_empty()) + }); + if let Some((request, deployments)) = team_owned { let routed = deployments.iter().flat_map(|deployment| { let public = deployment.model_info.team_public_model_name.as_deref(); - let pattern = public.unwrap_or(model); - routed_wildcard(pattern, model, std::slice::from_ref(deployment)) + let pattern = public.unwrap_or(request); + routed_wildcard(pattern, request, std::slice::from_ref(deployment)) }); let kept = reserved(model, routed.collect(), caller, now)?; return resolved(ResolveKind::Wildcard, model, kept); @@ -428,6 +440,7 @@ pub fn resolve_request( mod tests { use std::collections::BTreeMap; + use rstest::rstest; use serde_json::{Value, json}; use super::*; @@ -544,8 +557,11 @@ mod tests { #[test] fn a_name_no_pattern_matches_is_tried_under_its_provider() { let registry = registry(wildcard_group()); - assert_eq!(wildcard_name(®istry, "gpt-4o", false), Some("gpt-4o")); - assert_eq!(wildcard_name(®istry, "openai/gpt-4o", false), None); + assert_eq!( + wildcard_name(®istry, "gpt-4o", false, None), + Some("gpt-4o") + ); + assert_eq!(wildcard_name(®istry, "openai/gpt-4o", false, None), None); let caller = Caller { wildcard_request: Some("openai/gpt-4o"), ..Caller::default() @@ -564,6 +580,48 @@ mod tests { assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o"); } + #[rstest] + fn a_bare_name_is_tried_under_its_provider_against_the_teams_own_wildcard() { + let registry = registry(json!([ + {"model_name": "model_name_tw_1", "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "tw", "team_id": "TW", "team_public_model_name": "openai/*"}}, + {"model_name": "g", "litellm_params": {"model": "openai/g"}, "model_info": {"id": "g"}}, + ])); + // No proxy-wide wildcard: only the team's patterns call for the provider spelling. + assert_eq!( + wildcard_name(®istry, "gpt-4o-mini", false, Some("TW")), + Some("gpt-4o-mini") + ); + assert_eq!(wildcard_name(®istry, "gpt-4o-mini", false, None), None); + assert_eq!( + wildcard_name(®istry, "openai/gpt-4o-mini", false, Some("TW")), + None + ); + let resolve = |team_id| { + let caller = Caller { + team_id, + wildcard_request: Some("openai/gpt-4o-mini"), + ..Caller::default() + }; + resolve_request( + ®istry, + "gpt-4o-mini", + false, + caller, + None, + None, + Timestamp::now(), + ) + .unwrap() + }; + let resolution = resolve(Some("TW")).unwrap(); + assert_eq!(resolution.kind, ResolveKind::Wildcard); + assert_eq!(resolution.group, "gpt-4o-mini"); + assert_eq!(ids(Some(resolution.clone())), ["tw"]); + assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o-mini"); + assert!(resolve(Some("OTHER")).is_none()); + } + #[test] fn a_deployment_named_by_id_refuses_the_endpoint_like_any_other() { let registry = registry(json!([ diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index 1972b2cbd90..c7f3b1f9c31 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -281,8 +281,12 @@ impl Core { // No deployment for the name: route to the first default fallback. let default_target = crate::resolve::default_target(&snapshot.settings.dynamic); // `get_pattern`: a name no pattern matches is tried again under its inferred provider. - let wildcard_name = - crate::resolve::wildcard_name(registry, name, request.specific_deployment); + let wildcard_name = crate::resolve::wildcard_name( + registry, + name, + request.specific_deployment, + request.team_id.as_deref(), + ); let wildcard_request = match wildcard_name { Some(model) => { let scope = StrategyScope::Default; @@ -308,7 +312,7 @@ impl Core { caller, default_target, snapshot.settings.dynamic.default_ref(), - jiff::Timestamp::now(), + self.health.timestamp(), )? .ok_or_else(|| Error::UnknownModel { // The name `_common_checks_available_deployment` looked for, past the alias. diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index 582dfc78ebf..ae416abf3a9 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -1808,3 +1808,57 @@ async fn a_listed_pass_through_pick_names_the_group_an_alias_resolves_to(config: Error::NoPassThroughDeployments { model } if model == "primary" )); } + +struct StoppedClock(u64); + +impl litellm_router_core::health::Clock for StoppedClock { + fn now(&self) -> std::time::SystemTime { + std::time::UNIX_EPOCH + Duration::from_secs(self.0) + } +} + +fn utc_hour(hour: u64) -> u64 { + 1_767_225_600 + hour * 3_600 +} + +#[rstest] +#[case::inside_the_window(utc_hour(10), vec!["open"])] +#[case::outside_the_window(utc_hour(20), vec!["open", "reserved"])] +#[tokio::test] +async fn access_windows_follow_the_routers_clock(#[case] now: u64, #[case] expected: Vec<&str>) { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "chat", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "open"}}, + {"model_name": "chat", "litellm_params": {"model": "openai/b"}, "model_info": { + "id": "reserved", + "access_windows": [{"start": "09:00", "end": "17:00", "timezone": "UTC", "team_ids": ["batch"]}], + }}, + ] + })) + .unwrap(); + let router = Router::builder( + config, + RouterDynamicConfig::default(), + Arc::new(ScriptedHost::default()), + ) + .randomness(Arc::new(SeededRandomness::new(7))) + .clock(Arc::new(StoppedClock(now))) + .build() + .unwrap(); + let mut picked = std::collections::BTreeSet::new(); + for _ in 0..40 { + let deployment = router + .select(&Op::from("acompletion"), &request_for("chat")) + .await + .unwrap(); + picked.insert(deployment.id.clone()); + } + assert_eq!(picked.into_iter().collect::>(), expected); +} + +fn request_for(model: &str) -> RouterRequest { + RouterRequest { + model: model.to_owned(), + ..RouterRequest::default() + } +} diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 0579eb15baf..385522c2435 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -898,6 +898,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def redacts_messages_itself(self) -> bool: return False + def exempt_from_message_redaction(self) -> bool: + """Router pre-call checks that key routing off the request and never log it return True, so + `turn_off_message_logging` leaves the success event they receive unredacted.""" + return False + def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict: """ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it. diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 0b8c128286a..d3cc04408b6 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -71,12 +71,12 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj _TERMINAL_STATUSES: Final = frozenset( - {"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action"} + {"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action", "requires_continuation"} ) _POLLABLE_STATUSES: Final = frozenset({"in_progress", "queued"}) -_STATUSES_THAT_PRODUCED_OUTPUT: Final = frozenset({"completed", "requires_action"}) +_STATUSES_THAT_PRODUCED_OUTPUT: Final = frozenset({"completed", "requires_action", "requires_continuation"}) SettlementOutcome: TypeAlias = Literal["billed", "released", "unsettled"] diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index acd43c9b830..cd144646f79 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -95,6 +95,7 @@ async def acreate( response_mime_type: str | None = None, # Continuation previous_interaction_id: str | None = None, + continuation_token: str | None = None, # Extra params extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, @@ -129,6 +130,7 @@ async def acreate( response_format: JSON schema for response format response_mime_type: MIME type of the response previous_interaction_id: ID of previous interaction for continuation + continuation_token: Token from a `requires_continuation` interaction to resume its long decode extra_headers: Additional headers extra_body: Additional body parameters timeout: Request timeout @@ -164,6 +166,7 @@ async def acreate( response_format=response_format, response_mime_type=response_mime_type, previous_interaction_id=previous_interaction_id, + continuation_token=continuation_token, extra_headers=extra_headers, extra_body=extra_body, timeout=timeout, @@ -224,6 +227,7 @@ def create( response_mime_type: str | None = None, # Continuation previous_interaction_id: str | None = None, + continuation_token: str | None = None, # Extra params extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, @@ -262,6 +266,7 @@ def create( response_format: JSON schema for response format response_mime_type: MIME type of the response previous_interaction_id: ID of previous interaction for continuation + continuation_token: Token from a `requires_continuation` interaction to resume its long decode extra_headers: Additional headers extra_body: Additional body parameters timeout: Request timeout diff --git a/litellm/interactions/utils.py b/litellm/interactions/utils.py index 3895a85061d..1d7465a99c1 100644 --- a/litellm/interactions/utils.py +++ b/litellm/interactions/utils.py @@ -20,6 +20,7 @@ INTERACTIONS_API_OPTIONAL_PARAMS: Final = { "response_format", "response_mime_type", "previous_interaction_id", + "continuation_token", "agent_config", } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9c49efc529f..9b6ec8d28d6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4405,7 +4405,11 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(callback, CustomLogger): # custom logger class from litellm.integrations.shadow_eval_logger import ShadowEvalLogger - model_call_details: dict = self.model_call_details + model_call_details: dict = ( + self._model_call_details_with_request_messages() + if callback.exempt_from_message_redaction() + else self.model_call_details + ) ################################## # call redaction hook for custom logger model_call_details = callback.redact_standard_logging_payload_from_model_call_details( @@ -4969,6 +4973,21 @@ class Logging(LiteLLMLoggingBaseClass): ) -> None: executor.submit(self.success_handler, result, start_time, end_time, cache_hit) + def _model_call_details_with_request_messages(self) -> dict[str, object]: # mutable-ok: loggers take a dict + """The success-event payload with the request messages put back where `turn_off_message_logging` + redacted them in place, for a callback that routes on the request and never logs it.""" + request_messages: Final = cast(object, self.messages) # cast-ok: the request messages are only put back as-is + standard_logging_object: Final = self.model_call_details.get("standard_logging_object") + return { + **self.model_call_details, + "messages": request_messages, + **( + {"standard_logging_object": {**standard_logging_object, "messages": request_messages}} + if isinstance(standard_logging_object, dict) + else {} + ), + } + def _get_completed_success_callbacks( self, *, asynchronous: bool ) -> tuple[str | Callable[..., object] | CustomLogger, ...]: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 4c3907dc52f..eda98ae2c67 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5879,7 +5879,7 @@ def _bedrock_tools_pt( description = name defs = parameters.pop("$defs", {}) - defs_copy = copy.deepcopy(defs) if isinstance(defs, dict) else {} + defs_copy = defs if isinstance(defs, dict) else {} # Expand $ref references in parameters using the definitions # Note: We don't pre-flatten defs as that causes exponential memory growth # with circular references (see issue #19098). unpack_defs handles nested diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 4f890ed922f..ac4addb0847 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -827,11 +827,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return cls._raise_invalid_reasoning_effort(model, value, llm_provider) def get_supported_openai_params(self, model: str): + sampling_params: Final = ("temperature", "top_p") if self._supports_sampling_params(model) else () params: Final = [ "stream", "stop", - "temperature", - "top_p", + *sampling_params, "max_tokens", "max_completion_tokens", "tools", diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index 1895a97af7d..50da73cc8b9 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -126,6 +126,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): "response_format", "response_mime_type", "previous_interaction_id", + "continuation_token", ] def validate_environment( @@ -206,6 +207,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): "environment", "response_modalities", "previous_interaction_id", + "continuation_token", ] dropped_keys: Final = ( frozenset({"agent_config"}) @@ -301,7 +303,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): ) -> InteractionsAPIStreamingResponse: """Parse streaming chunk.""" verbose_logger.debug("Google AI Interactions streaming chunk: %s", parsed_chunk) - return InteractionsAPIStreamingResponse(**parsed_chunk) + return InteractionsAPIStreamingResponse.model_validate(parsed_chunk) # GET / DELETE / CANCEL - just build URLs, responses match spec directly diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a5b9a5cb457..a29f97ee9ea 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4046,6 +4046,8 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 2.5e-07, + "deprecation_date": "2027-09-01", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "litellm_provider": "azure_ai", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -4068,7 +4070,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, "prompt_cache_min_tokens": 512, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models" + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" }, "azure_ai/claude-opus-5": { "deprecation_date": "2027-07-08", @@ -13458,13 +13460,16 @@ "supports_reasoning": true }, "azure_ai/mistral-document-ai-2512": { + "deprecation_date": "2027-02-15", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.003, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" ], - "source": "https://ai.azure.com/catalog/models/mistral-document-ai-2512" + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", + "successor_model": "azure_ai/mistral-ocr-4-0" }, "azure_ai/mistral-ocr-4-0": { "litellm_provider": "azure_ai", @@ -40893,6 +40898,22 @@ "output_cost_per_token": 4e-07, "supports_tool_choice": true }, + "jina-reranker-v3.5": { + "input_cost_per_token": 5e-08, + "litellm_provider": "jina_ai", + "max_input_tokens": 131072, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" + }, + "jina-reranker-m0": { + "input_cost_per_token": 5e-08, + "litellm_provider": "jina_ai", + "max_input_tokens": 10240, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" + }, "jina-reranker-v2-base-multilingual": { "input_cost_per_token": 5e-08, "litellm_provider": "jina_ai", @@ -43771,7 +43792,8 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/zai-glm-5-3": { "cache_read_input_token_cost": 1.4e-07, @@ -45270,6 +45292,18 @@ "source": "https://tokenfactory.nebius.com/api/public/models_info", "cache_read_input_token_cost": 1.5e-07 }, + "nebius/moonshotai/Kimi-K3.1-Pro": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "nebius", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "supports_prompt_caching": true, + "source": "https://tokenfactory.nebius.com/api/public/models_info" + }, "nebius/moonshotai/Kimi-K3": { "max_input_tokens": 1048576, "max_output_tokens": 1024000, @@ -71189,6 +71223,8 @@ }, "gpt-5.5-cyber-preview": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-02-01", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "input_cost_per_token": 1.25e-05, "litellm_provider": "openai", "mode": "chat", @@ -72148,6 +72184,7 @@ "output_cost_per_token": 4.5e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_audio_output": true, "tpm": 250000, "rpm": 10, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, @@ -72170,6 +72207,7 @@ "output_cost_per_token": 4.5e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_audio_output": true, "tpm": 250000, "rpm": 10, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, @@ -76469,6 +76507,33 @@ "supports_tool_choice": true, "supports_vision": true }, + "us-gov.xai.grok-4.7": { + "supports_regex_lookaround": false, + "supported_endpoints": ["/v1/chat/completions"], + "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_response_format": true, + "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_batches": 1.32e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 7.92e-06, + "output_cost_per_token_batches": 3.96e-06, + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "baseten/deepseek-ai/DeepSeek-V4.1-Flash-Fast": { "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 6e-07, @@ -78437,7 +78502,7 @@ "cache_creation_input_token_cost": 1e-05, "cache_creation_input_token_cost_above_1hr": 1.6e-05, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2027-10-01", + "deprecation_date": "2027-01-15", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "input_cost_per_token": 8e-06, "litellm_provider": "azure_ai", @@ -78447,7 +78512,7 @@ "output_cost_per_token": 4e-05, "prompt_cache_min_tokens": 512.0, "search_context_cost_per_query": {"search_context_size_high":0.02,"search_context_size_low":0.02,"search_context_size_medium":0.02}, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models", + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -82782,6 +82847,39 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true }, + "gpt-6-cyber": { + "cache_creation_input_token_cost": 2.5e-05, + "cache_creation_input_token_cost_above_12hr": 5e-05, + "cache_read_input_token_cost": 2e-06, + "input_cost_per_token": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0001, + "source": "https://developers.openai.com/api/docs/models/gpt-6-cyber", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true + }, "gpt-6-astra-law": { "cache_read_input_token_cost": 1e-06, "input_cost_per_token": 1e-05, @@ -83380,6 +83478,28 @@ "supports_response_schema": true, "source": "https://empiriolabs.ai/models/kimi-k3" }, + "empiriolabs/kimi-k3-1-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1.8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/kimi-k3-1-pro" + }, "empiriolabs/mimo-v2-6-flash": { "litellm_provider": "empiriolabs", "mode": "chat", @@ -86383,5 +86503,32 @@ "supports_system_messages": false, "supports_video_input": true, "supports_vision": true + }, + "databricks/databricks-kimi-k3-1-pro": { + "cache_creation_input_token_cost": 1.99997e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "metadata": {"notes":"Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference."}, + "mode": "chat", + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.000142857, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 100.0 } } diff --git a/litellm/models/plugin_upload.py b/litellm/models/plugin_upload.py index bd5e85c79cc..95db82dec79 100644 --- a/litellm/models/plugin_upload.py +++ b/litellm/models/plugin_upload.py @@ -41,6 +41,14 @@ class PluginArchiveMetadata: description: str | None = None has_agent_skill_metadata: bool = False updated_at: datetime | None = None + created_at: datetime | None = None + created_by: str | None = None + + +@dataclass(frozen=True, slots=True) +class PluginReview: + approved_by: str + security_scan_passed: bool @dataclass(frozen=True, slots=True) @@ -51,6 +59,7 @@ class PluginPublication: manifest: Mapping[str, object] created_by: str | None content: bytes | None = None + review: PluginReview | None = None @dataclass(frozen=True, slots=True) @@ -66,6 +75,10 @@ class PublishedPlugin: version: str enabled: bool updated_at: datetime + created_at: datetime | None = None + created_by: str | None = None + review: PluginReview | None = None + published_by: str | None = None @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7d9e230f8ef..5cb63c98d56 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -7683,6 +7683,17 @@ "PluginListItem": { "description": "Plugin item in list responses.", "properties": { + "approved_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Approved By" + }, "author": { "anyOf": [ { @@ -7715,6 +7726,17 @@ ], "title": "Created At" }, + "created_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Created By" + }, "description": { "anyOf": [ { @@ -7737,6 +7759,17 @@ ], "title": "Domain" }, + "download_count": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Download Count" + }, "enabled": { "title": "Enabled", "type": "boolean" @@ -7801,6 +7834,22 @@ ], "title": "Namespace" }, + "published_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Published By" + }, + "security_scan_passed": { + "default": false, + "title": "Security Scan Passed", + "type": "boolean" + }, "source": { "additionalProperties": { "type": "string" @@ -59415,6 +59464,14 @@ }, "title": "Skills", "type": "array" + }, + "tags": { + "default": [], + "items": { + "type": "string" + }, + "title": "Tags", + "type": "array" } }, "required": [ @@ -59690,6 +59747,17 @@ }, "SkillHubSkill": { "properties": { + "approved_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Approved By" + }, "content_url": { "anyOf": [ { @@ -59701,6 +59769,29 @@ ], "title": "Content Url" }, + "created_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Created At" + }, + "created_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Created By" + }, "description": { "anyOf": [ { @@ -59732,7 +59823,7 @@ "type": "null" } ], - "description": "Gateway ZIP downloads, included in proxy-admin catalog listings", + "description": "Gateway ZIP downloads, including repeat downloads", "title": "Download Count" }, "download_url": { @@ -59779,6 +59870,22 @@ ], "title": "Namespace" }, + "published_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Published By" + }, + "security_scan_passed": { + "default": false, + "title": "Security Scan Passed", + "type": "boolean" + }, "source": { "additionalProperties": { "type": "string" @@ -60434,6 +60541,66 @@ "title": "Limit", "type": "integer" } + }, + { + "in": "query", + "name": "tag", + "required": false, + "schema": { + "default": "", + "title": "Tag", + "type": "string" + } + }, + { + "in": "query", + "name": "security_scan_passed", + "required": false, + "schema": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Security Scan Passed" + } + }, + { + "in": "query", + "name": "sort_by", + "required": false, + "schema": { + "default": "name", + "enum": [ + "name", + "download_count", + "updated_at", + "created_at", + "created_by", + "published_by", + "approved_by", + "security_scan_passed" + ], + "title": "Sort By", + "type": "string" + } + }, + { + "in": "query", + "name": "sort_order", + "required": false, + "schema": { + "default": "asc", + "enum": [ + "asc", + "desc" + ], + "title": "Sort Order", + "type": "string" + } } ], "responses": { diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 26cb0a2f3e6..3534dd167de 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -60,6 +60,7 @@ from litellm.proxy.common_utils.claude_code_marketplace_auth import ( ) from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.deps.repositories import PluginUploadRepositoryDep +from litellm.repositories.plugin_publication import read_plugin_publisher, read_plugin_review from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.proxy.auth.special_headers import SpecialHeaders from litellm.types.proxy.claude_code_endpoints import ( @@ -666,6 +667,7 @@ async def register_plugin( response_model=ListPluginsResponse, ) async def list_plugins( + uploads: PluginUploadRepositoryDep, enabled_only: bool = False, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> ListPluginsResponse: @@ -678,6 +680,17 @@ async def list_plugins( Returns: List of plugins with their metadata. """ + catalog: Final = await list_visible_plugins(enabled_only, user_api_key_dict) + if uploads is None: + raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) + counts: Final = await uploads.download_counts(tuple(plugin.id for plugin in catalog.plugins)) + return ListPluginsResponse( + plugins=[plugin.model_copy(update={"download_count": counts.get(plugin.id, 0)}) for plugin in catalog.plugins], + count=catalog.count, + ) + + +async def list_visible_plugins(enabled_only: bool, user_api_key_dict: UserAPIKeyAuth) -> ListPluginsResponse: try: prisma_client: Final = await get_prisma_client() @@ -693,9 +706,14 @@ async def list_plugins( # Parse manifest to get additional fields manifest = json.loads(p.manifest_json) if p.manifest_json else {} + review = read_plugin_review(MANIFEST_FIELDS.validate_python(manifest)) plugin_list.append( PluginListItem( id=p.id, + created_by=p.created_by, + approved_by=review.approved_by if review else None, + published_by=read_plugin_publisher(MANIFEST_FIELDS.validate_python(manifest)), + security_scan_passed=review.security_scan_passed if review else False, name=p.name, version=p.version, description=p.description, diff --git a/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py b/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py index 93d1ace435d..ffddbd31048 100644 --- a/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py +++ b/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py @@ -3,7 +3,8 @@ import hashlib from collections.abc import Mapping from dataclasses import dataclass, replace from datetime import datetime -from typing import Annotated, Final +from itertools import chain +from typing import Annotated, Final, Literal from urllib.parse import quote import httpx @@ -33,7 +34,7 @@ from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketp enable_plugin, get_plugin_downloads, header_key_caller, - list_plugins, + list_visible_plugins, plugin_archive_source_matches, plugin_archive_url, plugin_downloader, @@ -65,6 +66,7 @@ from litellm.proxy.discovery_endpoints.skill_import import ( parse_skill_source, select_imported_skill, ) +from litellm.repositories.plugin_publication import read_plugin_publisher, read_plugin_review from litellm.types.proxy.claude_code_endpoints import PluginListItem from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import AgentSkillsIndexEntry from litellm.types.proxy.skill_hub import ( @@ -78,6 +80,7 @@ from litellm.types.proxy.skill_hub import ( SkillHubSkill, SkillHubVersion, SkillHubVersions, + SkillSortField, ) router: Final = APIRouter(prefix="/skill-hub/skills", tags=["Skills Hub"]) @@ -121,6 +124,11 @@ def _skill_response(request: Request, plugin: PluginListItem, download_count: in namespace=plugin.namespace, keywords=None if plugin.keywords is None else tuple(plugin.keywords), updated_at=None if plugin.updated_at is None else datetime.fromisoformat(plugin.updated_at), + created_at=None if plugin.created_at is None else datetime.fromisoformat(plugin.created_at), + created_by=plugin.created_by, + approved_by=plugin.approved_by, + published_by=plugin.published_by, + security_scan_passed=plugin.security_scan_passed, visibility="public" if plugin.enabled else "private", source={"source": "archive", "url": _download_url(request, plugin.name)} if local_archive else plugin.source, download_url=_download_url(request, plugin.name) if local_archive else None, @@ -137,28 +145,74 @@ async def list_skills( query: str = "", offset: Annotated[int, Query(ge=0)] = 0, limit: Annotated[int, Query(ge=1, le=100)] = 50, + tag: str = "", + security_scan_passed: bool | None = None, + sort_by: SkillSortField = "name", + sort_order: Literal["asc", "desc"] = "asc", ) -> SkillHubList: """Browse the caller's public and granted private skills.""" - catalog: Final = await list_plugins(enabled_only=False, user_api_key_dict=caller) - needle: Final = query.casefold() + catalog: Final = await list_visible_plugins(enabled_only=False, user_api_key_dict=caller) + needle: Final = query.strip().casefold() + selected_tag: Final = tag.strip().casefold() matched: Final = tuple( plugin for plugin in sorted(catalog.plugins, key=lambda item: item.name) - if needle in " ".join((plugin.name, plugin.description or "", *(plugin.keywords or ()))).casefold() + if needle + in " ".join( + ( + plugin.name, + plugin.description or "", + plugin.domain or "", + plugin.namespace or "", + plugin.created_by or "", + plugin.approved_by or "", + plugin.published_by or "", + *(plugin.keywords or ()), + ) + ).casefold() + and (not selected_tag or selected_tag in {word.casefold() for word in plugin.keywords or ()}) + and (security_scan_passed is None or plugin.security_scan_passed == security_scan_passed) ) - page: Final = matched[offset : offset + limit] counts: Final[Mapping[str, int]] = ( - await _require_store(uploads).download_counts(tuple(plugin.id for plugin in page)) - if is_proxy_admin(caller) + await _require_store(uploads).download_counts(tuple(plugin.id for plugin in matched)) + if sort_by == "download_count" else {} ) + ordered: Final = sorted( + matched, key=lambda plugin: _skill_sort_key(plugin, sort_by, counts), reverse=sort_order == "desc" + ) + page: Final = ordered[offset : offset + limit] + page_counts: Final = ( + await _require_store(uploads).download_counts(tuple(plugin.id for plugin in page)) + if sort_by != "download_count" + else counts + ) return SkillHubList( - skills=tuple(_skill_response(request, plugin, counts.get(plugin.id)) for plugin in page), + skills=tuple(_skill_response(request, plugin, page_counts.get(plugin.id)) for plugin in page), count=len(matched), next_offset=offset + limit if offset + limit < len(matched) else None, + tags=tuple( + sorted(set(chain.from_iterable(plugin.keywords or () for plugin in catalog.plugins)), key=str.casefold) + ), ) +def _skill_sort_key(plugin: PluginListItem, field: SkillSortField, counts: Mapping[str, int]) -> str | float: + if field == "download_count": + return counts.get(plugin.id, 0) + if field == "security_scan_passed": + return int(plugin.security_scan_passed) + if field in {"created_at", "updated_at"}: + value: Final = plugin.created_at if field == "created_at" else plugin.updated_at + return datetime.fromisoformat(value).timestamp() if value else 0 + return { + "name": plugin.name, + "created_by": plugin.created_by or "", + "published_by": plugin.published_by or "", + "approved_by": plugin.approved_by or "", + }[field].casefold() + + @router.get("/{name}", response_model=SkillHubSkill) async def get_skill( request: Request, @@ -176,6 +230,7 @@ async def get_skill( def skill_metadata_response(request: Request, metadata: PluginArchiveMetadata) -> SkillHubSkill: manifest: Final = MANIFEST_FIELDS.validate_json(metadata.manifest_json or "{}") + review: Final = read_plugin_review(manifest) return _skill_response( request, PluginListItem.model_validate( @@ -187,7 +242,11 @@ def skill_metadata_response(request: Request, metadata: PluginArchiveMetadata) - "source": dict(metadata.source), "has_agent_skill_metadata": metadata.has_agent_skill_metadata, "enabled": metadata.enabled, - "created_at": None, + "created_at": metadata.created_at.isoformat() if metadata.created_at is not None else None, + "created_by": metadata.created_by, + "approved_by": review.approved_by if review else None, + "published_by": read_plugin_publisher(manifest), + "security_scan_passed": review.security_scan_passed if review else False, "updated_at": metadata.updated_at.isoformat() if metadata.updated_at is not None else None, "domain": manifest.get("domain"), "namespace": manifest.get("namespace"), @@ -416,7 +475,11 @@ def skill_publication_response( "version": published.version, "enabled": published.enabled, "has_agent_skill_metadata": has_agent_skill_metadata(publication.manifest), - "created_at": None, + "created_at": published.created_at.isoformat() if published.created_at is not None else None, + "created_by": published.created_by, + "approved_by": published.review.approved_by if published.review else None, + "published_by": published.published_by, + "security_scan_passed": published.review.security_scan_passed if published.review else False, "updated_at": published.updated_at.isoformat(), "domain": publication.manifest.get("domain"), "namespace": publication.manifest.get("namespace"), diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index e8f0dc68d83..326f6400ba3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -4,7 +4,8 @@ # # +-------------------------------------------------------------+ import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from itertools import chain, dropwhile, takewhile from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException @@ -14,6 +15,10 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -218,15 +223,14 @@ class ZscalerAIGuard(CustomGuardrail): zscaler_ai_guard_result = None direction: Final = "OUT" if input_type == "response" else "IN" verbose_proxy_logger.debug("direction: %s", direction) - # Concatenate all texts and send to Zscaler AI Guard - if texts: - concatenated_text: Final = " ".join(texts) + content: Final = self._content_to_scan(inputs, input_type) + if content: zscaler_ai_guard_result = await self.make_zscaler_ai_guard_api_call( zscaler_ai_guard_url=self.zscaler_ai_guard_url, api_key=self.api_key, policy_id=policy_id, direction=direction, - content=concatenated_text, + content=content, user=user_for_payload, **kwargs, ) @@ -244,6 +248,37 @@ class ZscalerAIGuard(CustomGuardrail): verbose_proxy_logger.debug("ZscalerAIGuard: Successfully applied guardrail.") return inputs + def _content_to_scan( + self, + inputs: GenericGuardrailAPIInputs, + input_type: Literal["request", "response"], + ) -> str | None: + structured_messages: Final = inputs.get("structured_messages") or [] + scans_last_user_turn: Final = ( + input_type == "request" + and not effective_scan_only_tool_results_for_guardrail(self) + and any("content" in message for message in structured_messages) + ) + texts: Final = ( + self._last_user_turn_texts(structured_messages) if scans_last_user_turn else inputs.get("texts", []) + ) + return " ".join(texts) if texts else None + + @staticmethod + def _last_user_turn_texts(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + from_last_user_message: Final = dropwhile(lambda message: message.get("role") != "user", reversed(messages)) + last_user_turn: Final = tuple(takewhile(lambda message: message.get("role") == "user", from_last_user_message)) + texts: Final = tuple(chain.from_iterable(map(ZscalerAIGuard._content_texts, reversed(last_user_turn)))) + return texts if any(text.strip() for text in texts) else () + + @staticmethod + def _content_texts(message: Mapping[str, object]) -> tuple[str, ...]: + content: Final = message.get("content") + if isinstance(content, str): + return (content,) + parts: Final = (as_str_mapping(part) for part in as_object_list(content) or ()) + return tuple(text for part in parts if part is not None and isinstance(text := part.get("text"), str)) + def extract_blocking_info(self, response): """ Extracts transaction ID and blocking detector details from a response. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c8895c01f7b..d28401524ed 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5482,7 +5482,9 @@ async def generate_key_helper_fn( if query_type == "insert_data": user_row = cast( # cast-ok: table_name="user" is the insert_data branch returning the user row "prisma_models.LiteLLM_UserTable | None", - await prisma_client.insert_data(data=user_data, table_name="user"), + await prisma_client.insert_data( + data=user_data, table_name="user", allow_existing_user=request_type != "user" + ), ) if user_row is None: diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 79d38ec2e53..013eb1ebff4 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -68,7 +68,7 @@ def _get_routing_strategies_from_router_class() -> list[str]: Dynamically extract routing strategies from the Router class __init__ method. """ # Get the __init__ signature - sig: Final = inspect.signature(Router.__init__) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped params, only routing_strategy is read + sig: Final = inspect.signature(Router) # Get the routing_strategy parameter routing_strategy_param: Final = sig.parameters.get("routing_strategy") diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index e48f68a2d0a..66c2e54f5c6 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1294,7 +1294,7 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou if litellm.default_internal_user_params: default_role = litellm.default_internal_user_params.get("user_role") - new_user_request: Final = NewUserRequest( + new_user_request: Final = SCIMNewUserRequest( user_id=user_id, user_email=user_id, # We don't have email from group membership user_alias=None, @@ -1873,7 +1873,7 @@ async def create_user( admin_group: Final = await _get_scim_admin_group() resolved_role: Final = _resolve_scim_user_role(user.groups or [], admin_group, default_role) - new_user_request: Final = NewUserRequest( + new_user_request: Final = SCIMNewUserRequest( user_id=user_id, user_email=user_data["user_email"], user_alias=user_data["user_alias"], diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 53079d84fdd..aaa261eb19b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -78,7 +78,6 @@ from litellm.proxy._types import ( LitellmUserRoles, Member, NewTeamRequest, - NewUserRequest, NewUserRequestTeam, NewUserResponse, ProxyErrorTypes, @@ -164,6 +163,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( MicrosoftGraphAPIUserGroupResponse, MicrosoftServicePrincipalTeam, RoleMappings, + SSONewUserRequest, TeamMappings, ) from litellm.types.proxy.ui_sso import ParsedOpenIDResult @@ -2915,7 +2915,7 @@ async def insert_sso_user( if user_defined_values["user_role"] is None: user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY - new_user_request: Final = NewUserRequest( + new_user_request: Final = SSONewUserRequest( user_id=user_defined_values["user_id"], user_email=normalize_email(user_defined_values["user_email"]), user_alias=_get_sso_user_alias(result_openid), @@ -2924,11 +2924,9 @@ async def insert_sso_user( budget_duration=user_defined_values["budget_duration"], sso_user_id=user_defined_values["user_id"], auto_create_key=False, + metadata={"auth_provider": getattr(result_openid, "provider")} if hasattr(result_openid, "provider") else {}, ) - if result_openid and hasattr(result_openid, "provider"): - new_user_request.metadata = {"auth_provider": getattr(result_openid, "provider")} - from litellm.proxy.proxy_server import premium_user with free_sso_user_creation(new_user_request.user_id, enabled=premium_user is not True): diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index b8cdea0dea4..cb1e51d0deb 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -24,7 +24,9 @@ from litellm.proxy._types import ( CommonProxyErrors, ) from litellm.proxy.db.health_check_latest import latest_by_model_name +from litellm.proxy.deps.repositories import PluginUploadRepositoryDep from litellm.proxy.utils import get_custom_url +from litellm.repositories.plugin_publication import read_plugin_publisher, read_plugin_review from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.router_strategy.complexity_router.fuse_presets import FusePresetCatalog, get_fuse_presets from litellm.types.agents import AgentCard @@ -71,6 +73,9 @@ class _EndpointEntry(TypedDict): class _PluginRow(Protocol): + @property + def created_by(self) -> str | None: ... + @property def id(self) -> str: ... @@ -366,7 +371,7 @@ async def get_mcp_servers(): "/public/skill_hub", tags=["public", "Claude Code Marketplace"], ) -async def public_skill_hub(): +async def public_skill_hub(uploads: PluginUploadRepositoryDep): """Return enabled (public) Claude Code skills, or none when the marketplace requires a key.""" from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( MANIFEST_FIELDS, @@ -385,13 +390,22 @@ async def public_skill_hub(): raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: plugins: Final = await _plugin_table(prisma_client).find_many(where={"enabled": True}) + counts: Final[Mapping[str, int]] = ( + await uploads.download_counts(tuple(plugin.id for plugin in plugins)) if uploads else {} + ) items: Final = [] for plugin in plugins: raw = plugin.manifest_json or {} manifest = json.loads(raw) if isinstance(raw, str) else raw + review = read_plugin_review(MANIFEST_FIELDS.validate_python(manifest)) items.append( PluginListItem( id=plugin.id, + download_count=counts.get(plugin.id), + created_by=plugin.created_by, + approved_by=review.approved_by if review else None, + published_by=read_plugin_publisher(MANIFEST_FIELDS.validate_python(manifest)), + security_scan_passed=review.security_scan_passed if review else False, name=plugin.name, enabled=plugin.enabled, created_at=str(plugin.created_at) if plugin.created_at else None, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 148d465099f..08942851af9 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -271,6 +271,10 @@ REQUIRED_ONE_OF_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, str]]] = Map {"acreate_interaction": ("model", "agent"), "acreate_agent": ("id", "name")} ) +BODY_PARAM_ALTERNATIVES_BY_ROUTE: Final[Mapping[str, Mapping[str, tuple[str, ...]]]] = MappingProxyType( + {"acreate_interaction": MappingProxyType({"input": ("continuation_token",)})} +) + JSON_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) _OPTIONAL_JSON_OBJECT_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter( Mapping[str, object] | None @@ -320,8 +324,11 @@ def _find_missing_required_body_param( ) if missing_merge_base_param is not None: return MissingBodyParam(name=missing_merge_base_param, model_deployments_loaded=True) + alternatives: Final = BODY_PARAM_ALTERNATIVES_BY_ROUTE.get(route_type, {}) missing_present_params: Final = tuple( - param for param in REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE.get(route_type, ()) if param not in data + param + for param in REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE.get(route_type, ()) + if param not in data and not any(alternative in data for alternative in alternatives.get(param, ())) ) if not missing_present_params: return None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 00187790901..3102674376a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -313,6 +313,7 @@ if TYPE_CHECKING: from prisma import models as prisma_models from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions from prisma.client import TransactionManager + from prisma.errors import UniqueViolationError from prisma.models import LiteLLM_DeprecatedVerificationToken from prisma.types import HttpConfig, LiteLLM_VerificationTokenInclude @@ -5559,6 +5560,25 @@ _LOCK_TEAMS_BY_ID_SQL: Final = ( ) +def _user_unique_conflict_message(metadata: object, user_id: object) -> str: + match metadata: + case {"target": ["user_id"]}: + return f"User with id {user_id} already exists" + case {"target": ["sso_user_id"]}: + return "User with this sso_user_id already exists" + case _: + return "User conflicts with an existing unique value" + + +def _raise_user_unique_conflict(error: "UniqueViolationError", user_id: object, *, allow_existing: bool) -> NoReturn: + if allow_existing: + raise error + raise HTTPException( + status_code=409, + detail={"error": _user_unique_conflict_message(error.meta, user_id)}, + ) from error + + class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() @@ -6446,12 +6466,15 @@ class PrismaClient: max_tries=3, # maximum number of retries max_time=10, # maximum total time to retry for on_backoff=on_backoff, # specifying the function to call on backoff + giveup=lambda exc: isinstance(exc, HTTPException) and exc.status_code == 409, ) @log_db_metrics async def insert_data( self, data: Mapping[str, object], table_name: Literal["user", "key", "config", "spend", "team", "user_notification"], + *, + allow_existing_user: bool = True, ): """ Add a key to the database. If it already exists, do nothing. @@ -6487,14 +6510,20 @@ class PrismaClient: verbose_proxy_logger.info("Data Inserted into Keys Table") return new_verification_token elif table_name == "user": + from prisma.errors import UniqueViolationError + db_data = self.jsonify_object(data=data) try: - new_user_row: Final = await UserRepository(self).upsert_user(db_data) + new_user_row: Final = await UserRepository(self).upsert_user( + db_data, allow_existing=allow_existing_user + ) if new_user_row is None: raise HTTPException( status_code=403, detail="SSO for more than 5 users requires a LiteLLM Enterprise license. Existing users can still sign in.", ) + except UniqueViolationError as e: + _raise_user_unique_conflict(e, db_data["user_id"], allow_existing=allow_existing_user) except Exception as e: if ( "Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`" diff --git a/litellm/repositories/plugin_publication.py b/litellm/repositories/plugin_publication.py index f5fab458c6b..6f94076c3e0 100644 --- a/litellm/repositories/plugin_publication.py +++ b/litellm/repositories/plugin_publication.py @@ -1,15 +1,18 @@ import hashlib import json +from collections.abc import Mapping +from dataclasses import asdict from datetime import datetime, timezone from typing import TYPE_CHECKING, Final, Protocol -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from semver import Version from litellm.models.plugin_upload import ( PluginPublication, PluginPublicationBase, PluginPublicationError, + PluginReview, PublishedPlugin, ) from litellm.repositories.prisma_protocols import TableActions @@ -23,6 +26,22 @@ if TYPE_CHECKING: _MANIFEST: Final = TypeAdapter(dict[str, object]) _SOURCE: Final = TypeAdapter(dict[str, str]) +_REVIEW: Final = TypeAdapter(PluginReview) +_PUBLISHER: Final = TypeAdapter[str | None](str | None) + + +def read_plugin_review(manifest: Mapping[str, object]) -> PluginReview | None: + try: + return _REVIEW.validate_python(manifest.get("_litellm_review")) + except ValidationError: + return None + + +def read_plugin_publisher(manifest: Mapping[str, object]) -> str | None: + try: + return _PUBLISHER.validate_python(manifest.get("_litellm_published_by")) + except ValidationError: + return None class PluginPublicationTransaction(Protocol): @@ -122,7 +141,22 @@ async def publish_plugin( ) if isinstance(version, PluginPublicationError): return version - manifest_json: Final = json.dumps({**publication.manifest, "version": version}) + current_manifest: Final = _MANIFEST.validate_json(current.manifest_json or "{}") if current is not None else {} + same_source: Final = current is not None and current_manifest.get("source") == publication.manifest.get("source") + review: Final = publication.review or (read_plugin_review(current_manifest) if same_source else None) + publisher: Final = publication.review.approved_by if publication.review is not None else publication.created_by + manifest_json: Final = json.dumps( + { + **{ + key: value + for key, value in publication.manifest.items() + if key not in {"_litellm_review", "_litellm_published_by"} + }, + "version": version, + **({"_litellm_published_by": publisher} if publisher is not None else {}), + **({"_litellm_review": asdict(review)} if review is not None else {}), + } + ) if publication.content is not None and verified_archive(publication.content, manifest_json) is None: return PluginPublicationError("conflict", "Uploaded archive does not match its published source digest") timestamp: Final = datetime.now(timezone.utc) @@ -141,8 +175,18 @@ async def publish_plugin( await tx.litellm_claudecodepluginupload.create( data={"plugin_name": publication.name, "content": Base64.encode(publication.content)} ) - await _snapshot(tx, created, publication.content, publication.created_by) - return PublishedPlugin(created.id, created.name, version, created.enabled, created.updated_at or timestamp) + await _snapshot(tx, created, publication.content, publisher) + return PublishedPlugin( + created.id, + created.name, + version, + created.enabled, + created.updated_at or timestamp, + created.created_at, + created.created_by, + review, + publisher, + ) previous_archive: Final = await tx.litellm_claudecodepluginupload.find_unique(where={"plugin_name": current.name}) previous_content: Final = ( verified_archive(previous_archive.content.decode(), current.manifest_json) @@ -152,11 +196,7 @@ async def publish_plugin( history_count: Final = await tx.litellm_claudecodepluginversion.count(where={"plugin_id": current.id}, take=1) if history_count == 0: await _snapshot(tx, current, previous_content, None) - unchanged_content: Final = ( - previous_content - if _MANIFEST.validate_json(current.manifest_json or "{}").get("source") == publication.manifest.get("source") - else None - ) + unchanged_content: Final = previous_content if same_source else None content: Final = publication.content if publication.content is not None else unchanged_content updated: Final = await tx.litellm_claudecodeplugintable.update(where={"id": current.id}, data=data) if updated is None: @@ -171,5 +211,15 @@ async def publish_plugin( "update": {"content": Base64.encode(content)}, }, ) - await _snapshot(tx, updated, content, publication.created_by) - return PublishedPlugin(updated.id, updated.name, version, updated.enabled, updated.updated_at or timestamp) + await _snapshot(tx, updated, content, publisher) + return PublishedPlugin( + updated.id, + updated.name, + version, + updated.enabled, + updated.updated_at or timestamp, + updated.created_at, + updated.created_by, + review, + publisher, + ) diff --git a/litellm/repositories/plugin_upload_repository.py b/litellm/repositories/plugin_upload_repository.py index 99e226e442b..8045abf9af9 100644 --- a/litellm/repositories/plugin_upload_repository.py +++ b/litellm/repositories/plugin_upload_repository.py @@ -287,4 +287,6 @@ def _archive_metadata(plugin: "prisma_models.LiteLLM_ClaudeCodePluginTable") -> description=plugin.description, has_agent_skill_metadata=has_agent_skill_metadata(manifest), updated_at=plugin.updated_at, + created_at=plugin.created_at, + created_by=plugin.created_by, ) diff --git a/litellm/repositories/skill_submission_repository.py b/litellm/repositories/skill_submission_repository.py index 829aa0aadbf..03821da042a 100644 --- a/litellm/repositories/skill_submission_repository.py +++ b/litellm/repositories/skill_submission_repository.py @@ -7,12 +7,13 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # Prisma's generated delegates are dynamic from urllib.parse import urlsplit -from pydantic import BaseModel, ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.models.plugin_upload import ( PluginPublication, PluginPublicationBase, PluginPublicationError, + PluginReview, PublishedPlugin, uploaded_plugin_source, ) @@ -312,7 +313,9 @@ class SkillSubmissionRepository: if approve and content is not None and publication_base_url else _SOURCE.validate_python(submission.proposal.manifest.get("source")) ) - publication: Final = await _publish_submission(transaction, row, source, content) if approve else None + publication: Final = ( + await _publish_submission(transaction, row, source, content, reviewer) if approve else None + ) if isinstance(publication, PluginPublicationError): return publication reviewed_at: Final = datetime.now(timezone.utc) @@ -346,6 +349,7 @@ async def _publish_submission( row: "prisma_models.LiteLLM_SkillSubmission", source: Mapping[str, str], content: bytes | None, + reviewer: str, ) -> PublishedPlugin | PluginPublicationError: stored: Final = _StoredProposal.model_validate_json(row.proposal_json) proposal: Final = stored.proposal(row.plugin_name) @@ -363,6 +367,24 @@ async def _publish_submission( manifest={**stored.manifest, "source": dict(source)}, content=content, created_by=row.submitted_by, + review=PluginReview(reviewer, scan_passed_for_archive(row.scan_json, source.get("sha256"))), ), base, ) + + +def scan_passed_for_archive(scan_json: str | None, digest: str | None) -> bool: + if not scan_json or not digest: + return False + try: + scan: Final = SkillScan.model_validate_json(scan_json) + except ValidationError: + return False + return ( + scan.status == "completed" + and scan.archive_sha256 == digest + and scan.report is not None + and scan.report.complete + and scan.report.risk_score == 0 + and not scan.report.findings + ) diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 1fce0717ea3..28fba2de564 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -168,7 +168,9 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): """ return await _count_billable_users(self.table) - async def upsert_user(self, data: Mapping[str, object]) -> "prisma_models.LiteLLM_UserTable | None": + async def upsert_user( + self, data: Mapping[str, object], *, allow_existing: bool = True + ) -> "prisma_models.LiteLLM_UserTable | None": """Return the user row, or None when a new free SSO user would exceed capacity.""" scope: Final = _FREE_SSO_CREATION.get() if ( @@ -176,6 +178,8 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): or (scope[0] is not None and scope[0] != data["user_id"]) or scope[1] is not asyncio.current_task() ): + if not allow_existing: + return await self.table.create(data=dict(data)) return await self.table.upsert( where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}} ) diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py index 5cb256a8ac0..be331a91475 100644 --- a/litellm/responses/mcp/stream_round_billing.py +++ b/litellm/responses/mcp/stream_round_billing.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Callable, Coroutine, Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass from datetime import datetime from typing import Final, TypeAlias @@ -297,12 +297,8 @@ class StreamedRoundBilling: self._close(parked) self._bank((self._dispatch(row, billed, self._logging_obj.attempt_scoped_copy()),), row) - def _dispatch( - self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj - ) -> BilledRoundsDispatch | Coroutine[object, object, None]: + def _dispatch(self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj) -> BilledRoundsDispatch: dispatch: Final = self._log_success(_SuccessLog(logger, row, billed.chat_usage)) - if billed.chat_usage is None: - return dispatch return BilledRoundsDispatch(dispatch, billed.spend) @staticmethod diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 249c6370b86..288e4a7119d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -53,7 +53,7 @@ from litellm.litellm_core_utils.internal_call_metadata import ( requests_background, without_budget_reservation, ) -from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -638,6 +638,24 @@ class BaseResponsesAPIStreamingIterator: self._handle_failure(e) raise + def _deferred_logging_dispatch( + self, logging_coroutine: Coroutine[object, object, None], logging_response: object + ) -> BilledRoundsDispatch | Coroutine[object, object, None]: + response: Final[object] = getattr(logging_response, "response", None) + if not isinstance(response, ResponsesAPIResponse): + return logging_coroutine + price: Final = ( + UncachedPrice(0.0, None) + if self._completed_response_cache_hit + else self.logging_obj.uncached_price(response) + ) + spend: Final = BilledRoundsSpend( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response.usage), + price.cost or 0.0, + price.cost_breakdown, + ) + return BilledRoundsDispatch(logging_coroutine, spend) + def _log_completed_response(self, *, is_async: bool) -> None: if self._completed_response_logged: return @@ -662,7 +680,9 @@ class BaseResponsesAPIStreamingIterator: if deferred_dispatch_armed: # End-of-stream guardrail scans write guardrail_information after # the terminal event; dispatching now would snapshot metadata early. - self.logging_obj._deferred_stream_complete_args = (logging_coroutine,) + self.logging_obj._deferred_stream_complete_args = ( + self._deferred_logging_dispatch(logging_coroutine, logging_response), + ) else: spawn_logging_task( asyncio.get_running_loop(), logging_coroutine, task_name="Responses stream success logging" diff --git a/litellm/router.py b/litellm/router.py index fa53cdf8bc6..4b51c16aa81 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -52,7 +52,7 @@ if TYPE_CHECKING: # ``Router(...)``'s parameters, their defaults and its docstring: ``router_arguments`` binds a call # against them and ``Router.__init__`` reports them. The body is never run. def _router_init( - self, + self: object, model_list: list[DeploymentTypedDict] | list[dict[str, Any]] | None = None, ## ASSISTANTS API ## assistants_config: AssistantsTypedDict | None = None, @@ -228,18 +228,16 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues, OpenAIFileObject from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, ModelResponse - _InitParams: Final = ParamSpec("_InitParams") + _P: Final = ParamSpec("_P") - def _signature_of( - source: Callable[_InitParams, None], - ) -> Callable[[Callable[..., None]], Callable[_InitParams, None]]: ... + def _takes_parameters_of(source: Callable[_P, None]) -> Callable[[Callable[..., None]], Callable[_P, None]]: ... class Router(RustCore, DeprecatedRouterMembers): """What type checkers see: the constructor's parameters and the endpoints' return types. The endpoints are served by the core at run time (litellm/router_rust/invoke/endpoints.py).""" - @_signature_of(_router_init) # pyright: ignore[reportUnknownArgumentType] # _router_init's untyped dict params - def __init__(self) -> None: ... + @_takes_parameters_of(_router_init) + def __init__(self) -> None: ... # pyright: ignore[reportMissingSuperCall] # a type-checking stub @classmethod def from_config(cls, config: RouterConfig, host: PythonHost = python_host) -> "Router": ... diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index 00dcd77c1e2..daff0226ee9 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -1911,10 +1911,9 @@ class RustEndpoints: if self._scheduler is None: from litellm.scheduler import Scheduler - self._scheduler = ( - Scheduler(polling_interval=self._polling_interval) # pyright: ignore[reportArgumentType] # as configured - if self._polling_interval is not None - else Scheduler() + self._scheduler = Scheduler( + polling_interval=self._polling_interval, # pyright: ignore[reportArgumentType] # as configured + redis_cache=self._core.routing_cache.redis_cache, ) return self._scheduler diff --git a/litellm/router_rust/invoke/walk.py b/litellm/router_rust/invoke/walk.py index 2c1481f068a..17e20119f43 100644 --- a/litellm/router_rust/invoke/walk.py +++ b/litellm/router_rust/invoke/walk.py @@ -264,10 +264,18 @@ class WalkMetadata: def hop_kwargs(self, event: Mapping[str, object]) -> dict[str, object]: # mutable-ok: a fresh dict for the loggers """The kwargs the level held after a hop (``kwargs["model"] = target``), as the - fallback loggers see them.""" + fallback loggers see them: the request's own metadata (the proxy's ``user_api_key_*`` + labels) with the level's walk stamps laid over it, as ``run_async_fallback``'s + ``_hop_metadata`` copies it.""" scope: Final = self._scope(event.get("scope")) model: Final = scope.metadata.get("model_group", event.get("to")) - return {**self._request, **scope.kwargs, "model": model, self._field: dict(scope.metadata)} + request_bucket: Final = as_str_mapping(self._request.get(self._field)) or {} + return { + **self._request, + **scope.kwargs, + "model": model, + self._field: {**request_bucket, **scope.metadata}, + } def log_retry(self, event: Mapping[str, object], error: BaseException | None) -> None: """A ``retry_logged`` event: the failure joins the scope's ``previous_models``.""" diff --git a/litellm/router_rust/surface/health.py b/litellm/router_rust/surface/health.py index b16adc38af9..0653518d0d7 100644 --- a/litellm/router_rust/surface/health.py +++ b/litellm/router_rust/surface/health.py @@ -279,6 +279,7 @@ class RustHealthCallbacks: is_caller_scoped_auth_failure, is_caller_timeout_408, is_failure_outside_selected_deployment, + is_response_lookup_not_found, is_video_lookup_not_found, ) @@ -288,6 +289,7 @@ class RustHealthCallbacks: is_failure_outside_selected_deployment(exception if isinstance(exception, BaseException) else None) or (isinstance(exception, Exception) and is_background_response_poll_not_found(exception, litellm_params)) or is_video_lookup_not_found(exception, logged) + or is_response_lookup_not_found(exception, logged) or is_caller_timeout_408(logged, status) # pyright: ignore[reportArgumentType] # the status as raised or is_caller_scoped_auth_failure(self._deployment_of(deployment_id), status) # pyright: ignore[reportArgumentType] # see above ) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 03485ebd5f8..4c918776223 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -108,6 +108,26 @@ def is_video_lookup_not_found(exception: object, failure_kwargs: Mapping[str, ob return getattr(exception, "status_code", None) == 404 and failure_kwargs.get("call_type") in VIDEO_ID_CALL_TYPES +# The calls on an existing response id. ``compact_responses`` is left out: it sends a model and input, +# so its 404 can mean a missing model rather than a missing response. +RESPONSE_ID_CALL_TYPES: Final = frozenset( + { + "aget_responses", + "get_responses", + "adelete_responses", + "delete_responses", + "acancel_responses", + "cancel_responses", + "alist_input_items", + } +) + + +def is_response_lookup_not_found(exception: object, failure_kwargs: Mapping[str, object]) -> bool: + """Whether a call on an existing response id failed with a provider 404: the response is not on that account.""" + return getattr(exception, "status_code", None) == 404 and failure_kwargs.get("call_type") in RESPONSE_ID_CALL_TYPES + + _EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = ( # ContentPolicyViolationError subclasses BadRequestError, so it must be checked first. (litellm.ContentPolicyViolationError, "ContentPolicyViolationErrorAllowedFails"), diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index 6033024123f..97bf4710fca 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -26,7 +26,7 @@ async def send_llm_exception_alert( ctx: "RoutingContext | None", request_kwargs: dict, error_traceback_str: str, - original_exception, + original_exception: Exception, ): """ Only runs if router.slack_alerting_logger is set @@ -34,7 +34,7 @@ async def send_llm_exception_alert( Parameters: ctx (RoutingContext): The router's routing context. - original_exception (Any): The original exception that occurred. + original_exception (Exception): The original exception that occurred. Returns: None diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 81012d38b65..e5d944b3aef 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -163,6 +163,9 @@ class PromptCachingDeploymentCheck(CustomLogger): return healthy_deployments + def exempt_from_message_redaction(self) -> bool: + return True + async def async_log_success_event( self, kwargs: _PromptCachingSuccessEventKwargs, diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index c4f7a2d8a7b..4d808742434 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -176,11 +176,16 @@ class Status1(Enum): incomplete = "incomplete" budget_exceeded = "budget_exceeded" queued = "queued" + requires_continuation = "requires_continuation" class InteractionStatusUpdate(LiteLLMBaseModel): interaction_id: str | None = None status: Status1 | None = None + continuation_token: str | None = Field( + None, + description="Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + ) event_type: Literal["interaction.status_update"] = "interaction.status_update" event_id: str | None = Field( None, @@ -343,6 +348,7 @@ class Status3(Enum): INCOMPLETE = "incomplete" BUDGET_EXCEEDED = "budget_exceeded" QUEUED = "queued" + REQUIRES_CONTINUATION = "requires_continuation" class ModelOption(RootModel[str]): @@ -604,6 +610,9 @@ class ToolChoice(RootModel[ToolChoiceType | ToolChoiceConfig]): class Usage(LiteLLMBaseModel): total_input_tokens: int | None = Field(None, description="Number of tokens in the prompt (context).") + initial_input_tokens: int | None = Field( + None, description="Number of tokens in the prompt of the first hop of a long decode." + ) input_tokens_by_modality: list[ModalityTokens] | None = Field( None, description="A breakdown of input token usage by modality." ) @@ -625,6 +634,7 @@ class Usage(LiteLLMBaseModel): None, description="A breakdown of tool-use token usage by modality." ) total_reasoning_tokens: int | None = Field(None, description="Number of tokens of thoughts for thinking models.") + total_thought_tokens: int | None = Field(None, description="Number of tokens of thoughts for thinking models.") total_tokens: int | None = Field( None, description="Total token count for the interaction request (prompt + responses + other\ninternal tokens).", @@ -787,6 +797,10 @@ class Interaction(LiteLLMBaseModel): description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: str | None = Field(None, description="The ID of the previous interaction, if any.") + continuation_token: str | None = Field( + None, + description="Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + ) input: str | list[Content] | list[Turn] | Content | None = Field( None, description="The inputs for the interaction." ) @@ -843,6 +857,10 @@ class CreateModelInteractionParams(LiteLLMBaseModel): description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: str | None = Field(None, description="The ID of the previous interaction, if any.") + continuation_token: str | None = Field( + None, + description="Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + ) input: str | list[Content] | list[Turn] | Content = Field(..., description="The inputs for the interaction.") generation_config: GenerationConfig | None = Field( None, @@ -894,6 +912,10 @@ class CreateAgentInteractionParams(LiteLLMBaseModel): description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: str | None = Field(None, description="The ID of the previous interaction, if any.") + continuation_token: str | None = Field( + None, + description="Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + ) input: str | list[Content] | list[Turn] | Content = Field(..., description="The inputs for the interaction.") agent_config: DynamicAgentConfig | DeepResearchAgentConfig | None = Field( None, description="Configuration for the agent.", discriminator="type" @@ -1080,6 +1102,7 @@ class InteractionsAPIResponse(BaseLiteLLMOpenAIResponseObject): # New schema field (Api-Revision: 2026-05-20). steps: list[dict[str, Any]] | None = None usage: dict[str, Any] | None = None + continuation_token: str | None = None _hidden_params: dict = PrivateAttr(default_factory=dict) @@ -1126,6 +1149,7 @@ class InteractionsAPIStreamingResponse(BaseLiteLLMOpenAIResponseObject): step: dict[str, Any] | None = None interaction_id: str | None = None interaction: dict[str, Any] | None = None + continuation_token: str | None = None _hidden_params: dict = PrivateAttr(default_factory=dict) diff --git a/litellm/types/proxy/claude_code_endpoints.py b/litellm/types/proxy/claude_code_endpoints.py index 5d8479b8ce1..dfbf1de17af 100644 --- a/litellm/types/proxy/claude_code_endpoints.py +++ b/litellm/types/proxy/claude_code_endpoints.py @@ -154,6 +154,11 @@ class PluginListItem(LiteLLMBaseModel): enabled: bool created_at: str | None updated_at: str | None + created_by: str | None = None + approved_by: str | None = None + published_by: str | None = None + security_scan_passed: bool = False + download_count: int | None = None has_agent_skill_metadata: bool = Field(default=False, exclude=True) diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index fc0b73f77b6..e79b509da1b 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -120,6 +120,7 @@ class BulkNewUserItem(NewUserRequest): model_config = ConfigDict(extra="forbid", protected_namespaces=()) + user_id: str | None = None auto_create_key: bool = False @field_validator("send_invite_email") diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index fadb309e0cd..72a39c61967 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -13,6 +13,7 @@ from pydantic import ( from pydantic_core.core_schema import SerializerFunctionWrapHandler from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.proxy.management_endpoints.user_requests import NewUserRequest SCIM_ENTERPRISE_USER_SCHEMA: Final = "urn:ietf:params:scim:schemas:extension:enterprise:2.0:User" SCIM_ENTERPRISE_METADATA_KEY: Final = "scim_enterprise" @@ -26,6 +27,14 @@ _SCIM_SERIALIZED_FIELDS_ADAPTER: Final = TypeAdapter( ) +class SCIMNewUserRequest(NewUserRequest): + """Preserve opaque directory IDs when provisioning SCIM users and group members.""" + + model_config = ConfigDict(frozen=True) + + user_id: str | None = None + + class LiteLLM_UserScimMetadata(LiteLLMBaseModel): """ Scim metadata stored in LiteLLM_UserTable.metadata diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index affb6d2e9ce..38ed799f38b 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -1,12 +1,19 @@ from typing import Literal -from pydantic import Field +from pydantic import ConfigDict, Field from typing_extensions import TypedDict from litellm.types.proxy.auth.user_roles import KeyManagementRoutes, LitellmUserRoles +from litellm.types.proxy.management_endpoints.user_requests import NewUserRequest from litellm.types.utils import LiteLLMPydanticObjectBase +class SSONewUserRequest(NewUserRequest): + model_config = ConfigDict(frozen=True) + + user_id: str | None = None + + class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase): """ Set default upperbound to max budget a key called via `/key/generate` can be. diff --git a/litellm/types/proxy/management_endpoints/user_requests.py b/litellm/types/proxy/management_endpoints/user_requests.py index e53dc4152c3..edda1f78430 100644 --- a/litellm/types/proxy/management_endpoints/user_requests.py +++ b/litellm/types/proxy/management_endpoints/user_requests.py @@ -1,7 +1,7 @@ from collections.abc import Mapping -from typing import Literal +from typing import Annotated, Final, Literal -from pydantic import Field, field_validator, model_validator +from pydantic import AfterValidator, Field, field_validator, model_validator from litellm.models.user import LiteLLM_UserTable from litellm.types.proxy.auth.user_roles import LitellmUserRoles @@ -15,7 +15,17 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): user_role: Literal["user", "admin"] = "user" +def _validate_custom_user_id(value: str) -> str: + normalized: Final = value.strip() + if not normalized: + raise ValueError("user_id must not be empty or contain only whitespace") + if "\x00" in normalized: + raise ValueError("user_id must not contain NUL characters") + return normalized + + class NewUserRequest(GenerateRequestBase): + user_id: Annotated[str, AfterValidator(_validate_custom_user_id)] | None = None tpm_limit: int | None = Field(default=None, ge=0) rpm_limit: int | None = Field(default=None, ge=0) max_budget: float | None = None diff --git a/litellm/types/proxy/skill_hub.py b/litellm/types/proxy/skill_hub.py index d575defcda7..050c3cb7eb8 100644 --- a/litellm/types/proxy/skill_hub.py +++ b/litellm/types/proxy/skill_hub.py @@ -7,6 +7,17 @@ from pydantic import ConfigDict, Field from litellm.models.skill_submission import SubmissionOperation, SubmissionStatus from litellm.types.llms.base import LiteLLMBaseModel +SkillSortField: TypeAlias = Literal[ + "name", + "download_count", + "updated_at", + "created_at", + "created_by", + "published_by", + "approved_by", + "security_scan_passed", +] + class SkillHubSkill(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) @@ -16,6 +27,11 @@ class SkillHubSkill(LiteLLMBaseModel): description: str | None = None version: str | None = None updated_at: datetime | None = None + created_at: datetime | None = None + created_by: str | None = None + approved_by: str | None = None + published_by: str | None = None + security_scan_passed: bool = False visibility: Literal["public", "private"] source: Mapping[str, str] download_url: str | None = None @@ -23,12 +39,11 @@ class SkillHubSkill(LiteLLMBaseModel): domain: str | None = None namespace: str | None = None keywords: tuple[str, ...] | None = None - download_count: int | None = Field( - None, description="Gateway ZIP downloads, included in proxy-admin catalog listings" - ) + download_count: int | None = Field(None, description="Gateway ZIP downloads, including repeat downloads") class SkillHubList(LiteLLMBaseModel): + tags: tuple[str, ...] = () model_config = ConfigDict(frozen=True) skills: tuple[SkillHubSkill, ...] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ba29a67a36c..ca434154f68 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -703,6 +703,8 @@ class CallTypes(str, Enum): adelete_responses = "adelete_responses" cancel_responses = "cancel_responses" acancel_responses = "acancel_responses" + compact_responses = "compact_responses" + acompact_responses = "acompact_responses" alist_input_items = "alist_input_items" llm_passthrough_route = "llm_passthrough_route" allm_passthrough_route = "allm_passthrough_route" @@ -806,6 +808,8 @@ CallTypesLiteral = Literal[ "send_message", "aresponses", "responses", + "compact_responses", + "acompact_responses", "acreate_skill", "acreate_realtime_client_secret", "arealtime_calls", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a5b9a5cb457..a29f97ee9ea 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4046,6 +4046,8 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 2.5e-07, + "deprecation_date": "2027-09-01", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "litellm_provider": "azure_ai", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -4068,7 +4070,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, "prompt_cache_min_tokens": 512, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models" + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" }, "azure_ai/claude-opus-5": { "deprecation_date": "2027-07-08", @@ -13458,13 +13460,16 @@ "supports_reasoning": true }, "azure_ai/mistral-document-ai-2512": { + "deprecation_date": "2027-02-15", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.003, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" ], - "source": "https://ai.azure.com/catalog/models/mistral-document-ai-2512" + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", + "successor_model": "azure_ai/mistral-ocr-4-0" }, "azure_ai/mistral-ocr-4-0": { "litellm_provider": "azure_ai", @@ -40893,6 +40898,22 @@ "output_cost_per_token": 4e-07, "supports_tool_choice": true }, + "jina-reranker-v3.5": { + "input_cost_per_token": 5e-08, + "litellm_provider": "jina_ai", + "max_input_tokens": 131072, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" + }, + "jina-reranker-m0": { + "input_cost_per_token": 5e-08, + "litellm_provider": "jina_ai", + "max_input_tokens": 10240, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" + }, "jina-reranker-v2-base-multilingual": { "input_cost_per_token": 5e-08, "litellm_provider": "jina_ai", @@ -43771,7 +43792,8 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/zai-glm-5-3": { "cache_read_input_token_cost": 1.4e-07, @@ -45270,6 +45292,18 @@ "source": "https://tokenfactory.nebius.com/api/public/models_info", "cache_read_input_token_cost": 1.5e-07 }, + "nebius/moonshotai/Kimi-K3.1-Pro": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "nebius", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "supports_prompt_caching": true, + "source": "https://tokenfactory.nebius.com/api/public/models_info" + }, "nebius/moonshotai/Kimi-K3": { "max_input_tokens": 1048576, "max_output_tokens": 1024000, @@ -71189,6 +71223,8 @@ }, "gpt-5.5-cyber-preview": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-02-01", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "input_cost_per_token": 1.25e-05, "litellm_provider": "openai", "mode": "chat", @@ -72148,6 +72184,7 @@ "output_cost_per_token": 4.5e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_audio_output": true, "tpm": 250000, "rpm": 10, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, @@ -72170,6 +72207,7 @@ "output_cost_per_token": 4.5e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_audio_output": true, "tpm": 250000, "rpm": 10, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, @@ -76469,6 +76507,33 @@ "supports_tool_choice": true, "supports_vision": true }, + "us-gov.xai.grok-4.7": { + "supports_regex_lookaround": false, + "supported_endpoints": ["/v1/chat/completions"], + "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_response_format": true, + "cache_read_input_token_cost": 6.6e-07, + "cache_read_input_token_cost_flex": 3.3e-07, + "cache_read_input_token_cost_priority": 1.155e-06, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_batches": 1.32e-06, + "input_cost_per_token_flex": 1.32e-06, + "input_cost_per_token_priority": 4.62e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 7.92e-06, + "output_cost_per_token_batches": 3.96e-06, + "output_cost_per_token_flex": 3.96e-06, + "output_cost_per_token_priority": 1.386e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json", + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "baseten/deepseek-ai/DeepSeek-V4.1-Flash-Fast": { "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 6e-07, @@ -78437,7 +78502,7 @@ "cache_creation_input_token_cost": 1e-05, "cache_creation_input_token_cost_above_1hr": 1.6e-05, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2027-10-01", + "deprecation_date": "2027-01-15", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "input_cost_per_token": 8e-06, "litellm_provider": "azure_ai", @@ -78447,7 +78512,7 @@ "output_cost_per_token": 4e-05, "prompt_cache_min_tokens": 512.0, "search_context_cost_per_query": {"search_context_size_high":0.02,"search_context_size_low":0.02,"search_context_size_medium":0.02}, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models", + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -82782,6 +82847,39 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true }, + "gpt-6-cyber": { + "cache_creation_input_token_cost": 2.5e-05, + "cache_creation_input_token_cost_above_12hr": 5e-05, + "cache_read_input_token_cost": 2e-06, + "input_cost_per_token": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0001, + "source": "https://developers.openai.com/api/docs/models/gpt-6-cyber", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true + }, "gpt-6-astra-law": { "cache_read_input_token_cost": 1e-06, "input_cost_per_token": 1e-05, @@ -83380,6 +83478,28 @@ "supports_response_schema": true, "source": "https://empiriolabs.ai/models/kimi-k3" }, + "empiriolabs/kimi-k3-1-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1.8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/kimi-k3-1-pro" + }, "empiriolabs/mimo-v2-6-flash": { "litellm_provider": "empiriolabs", "mode": "chat", @@ -86383,5 +86503,32 @@ "supports_system_messages": false, "supports_video_input": true, "supports_vision": true + }, + "databricks/databricks-kimi-k3-1-pro": { + "cache_creation_input_token_cost": 1.99997e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "metadata": {"notes":"Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference."}, + "mode": "chat", + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.000142857, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 100.0 } } diff --git a/osv-scanner.toml b/osv-scanner.toml index 45fb7635088..11c2d4fbd87 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -83,11 +83,6 @@ id = "GHSA-45c3-73f7-pv4m" ignoreUntil = 2026-11-27 reason = "dnspython 2.9.0 (published 2026-10-09) is the only fix and requires Python >= 3.11, so the Python 3.10 resolution stays on 2.8.0 while 3.11+ locks 2.9.0; the advisory is CPU use in dns.name.from_wire_parser and litellm only reaches dnspython through pydantic EmailStr, which never does DNS lookups; remove this entry once 3.10 is dropped or a 3.10-compatible fix ships" -[[IgnoredVulns]] -id = "GHSA-xcx6-jx8p-x3h5" -ignoreUntil = 2026-11-24 -reason = "@azure/msal-node is fixed only in 7.1.0 (published 2026-10-13), but it is dev-only in vscode-extension via @vscode/vsce > @azure/identity, and the latest @azure/identity (4.14.0) still requires msal-node ^6.0.0, so the lock cannot move without forcing a major override; the advisory is in the auth_code flow, which extension packaging does not run; bump and remove this entry once @azure/identity accepts msal-node >= 7.1.0" - [[PackageOverrides]] name = "pillow" version = "12.3.0" diff --git a/packaging/litellm-proxy/pyproject.toml b/packaging/litellm-proxy/pyproject.toml new file mode 100644 index 00000000000..172f40a7228 --- /dev/null +++ b/packaging/litellm-proxy/pyproject.toml @@ -0,0 +1,22 @@ +[project] +name = "litellm-proxy" +version = "1.109.0" +description = "Opt-in Rust HTTP forwarding server for the LiteLLM Python proxy" +requires-python = ">=3.10, <3.16" +license = "MIT" + +[build-system] +requires = ["maturin==1.15.0"] +build-backend = "maturin" + +[tool.maturin] +manifest-path = "../../litellm-rust/crates/gateway/Cargo.toml" +bindings = "bin" +targets = [{ name = "litellm-gateway", kind = "bin" }] +profile = "release" +editable-profile = "dev" +include = [ + { path = "rust-toolchain.toml", format = "sdist" }, + { path = ".cargo/config.toml", format = "sdist" }, + { path = "model_prices_and_context_window.json", format = "sdist" }, +] diff --git a/pyproject.toml b/pyproject.toml index ee6cb6c1dcc..0b81618355d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -359,6 +359,7 @@ exclude = [ ] [tool.uv] +concurrent-builds = 1 cache-keys = [ { file = "pyproject.toml" }, { file = "rust-toolchain.toml" }, @@ -395,11 +396,12 @@ brotli = "2026-10-22T14:01:00Z" pypdf = "2026-10-28T17:01:00Z" [tool.uv.sources] +litellm-proxy = { workspace = true } litellm-proxy-extras = { workspace = true } litellm-enterprise = { workspace = true } [tool.uv.workspace] -members = ["enterprise", "litellm-proxy-extras"] +members = ["enterprise", "litellm-proxy-extras", "packaging/litellm-proxy"] [tool.commitizen] version = "1.109.0" diff --git a/scripts/build_core_distribution.py b/scripts/build_core_distribution.py index 41d7208d2c3..c458f6cb44f 100644 --- a/scripts/build_core_distribution.py +++ b/scripts/build_core_distribution.py @@ -1,4 +1,4 @@ -"""Build the independent core distribution from shared repository sources.""" +"""Build standalone core or gateway distributions from shared repository sources.""" import argparse import shutil @@ -19,8 +19,16 @@ SOURCES: Final = ( "LICENSE", ) +IN_TREE_RUST_PATH: Final = "../../litellm-rust/" -def stage_core_distribution(source: Path, destination: Path) -> None: + +class BuildOptions(argparse.Namespace): + package: str + out_dir: Path + sdist_only: bool + + +def stage_core_distribution(source: Path, destination: Path, *, package: str = "litellm-core") -> None: version: Final = subprocess.run( ["uv", "version", "--short"], cwd=source, check=True, capture_output=True, text=True ).stdout.strip() @@ -53,8 +61,7 @@ def stage_core_distribution(source: Path, destination: Path) -> None: )(directory, names) destination.mkdir(parents=True, exist_ok=True) - for name in SOURCES: - path: Final = source / name + for name, path in ((name, source / name) for name in SOURCES if package == "litellm-core" or name != "litellm"): if path.is_dir(): shutil.copytree( path, @@ -63,15 +70,19 @@ def stage_core_distribution(source: Path, destination: Path) -> None: ) else: shutil.copy2(path, destination / name) - shutil.copy2(source / "packaging/litellm-core/pyproject.toml", destination / "pyproject.toml") + template: Final = (source / "packaging" / package / "pyproject.toml").read_text() + if package == "litellm-proxy" and IN_TREE_RUST_PATH not in template: + raise SystemExit(f"packaging/{package}/pyproject.toml no longer points at {IN_TREE_RUST_PATH}") + metadata: Final = template.replace(IN_TREE_RUST_PATH, "litellm-rust/") + (destination / "pyproject.toml").write_text(metadata) subprocess.run(["uv", "version", version, "--frozen"], cwd=destination, check=True) -def build_core_distribution(output: Path, *, sdist_only: bool = False) -> None: +def build_core_distribution(output: Path, *, sdist_only: bool = False, package: str = "litellm-core") -> None: output_path: Final = output.resolve() with tempfile.TemporaryDirectory(prefix="litellm-core-") as temporary: stage: Final = Path(temporary) - stage_core_distribution(ROOT, stage) + stage_core_distribution(ROOT, stage, package=package) subprocess.run( ["uv", "build", "--python", sys.executable, "--out-dir", str(output_path)] + (["--sdist"] if sdist_only else []), @@ -82,12 +93,13 @@ def build_core_distribution(output: Path, *, sdist_only: bool = False) -> None: def main() -> None: parser: Final = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--package", choices=("litellm-core", "litellm-proxy"), default="litellm-core") parser.add_argument("--out-dir", type=Path, default=ROOT / "dist/core") parser.add_argument( "--sdist-only", action="store_true", help="Build only the source archive, without compiling a wheel" ) - args: Final = parser.parse_args() - build_core_distribution(Path(args.out_dir), sdist_only=args.sdist_only) + args: Final = parser.parse_args(namespace=BuildOptions()) + build_core_distribution(Path(args.out_dir), sdist_only=args.sdist_only, package=args.package) if __name__ == "__main__": diff --git a/tests/base_sdk_tests/test_core_distribution.py b/tests/base_sdk_tests/test_core_distribution.py index c22fa4e54fc..04536630200 100644 --- a/tests/base_sdk_tests/test_core_distribution.py +++ b/tests/base_sdk_tests/test_core_distribution.py @@ -20,12 +20,21 @@ else: ROOT: Final = Path(__file__).resolve().parents[2] +@pytest.mark.parametrize( + ("package_cli_args", "package"), + [ + (("--package", "litellm-proxy"), "litellm-proxy"), + ((), "litellm-core"), + ], +) @pytest.mark.parametrize("sdist_only", [False, True]) -def test_build_command_selects_requested_distributions(tmp_path: Path, sdist_only: bool) -> None: +def test_build_command_selects_requested_distributions( + tmp_path: Path, sdist_only: bool, package_cli_args: tuple[str, ...], package: str +) -> None: from scripts.build_core_distribution import main output: Final = tmp_path / "dist" - arguments: Final = ["build_core_distribution.py", "--out-dir", str(output)] + ( + arguments: Final = ["build_core_distribution.py", *package_cli_args, "--out-dir", str(output)] + ( ["--sdist-only"] if sdist_only else [] ) with ( @@ -42,6 +51,7 @@ def test_build_command_selects_requested_distributions(tmp_path: Path, sdist_onl assert command[command.index("--out-dir") + 1] == str(output) assert run.call_args.kwargs["check"] is True assert run.call_args.kwargs["cwd"] == stage.call_args.args[1] + assert stage.call_args.kwargs["package"] == package assert not stage.call_args.args[1].exists() @@ -53,7 +63,10 @@ def source_repository(tmp_path: Path) -> Path: return source -def test_staging_stamps_release_version_without_modifying_sources(tmp_path: Path, source_repository: Path) -> None: +@pytest.mark.parametrize("package", ["litellm-core", "litellm-proxy"]) +def test_staging_stamps_release_version_without_modifying_sources( + tmp_path: Path, source_repository: Path, package: str +) -> None: from scripts.build_core_distribution import SOURCES, stage_core_distribution source: Final = source_repository @@ -61,16 +74,44 @@ def test_staging_stamps_release_version_without_modifying_sources(tmp_path: Path path: Final = source / name path.write_text(f"shared {name}") (source / "pyproject.toml").write_text('[project]\nname = "litellm"\nversion = "9.8.7rc1"\n') - manifest: Final = source / "packaging/litellm-core/pyproject.toml" + manifest: Final = source / "packaging" / package / "pyproject.toml" manifest.parent.mkdir(parents=True) - manifest.write_bytes((ROOT / "packaging/litellm-core/pyproject.toml").read_bytes()) + manifest.write_bytes((ROOT / "packaging" / package / "pyproject.toml").read_bytes()) original: Final = manifest.read_bytes() stage: Final = tmp_path / "stage" - stage_core_distribution(source, stage) - assert tomllib.loads((stage / "pyproject.toml").read_text())["project"]["version"] == "9.8.7rc1" + stage_core_distribution(source, stage, package=package) + staged_project: Final = tomllib.loads((stage / "pyproject.toml").read_text()) + assert staged_project["project"]["name"] == package + assert staged_project["project"]["version"] == "9.8.7rc1" + if package == "litellm-proxy": + assert staged_project["tool"]["maturin"]["targets"] == [{"name": "litellm-gateway", "kind": "bin"}] + assert not (stage / "litellm").exists() + assert all( + (stage / name).exists() + for name in ("litellm-rust", ".cargo", "rust-toolchain.toml", "model_prices_and_context_window.json") + ) assert manifest.read_bytes() == original assert (source / "pyproject.toml").read_text() == '[project]\nname = "litellm"\nversion = "9.8.7rc1"\n' - assert all((stage / name).read_bytes() == (source / name).read_bytes() for name in SOURCES) + expected: Final = tuple(name for name in SOURCES if package == "litellm-core" or name != "litellm") + assert all((stage / name).read_bytes() == (source / name).read_bytes() for name in expected) + + +def test_gateway_staging_refuses_manifest_without_in_tree_rust_path(tmp_path: Path, source_repository: Path) -> None: + from scripts.build_core_distribution import SOURCES, stage_core_distribution + + source: Final = source_repository + for name in SOURCES: + (source / name).write_text(f"shared {name}") + (source / "pyproject.toml").write_text('[project]\nname = "litellm"\nversion = "9.8.7"\n') + manifest: Final = source / "packaging" / "litellm-proxy" / "pyproject.toml" + manifest.parent.mkdir(parents=True) + manifest.write_text( + (ROOT / "packaging" / "litellm-proxy" / "pyproject.toml") + .read_text() + .replace("../../litellm-rust/", "../../rust/") + ) + with pytest.raises(SystemExit, match="no longer points at"): + stage_core_distribution(source, tmp_path / "stage", package="litellm-proxy") @pytest.fixture(scope="module") @@ -99,7 +140,9 @@ def test_core_wheel_metadata_and_resources(distributions: tuple[Path, Path]) -> assert metadata["Version"] == tomllib.loads((ROOT / "pyproject.toml").read_text())["project"]["version"] assert not metadata.get_all("Provides-Extra") assert not any(n.endswith(".dist-info/entry_points.txt") for n in names) - requirements: Final = tomllib.loads((ROOT / "packaging/litellm-core/pyproject.toml").read_text())["project"]["dependencies"] + requirements: Final = tomllib.loads((ROOT / "packaging/litellm-core/pyproject.toml").read_text())["project"][ + "dependencies" + ] for python_version in ("3.10", "3.11", "3.12", "3.13", "3.14"): environment: Final = {"python_version": python_version, "python_full_version": python_version + ".0"} assert { diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index cd951c6cd28..d89badfb33e 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -89,6 +89,7 @@ importlib-metadata: >=6.8.0 # Apache 2.0 License tokenizers: >=0.20.2 # Apache 2.0 License jinja2: >=3.1.4 # BSD 3-Clause License litellm-proxy-extras: >=0.1.1 # MIT License +litellm-proxy: >=1.109.0 # MIT, verified in repository LICENSE for litellm-rust and packaging/litellm-proxy litellm-enterprise: >=0.1.1 # LiteLLM Enterprise License a2a-sdk: >=0.3.22 # Apache 2.0 license anyio: >=4.5.0 # Unknown license diff --git a/tests/integration/_support/daily_activity.py b/tests/integration/_support/daily_activity.py index 346fc156ea9..cedcb80843a 100644 --- a/tests/integration/_support/daily_activity.py +++ b/tests/integration/_support/daily_activity.py @@ -272,16 +272,32 @@ def assert_key_reported( date: str, metadata: Mapping[str, JsonValue], metrics: Mapping[str, float], + *, + entity_rows: Mapping[str, int] | None = None, ) -> None: assert response.status_code == 200, response.text body: Final = object_value(response.json()) records: Final = tuple(object_value(record) for record in records_of_key(body, api_key)) assert records, response.text assert all(record["metadata"] == metadata for record in records), response.text - assert all(counted(record["metrics"]) == pytest.approx(metrics) for record in records), response.text days: Final = body["results"] assert isinstance(days, list) and len(days) == 1, response.text day: Final = object_value(days[0]) + breakdown: Final = object_value(day["breakdown"]) + if entity_rows is None: + assert all(counted(record["metrics"]) == pytest.approx(metrics) for record in records), response.text + else: + totals: Final = records_of_key({name: value for name, value in breakdown.items() if name != "entities"}, api_key) + assert totals, response.text + assert all(counted(object_value(record)["metrics"]) == pytest.approx(metrics) for record in totals), response.text + entities: Final = object_value(breakdown["entities"]) + assert set(entities) == set(entity_rows), response.text + for entity, rows in entity_rows.items(): + bucket: Final = object_value(entities[entity]) + assert counted(bucket["metrics"]) == pytest.approx(seeded_metrics(rows)), response.text + keys: Final = object_value(bucket["api_key_breakdown"]) + assert set(keys) == {api_key}, response.text + assert counted(object_value(keys[api_key])["metrics"]) == pytest.approx(seeded_metrics(rows)), response.text assert day["date"] == date, response.text assert counted(day["metrics"]) == pytest.approx(metrics), response.text assert object_value(body["metadata"])["total_spend"] == pytest.approx(metrics["spend"]), response.text diff --git a/tests/integration/_support/device_login.py b/tests/integration/_support/device_login.py index 55519678f40..55eb5a753bb 100644 --- a/tests/integration/_support/device_login.py +++ b/tests/integration/_support/device_login.py @@ -25,6 +25,7 @@ from pydantic import JsonValue CHATGPT_AUTH_HOST: Final = "auth.openai.com" GITHUB_HOST: Final = "github.com" GITHUB_API_HOST: Final = "api.github.com" +COPILOT_API_HOST: Final = "api.githubcopilot.com" AUTH_HOSTS: Final = (CHATGPT_AUTH_HOST, GITHUB_HOST, GITHUB_API_HOST) CHATGPT_ACCOUNT: Final = "acct-device-login" _FAR_FUTURE: Final = 4102444800 @@ -79,7 +80,17 @@ COPILOT_MINTED_KEY: Final = "tid=device-login;minted" COPILOT_REJECTED_KEY: Final = "tid=device-login;rejected" CONTROL_KEY: Final = "device-login-control-key" ACCEPTED_BEARERS: Final = frozenset( - {CHATGPT_STORED, CHATGPT_REFRESHED, CHATGPT_FIRST_LOGIN, COPILOT_STORED_KEY, COPILOT_MINTED_KEY, CONTROL_KEY} + { + CHATGPT_STORED, + CHATGPT_REFRESHED, + CHATGPT_FIRST_LOGIN, + COPILOT_ACCESS, + COPILOT_FIRST_LOGIN_ACCESS, + COPILOT_UNSAVED_KEY_ACCESS, + COPILOT_STORED_KEY, + COPILOT_MINTED_KEY, + CONTROL_KEY, + } ) REJECTED_DETAIL: Final = "the scripted provider rejected this bearer token" _DEVICE_AUTH_ID: Final = "device-auth-synthetic" @@ -180,6 +191,10 @@ def _auth_hosts(api_url: str, grant: threading.Event) -> Callable[[Request], Rep return _json(200, {"access_token": COPILOT_FIRST_LOGIN_ACCESS, "token_type": "bearer"}) case ("api.github.com", "GET", "/copilot_internal/v2/token"): return _copilot_api_key(request, api_url) + case ("api.github.com", "GET", "/copilot_internal/user"): + if request.headers.get("authorization") == f"token {COPILOT_REJECTED_ACCESS}": + return _json(401, {"message": "Bad credentials"}) + return _json(200, {"endpoints": {"api": f"https://{COPILOT_API_HOST}"}}) case _: return _switched_off() @@ -204,9 +219,22 @@ class Switches: switch.clear() +@dataclass(frozen=True, slots=True) +class _ProviderAPI: + plain: Wire + secure: Wire + + @property + def url(self) -> str: + return self.plain.url + + def drain(self) -> tuple[Request, ...]: + return (*self.plain.drain(), *self.secure.drain()) + + @dataclass(frozen=True, slots=True) class Peers: - api: Wire + api: _ProviderAPI auth: Wire tunnel: str cert: Path @@ -228,6 +256,14 @@ class Peers: def auth_connections(self) -> tuple[str, ...]: return tuple(self.authorities.get_nowait() for _ in range(self.authorities.qsize())) + def assert_only_copilot_discovery(self, access_token: str) -> None: + assert set(self.auth_connections()) <= {f"{GITHUB_API_HOST}:443"} + requests: Final = self.auth.drain() + assert { + (request.method, request.target, request.headers.get("host"), request.headers.get("authorization")) + for request in requests + } <= {("GET", "/copilot_internal/user", GITHUB_API_HOST, f"token {access_token}")}, requests + def dropped(self) -> tuple[Hangup, ...]: return tuple(self.hangups.get_nowait() for _ in range(self.hangups.qsize())) @@ -279,9 +315,8 @@ class _TunnelServer(socketserver.ThreadingTCPServer): @contextmanager def _tunnel( - destination: Wire, switches: Switches, authorities: SimpleQueue[str], hangups: SimpleQueue[Hangup] + destination: Wire, provider: Wire, switches: Switches, authorities: SimpleQueue[str], hangups: SimpleQueue[Hangup] ) -> Generator[str]: - destination_port: Final = int(destination.url.rsplit(":", 1)[1]) class Tunnel(socketserver.StreamRequestHandler): rbufsize = 0 @@ -306,20 +341,24 @@ def _tunnel( if len(request_line) < 2: return authority: Final = request_line[1] - if authority.rsplit(":", 1)[0] not in AUTH_HOSTS: + host: Final = authority.rsplit(":", 1)[0] + if host not in (*AUTH_HOSTS, COPILOT_API_HOST): self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n") return - authorities.put(authority) - if switches.refuse_connect.is_set(): + is_auth: Final = host in AUTH_HOSTS + if is_auth: + authorities.put(authority) + if is_auth and switches.refuse_connect.is_set(): self.wfile.write(b"HTTP/1.1 502 Bad Gateway\r\ncontent-length: 0\r\n\r\n") return - if self.hold(authority, opened, switches.hold_connect) is None: + if is_auth and self.hold(authority, opened, switches.hold_connect) is None: return self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n") - client_hello: Final = self.hold(authority, opened, switches.stall_tls) + client_hello: Final = self.hold(authority, opened, switches.stall_tls) if is_auth else b"" if client_hello is None: return self.request.settimeout(10) + destination_port: Final = int((destination if is_auth else provider).url.rsplit(":", 1)[1]) with socket.create_connection(("127.0.0.1", destination_port), timeout=10) as upstream: upstream.sendall(client_hello) outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream)) @@ -340,16 +379,17 @@ def _tunnel( @contextmanager def device_login_peers(directory: Path) -> Generator[Peers]: - cert, key = write_self_signed_cert(directory, AUTH_HOSTS) + cert, key = write_self_signed_cert(directory, (*AUTH_HOSTS, COPILOT_API_HOST)) switches: Final = Switches(threading.Event(), threading.Event(), threading.Event(), threading.Event()) authorities: Final[SimpleQueue[str]] = SimpleQueue() hangups: Final[SimpleQueue[Hangup]] = SimpleQueue() with ( wire_server(_provider_api(rv.ResponsesVendor())) as api, + wire_server(_provider_api(rv.ResponsesVendor()), tls=server_context(cert, key)) as copilot_api, wire_server(_auth_hosts(api.url, switches.grant), tls=server_context(cert, key)) as auth, - _tunnel(auth, switches, authorities, hangups) as tunnel, + _tunnel(auth, copilot_api, switches, authorities, hangups) as tunnel, ): - yield Peers(api, auth, tunnel, cert, authorities, hangups, switches) + yield Peers(_ProviderAPI(api, copilot_api), auth, tunnel, cert, authorities, hangups, switches) def chatgpt_record( diff --git a/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py b/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py new file mode 100644 index 00000000000..3227caaa7f5 --- /dev/null +++ b/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py @@ -0,0 +1,114 @@ +"""GET, compact, DELETE and cancel on /v1/responses answer 200 when a during_call guardrail is on by default +(LIT-9770): the hook used to raise `'acompact_responses' is not a valid CallTypes` and the proxy returned 500""" + +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, string_value +from integration._support.process import owned_proxy +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.shares_worker_state + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MODEL: Final = "openai/gpt-5.4-mini" + + +def _input() -> list[JsonValue]: + return [ + {"role": "user", "content": f"Remember the number {uuid.uuid4().hex}."}, + {"role": "assistant", "content": "Got it."}, + ] + + +def _response(response_id: str) -> Reply: + body: Final[dict[str, JsonValue]] = { + "id": response_id, + "object": "response", + "created_at": 1, + "status": "completed", + "model": _MODEL.removeprefix("openai/"), + "output": [ + { + "type": "message", + "id": f"msg_{uuid.uuid4().hex}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Got it.", "annotations": []}], + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}, + } + return Reply(body=json.dumps(body).encode()) + + +def _deleted(response_id: str) -> Reply: + return Reply(body=json.dumps({"id": response_id, "object": "response.deleted", "deleted": True}).encode()) + + +def _guarded_config(path: Path) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + guardrail: Final[dict[str, JsonValue]] = { + "guardrail_name": f"content-filter-{uuid.uuid4().hex}", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "during_call", + "default_on": True, + "blocked_words": [{"keyword": "forbiddenword", "action": "BLOCK"}], + }, + } + path.write_text(yaml.safe_dump({**base, "guardrails": [guardrail]})) + return path + + +@pytest.fixture(scope="module") +def guarded(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("during-call-guardrail") + with owned_proxy(session_gateway, directory, {}, config=_guarded_config(directory / "config.yaml")) as candidate: + yield candidate + + +def _create(guarded: Gateway, provider: SharedProvider, upstream_id: str) -> str: + provider.expect(_response(upstream_id)) + created: Final = guarded.request("POST", "/v1/responses", {"model": _MODEL, "input": _input(), "store": True}) + assert created.status_code == 200, created.text + (upstream,) = provider.received() + assert upstream.target.endswith("/responses"), upstream.target + return string_value(_JSON_OBJECT.validate_json(created.content)["id"]) + + +def _lifecycle(guarded: Gateway, provider: SharedProvider, operation: str) -> tuple[httpx.Response, str]: + if operation == "compact": + provider.expect(_response(f"resp_{uuid.uuid4().hex}")) + return guarded.request( + "POST", "/v1/responses/compact", {"model": _MODEL, "input": _input()} + ), "/responses/compact" + upstream_id: Final = f"resp_{uuid.uuid4().hex}" + managed_id: Final = _create(guarded, provider, upstream_id) + provider.expect(_deleted(upstream_id) if operation == "delete" else _response(upstream_id)) + if operation == "get": + return guarded.request("GET", f"/v1/responses/{managed_id}"), f"/responses/{upstream_id}" + if operation == "delete": + return guarded.request("DELETE", f"/v1/responses/{managed_id}"), f"/responses/{upstream_id}" + return guarded.request("POST", f"/v1/responses/{managed_id}/cancel"), f"/responses/{upstream_id}/cancel" + + +@pytest.mark.parametrize("operation", ["get", "compact", "delete", "cancel"]) +def test_responses_lifecycle_answers_200_with_a_during_call_guardrail_on( + guarded: Gateway, provider: SharedProvider, operation: str +) -> None: + response, upstream_target = _lifecycle(guarded, provider, operation) + assert response.status_code == 200, response.text + (upstream,) = provider.received() + assert upstream.target.endswith(upstream_target), upstream.target diff --git a/tests/integration/database/test_skill_submissions.py b/tests/integration/database/test_skill_submissions.py index cc00bf6e359..e55e19fb110 100644 --- a/tests/integration/database/test_skill_submissions.py +++ b/tests/integration/database/test_skill_submissions.py @@ -134,6 +134,8 @@ async def test_approval_commits_version_archive_history_and_review(pending_updat assert current is not None published_manifest: Final = { **case.proposal.manifest, + "_litellm_review": {"approved_by": "integration-reviewer", "security_scan_passed": False}, + "_litellm_published_by": "integration-reviewer", "source": { **TypeAdapter(dict[str, str]).validate_python(case.submission.proposal.manifest["source"]), "url": f"https://skills.example.test/claude-code/plugins/{case.original.name}/archive/" @@ -148,17 +150,22 @@ async def test_approval_commits_version_archive_history_and_review(pending_updat ) metadata: Final = await case.publications.get_metadata(case.original.name) assert metadata is not None + assert metadata.created_by == before.created_by == case.original.created_by + assert metadata.created_at == before.created_at + assert metadata.updated_at is not None and before.updated_at is not None + assert metadata.updated_at > before.updated_at assert await case.publications.get_content(metadata) == case.proposal.content history: Final = await case.publications.versions(metadata, 10, 0) assert tuple((version.version, version.description, version.archive_available) for version in history) == ( ("1.0.1", case.proposal.description, True), ("1.0.0", case.original.description, True), ) + assert history[0].published_by == "integration-reviewer" + assert history[1].published_by == case.original.created_by for version, publication, manifest in zip( - history, (case.proposal, case.original), (published_manifest, case.original.manifest), strict=True + history, (case.proposal, case.original), (published_manifest, {**case.original.manifest, "_litellm_published_by": case.original.created_by}), strict=True ): assert version.manifest == {**manifest, "version": version.version} - assert version.published_by == publication.created_by assert await case.publications.version_content(metadata, version.id) == publication.content diff --git a/tests/integration/management/test_claude_code_plugin_installation_preference.py b/tests/integration/management/test_claude_code_plugin_installation_preference.py index aaf6f02ff4e..6d41c690709 100644 --- a/tests/integration/management/test_claude_code_plugin_installation_preference.py +++ b/tests/integration/management/test_claude_code_plugin_installation_preference.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from functools import partial from pathlib import Path +from threading import Lock from typing import Final import httpx @@ -248,7 +249,7 @@ def test_duplicate_json_key_keeps_the_last_value(gateway: Gateway) -> None: assert stored_manifest(name)["installation_preference"] == "available" -def test_repeated_identical_update_is_idempotent(gateway: Gateway) -> None: +def test_repeated_identical_updates_publish_patch_versions_without_changing_preference(gateway: Gateway) -> None: name: Final = plugin_name() with gateway.scenario() as scenario: registered(scenario, name, installation_preference="required") @@ -257,7 +258,9 @@ def test_repeated_identical_update_is_idempotent(gateway: Gateway) -> None: replaced(gateway, name, installation_preference="required") second: Final = stored_rows(name) assert len(first) == 1, first - assert second == first + first_manifest: Final = JSON_OBJECT.validate_json(string_value(first[0]["manifest_json"])) + assert first_manifest["version"] == "1.0.1", first + assert second == [{**first[0], "manifest_json": json.dumps({**first_manifest, "version": "1.0.2"})}] assert_preference_served(gateway, name, "required") @@ -442,14 +445,16 @@ def test_burst_reads_across_processes_while_the_preference_flips(gateway: Gatewa name: Final = plugin_name() targets: Final = (gateway, peer) flips: Final = tuple(PREFERENCES[(index + 1) % len(PREFERENCES)] for index in range(13)) + writer: Final = Lock() with gateway.scenario() as scenario: registered(scenario, name, installation_preference="available") def job(index: int) -> tuple[str, JsonValue]: if index % 4 == 3: - return "flip", replace( - gateway, name, replacement(name, installation_preference=flips[index // 4]) - ).status_code + with writer: + return "flip", replace( + gateway, name, replacement(name, installation_preference=flips[index // 4]) + ).status_code return "read", marketplace_entry(targets[index % 2], name).get("installationPreference") with ThreadPoolExecutor(max_workers=8) as pool: diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index f6743a8751a..0c97ae62d07 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -42,7 +42,6 @@ from openai.types.chat import ChatCompletionMessageParam from openai.types.responses.tool_param import Mcp from pydantic import JsonValue, TypeAdapter - pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state] @@ -589,7 +588,11 @@ def _upstream_model(surface: Surface) -> str: }[surface] -PRICES: Final[Mapping[str, JsonValue]] = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "num_retries": 0} +PRICES: Final[Mapping[str, JsonValue]] = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "num_retries": 0, +} @contextmanager @@ -917,7 +920,9 @@ def test_an_open_server_in_another_gateway_cannot_widen_this_gateways_tools(gate register_mcp(scenario, local, "open" + uuid.uuid4().hex[:8], allow_all_keys=True) listed, called = official_client_outcomes(shared, key, "/mcp", f"{foreign_alias}-add", ADD) assert set(listed.tools) == { - f"{foreign_alias}-add", f"{foreign_alias}-multiply", f"{foreign_alias}-fail" + f"{foreign_alias}-add", + f"{foreign_alias}-multiply", + f"{foreign_alias}-fail", }, listed.tools assert called.ok and called.text == str(ADD["a"] + ADD["b"]), called assert [call["body"]["params"]["name"] for call in _peer_add_calls(foreign)] == ["add"] @@ -1528,6 +1533,7 @@ def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_ one_round_rows: Final = _model_rows( one_round.headers["x-litellm-call-id"], _logs_a_failure_billing(_price(FIRST_ROUND) + FIXED_MARGIN) ) + assert _statuses(one_round_rows) == ["failure"], one_round_rows rig.wire.drain() response: Final = rig.send(key, [AUTO], stream=True, guardrails=[STREAM_BLOCK]) assert STREAM_BLOCK_DETAIL in response.text, response.text diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 91f502179f3..384d6635bb0 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -67,9 +67,9 @@ class _TokenDirs: def environment(self) -> Mapping[str, str]: return MappingProxyType({"CHATGPT_TOKEN_DIR": str(self.chatgpt), "GITHUB_COPILOT_TOKEN_DIR": str(self.copilot)}) - def store_valid_tokens(self, api_url: str) -> None: + def store_valid_tokens(self) -> None: dl.write_chatgpt(self.chatgpt, dl.chatgpt_record(dl.CHATGPT_STORED)) - dl.write_copilot_key(self.copilot, dl.COPILOT_STORED_KEY, api_url) + dl.write_copilot_access(self.copilot, dl.COPILOT_ACCESS) def _token_dirs(directory: Path) -> _TokenDirs: @@ -242,16 +242,16 @@ def test_booting_without_tokens_lists_login_deployments_and_serves_only_after_to assert peers.auth_connections() == () assert dl.bearer(peers.api.drain()[-1]) == dl.CONTROL_KEY - dirs.store_valid_tokens(peers.api.url) + dirs.store_valid_tokens() _reload(owned) _every_worker_lists(owned, workers) for served_logins, name in enumerate(_LOGIN_NAMES, start=1): served: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert served.status == 200, served.text _successes_logged(key, at_least=_CONTROL_CALLS_BEFORE_THE_RELOAD + served_logins) - assert peers.auth_connections() == () + peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) forwarded: Final = peers.api.drain() - assert {dl.bearer(request) for request in forwarded} == {dl.CHATGPT_STORED, dl.COPILOT_STORED_KEY} + assert {dl.bearer(request) for request in forwarded} == {dl.CHATGPT_STORED, dl.COPILOT_ACCESS} dl.write_chatgpt(dirs.chatgpt, dl.chatgpt_record(dl.CHATGPT_EXPIRED, refresh_token=dl.GOOD_REFRESH)) peers.switches.hold_connect.set() @@ -287,7 +287,7 @@ def test_a_token_removed_after_boot_turns_every_request_into_the_refusal( ) -> None: dirs: Final = _token_dirs(tmp_path) with dl.device_login_peers(tmp_path) as peers: - dirs.store_valid_tokens(peers.api.url) + dirs.store_valid_tokens() with owned_proxy_process( gateway, tmp_path, @@ -310,7 +310,7 @@ def test_a_token_removed_after_boot_turns_every_request_into_the_refusal( refused: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert refused.status == 400, refused.text assert refusal in refused.message(), refused.text - assert peers.auth_connections() == () + peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) tail: Final = owned.log.read_text() assert dl.CHATGPT_USER_CODE not in tail and dl.COPILOT_USER_CODE not in tail, tail[-3000:] control: Final = _chat(_base_url(owned), key, _CONTROL, uuid.uuid4().hex) diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index 39e36fd95fa..9074ea83189 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -2,7 +2,6 @@ import asyncio import hashlib import itertools import json -import re import threading import time import uuid @@ -26,9 +25,10 @@ from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.timing import assert_cut_short, timed from integration._support.wire import Request -from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS from pydantic import JsonValue +from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS + pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] Provider: TypeAlias = Literal["chatgpt", "copilot"] @@ -52,7 +52,7 @@ _REFUSALS: Final[Mapping[Provider, str]] = MappingProxyType( {"chatgpt": dl.CHATGPT_REFUSAL, "copilot": dl.COPILOT_REFUSAL} ) _STORED_BEARERS: Final[Mapping[Provider, str]] = MappingProxyType( - {"chatgpt": dl.CHATGPT_STORED, "copilot": dl.COPILOT_STORED_KEY} + {"chatgpt": dl.CHATGPT_STORED, "copilot": dl.COPILOT_ACCESS} ) _CONTROL_MODEL: Final = "device-login-control" _EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0, "cache": {"no-cache": True}}) @@ -154,7 +154,7 @@ class _Scene: if provider == "chatgpt": dl.write_chatgpt(self.rig.chatgpt_dir, dl.chatgpt_record(dl.CHATGPT_STORED)) return - dl.write_copilot_key(self.rig.copilot_dir, dl.COPILOT_STORED_KEY, self.rig.peers.api.url) + dl.write_copilot_access(self.rig.copilot_dir, dl.COPILOT_ACCESS) def _config(peers: dl.Peers, directory: Path) -> Path: @@ -403,7 +403,7 @@ def test_missing_token_is_refused_with_the_login_instructions_before_any_auth_ho @pytest.mark.parametrize("cell", _CELLS, ids=_Cell.name) -def test_stored_token_serves_the_request_without_an_auth_host_call(scene: _Scene, cell: _Cell) -> None: +def test_stored_token_serves_the_request_without_a_device_login(scene: _Scene, cell: _Cell) -> None: scene.store_valid_token(cell.provider) call: Final = scene.call(cell) served: Final = _serve(call) @@ -412,7 +412,10 @@ def test_stored_token_serves_the_request_without_an_auth_host_call(scene: _Scene received: Final = scene.rig.peers.api.drain() (forwarded,) = _carrying(received, call.marker) assert dl.bearer(forwarded) == _STORED_BEARERS[cell.provider], _describe(received) - assert scene.rig.peers.auth_connections() == () + if cell.provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () (row,) = _spend_rows(scene.key) assert row["status"] == "success", row @@ -487,7 +490,7 @@ def test_missing_token_is_refused_on_the_other_model_routes(scene: _Scene, route assert _describe(scene.rig.peers.api.drain()) == () -def test_copilot_embeddings_are_served_with_the_stored_api_key(scene: _Scene) -> None: +def test_copilot_embeddings_are_served_with_the_stored_access_token(scene: _Scene) -> None: scene.store_valid_token("copilot") served: Final = scene.post( "/v1/embeddings", {"model": "github_copilot/text-embedding-3-small", "input": "login", **_EXTRA} @@ -496,9 +499,9 @@ def test_copilot_embeddings_are_served_with_the_stored_api_key(scene: _Scene) -> (item,) = rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(served.text)["data"]) assert item["embedding"] == [0.25, 0.5, 0.75], served.text (forwarded,) = scene.rig.peers.api.drain() - assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_STORED_KEY), forwarded.target + assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_ACCESS), forwarded.target assert forwarded.target.endswith("/embeddings"), forwarded.target - assert scene.rig.peers.auth_connections() == () + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) @pytest.mark.parametrize("provider", _PROVIDERS) @@ -562,71 +565,67 @@ def test_stored_token_the_provider_rejects_surfaces_its_401_without_a_login( if provider == "chatgpt": dl.write_chatgpt(scene.rig.chatgpt_dir, dl.chatgpt_record(dl.CHATGPT_REJECTED)) else: - dl.write_copilot_key(scene.rig.copilot_dir, dl.COPILOT_REJECTED_KEY, scene.rig.peers.api.url) + dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_REJECTED_ACCESS) call: Final = scene.call(_Cell(provider, endpoint, False, "httpx")) served: Final = _serve(call) assert served.status == 401, served.text assert dl.REJECTED_DETAIL in served.text, served.text received: Final = scene.rig.peers.api.drain() assert {dl.bearer(request) for request in _carrying(received, call.marker)} == { - dl.CHATGPT_REJECTED if provider == "chatgpt" else dl.COPILOT_REJECTED_KEY + dl.CHATGPT_REJECTED if provider == "chatgpt" else dl.COPILOT_REJECTED_ACCESS }, _describe(received) - assert scene.rig.peers.auth_connections() == () + if provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_REJECTED_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () @pytest.mark.parametrize("stale_key", (False, True), ids=("missing-api-key", "expired-api-key")) -def test_copilot_api_key_is_minted_once_from_the_stored_access_token(scene: _Scene, stale_key: bool) -> None: +def test_copilot_access_token_is_forwarded_without_minting_an_api_key(scene: _Scene, stale_key: bool) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_ACCESS) if stale_key: dl.write_copilot_key(scene.rig.copilot_dir, dl.COPILOT_REJECTED_KEY, scene.rig.peers.api.url, -3600) + legacy_key: Final = scene.rig.copilot_dir / "api-key.json" + stored_key: Final = legacy_key.read_bytes() if stale_key else None call: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) served: Final = _serve(call) assert served.status == 200, served.text assert set(rv.MARKER.findall(served.text)) == {call.marker}, served.text - (minted,) = scene.rig.peers.auth.drain() - assert (minted.method, minted.target, minted.headers.get("host"), minted.headers.get("authorization")) == ( - "GET", - "/copilot_internal/v2/token", - dl.GITHUB_API_HOST, - f"token {dl.COPILOT_ACCESS}", - ) + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) (forwarded,) = _carrying(scene.rig.peers.api.drain(), call.marker) - assert dl.bearer(forwarded) == dl.COPILOT_MINTED_KEY - assert json.loads((scene.rig.copilot_dir / "api-key.json").read_text())["token"] == dl.COPILOT_MINTED_KEY + assert dl.bearer(forwarded) == dl.COPILOT_ACCESS + assert (legacy_key.read_bytes() if legacy_key.exists() else None) == stored_key -def test_copilot_api_key_that_cannot_be_saved_is_minted_at_most_once_per_worker(scene: _Scene) -> None: +def test_copilot_access_token_serves_when_the_legacy_api_key_path_is_a_directory(scene: _Scene) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_UNSAVED_KEY_ACCESS) unwritable_key_file: Final = scene.rig.copilot_dir / "api-key.json" unwritable_key_file.mkdir() first: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) first_served: Final = _serve(first) assert first_served.status == 200, first_served.text - (minted,) = scene.rig.peers.auth.drain() - assert minted.headers.get("authorization") == f"token {dl.COPILOT_UNSAVED_KEY_ACCESS}" later: Final = tuple(scene.call(_Cell("copilot", "chat", False, "httpx")) for _ in range(3)) later_served: Final = tuple(_serve(call) for call in later) assert [served.status for served in later_served] == [200, 200, 200], [served.text for served in later_served] - assert len(scene.rig.peers.auth.drain()) <= 1 + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_UNSAVED_KEY_ACCESS) received: Final = scene.rig.peers.api.drain() forwarded: Final = tuple(_carrying(received, call.marker) for call in (first, *later)) assert [tuple(dl.bearer(request) for request in requests) for requests in forwarded] == [ - (dl.COPILOT_MINTED_KEY,) + (dl.COPILOT_UNSAVED_KEY_ACCESS,) ] * 4, _describe(received) assert unwritable_key_file.is_dir() def test_copilot_access_token_github_rejects_fails_the_request_without_a_device_login(scene: _Scene) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_REJECTED_ACCESS) - served: Final = _serve(scene.call(_Cell("copilot", "chat", False, "httpx"))) - assert served.status == 400, served.text - assert "Failed to refresh API key" in _error_message(served.text), served.text + call: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) + served: Final = _serve(call) + assert served.status == 401, served.text + assert dl.REJECTED_DETAIL in _error_message(served.text), served.text _failure_logged(served) - reached: Final = scene.rig.peers.auth.drain() - assert {(request.method, request.target) for request in reached} == {("GET", "/copilot_internal/v2/token")}, ( - _describe(reached) - ) - assert _describe(scene.rig.peers.api.drain()) == () + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_REJECTED_ACCESS) + (forwarded,) = _carrying(scene.rig.peers.api.drain(), call.marker) + assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_REJECTED_ACCESS) def test_expired_copilot_api_key_with_no_access_token_is_refused_before_any_auth_host_call(scene: _Scene) -> None: @@ -727,7 +726,10 @@ def test_a_token_mounted_after_a_refusal_serves_without_a_restart_and_removing_i after: Final = _serve(scene.call(cell)) assert after.status == 400, after.text assert _REFUSALS[provider] in _error_message(after.text), after.text - assert scene.rig.peers.auth_connections() == () + if provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () _COOLDOWN_SECONDS_LEFT: Final = 45 diff --git a/tests/integration/providers/test_gemini_interactions_wire.py b/tests/integration/providers/test_gemini_interactions_wire.py index 58a7ccb31c3..ded0718bd2b 100644 --- a/tests/integration/providers/test_gemini_interactions_wire.py +++ b/tests/integration/providers/test_gemini_interactions_wire.py @@ -1,5 +1,4 @@ import json -import uuid from pathlib import Path from typing import Final @@ -147,6 +146,7 @@ def test_interaction_forwards_native_body_to_gemini(gateway: Gateway, prefix: st "agent": None, "updated": None, "outputs": None, + "continuation_token": None, }, response.text assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text @@ -179,6 +179,7 @@ def test_interaction_authenticates_google_key_header(gateway: Gateway, prefix: s "agent": None, "updated": None, "outputs": None, + "continuation_token": None, }, response.text assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text @@ -218,6 +219,7 @@ def test_agent_interaction_uses_env_gemini_credentials(gateway: Gateway, tmp_pat "agent": "deep-research-pro-preview-12-2025", "updated": None, "outputs": None, + "continuation_token": None, }, response.text assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text @@ -277,3 +279,78 @@ def test_stream_interaction_ends_without_openai_done_terminator(gateway: Gateway expected: Final = b"".join(_proxy_stream_event(payload, model) for _event_type, payload in _STREAM_EVENTS) assert response.content == expected, response.text assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions?alt=sse")] + + +_LONG_DECODE_MODEL: Final = "gemini/gemini-4-argon" +_LONG_DECODE_BODY: Final = { + "input": "prove the prime number theorem", + "generation_config": {"max_output_tokens": 120000, "thinking_level": "high"}, +} + + +def _hop(hop_id: str, status: str, text: str, token: str | None, input_tokens: int, output_tokens: int) -> dict: + return { + "id": hop_id, + "object": "interaction", + "model": "gemini-4-argon", + "status": status, + "created": "2026-11-05T00:00:00Z", + "steps": [{"type": "model_output", "content": [{"type": "text", "text": text}]}], + "usage": { + "total_input_tokens": input_tokens, + "total_output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + "initial_input_tokens": 79, + }, + "continuation_token": token, + } + + +_FIRST_HOP: Final = _hop("interaction-hop-1", "requires_continuation", "Lemma 1.", "tok-1", 79, 32768) +_SECOND_HOP: Final = _hop("interaction-hop-2", "completed", " QED.", None, 32847, 1533) + + +def test_interaction_keeps_requires_continuation_status_and_token_for_the_caller(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.target == "/v1beta/interactions" + return Reply(body=json.dumps(_FIRST_HOP).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_LONG_DECODE_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY) + key: Final = scenario.key(models=[model]) + response: Final = gateway.request( + "POST", "/v1beta/interactions", {"model": model, **_LONG_DECODE_BODY}, key=key + ) + assert response.status_code == 200, response.text + sent: Final = [json.loads(r.body) for r in wire.drain()] + + assert sent == [{"model": "gemini-4-argon", **_LONG_DECODE_BODY}], sent + payload: Final = response.json() + assert payload["id"] == "interaction-hop-1", response.text + assert payload["status"] == "requires_continuation", response.text + assert payload["continuation_token"] == "tok-1", response.text + assert payload["usage"]["initial_input_tokens"] == 79, response.text + + +def test_interaction_forwards_a_caller_continuation_token_to_resume_a_long_decode(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.target == "/v1beta/interactions" + return Reply(body=json.dumps(_SECOND_HOP).encode()) + + resume: Final = { + "generation_config": _LONG_DECODE_BODY["generation_config"], + "previous_interaction_id": "interaction-hop-1", + "continuation_token": "tok-1", + } + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_LONG_DECODE_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY) + key: Final = scenario.key(models=[model]) + response: Final = gateway.request("POST", "/v1beta/interactions", {"model": model, **resume}, key=key) + assert response.status_code == 200, response.text + sent: Final = [json.loads(r.body) for r in wire.drain()] + + assert sent == [{"model": "gemini-4-argon", **resume}], sent + payload: Final = response.json() + assert payload["id"] == "interaction-hop-2", response.text + assert payload["status"] == "completed", response.text + assert payload["continuation_token"] is None, response.text diff --git a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py index 018c6a653d1..696253f6a98 100644 --- a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py +++ b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py @@ -299,14 +299,14 @@ def cooldown_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Cooldown pytest.param("chat", _BRIDGED, "pause_after_content", id="chat-bridged-to-responses"), ), ) -def test_a_deployment_whose_stream_stalls_stays_out_of_cooldown( +def test_a_deployment_whose_stream_stalls_is_cooled_down( cooldown_rig: _CooldownRig, endpoint: Endpoint, backend: str, shape: Shape ) -> None: with _peer(shape) as wire, cooldown_rig.gateway.scenario() as scenario: group, identity = _deployment(scenario, wire, backend) response: Final = _streamed(cooldown_rig.gateway, endpoint, group) error: Final = _failure_row(cooldown_rig.gateway, _ROWS, (group,)) - assert not cooled(cooldown_rig.cache, identity), (error, response.text) + assert eventually(lambda: cooled(cooldown_rig.cache, identity), bool, _ROW_SECONDS), (error, response.text) @_OWNED_PROXY diff --git a/tests/integration/routing/test_include_fallback_errors_wire.py b/tests/integration/routing/test_include_fallback_errors_wire.py index 9075f5e41d3..2f48bc43a22 100644 --- a/tests/integration/routing/test_include_fallback_errors_wire.py +++ b/tests/integration/routing/test_include_fallback_errors_wire.py @@ -61,10 +61,10 @@ _PROVIDERS: Final = MappingProxyType( ) _STREAMING: Final = (pytest.param(False, id="non-stream"), pytest.param(True, id="stream")) _NON_BOOLEAN_FLAGS: Final = ( - pytest.param(1, True, id="int"), - pytest.param("", False, id="empty-string"), - pytest.param([], False, id="list"), - pytest.param("x" * 5120, True, id="five-kilobytes"), + pytest.param(1, id="int"), + pytest.param("", id="empty-string"), + pytest.param([], id="list"), + pytest.param("x" * 5120, id="five-kilobytes"), ) _OPENAI_UNAUTHORIZED: Final = Reply( status=401, @@ -326,10 +326,8 @@ def test_default_gateway_keeps_the_flag_off_the_wire_on_a_fallback( assert observed.attempted == "1", observed -@pytest.mark.parametrize(("value", "errors_reported"), _NON_BOOLEAN_FLAGS) -def test_default_gateway_keeps_a_non_boolean_flag_off_the_wire( - registered: _Registered, value: JsonValue, errors_reported: bool -) -> None: +@pytest.mark.parametrize("value", _NON_BOOLEAN_FLAGS) +def test_default_gateway_keeps_a_non_boolean_flag_off_the_wire(registered: _Registered, value: JsonValue) -> None: sent: Final = _send( registered.gateway, _PATHS["chat"], @@ -497,9 +495,9 @@ def test_exposed_chat_without_a_fallback_keeps_the_flag_off_the_wire_and_the_err assert observed.errors_header is None, observed -@pytest.mark.parametrize(("value", "errors_reported"), _NON_BOOLEAN_FLAGS) -def test_exposed_chat_keeps_a_non_boolean_flag_off_the_wire_and_the_errors_header_follows_its_truthiness( - exposed: _Exposed, value: JsonValue, errors_reported: bool +@pytest.mark.parametrize("value", _NON_BOOLEAN_FLAGS) +def test_exposed_chat_keeps_a_non_boolean_flag_off_the_wire_and_reports_no_errors( + exposed: _Exposed, value: JsonValue ) -> None: sent: Final = _send( exposed.proxy, @@ -515,7 +513,7 @@ def test_exposed_chat_keeps_a_non_boolean_flag_off_the_wire_and_the_errors_heade assert observed.leaks == (), observed assert observed.hit == (_EXPOSED_OPENAI_PRIMARY, _EXPOSED_OPENAI_BACKUP), observed assert observed.attempted == "1", observed - assert _errors_mention_the_outage(observed.errors_header) is errors_reported, observed + assert observed.errors_header is None, observed def test_exposed_proxy_rejects_an_unauthenticated_flagged_request_before_the_wire(exposed: _Exposed) -> None: diff --git a/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py new file mode 100644 index 00000000000..13bf67582b3 --- /dev/null +++ b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py @@ -0,0 +1,95 @@ +import collections +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import eventually, gateway_from_environment, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.wire import Reply, Wire, wire_server +from integration.routing._prompt_cache_pin_support import ( + DEFAULT_MARK, + anthropic_deployment, + ask_turn, + chat_body, + claude_peer, + new_marker, + one_new_pin, + pin_config, + post_chat, + read_pins, + system_turn, +) +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +from litellm.constants import REDACTED_BY_LITELLM + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +_GROUP: Final = "pin-logging-off" +_LEFT: Final = "pin-logging-off-left" +_RIGHT: Final = "pin-logging-off-right" +_FOLLOW_UPS: Final = 20 +_REDACTED_MESSAGES: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}] +_LOGGED_EVENT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + redis: "Redis[bytes]" + logger_sink: Wire + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("prompt-cache-pin-message-logging-off") + with ( + gateway_from_environment() as gateway, + owned_redis(directory) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + wire_server(lambda request: Reply()) as logger_sink, + ): + config: Final = pin_config( + directory, + [anthropic_deployment(_GROUP, _LEFT, left.url), anthropic_deployment(_GROUP, _RIGHT, right.url)], + litellm_settings={ + "turn_off_message_logging": True, + "callbacks": ["integration._support.redaction_probe.probe"], + "cache": False, + "enable_redis_auth_cache": False, + }, + ) + environment: Final = { + "REDIS_HOST": cache.host, + "REDIS_PORT": str(cache.port), + "REDACTION_PROBE_URL": logger_sink.url, + } + with ( + owned_proxy_process(gateway, directory, environment, config=config, workers=1) as owned, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + yield _Rig(owned, redis, logger_sink) + + +def test_follow_ups_stay_on_the_first_deployment_while_the_logger_sees_redacted_messages(rig: _Rig) -> None: + before: Final = frozenset(read_pins(rig.redis)) + body: Final = chat_body(_GROUP, [system_turn(new_marker(), DEFAULT_MARK), ask_turn(new_marker())]) + + _, first_served = post_chat(rig.owned.gateway, body) + _, pin = one_new_pin(rig.redis, before) + assert pin.model_id == first_served, (pin, first_served) + + served: Final = tuple(post_chat(rig.owned.gateway, body)[1] for _ in range(_FOLLOW_UPS)) + assert served == (first_served,) * _FOLLOW_UPS, collections.Counter(served) + + eventually(rig.logger_sink.received.qsize, lambda count: count >= 1 + _FOLLOW_UPS, seconds=30) + logged: Final = tuple( + object_value(_LOGGED_EVENT.validate_json(request.body)["standard"])["messages"] + for request in rig.logger_sink.drain() + ) + assert len(logged) == 1 + _FOLLOW_UPS, len(logged) + assert all(messages == _REDACTED_MESSAGES for messages in logged), logged diff --git a/tests/integration/sdk/test_device_code_login_guard_sdk.py b/tests/integration/sdk/test_device_code_login_guard_sdk.py index 7640fe13d5d..256871faf7c 100644 --- a/tests/integration/sdk/test_device_code_login_guard_sdk.py +++ b/tests/integration/sdk/test_device_code_login_guard_sdk.py @@ -184,7 +184,7 @@ def clean(bench: _Bench) -> Iterator[_Bench]: _FIRST_LOGIN_BEARERS: Final[Mapping[Provider, str]] = MappingProxyType( - {"chatgpt": dl.CHATGPT_FIRST_LOGIN, "copilot": dl.COPILOT_MINTED_KEY} + {"chatgpt": dl.CHATGPT_FIRST_LOGIN, "copilot": dl.COPILOT_FIRST_LOGIN_ACCESS} ) @@ -205,7 +205,7 @@ def test_sync_call_on_the_main_thread_keeps_the_interactive_first_login( assert {dl.bearer(request) for request in forwarded} == {_FIRST_LOGIN_BEARERS[provider]}, [ (request.method, request.target) for request in forwarded ] - expected_files: Final = ("auth.json",) if provider == "chatgpt" else ("access-token", "api-key.json") + expected_files: Final = ("auth.json",) if provider == "chatgpt" else ("access-token",) assert clean.token_files() == expected_files diff --git a/tests/integration/sdk/test_router_include_fallback_errors_wire.py b/tests/integration/sdk/test_router_include_fallback_errors_wire.py index fb17e05b02a..a8952c8c2b5 100644 --- a/tests/integration/sdk/test_router_include_fallback_errors_wire.py +++ b/tests/integration/sdk/test_router_include_fallback_errors_wire.py @@ -36,11 +36,11 @@ _BURST: Final = 12 _CALLBACK_WINDOW: Final = 15.0 _STREAMING: Final = (pytest.param(False, id="non-stream"), pytest.param(True, id="stream")) _NON_BOOLEAN_FLAGS: Final = ( - pytest.param(1, True, id="int"), - pytest.param("", False, id="empty-string"), - pytest.param([], False, id="list"), - pytest.param("x" * 5120, True, id="five-kilobytes"), - pytest.param(False, False, id="false"), + pytest.param(1, id="int"), + pytest.param("", id="empty-string"), + pytest.param([], id="list"), + pytest.param("x" * 5120, id="five-kilobytes"), + pytest.param(False, id="false"), ) _UNAUTHORIZED: Final = Reply( status=401, @@ -263,10 +263,8 @@ def test_without_a_fallback_the_flag_still_stays_off_the_wire() -> None: assert outcome.errors == (), outcome -@pytest.mark.parametrize(("value", "errors_reported"), _NON_BOOLEAN_FLAGS) -def test_a_non_boolean_flag_stays_off_the_wire_and_the_errors_header_follows_its_truthiness( - value: object, errors_reported: bool -) -> None: +@pytest.mark.parametrize("value", _NON_BOOLEAN_FLAGS) +def test_a_flag_other_than_true_stays_off_the_wire_and_reports_no_errors(value: object) -> None: with wire_server(_peer) as wire: router: Final = _router(wire) outcome: Final = _outcome(wire, _complete(router, "primary", stream=False, include_fallback_errors=value)) @@ -274,7 +272,7 @@ def test_a_non_boolean_flag_stays_off_the_wire_and_the_errors_header_follows_its assert outcome.hit == ("primary", "backup"), outcome assert outcome.text == f"{_ANSWER} backup", outcome assert outcome.attempted == 1, outcome - assert (len(outcome.errors) == 1) is errors_reported, outcome + assert outcome.errors == (), outcome class _Recorder(CustomLogger): diff --git a/tests/integration/spend/fixtures/daily_activity_user.json b/tests/integration/spend/fixtures/daily_activity_user.json index ae05e661071..058e12a5044 100644 --- a/tests/integration/spend/fixtures/daily_activity_user.json +++ b/tests/integration/spend/fixtures/daily_activity_user.json @@ -1 +1 @@ -{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"user-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}},"guardrails":{}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"user-1":5}}} +{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"user-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{"user_email":"user@example.com","user_alias":null},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}},"guardrails":{}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"user-1":5}}} diff --git a/tests/integration/spend/test_daily_activity_key_owner.py b/tests/integration/spend/test_daily_activity_key_owner.py index cec19ce5ea0..5a9d8c7467c 100644 --- a/tests/integration/spend/test_daily_activity_key_owner.py +++ b/tests/integration/spend/test_daily_activity_key_owner.py @@ -57,6 +57,7 @@ def test_key_whose_daily_spend_names_two_users_is_reported_with_no_owner(gateway DAY, key_metadata(), seeded_metrics(2), + entity_rows={first: 1, second: 1}, ) @@ -74,6 +75,7 @@ def test_daily_spend_rows_naming_no_user_do_not_hide_the_one_user_the_others_nam DAY, key_metadata(user=owner, email=email), seeded_metrics(2), + entity_rows={owner: 1, "Unassigned": 1}, ) diff --git a/tests/integration/spend/test_daily_activity_key_owner_faults.py b/tests/integration/spend/test_daily_activity_key_owner_faults.py index 1dba97d6ae1..2eeca3c1588 100644 --- a/tests/integration/spend/test_daily_activity_key_owner_faults.py +++ b/tests/integration/spend/test_daily_activity_key_owner_faults.py @@ -199,7 +199,7 @@ def test_key_stops_being_reported_with_an_owner_once_a_second_user_spends_with_i with daily_rows((user_row(second, api_key, DAY),)): shared: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key) assert_key_reported(alone, api_key, DAY, key_metadata(user=first, email=email), seeded_metrics(1)) - assert_key_reported(shared, api_key, DAY, key_metadata(), seeded_metrics(2)) + assert_key_reported(shared, api_key, DAY, key_metadata(), seeded_metrics(2), entity_rows={first: 1, second: 1}) @pytest.mark.timeout(300) diff --git a/tests/unit/interactions/interactions.openapi.json b/tests/unit/interactions/interactions.openapi.json index 3830f5093cb..b67ad0aceaa 100644 --- a/tests/unit/interactions/interactions.openapi.json +++ b/tests/unit/interactions/interactions.openapi.json @@ -756,7 +756,7 @@ { "label": "list", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n res, err := client.Environments.ListEnvironments(ctx, operations.ListEnvironmentsRequest{})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.ListEnvironmentsResponse != nil {\n for _, env := range res.ListEnvironmentsResponse.Environments {\n if env.ID != \"\" {\n fmt.Println(env.ID)\n }\n }\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n res, err := client.Environments.ListEnvironments(ctx, operations.ListEnvironmentsRequest{})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.ListEnvironmentsResponse != nil {\n for _, env := range res.ListEnvironmentsResponse.Environments {\n if env.ID != \"\" {\n fmt.Println(env.ID)\n }\n }\n }\n}\n" } ] }, @@ -834,12 +834,12 @@ { "label": "create", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n req := operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n Sources: []gaos_interactions.Source{\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeInline),\n Target: genai.Ptr(\"main.py\"),\n Content: genai.Ptr(\"print('Hello, World!')\"),\n },\n },\n },\n }\n\n res, err := client.Environments.CreateEnvironment(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n req := operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n Sources: []gaos_interactions.Source{\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeInline),\n Target: genai.Ptr(\"main.py\"),\n Content: genai.Ptr(\"print('Hello, World!')\"),\n },\n },\n },\n }\n\n res, err := client.Environments.CreateEnvironment(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" }, { "label": "copy", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.CreateEnvironment(ctx, operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n FromEnvironment: genai.Ptr(sourceEnvID),\n },\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Environments.CreateEnvironment(ctx, operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n FromEnvironment: genai.Ptr(sourceEnvID),\n },\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" } ] } @@ -967,12 +967,12 @@ { "label": "list_files", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"src\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n fileType := \"\"\n if file.Type != nil {\n fileType = string(*file.Type)\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, fileType, size)\n }\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"src\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n fileType := \"\"\n if file.Type != nil {\n fileType = string(*file.Type)\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, fileType, size)\n }\n }\n}\n" }, { "label": "get_file", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"main.py\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, size)\n }\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"main.py\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, size)\n }\n }\n}\n" } ] } @@ -1036,7 +1036,7 @@ { "label": "delete", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n _, err = client.Environments.DeleteEnvironment(ctx, operations.DeleteEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n fmt.Println(\"Environment deleted successfully.\")\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n _, err = client.Environments.DeleteEnvironment(ctx, operations.DeleteEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n fmt.Println(\"Environment deleted successfully.\")\n}\n" } ] }, @@ -1093,7 +1093,7 @@ { "label": "get", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.GetEnvironment(ctx, operations.GetEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Environments.GetEnvironment(ctx, operations.GetEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n }\n}\n" } ] } @@ -1171,7 +1171,7 @@ "deep_research": { "summary": "Deep Research", "value": { - "agent": "deep-research-pro-preview-12-2025", + "agent": "deep-research-preview-04-2026", "created": "2025-11-26T12:22:47Z", "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", "object": "interaction", @@ -1472,7 +1472,7 @@ { "label": "deep_research", "lang": "sh", - "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"deep-research-pro-preview-12-2025\",\n \"input\": \"Find a cure to cancer\",\n \"background\": true\n }'\n" + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"deep-research-preview-04-2026\",\n \"input\": \"Find a cure to cancer\",\n \"background\": true\n }'\n" }, { "label": "antigravity", @@ -1517,7 +1517,7 @@ { "label": "deep_research", "lang": "python", - "source": "from google import genai\n\nclient = genai.Client()\ninteraction = client.interactions.create(\n agent=\"deep-research-pro-preview-12-2025\",\n input=\"find a cure to cancer\",\n background=True,\n)\nprint(interaction.status)\n" + "source": "from google import genai\n\nclient = genai.Client()\ninteraction = client.interactions.create(\n agent=\"deep-research-preview-04-2026\",\n input=\"find a cure to cancer\",\n background=True,\n)\nprint(interaction.status)\n" }, { "label": "antigravity", @@ -1562,7 +1562,7 @@ { "label": "deep_research", "lang": "javascript", - "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n agent: 'deep-research-pro-preview-12-2025',\n input: 'find a cure to cancer',\n background: true,\n});\nconsole.log(interaction.status);\n" + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n agent: 'deep-research-preview-04-2026',\n input: 'find a cure to cancer',\n background: true,\n});\nconsole.log(interaction.status);\n" }, { "label": "antigravity", @@ -1607,7 +1607,7 @@ { "label": "deep_research", "lang": "java", - "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionStatus;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"deep-research-pro-preview-12-2025\"))\n .input(InteractionsInput.of(\"find a cure to cancer\"))\n .background(true)\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.status().map(InteractionStatus::value).orElse(\"\"));\n" + "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionStatus;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"deep-research-preview-04-2026\"))\n .input(InteractionsInput.of(\"find a cure to cancer\"))\n .background(true)\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.status().map(InteractionStatus::value).orElse(\"\"));\n" }, { "label": "antigravity", @@ -1652,7 +1652,7 @@ { "label": "deep_research", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"deep-research-pro-preview-12-2025\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"find a cure to cancer\")),\n Background: genai.Ptr(true),\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil {\n fmt.Println(res.Interaction.Status)\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"deep-research-preview-04-2026\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"find a cure to cancer\")),\n Background: genai.Ptr(true),\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil {\n fmt.Println(res.Interaction.Status)\n }\n}\n" }, { "label": "antigravity", @@ -3041,6 +3041,22 @@ "description": "An agent definition for the CreateAgent API.\nThis message is the target for annotation-parser-based JSON parsing.\nNew format:\n {\n \"id\": \"customer-sentinel\",\n \"base_agent\": \"\",\n \"system_instruction\": \"...\",\n \"base_environment\": { \"type\": \"remote\", \"sources\": [...] },\n \"tools\": [ {\"type\": \"code_execution\"} ]\n }", "type": "object", "properties": { + "id": { + "description": "The unique identifier for the agent.", + "type": "string" + }, + "base_agent": { + "description": "The base agent to extend.", + "type": "string" + }, + "system_instruction": { + "description": "System instruction for the agent.", + "type": "string" + }, + "description": { + "description": "Agent description for developers to quickly read and understand.", + "type": "string" + }, "agent_config": { "description": "Configuration parameters for the agent.", "oneOf": [ @@ -3060,22 +3076,6 @@ } ] }, - "id": { - "description": "The unique identifier for the agent.", - "type": "string" - }, - "base_agent": { - "description": "The base agent to extend.", - "type": "string" - }, - "system_instruction": { - "description": "System instruction for the agent.", - "type": "string" - }, - "description": { - "description": "Agent description for developers to quickly read and understand.", - "type": "string" - }, "tools": { "description": "The tools available to the agent.", "type": "array", @@ -3091,13 +3091,11 @@ "description": "The agent to interact with.", "type": "string", "enum": [ - "deep-research-pro-preview-12-2025", "deep-research-preview-04-2026", "deep-research-max-preview-04-2026", "antigravity-preview-05-2026" ], "x-speakeasy-enum-descriptions": [ - "Gemini Deep Research Agent", "Gemini Deep Research Agent", "Gemini Deep Research Max Agent", "Use the Antigravity managed agent to perform multi-step tasks that require reasoning, file operations, and tool use." @@ -3330,14 +3328,6 @@ "type": { "const": "audio" }, - "data": { - "type": "string", - "format": "byte", - "x-speakeasy-base64-input-mode": "file" - }, - "uri": { - "type": "string" - }, "mime_type": { "type": "string", "enum": [ @@ -3371,6 +3361,14 @@ "WEBM audio format" ] }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, "rate": { "description": "Deprecated. Use sample_rate instead. The value is ignored.", "type": "integer", @@ -3817,32 +3815,6 @@ } ] }, - "response_format": { - "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", - "oneOf": [ - { - "$ref": "#/components/schemas/ResponseFormat" - }, - { - "title": "ResponseFormatList", - "type": "array", - "items": { - "$ref": "#/components/schemas/ResponseFormat" - } - } - ] - }, - "environment": { - "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", - "oneOf": [ - { - "$ref": "#/components/schemas/EnvironmentConfig" - }, - { - "type": "string" - } - ] - }, "id": { "description": "Required. Output only. A unique identifier for the interaction completion.", "type": "string", @@ -3858,7 +3830,8 @@ "failed", "cancelled", "incomplete", - "budget_exceeded" + "budget_exceeded", + "requires_continuation" ], "readOnly": true, "x-google-enum-deprecated": [ @@ -3868,7 +3841,8 @@ false, false, false, - true + true, + false ], "x-google-enum-descriptions": [ "The interaction is in progress.", @@ -3877,7 +3851,8 @@ "The interaction failed.", "The interaction was cancelled.", "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", - "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11).", + "The interaction requires at least one more request for it to be\ncompleted. The next request should set the continuation_token to the one\nreturned by this request." ] }, "created": { @@ -3909,6 +3884,21 @@ "$ref": "#/components/schemas/ResponseModality" } }, + "response_format": { + "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", + "oneOf": [ + { + "$ref": "#/components/schemas/ResponseFormat" + }, + { + "title": "ResponseFormatList", + "type": "array", + "items": { + "$ref": "#/components/schemas/ResponseFormat" + } + } + ] + }, "response_mime_type": { "description": "The mime type of the response. This is required if response_format is set.", "type": "string", @@ -3918,6 +3908,17 @@ "description": "The ID of the previous interaction, if any.", "type": "string" }, + "environment": { + "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentConfig" + }, + { + "type": "string" + } + ] + }, "environment_id": { "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", "type": "string", @@ -3945,6 +3946,12 @@ "type": "string" } }, + "continuation_token": { + "description": "Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, "stream": { "description": "Input only. Whether the interaction will be streamed.", "type": "boolean", @@ -4018,32 +4025,6 @@ "$ref": "#/components/schemas/GenerationConfig", "writeOnly": true }, - "response_format": { - "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", - "oneOf": [ - { - "$ref": "#/components/schemas/ResponseFormat" - }, - { - "title": "ResponseFormatList", - "type": "array", - "items": { - "$ref": "#/components/schemas/ResponseFormat" - } - } - ] - }, - "environment": { - "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", - "oneOf": [ - { - "$ref": "#/components/schemas/EnvironmentConfig" - }, - { - "type": "string" - } - ] - }, "id": { "description": "Required. Output only. A unique identifier for the interaction completion.", "type": "string", @@ -4059,7 +4040,8 @@ "failed", "cancelled", "incomplete", - "budget_exceeded" + "budget_exceeded", + "requires_continuation" ], "readOnly": true, "x-google-enum-deprecated": [ @@ -4069,7 +4051,8 @@ false, false, false, - true + true, + false ], "x-google-enum-descriptions": [ "The interaction is in progress.", @@ -4078,7 +4061,8 @@ "The interaction failed.", "The interaction was cancelled.", "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", - "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11).", + "The interaction requires at least one more request for it to be\ncompleted. The next request should set the continuation_token to the one\nreturned by this request." ] }, "created": { @@ -4110,6 +4094,21 @@ "$ref": "#/components/schemas/ResponseModality" } }, + "response_format": { + "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", + "oneOf": [ + { + "$ref": "#/components/schemas/ResponseFormat" + }, + { + "title": "ResponseFormatList", + "type": "array", + "items": { + "$ref": "#/components/schemas/ResponseFormat" + } + } + ] + }, "response_mime_type": { "description": "The mime type of the response. This is required if response_format is set.", "type": "string", @@ -4119,6 +4118,17 @@ "description": "The ID of the previous interaction, if any.", "type": "string" }, + "environment": { + "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentConfig" + }, + { + "type": "string" + } + ] + }, "environment_id": { "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", "type": "string", @@ -4146,6 +4156,12 @@ "type": "string" } }, + "continuation_token": { + "description": "Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, "stream": { "description": "Input only. Whether the interaction will be streamed.", "type": "boolean", @@ -4348,11 +4364,31 @@ "type": "string", "enum": [ "application/pdf", - "text/csv" + "text/csv", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + "application/vnd.ms-powerpoint", + "application/vnd.oasis.opendocument.presentation", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "application/msword", + "application/vnd.openxmlformats-officedocument.wordprocessingml.template", + "application/vnd.oasis.opendocument.text", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "application/vnd.ms-excel", + "application/vnd.oasis.opendocument.spreadsheet" ], "x-google-enum-descriptions": [ "PDF document format", - "CSV document format" + "CSV document format", + "PPTX presentation format", + "PPT presentation format", + "ODP presentation format", + "DOCX document format", + "DOC document format", + "DOTX document template format", + "ODT document format", + "XLSX spreadsheet format", + "XLS spreadsheet format", + "ODS spreadsheet format" ] } }, @@ -4379,6 +4415,37 @@ "type": { "const": "document" }, + "mime_type": { + "type": "string", + "enum": [ + "application/pdf", + "text/csv", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + "application/vnd.ms-powerpoint", + "application/vnd.oasis.opendocument.presentation", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "application/msword", + "application/vnd.openxmlformats-officedocument.wordprocessingml.template", + "application/vnd.oasis.opendocument.text", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "application/vnd.ms-excel", + "application/vnd.oasis.opendocument.spreadsheet" + ], + "x-google-enum-descriptions": [ + "PDF document format", + "CSV document format", + "PPTX presentation format", + "PPT presentation format", + "ODP presentation format", + "DOCX document format", + "DOC document format", + "DOTX document template format", + "ODT document format", + "XLSX spreadsheet format", + "XLS spreadsheet format", + "ODS spreadsheet format" + ] + }, "data": { "type": "string", "format": "byte", @@ -4386,17 +4453,6 @@ }, "uri": { "type": "string" - }, - "mime_type": { - "type": "string", - "enum": [ - "application/pdf", - "text/csv" - ], - "x-google-enum-descriptions": [ - "PDF document format", - "CSV document format" - ] } }, "required": [ @@ -4474,23 +4530,6 @@ "type": "string", "readOnly": true }, - "network": { - "description": "Network configuration for the environment.", - "oneOf": [ - { - "$ref": "#/components/schemas/EnvironmentNetworkEgressAllowlist" - }, - { - "type": "string", - "enum": [ - "disabled" - ], - "x-google-enum-descriptions": [ - "All network egress is blocked." - ] - } - ] - }, "sources": { "description": "Sources to be mounted into the environment.", "type": "array", @@ -4526,6 +4565,23 @@ "" ] }, + "network": { + "description": "Network configuration for the environment.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentNetworkEgressAllowlist" + }, + { + "type": "string", + "enum": [ + "disabled" + ], + "x-google-enum-descriptions": [ + "All network egress is blocked." + ] + } + ] + }, "file_count": { "description": "Output only. The number of files in the environment, output only.", "type": "string", @@ -4551,6 +4607,12 @@ "type": { "const": "remote" }, + "sources": { + "type": "array", + "items": { + "$ref": "#/components/schemas/Source" + } + }, "network": { "description": "Network configuration for the environment.", "oneOf": [ @@ -4568,12 +4630,6 @@ } ] }, - "sources": { - "type": "array", - "items": { - "$ref": "#/components/schemas/Source" - } - }, "environment_id": { "description": "Optional. The environment ID for the interaction. If specified, the request will\nupdate the existing environment instead of creating a new one.", "type": "string" @@ -5347,30 +5403,6 @@ "description": "Configuration parameters for model interactions.", "type": "object", "properties": { - "tool_choice": { - "description": "The tool choice configuration.", - "oneOf": [ - { - "$ref": "#/components/schemas/ToolChoiceConfig" - }, - { - "type": "string", - "enum": [ - "auto", - "any", - "none", - "validated" - ], - "x-google-enum-descriptions": [ - "Auto tool choice.", - "Any tool choice.", - "No tool choice.", - "Validated tool choice." - ], - "x-speakeasy-model-namespace": "interactions" - } - ] - }, "temperature": { "description": "Controls the randomness of the output.", "type": "number", @@ -5395,6 +5427,30 @@ "type": "string" } }, + "tool_choice": { + "description": "The tool choice configuration.", + "oneOf": [ + { + "$ref": "#/components/schemas/ToolChoiceConfig" + }, + { + "type": "string", + "enum": [ + "auto", + "any", + "none", + "validated" + ], + "x-google-enum-descriptions": [ + "Auto tool choice.", + "Any tool choice.", + "No tool choice.", + "Validated tool choice." + ], + "x-speakeasy-model-namespace": "interactions" + } + ] + }, "thinking_level": { "description": "The level of thought tokens that the model should generate.", "$ref": "#/components/schemas/ThinkingLevel" @@ -6211,14 +6267,6 @@ "type": { "const": "image" }, - "data": { - "type": "string", - "format": "byte", - "x-speakeasy-base64-input-mode": "file" - }, - "uri": { - "type": "string" - }, "mime_type": { "type": "string", "enum": [ @@ -6242,6 +6290,14 @@ "TIFF image format" ] }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, "resolution": { "description": "The resolution of the media.", "$ref": "#/components/schemas/MediaResolution" @@ -6385,7 +6441,8 @@ "failed", "cancelled", "incomplete", - "budget_exceeded" + "budget_exceeded", + "requires_continuation" ], "readOnly": true, "x-google-enum-deprecated": [ @@ -6395,7 +6452,8 @@ false, false, false, - true + true, + false ], "x-google-enum-descriptions": [ "The interaction is in progress.", @@ -6404,7 +6462,8 @@ "The interaction failed.", "The interaction was cancelled.", "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", - "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11).", + "The interaction requires at least one more request for it to be\ncompleted. The next request should set the continuation_token to the one\nreturned by this request." ] }, "created": { @@ -6434,9 +6493,33 @@ "$ref": "#/components/schemas/GenerationConfig", "writeOnly": true }, + "system_instruction": { + "description": "System instruction for the interaction.", + "type": "string" + }, + "tools": { + "description": "A list of tool declarations the model may call during interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Tool" + } + }, "input": { "$ref": "#/components/schemas/InteractionsInput" }, + "usage": { + "description": "Output only. Statistics on the interaction request's token usage.", + "$ref": "#/components/schemas/Usage", + "readOnly": true + }, + "response_modalities": { + "description": "The requested modalities of the response (TEXT, IMAGE, AUDIO).", + "type": "array", + "deprecated": true, + "items": { + "$ref": "#/components/schemas/ResponseModality" + } + }, "response_format": { "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", "oneOf": [ @@ -6452,6 +6535,15 @@ } ] }, + "response_mime_type": { + "description": "The mime type of the response. This is required if response_format is set.", + "type": "string", + "deprecated": true + }, + "previous_interaction_id": { + "description": "The ID of the previous interaction, if any.", + "type": "string" + }, "environment": { "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", "oneOf": [ @@ -6463,39 +6555,6 @@ } ] }, - "system_instruction": { - "description": "System instruction for the interaction.", - "type": "string" - }, - "tools": { - "description": "A list of tool declarations the model may call during interaction.", - "type": "array", - "items": { - "$ref": "#/components/schemas/Tool" - } - }, - "usage": { - "description": "Output only. Statistics on the interaction request's token usage.", - "$ref": "#/components/schemas/Usage", - "readOnly": true - }, - "response_modalities": { - "description": "The requested modalities of the response (TEXT, IMAGE, AUDIO).", - "type": "array", - "deprecated": true, - "items": { - "$ref": "#/components/schemas/ResponseModality" - } - }, - "response_mime_type": { - "description": "The mime type of the response. This is required if response_format is set.", - "type": "string", - "deprecated": true - }, - "previous_interaction_id": { - "description": "The ID of the previous interaction, if any.", - "type": "string" - }, "environment_id": { "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", "type": "string", @@ -6531,6 +6590,12 @@ }, "readOnly": true }, + "continuation_token": { + "description": "Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, "stream": { "description": "Input only. Whether the interaction will be streamed.", "type": "boolean", @@ -6813,16 +6878,16 @@ }, "readOnly": true }, - "agent": { - "description": "The agent to interact with.", - "type": "string" - }, "continuation_token": { "description": "Opaque token to resume a long decode.\nOutput: set when status is INCOMPLETE and decoding can be resumed.\nInput: pass the latest token back unchanged in CreateInteraction to\ncontinue decoding.", "type": "string", "format": "byte", "x-speakeasy-base64-input-mode": "file" }, + "agent": { + "description": "The agent to interact with.", + "type": "string" + }, "object": { "description": "Output only. The resource type.", "type": "string", @@ -6863,7 +6928,8 @@ "failed", "cancelled", "incomplete", - "budget_exceeded" + "budget_exceeded", + "requires_continuation" ], "x-google-enum-deprecated": [ false, @@ -6872,7 +6938,8 @@ false, false, false, - true + true, + false ], "x-google-enum-descriptions": [ "The interaction is in progress.", @@ -6881,7 +6948,8 @@ "The interaction failed.", "The interaction was cancelled.", "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", - "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11).", + "The interaction requires at least one more request for it to be\ncompleted. The next request should set the continuation_token to the one\nreturned by this request." ] }, "event_type": { @@ -7076,14 +7144,14 @@ "type": { "const": "mcp_server" }, - "url": { - "description": "The full URL for the MCPServer endpoint.\nExample: \"https://api.example.com/mcp\"", - "type": "string" - }, "name": { "description": "The name of the MCPServer.", "type": "string" }, + "url": { + "description": "The full URL for the MCPServer endpoint.\nExample: \"https://api.example.com/mcp\"", + "type": "string" + }, "headers": { "description": "Optional: Fields for authentication headers, timeouts, etc., if needed.", "type": "object", @@ -7362,7 +7430,7 @@ }, "ModelOption": { "title": "Model", - "description": "The model that will complete your prompt.\\n\\nSee [models](https://ai.google.dev/gemini-api/docs/models) for additional details.", + "description": "The model that will complete your prompt.\\n\\nSee [models](https://aistudio.google.com/docs/models) for additional details.", "type": "string", "enum": [ "gemma-4-26b-a4b-it", @@ -8695,6 +8763,20 @@ "description": "Configuration for speech recognition (transcription).", "type": "object", "properties": { + "language_codes": { + "description": "Optional. BCP-47 language codes providing hints about the languages present in the\naudio. If omitted or empty, defaults to automatic language detection.", + "type": "array", + "items": { + "type": "string" + } + }, + "custom_vocabulary": { + "description": "Optional. A list of custom vocabulary phrases to bias the speech recognition model\ntoward recognizing specific terms.", + "type": "array", + "items": { + "type": "string" + } + }, "mode": { "description": "Discriminated transcription mode options or enum.", "oneOf": [ @@ -8714,20 +8796,6 @@ } ] }, - "language_codes": { - "description": "Optional. BCP-47 language codes providing hints about the languages present in the\naudio. If omitted or empty, defaults to automatic language detection.", - "type": "array", - "items": { - "type": "string" - } - }, - "custom_vocabulary": { - "description": "Optional. A list of custom vocabulary phrases to bias the speech recognition model\ntoward recognizing specific terms.", - "type": "array", - "items": { - "type": "string" - } - }, "timestamp_granularities": { "description": "Optional. The granularity of timestamps to include in the transcription output.\nSupported values: \"word\". If empty, no timestamps are generated.", "type": "array", @@ -9332,6 +9400,11 @@ "type": "integer", "format": "int32" }, + "initial_input_tokens": { + "description": "Number of input tokens used in the first turn of this interaction.", + "type": "integer", + "format": "int32" + }, "input_tokens_by_modality": { "description": "A breakdown of input token usage by modality.", "type": "array", @@ -9497,6 +9570,10 @@ "description": "The URI of the video.", "type": "string" }, + "resolution": { + "description": "The resolution of the media.", + "$ref": "#/components/schemas/MediaResolution" + }, "processing": { "description": "How the model processes this video for understanding.", "oneOf": [ @@ -9516,10 +9593,6 @@ } ] }, - "resolution": { - "description": "The resolution of the media.", - "$ref": "#/components/schemas/MediaResolution" - }, "name": { "description": "A user-defined name for this content block. Can be referenced by the model\nin the final response.", "type": "string" @@ -9573,14 +9646,6 @@ "type": { "const": "video" }, - "data": { - "type": "string", - "format": "byte", - "x-speakeasy-base64-input-mode": "file" - }, - "uri": { - "type": "string" - }, "mime_type": { "type": "string", "enum": [ @@ -9608,6 +9673,14 @@ "JPEG 2000 video format" ] }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, "resolution": { "description": "The resolution of the media.", "$ref": "#/components/schemas/MediaResolution" diff --git a/tests/unit/interactions/test_main.py b/tests/unit/interactions/test_main.py index 5af5dd7d82a..cd32eb11485 100644 --- a/tests/unit/interactions/test_main.py +++ b/tests/unit/interactions/test_main.py @@ -559,28 +559,34 @@ class TestInteractionsAcreateOffline: @pytest.mark.usefixtures("fake_provider_credentials") @pytest.mark.asyncio - async def test_acreate_gemini_keeps_continuation_token_out_of_the_request(self, respx_mock: MockRouter) -> None: + async def test_acreate_gemini_sends_a_caller_continuation_token_to_resume_a_decode( + self, respx_mock: MockRouter + ) -> None: route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( return_value=httpx.Response( 200, json={ "id": "interaction-offline", "object": "interaction", - "model": "gemini-3.8-flash", + "model": "gemini-4-argon", "status": "completed", "steps": [{"type": "model_output", "content": [{"type": "text", "text": "Hi"}]}], - "usage": {"input_tokens": 6, "output_tokens": 3}, + "usage": {"total_input_tokens": 6, "total_output_tokens": 3}, }, ) ) response: Final = await interactions.acreate( - model="gemini/gemini-3.8-flash", - input="Prove it", + model="gemini/gemini-4-argon", + previous_interaction_id="interaction-hop-1", continuation_token="dG9rLTE=", api_key="gemini-offline", ) - # The Interactions API answered 400 "Unknown parameter 'continuation_token'." on 2026-10-13 - assert json.loads(route.calls.last.request.content) == {"model": "gemini-3.8-flash", "input": "Prove it"} + # Google's spec (2026-11-05) takes continuation_token on CreateModelInteractionParams; a 2026-10-13 400 is gone + assert json.loads(route.calls.last.request.content) == { + "model": "gemini-4-argon", + "previous_interaction_id": "interaction-hop-1", + "continuation_token": "dG9rLTE=", + } assert response.status == "completed" @pytest.mark.usefixtures("fake_provider_credentials") @@ -746,9 +752,7 @@ def test_create_with_system_instruction_for_google_and_responses( def test_create_streaming_returns_interaction_events( respx_mock: MockRouter, ) -> None: - route: Final = respx_mock.post( - "https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse" - ).mock( + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse").mock( return_value=httpx.Response( 200, headers={"content-type": "text/event-stream"}, @@ -856,3 +860,65 @@ def test_create_streaming_hands_completed_interaction_to_success_callback( ) assert recorder.logged == (_GEMINI_COMPLETED_INTERACTION,) + + +_GEMINI_URL: Final = "https://generativelanguage.googleapis.com/v1beta/interactions" +_LONG_DECODE_REQUEST: Final = { + "model": "gemini/gemini-4-argon", + "input": "Prove the prime number theorem.", + "generation_config": {"max_output_tokens": 120000, "thinking_level": "high"}, +} +_FIRST_HOP: Final = { + "id": "hop-1", + "object": "interaction", + "model": "gemini-4-argon", + "status": "requires_continuation", + "steps": [{"type": "model_output", "content": [{"type": "text", "text": "Lemma 1."}]}], + "usage": { + "total_input_tokens": 79, + "total_output_tokens": 32668, + "total_thought_tokens": 100, + "total_tokens": 32847, + "initial_input_tokens": 79, + }, + "continuation_token": "tok-1", +} + + +@pytest.mark.usefixtures("fake_provider_credentials") +def test_create_keeps_requires_continuation_status_and_token_for_the_caller(respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(_GEMINI_URL).mock(return_value=httpx.Response(200, json=_FIRST_HOP)) + + response: Final = interactions.create(**_LONG_DECODE_REQUEST, api_key="gemini-offline") + + assert route.call_count == 1 + assert response.status == "requires_continuation" + assert response.continuation_token == "tok-1" + assert response.usage == _FIRST_HOP["usage"] + + +@pytest.mark.usefixtures("fake_provider_credentials") +def test_create_sends_a_caller_continuation_token_to_resume_a_long_decode(respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(_GEMINI_URL).mock( + return_value=httpx.Response( + 200, json={**_FIRST_HOP, "id": "hop-2", "status": "completed", "continuation_token": None} + ) + ) + + response: Final = interactions.create( + model="gemini/gemini-4-argon", + previous_interaction_id="hop-1", + continuation_token="tok-1", + generation_config=_LONG_DECODE_REQUEST["generation_config"], + api_key="gemini-offline", + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert sent == { + "model": "gemini-4-argon", + "previous_interaction_id": "hop-1", + "continuation_token": "tok-1", + "generation_config": _LONG_DECODE_REQUEST["generation_config"], + }, sent + assert response.status == "completed" + assert response.continuation_token is None diff --git a/tests/unit/interactions/test_openapi_compliance.py b/tests/unit/interactions/test_openapi_compliance.py index b6cc92e0adc..71db2cd24e3 100644 --- a/tests/unit/interactions/test_openapi_compliance.py +++ b/tests/unit/interactions/test_openapi_compliance.py @@ -2,7 +2,7 @@ OpenAPI compliance tests for Google Interactions API. Validates that our SDK requests/responses match the OpenAPI spec snapshot in -interactions.openapi.json, taken from https://ai.google.dev/static/api/interactions.openapi.json +interactions.openapi.json, taken from https://aistudio.google.com/api-reference/interactions.openapi.json (litellm/types/interactions/README.md says how to refresh it). The scheduled interactions-spec-canary workflow runs the same tests against the live spec by passing --interactions-spec-url, so a change Google makes shows up there, not in every PR. @@ -23,7 +23,7 @@ import pytest from openapi_core import OpenAPI from pydantic import BaseModel, TypeAdapter -OPENAPI_SPEC_URL = "https://ai.google.dev/static/api/interactions.openapi.json" +OPENAPI_SPEC_URL = "https://aistudio.google.com/api-reference/interactions.openapi.json" OPENAPI_SPEC_SNAPSHOT: Final = Path(__file__).with_name("interactions.openapi.json") SPEC_URL_OPTION: Final = TypeAdapter(str | None) @@ -47,10 +47,10 @@ def _requested_spec_url(request: pytest.FixtureRequest) -> str | None: def _load_openapi_spec(live_url: str | None) -> LoadedSpec: if live_url is None: return LoadedSpec(source=str(OPENAPI_SPEC_SNAPSHOT), spec=json.loads(OPENAPI_SPEC_SNAPSHOT.read_text())) - with httpx.Client(transport=httpx.HTTPTransport(retries=3), timeout=30.0) as client: + with httpx.Client(transport=httpx.HTTPTransport(retries=3), timeout=30.0, follow_redirects=True) as client: response: Final = client.get(live_url) response.raise_for_status() - return LoadedSpec(source=str(response.url), spec=response.json()) + return LoadedSpec(source=live_url, spec=response.json()) def _model_create_request_schema(spec_dict: Dict[str, Any]) -> Dict[str, Any]: diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 80845b1f12d..3a4231443dd 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -11047,3 +11047,24 @@ def test_claude_cache_control_cap_does_not_count_a_message_level_hint_that_list_ assert sorted(_cached_texts(body)) == ["system prompt", "turn 1", "turn 2", "turn 3"] assert json.dumps(body).count('"cache_control"') == 4 + + +@pytest.mark.parametrize("model", ["claude-opus-4-7", "claude-haiku-4-5-20251001"]) +def test_get_supported_params_reports_sampling_params_only_when_the_request_keeps_them( + local_model_cost_map, model: str +): + """What get_supported_openai_params says about temperature/top_p must match what + map_openai_params does with them under drop_params, so a caller that pre-flights + from the list is never told a param is fine that the request path then drops.""" + config = AnthropicConfig() + + supported = config.get_supported_openai_params(model=model) + sent = config.map_openai_params( + non_default_params={"temperature": 0.1, "top_p": 0.5}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert ("temperature" in supported) == ("temperature" in sent), (model, supported, sent) + assert ("top_p" in supported) == ("top_p" in sent), (model, supported, sent) diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index d66ba3f324e..65f8400e53c 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -477,6 +477,7 @@ def test_default_max_retries_env_var_reaches_azure_sdk_client(): "aget_responses", "adelete_responses", "acancel_responses", + "acompact_responses", "aresponses_websocket", "alist_input_items", "acreate_fine_tuning_job", diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 98e204dd44e..842500a7b27 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -2161,7 +2161,7 @@ class TestVertexClaudeBatch: monkeypatch.setattr(litellm, "drop_params", False) entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5) - with pytest.raises(VertexAIError, match=r"line 1 .*top_p=0.5.*drop_params") as exc_info: + with pytest.raises(VertexAIError, match=r"line 1 .*top_p.*drop_params") as exc_info: _uploaded_rows(config, {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9"}, [entry]) assert exc_info.value.status_code == 400 diff --git a/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py index 7e5b3a7abd7..a1be416f4cb 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -74,6 +74,7 @@ def _make_mock_prisma(): record.description = data.get("description") record.manifest_json = data.get("manifest_json", "{}") record.enabled = data.get("enabled", True) + record.created_by = data.get("created_by") record.created_at = data.get("created_at") record.updated_at = data.get("updated_at") store[data["name"]] = record @@ -221,7 +222,7 @@ async def test_metadata_edits_preserve_uploaded_source_and_reject_a_changed_http await table.create(data={"name": _UPLOAD_NAME, "manifest_json": json.dumps({"source": _UPLOAD_SOURCE})}) request: Final = UpdatePluginRequest(source=_UPLOAD_SOURCE, description="Updated description", version="1.0.0") response: Final = await update_plugin(uploads=_uploads(), plugin_name=_UPLOAD_NAME, request=request, user_api_key_dict=_USER) - expected: Final = {"name": _UPLOAD_NAME, "source": _UPLOAD_SOURCE, "description": "Updated description", "version": "1.0.0"} + expected: Final = {"name": _UPLOAD_NAME, "source": _UPLOAD_SOURCE, "description": "Updated description", "version": "1.0.0", "_litellm_published_by": _USER.user_id} assert response.plugin.description == "Updated description" assert await _read_stored_manifest(_UPLOAD_NAME) == expected with pytest.raises(HTTPException) as error: @@ -354,7 +355,7 @@ async def _register_public_and_private_plugins() -> None: async def _listed_names(user: UserAPIKeyAuth) -> set[str]: - response = await list_plugins(user_api_key_dict=user) + response = await list_plugins(uploads=_uploads(), user_api_key_dict=user) return {plugin.name for plugin in response.plugins} @@ -474,7 +475,7 @@ async def test_update_plugin_propagates_installation_preference_to_get_and_list( ) plugin = await get_plugin(plugin_name=name, user_api_key_dict=_USER) - listed_plugin = (await list_plugins(user_api_key_dict=_USER)).plugins[0] + listed_plugin = (await list_plugins(uploads=_uploads(), user_api_key_dict=_USER)).plugins[0] assert response.plugin.installation_preference == "auto_install" assert plugin["installation_preference"] == "auto_install" @@ -514,7 +515,7 @@ async def test_stored_unknown_installation_preference_is_not_served(stored_value ) marketplace = json.loads((await get_marketplace(request=MagicMock())).body) - listed_plugin = (await list_plugins(user_api_key_dict=_USER)).plugins[0] + listed_plugin = (await list_plugins(uploads=_uploads(), user_api_key_dict=_USER)).plugins[0] plugin = await get_plugin(plugin_name=name, user_api_key_dict=_USER) assert marketplace["plugins"] == [{"name": name, "source": _ARCHIVE_SOURCE}] @@ -1288,7 +1289,7 @@ async def test_skill_upload_creates_metadata_and_bytes_atomically_and_rejects_du "url": f"http://test/claude-code/plugins/{_UPLOAD_NAME}/archive/{digest}.zip", "sha256": digest, } - assert json.loads(stored["manifest_json"]) == {**metadata, "source": source} + assert json.loads(stored["manifest_json"]) == {**metadata, "source": source, "_litellm_published_by": stored["created_by"]} assert stored == { "name": _UPLOAD_NAME, "description": "My notes", "version": "2.0.0", "manifest_json": stored["manifest_json"], "created_by": litellm.proxy.proxy_server.litellm_proxy_admin_name, diff --git a/tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py index 8127fc80bab..7912df3939b 100644 --- a/tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py +++ b/tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py @@ -196,25 +196,25 @@ def _assert_portable(response: httpx.Response) -> None: assert stat.S_IMODE(archive.getinfo("scripts/check.sh").external_attr >> 16) == 0o755 -@pytest.mark.parametrize("key, expected", ((MASTER_KEY, 7), (PLAIN_KEY, None))) -async def test_catalog_counts_are_admin_only_and_batched_for_the_page( - hub: HubClient, key: str, expected: int | None, +@pytest.mark.parametrize("route", (f"{HUB}?limit=1", "/claude-code/plugins")) +@pytest.mark.parametrize("key", (MASTER_KEY, PLAIN_KEY)) +async def test_catalog_counts_are_batched_for_the_page( + hub: HubClient, key: str, route: str, ) -> None: assert (await _upload(hub, _zip(FILES))).status_code == 200 record: Final = _stored_upload(hub) hub.plugins.find_many.return_value = [record, record.model_copy(update={"id": "other", "name": "zzz-other"})] hub.database.query_raw.reset_mock() hub.database.query_raw.return_value = [{"id": record.id, "download_count": 7}] - response: Final = await hub.client.get(f"{HUB}?limit=1", headers=_headers(key)) + response: Final = await hub.client.get(route, headers=_headers(key)) assert response.status_code == 200, response.text - assert [(item["id"], item["download_count"]) for item in response.json()["skills"]] == [(record.id, expected)] - if expected is None: - hub.database.query_raw.assert_not_awaited() - else: - assert hub.database.query_raw.await_count == 1 - assert hub.database.query_raw.await_args.args[1] == (record.id,) - assert (await hub.client.get(f"{HUB}?offset=2", headers=_headers(key))).json()["skills"] == [] - assert hub.database.query_raw.await_count == 1 + items: Final = response.json()["plugins" if route == "/claude-code/plugins" else "skills"] + expected: Final = {record.id: 7, "other": 0} if route == "/claude-code/plugins" else {record.id: 7} + assert {item["id"]: item["download_count"] for item in items} == expected + assert hub.database.query_raw.await_count == 1 + assert set(hub.database.query_raw.await_args.args[1]) == set(expected) + assert (await hub.client.get(f"{HUB}?offset=2", headers=_headers(key))).json()["skills"] == [] + assert hub.database.query_raw.await_count == 1 @pytest.mark.parametrize("operation", ("upload", "import")) @@ -806,3 +806,98 @@ async def test_configured_scan_persists_result_without_changing_manual_approval( assert (await hub.client.get(path, headers=_headers(MASTER_KEY))).json() == before approved: Final = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY)) assert approved.status_code == 200 and approved.json()["submission"]["status"] == "approved" + + +async def test_catalog_filters_metadata_before_pagination_and_preserves_tag_facets(hub: HubClient) -> None: + hub.plugins.find_many.return_value = [ + prisma_models.LiteLLM_ClaudeCodePluginTable( + id=name, name=name, enabled=True, anonymous_download_count=0, created_by="author@example.test", + created_at=datetime(2026, 1, 2), updated_at=datetime(2026, 2, 3), + manifest_json=json.dumps({ + "source": {"source": "github", "repo": "example/skills"}, + "keywords": tags, + "_litellm_published_by": "publisher@example.test", + **({"_litellm_review": {"approved_by": "reviewer@example.test", "security_scan_passed": scanned}} if approved else {}), + }), + ) + for name, tags, approved, scanned in ( + ("a-other", ["support"], False, False), + ("b-review", ["engineering"], True, False), + ("c-review", ["engineering"], True, True), + ) + ] + response = await hub.client.get(HUB, headers=_headers(MASTER_KEY), params={ + "query": " REVIEWER ", "tag": "ENGINEERING", "limit": 1, "offset": 1, + }) + assert response.status_code == 200 + body = response.json() + assert body["count"] == 2 + assert body["tags"] == ["engineering", "support"] + assert [skill["name"] for skill in body["skills"]] == ["c-review"] + skill = body["skills"][0] + assert skill["created_by"] == "author@example.test" + assert skill["published_by"] == "publisher@example.test" + assert skill["approved_by"] == "reviewer@example.test" + assert skill["created_at"].startswith("2026-01-02") + assert skill["updated_at"].startswith("2026-02-03") + assert skill["security_scan_passed"] is True + scanned = await hub.client.get(HUB, headers=_headers(MASTER_KEY), params={"security_scan_passed": "true"}) + assert [skill["name"] for skill in scanned.json()["skills"]] == ["c-review"] + unscanned = await hub.client.get(HUB, headers=_headers(MASTER_KEY), params={"security_scan_passed": "false"}) + assert [skill["name"] for skill in unscanned.json()["skills"]] == ["a-other", "b-review"] + assert unscanned.json()["skills"][0]["approved_by"] is None + + +@pytest.mark.parametrize("has_issues", (False, True)) +async def test_admin_approval_does_not_override_scan_result(pending_hub: HubClient, has_issues: bool) -> None: + hub: Final = pending_hub + pending: Final = await _upload(hub, _zip(FILES), MEMBER_KEY, version=None) + receipt: Final = pending.json() + submission: Final = hub.database.litellm_skillsubmission.rows[0] + digest: Final = hashlib.sha256(submission.archive_content.decode()).hexdigest() + await hub.database.litellm_skillsubmission.update( + where={"id": submission.id}, data={"scan_json": json.dumps({ + "id": "scan", "requested_by": "admin", "started_at": "2026-01-01T00:00:00Z", + "expires_at": "2026-01-01T00:03:00Z", "status": "completed", "archive_sha256": digest, + "report": {"scanner_version": "test", "scan_mode": "static-only", "complete": True, + "risk_score": 0, "findings": [{"rule_id": "rule", "severity": "LOW", + "message": "Issue", "file": "SKILL.md"}] if has_issues else []}, + })}, + ) + approved: Final = await hub.client.put( + f"/skill-hub/submissions/{receipt['submission_id']}/approve", headers=_headers(MASTER_KEY), + ) + assert approved.status_code == 200 + skill: Final = approved.json()["skill"] + assert skill["security_scan_passed"] is not has_issues + assert skill["created_by"] == MEMBER_KEY + assert skill["approved_by"] == approved.json()["submission"]["reviewed_by"] + assert skill["published_by"] == skill["approved_by"] + _stored_upload(hub) + listed: Final = (await hub.client.get(HUB, headers=_headers(MASTER_KEY))).json()["skills"][0] + assert listed["security_scan_passed"] is not has_issues + assert listed["published_by"] == skill["published_by"] + + +@pytest.mark.parametrize("field,order,expected", [ + ("name", "asc", "a"), ("name", "desc", "c"), + ("download_count", "desc", "a"), ("download_count", "asc", "b"), + ("updated_at", "desc", "c"), ("updated_at", "asc", "b"), +]) +async def test_catalog_sorts_all_matches_before_pagination( + hub: HubClient, field: str, order: str, expected: str, +) -> None: + hub.plugins.find_many.return_value = [ + prisma_models.LiteLLM_ClaudeCodePluginTable( + id=name, name=name, enabled=True, anonymous_download_count=0, updated_at=datetime(2026, 1, day), + manifest_json=json.dumps({"source": {"source": "github", "repo": "example/skills"}}), + ) for name, day in (("a", 2), ("b", 1), ("c", 3)) + ] + hub.database.query_raw.return_value = [ + {"id": name, "download_count": count} for name, count in (("a", 8), ("b", 0), ("c", 3)) + ] + response = await hub.client.get(HUB, headers=_headers(PLAIN_KEY), params={ + "sort_by": field, "sort_order": order, "limit": 1, "offset": 0, + }) + assert response.status_code == 200, response.text + assert [skill["name"] for skill in response.json()["skills"]] == [expected] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 84f4d7a4173..de3ced0c583 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3291,3 +3291,21 @@ async def test_pass_through_guardrail_text_is_the_client_body_without_the_inject texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls] assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"] + + +@pytest.mark.asyncio +async def test_moderation_hook_lets_responses_compact_through_without_a_translation() -> None: + handler: Final = UnifiedLLMGuardrails() + guardrail: Final = RecordingGuardrail() + data: Final = { + "guardrail_to_apply": guardrail, + "model": "gpt-5.4-mini", + "input": [{"role": "user", "content": "Remember the number 42."}], + } + + result: Final = await handler.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), call_type="acompact_responses" + ) + + assert result is data + assert guardrail.apply_calls == [] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py index bccc531793f..712750f702e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py @@ -23,6 +23,8 @@ from fastapi import HTTPException import litellm from litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import ZscalerAIGuard +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.types.utils import GenericGuardrailAPIInputs from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -141,6 +143,216 @@ async def test_apply_guardrail_text_concatenation(mock_api_call): assert call_args.kwargs["content"] == "Hello world" +_CONVERSATION: Final = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is two plus two?"}, + {"role": "assistant", "content": "Four."}, +] + + +async def _execute_policy_bodies( + inputs: GenericGuardrailAPIInputs, + input_type: Literal["request", "response"], + scan_only_tool_results: bool = False, +) -> tuple[object, ...]: + received: Final = asyncio.Queue[object]() + + async def execute_policy(request: web.Request) -> web.Response: + await received.put(await request.json()) + return web.json_response({"statusCode": 200, "action": "ALLOW"}) + + app: Final = web.Application() + app.router.add_post("/v1/detection/execute-policy", execute_policy) + handler: Final = InMemoryGuardrailHandler() + async with TestServer(app) as server: + guardrail_id: Final = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "zscaler-last-user-turn", + "litellm_params": { + "guardrail": "zscaler_ai_guard", + "mode": "pre_call", + "api_key": "test_key", + "api_base": str(server.make_url("/v1/detection/execute-policy")), + "policy_id": 1, + "scan_only_tool_results": scan_only_tool_results, + }, + }, + )["guardrail_id"] + try: + guardrail: Final = handler.guardrail_id_to_custom_guardrail[guardrail_id] + assert guardrail is not None + returned: Final = await guardrail.apply_guardrail(inputs, {}, input_type) + finally: + handler.delete_in_memory_guardrail(guardrail_id) + + assert returned is inputs + return tuple(received.get_nowait() for _ in range(received.qsize())) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inputs", "expected_content"), + [ + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [*_CONVERSATION, {"role": "user", "content": "And three plus three?"}], + }, + "And three plus three?", + ), + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [ + *_CONVERSATION, + { + "role": "user", + "content": [{"type": "text", "text": "And three"}, {"type": "text", "text": "plus three?"}], + }, + ], + }, + "And three plus three?", + ), + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [ + *_CONVERSATION, + {"role": "user", "content": "And three plus three?"}, + {"role": "user", "content": "Answer in words."}, + ], + }, + "And three plus three? Answer in words.", + ), + ( + { + "texts": ["Summarize", "this report"], + "structured_messages": [ + { + "role": "user", + "parts": [{"kind": "text", "text": "Summarize"}, {"kind": "text", "text": "this report"}], + "messageId": "msg-1", + } + ], + }, + "Summarize this report", + ), + ({"texts": ["first input", "second input"]}, "first input second input"), + ( + { + "texts": ["every", "message"], + "structured_messages": [*_CONVERSATION, {"role": "assistant", "content": "Prefilled answer:"}], + }, + "What is two plus two?", + ), + ( + { + "texts": ["every", "message"], + "structured_messages": [ + *_CONVERSATION, + {"role": "tool", "tool_call_id": "call_1", "content": "tool output"}, + ], + }, + "What is two plus two?", + ), + ( + { + "texts": ["every", "message"], + "structured_messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "assistant", "content": "Prefilled answer:"}, + ], + }, + None, + ), + ( + { + "texts": ["You are a helpful assistant.", " "], + "structured_messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": " "}, + ], + }, + None, + ), + ( + { + "texts": ["every", "message", "\n"], + "structured_messages": [ + *_CONVERSATION, + { + "role": "user", + "content": [ + {"type": "text", "text": "\n"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + ], + }, + ], + }, + None, + ), + ], + ids=[ + "string-content", + "content-parts", + "consecutive-user-messages", + "a2a-message-parts", + "texts-only-route", + "ends-on-assistant", + "ends-on-tool", + "no-user-message", + "whitespace-only-latest-turn", + "whitespace-text-beside-image", + ], +) +async def test_request_scan_sends_only_the_last_user_turn( + inputs: GenericGuardrailAPIInputs, expected_content: str | None +) -> None: + bodies: Final = await _execute_policy_bodies(inputs, "request") + + assert bodies == ( + () if expected_content is None else ({"direction": "IN", "content": expected_content, "policyId": 1},) + ) + + +@pytest.mark.asyncio +async def test_response_scan_sends_every_text() -> None: + inputs: Final = GenericGuardrailAPIInputs( + texts=["First choice.", "Second choice."], + structured_messages=[*_CONVERSATION, {"role": "user", "content": "And three plus three?"}], + ) + + bodies: Final = await _execute_policy_bodies(inputs, "response") + + assert bodies == ({"direction": "OUT", "content": "First choice. Second choice.", "policyId": 1},) + + +@pytest.mark.asyncio +async def test_scan_only_tool_results_sends_the_tool_results() -> None: + inputs: Final = GenericGuardrailAPIInputs( + texts=["Sunny, 22C"], + structured_messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Sunny, 22C"}, + ], + ) + + bodies: Final = await _execute_policy_bodies(inputs, "request", scan_only_tool_results=True) + + assert bodies == ({"direction": "IN", "content": "Sunny, 22C", "policyId": 1},) + + @pytest.mark.asyncio @patch( "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call", diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 0c7292674a6..892c891f20c 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1939,6 +1939,7 @@ async def test_callback_handles_every_status_the_interactions_api_can_return(): assert released == { "completed", "requires_action", + "requires_continuation", "failed", "cancelled", "incomplete", diff --git a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 180362f45f1..94e0f91ecd1 100644 --- a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -411,10 +411,11 @@ async def test_create_user_without_groups_defers_to_default_team(mocker: MockerF @pytest.mark.asyncio -async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch): +@pytest.mark.parametrize("user_id", ["new-user", " employee-42 "]) +async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch, user_id: str): scim_user = SCIMUser( schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], - userName="new-user", + userName=user_id, emails=[SCIMUserEmail(value="new@example.com")], groups=[SCIMUserGroup(value="idp-team")], ) @@ -428,10 +429,12 @@ async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, mo await create_user(user=scim_user) assert new_user_mock.call_args.kwargs["data"].teams == ["idp-team"] + assert new_user_mock.call_args.kwargs["data"].user_id == user_id @pytest.mark.asyncio -async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch): +@pytest.mark.parametrize("user_id", ["group-user", " employee-42 "]) +async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch, user_id: str): monkeypatch.setattr( "litellm.default_internal_user_params", {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, @@ -439,14 +442,15 @@ async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFi ) new_user_mock = mocker.patch( # test-quality-ok: new_user is imported inside the helper, not injectable "litellm.proxy.management_endpoints.internal_user_endpoints.new_user", - AsyncMock(return_value=NewUserResponse(user_id="group-user", key="k")), + AsyncMock(return_value=NewUserResponse(user_id=user_id, key="k")), ) - created = await _create_user_if_not_exists(user_id="group-user") + created = await _create_user_if_not_exists(user_id=user_id) assert created is not None assert new_user_mock.call_args.kwargs["data"].teams is None assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + assert new_user_mock.call_args.kwargs["data"].user_id == user_id @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index cee0c90f80c..b96a5b606eb 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -13,6 +13,7 @@ import httpx import pytest from fastapi import HTTPException from fastapi.testclient import TestClient +from pydantic import ValidationError from pytest_mock import MockerFixture from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -48,6 +49,51 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( client = TestClient(app) + +@pytest.mark.parametrize("user_id", ["", " \t\n", "user\x00id", 123, True, [], {}]) +def test_new_user_rejects_invalid_custom_id(user_id: object) -> None: + with pytest.raises(ValidationError) as exc: + NewUserRequest.model_validate({"user_id": user_id}) + assert exc.value.errors()[0]["loc"] == ("user_id",) + + +@pytest.mark.parametrize("user_id", [None, "employee-42", "社員@example.test", " employee-42 "]) +def test_new_user_normalizes_custom_id_and_preserves_auto_generation(user_id: str | None) -> None: + request: Final = NewUserRequest(user_id=user_id) + assert request.user_id == (user_id.strip() if user_id is not None else None) + + +@pytest.mark.asyncio +async def test_new_user_returns_conflict_when_id_is_taken_after_duplicate_check(mocker: MockerFixture) -> None: + from prisma.errors import UniqueViolationError + + from litellm.proxy import proxy_server + from litellm.proxy.utils import PrismaClient + + table: Final = SimpleNamespace( + find_first=AsyncMock(return_value=None), + count=AsyncMock(return_value=0), + create=AsyncMock(side_effect=UniqueViolationError( + {"user_facing_error": {"meta": {"target": ["user_id"]}}}, + message="Unique constraint failed on user_id", + )), + upsert=AsyncMock(return_value=LiteLLM_UserTable(user_id="employee-42")), + ) + mocker.patch("prisma.Prisma", return_value=SimpleNamespace(litellm_usertable=table)) + database: Final = PrismaClient( + database_url="postgresql://unused:unused@localhost/unused", + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", database) + with pytest.raises(ProxyException) as exc: + await new_user( + data=NewUserRequest(user_id="employee-42", auto_create_key=False, send_invite_email=False), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert str(exc.value.code) == "409" + assert "employee-42" in exc.value.message + + if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache from litellm.proxy.utils import PrismaClient @@ -4862,7 +4908,9 @@ async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget, models = [] class _FakePrisma: - async def insert_data(self, data, table_name): + async def insert_data( + self, data: Mapping[str, object], table_name: str, *, allow_existing_user: bool = True + ) -> _FakeUserRow | SimpleNamespace: if table_name == "user": captured["user_data"] = dict(data) return _FakeUserRow() diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index ba78fefd653..6ccb019b937 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_UserTable, NewUserResponse, + SSOUserDefinedValues, ProxyException, ) from litellm.proxy.auth.handle_jwt import JWTHandler @@ -52,6 +53,62 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( from tests._master_key import MASTER_KEY +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("billable_users", "premium", "admitted"), + [(4, False, True), (5, False, False), (26, False, False), (26, True, True)], +) +async def test_sso_auto_creation_enforces_free_user_cap_through_new_user( + monkeypatch: pytest.MonkeyPatch, billable_users: int, premium: bool, admitted: bool +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.utils import PrismaClient + + row: Final = LiteLLM_UserTable(user_id="new-sso-user", user_email="new-sso-user@example.test") + table: Final = SimpleNamespace( + find_first=AsyncMock(return_value=None), + find_unique=AsyncMock(return_value=None), + count=AsyncMock(side_effect=lambda where=None: billable_users if where is None else 0), + create=AsyncMock(return_value=row), + upsert=AsyncMock(return_value=row), + ) + transaction: Final = MagicMock() + transaction.__aenter__ = AsyncMock( + return_value=SimpleNamespace(litellm_usertable=table, execute_raw=AsyncMock(), query_raw=AsyncMock()) + ) + transaction.__aexit__ = AsyncMock(return_value=None) + monkeypatch.setattr( + "prisma.Prisma", + lambda **kwargs: SimpleNamespace(litellm_usertable=table, tx=MagicMock(return_value=transaction)), + ) + database: Final = PrismaClient( + database_url="postgresql://unused:unused@localhost/unused", + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "premium_user", premium) + monkeypatch.setattr("asyncio.sleep", AsyncMock()) + values: Final[SSOUserDefinedValues] = { + "models": [], + "user_id": row.user_id, + "user_email": row.user_email, + "user_role": "internal_user_viewer", + "max_budget": None, + "budget_duration": None, + } + if admitted: + response: Final = await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values) + assert response.user_id == row.user_id + assert table.create.await_count + table.upsert.await_count == 1 + else: + with pytest.raises(ProxyException) as error: + await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values) + assert str(error.value.code) == "403" + assert "SSO for more than 5 users" in error.value.message + table.create.assert_not_awaited() + table.upsert.assert_not_awaited() + + def _sso_http_handler(route: Callable[[httpx.Request], httpx.Response] | None = None) -> AsyncHTTPHandler: """litellm's SSO httpx client answered in-process, by ``route`` or by the active respx mock.""" handler = AsyncHTTPHandler() @@ -1281,7 +1338,8 @@ async def test_upsert_sso_user_fills_user_alias_for_existing_user(): @pytest.mark.asyncio -async def test_insert_sso_user_sets_user_alias_from_display_name(): +@pytest.mark.parametrize("user_id", ["S-1-5-21-adfs-user", " employee-42 "]) +async def test_insert_sso_user_sets_user_alias_from_display_name(user_id: str): """ A newly created SSO user is inserted with the IdP display name as user_alias. """ @@ -1290,7 +1348,7 @@ async def test_insert_sso_user_sets_user_alias_from_display_name(): from litellm.proxy.management_endpoints.ui_sso import insert_sso_user sso_result = CustomOpenID( - id="S-1-5-21-adfs-user", + id=user_id, email="jane.doe@example.com", first_name="Jane", last_name="Doe", @@ -1300,7 +1358,7 @@ async def test_insert_sso_user_sets_user_alias_from_display_name(): ) user_defined_values: SSOUserDefinedValues = { "models": [], - "user_id": "S-1-5-21-adfs-user", + "user_id": user_id, "user_email": "jane.doe@example.com", "max_budget": None, "user_role": "internal_user", @@ -1309,12 +1367,14 @@ async def test_insert_sso_user_sets_user_alias_from_display_name(): with patch( "litellm.proxy.management_endpoints.ui_sso.new_user", - return_value=NewUserResponse(user_id="S-1-5-21-adfs-user", key="sk-xxxxx", teams=None), + return_value=NewUserResponse(user_id=user_id, key="sk-xxxxx", teams=None), ) as mock_new_user: await insert_sso_user(result_openid=sso_result, user_defined_values=user_defined_values) new_user_request = mock_new_user.call_args.kwargs["data"] - assert new_user_request.user_id == "S-1-5-21-adfs-user" + assert new_user_request.user_id == user_id + assert new_user_request.sso_user_id == user_id + assert new_user_request.metadata == {"auth_provider": "generic"} assert new_user_request.user_email == "jane.doe@example.com" assert new_user_request.user_alias == "Doe, Jane" diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py index cf1ac61c93a..2b966a8ebeb 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py @@ -191,6 +191,19 @@ async def _run(prisma, users, caller=ADMIN, license=None, generate_key=_no_keys) ) +@pytest.mark.asyncio +@pytest.mark.parametrize("user_id", [None, "", " employee-42 ", " \t"]) +async def test_bulk_keeps_existing_user_id_generation_and_whitespace_behavior(user_id: str | None) -> None: + prisma: Final = _FakePrisma() + response: Final = await _run(prisma, [{"user_id": user_id}]) + result: Final = response.data[0] + assert result.success + assert result.user_id + assert result.user_id in prisma.db.litellm_usertable.rows + if user_id: + assert result.user_id == user_id + + @pytest.mark.asyncio async def test_creates_users_and_team_membership_in_every_store(): prisma = _FakePrisma(teams=[_team("t1", [Member(user_id="existing", role="admin")]), _team("t2")]) diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index dd3ffad4727..d1afcf80ec9 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -1,3 +1,5 @@ +from litellm.proxy.deps.repositories import get_plugin_upload_repository +from litellm.repositories.plugin_upload_repository import PluginUploadRepository import asyncio import json import re @@ -1804,7 +1806,6 @@ def test_public_skill_hub_echoes_registered_installation_preference(preference: from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import register_plugin - from litellm.repositories.plugin_upload_repository import PluginUploadRepository from litellm.types.proxy.claude_code_endpoints import RegisterPluginRequest def create(data: Mapping[str, object]) -> LiteLLM_ClaudeCodePluginTable: @@ -1826,6 +1827,7 @@ def test_public_skill_hub_echoes_registered_installation_preference(preference: prisma.db.litellm_claudecodepluginversion.create = AsyncMock() app: Final = FastAPI() app.include_router(router) + app.dependency_overrides[get_plugin_upload_repository] = lambda: PluginUploadRepository(prisma) with patch("litellm.proxy.proxy_server.prisma_client", prisma): asyncio.run( @@ -1843,8 +1845,10 @@ def test_public_skill_hub_echoes_registered_installation_preference(preference: user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), ) ) + prisma.db.query_raw.return_value = [{"id": "plugin-id", "download_count": 12}] response: Final = TestClient(app).get("/public/skill_hub") + assert response.json()["plugins"][0]["download_count"] == 12 assert response.status_code == 200 assert [(p["name"], p["installation_preference"]) for p in response.json()["plugins"]] == [ ("security-review", preference) @@ -1854,6 +1858,7 @@ def test_public_skill_hub_echoes_registered_installation_preference(preference: _INTERNAL_PLUGIN: Final = SimpleNamespace( id="plugin-id", name="deploy-helper", + created_by=None, enabled=True, created_at=None, updated_at=None, @@ -1874,8 +1879,11 @@ def _skill_hub_response(find_many: AsyncMock, general_settings: Mapping[str, obj table.find_many = find_many prisma: Final = MagicMock() prisma.db.litellm_claudecodeplugintable = table + prisma.writer_db = prisma.db + prisma.db.query_raw = AsyncMock(return_value=[]) app: Final = FastAPI() app.include_router(router) + app.dependency_overrides[get_plugin_upload_repository] = lambda: PluginUploadRepository(prisma) with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), @@ -1909,6 +1917,7 @@ def test_public_skill_hub_500_body_carries_no_exception_text() -> None: def test_public_skill_hub_without_a_database_answers_db_not_connected() -> None: app: Final = FastAPI() app.include_router(router) + app.dependency_overrides[get_plugin_upload_repository] = lambda: None with ( patch("litellm.proxy.proxy_server.prisma_client", None), diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 71f816ea183..cb903fb134a 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1522,6 +1522,37 @@ def test_raise_if_required_body_param_missing_rejects_missing_param(route_type, assert exc_info.value.message == f"{route}: Missing required parameter: '{param}'." +@pytest.mark.parametrize( + "data", + [ + {"model": "gemini-4-argon", "input": "prove the prime number theorem"}, + {"model": "gemini-4-argon", "previous_interaction_id": "v1_abc", "continuation_token": "tok-1"}, + {"model": "gemini-4-argon", "input": "prove it again", "continuation_token": "tok-1"}, + ], +) +def test_create_interaction_accepts_input_or_a_continuation_token(data: dict[str, object]) -> None: + from litellm.proxy.route_llm_request import raise_if_required_body_param_missing + + assert raise_if_required_body_param_missing(route_type="acreate_interaction", data=data, llm_router=None) is None + + +def test_create_interaction_without_input_or_continuation_token_names_input() -> None: + from litellm.proxy.route_llm_request import ( + ProxyMissingRequiredParamError, + raise_if_required_body_param_missing, + ) + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + raise_if_required_body_param_missing( + route_type="acreate_interaction", + data={"model": "gemini-4-argon", "previous_interaction_id": "v1_abc"}, + llm_router=None, + ) + + assert exc_info.value.param == "input" + assert exc_info.value.message == "/interactions: Missing required parameter: 'input'." + + @pytest.mark.parametrize( "data, param", [ diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py index e335b08f040..8603de0c3a1 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py @@ -18,8 +18,10 @@ from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException +from prisma.errors import UniqueViolationError from litellm._service_logger import ServiceTypes from litellm.models.user import LiteLLM_UserTable @@ -192,6 +194,62 @@ async def test_insert_data_team_serializes_members(prisma_client: PrismaClient) assert create_payload["team_id"] == "t1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "message"), + [ + ({"target": ["user_id"]}, "User with id requested-user already exists"), + ({"target": ["sso_user_id"]}, "User with this sso_user_id already exists"), + (None, "User conflicts with an existing unique value"), + ({"target": ["unknown_field"]}, "User conflicts with an existing unique value"), + ], +) +async def test_user_conflict_names_constraint_and_does_not_retry( + prisma_client: PrismaClient, + metadata: Mapping[str, object] | None, + message: str, +) -> None: + create: Final = AsyncMock(side_effect=UniqueViolationError({"user_facing_error": {"meta": metadata}})) + prisma_client.db.litellm_usertable.create = create + with pytest.raises(HTTPException) as error: + await prisma_client.insert_data( + data={"user_id": "requested-user", "sso_user_id": "sso-subject"}, + table_name="user", allow_existing_user=False, + ) + assert error.value.status_code == 409 + assert error.value.detail == {"error": message} + assert create.await_count == 1 + + +@pytest.mark.asyncio +async def test_user_upsert_recovers_from_concurrent_creation( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + row: Final = LiteLLM_UserTable(user_id="requested-user") + conflict: Final = UniqueViolationError({"user_facing_error": {"meta": {"target": ["user_id"]}}}) + prisma_client.db.litellm_usertable.upsert = AsyncMock(side_effect=[conflict, row]) + monkeypatch.setattr("asyncio.sleep", AsyncMock()) + + result: Final = await prisma_client.insert_data(data={"user_id": row.user_id}, table_name="user") + + assert result is row + + +@pytest.mark.asyncio +async def test_user_insert_still_retries_transient_connection_failure( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + row: Final = LiteLLM_UserTable(user_id="requested-user") + create: Final = AsyncMock(side_effect=[httpx.ConnectError("connection unavailable"), row]) + prisma_client.db.litellm_usertable.create = create + monkeypatch.setattr("asyncio.sleep", AsyncMock()) + result: Final = await prisma_client.insert_data( + data={"user_id": row.user_id}, table_name="user", allow_existing_user=False + ) + assert result is row + assert create.await_count == 2 + + @pytest.mark.asyncio async def test_insert_data_user_organization_fk_raises_400( prisma_client: PrismaClient, diff --git a/tests/unit/repositories/test_plugin_publication.py b/tests/unit/repositories/test_plugin_publication.py index 0a784553d09..7a96b0fd5a7 100644 --- a/tests/unit/repositories/test_plugin_publication.py +++ b/tests/unit/repositories/test_plugin_publication.py @@ -16,6 +16,7 @@ from litellm.models.plugin_upload import ( PluginPublication, PluginPublicationBase, PluginPublicationError, + PluginReview, PublishedPlugin, ) from litellm.repositories.plugin_publication import ( @@ -144,6 +145,44 @@ async def test_first_legacy_update_preserves_its_original_snapshot( } assert replacement.kwargs["data"] == { "plugin_id": current.id, "version": result.version, "description": "replacement", - "manifest_json": json.dumps({"source": source, "version": result.version}), + "manifest_json": json.dumps({"source": source, "version": result.version, "_litellm_published_by": "admin"}), "published_at": result.updated_at, "published_by": "admin", "content": None, "archive_available": False, } + + +@pytest.mark.parametrize("same_source", (True, False)) +@pytest.mark.parametrize("review", [None, PluginReview("reviewer@example.test", False)]) +@pytest.mark.asyncio +async def test_publication_preserves_only_trusted_review_for_unchanged_source( + publication_tx: MagicMock, review: PluginReview | None, same_source: bool, +) -> None: + original = publication_tx.litellm_claudecodeplugintable.find_unique.return_value + stored_review = {"approved_by": "original-reviewer", "security_scan_passed": True} + stored_manifest = {**json.loads(original.manifest_json), "_litellm_review": stored_review} + current = original.model_copy(update={"manifest_json": json.dumps(stored_manifest)}) + publication_tx.litellm_claudecodeplugintable.find_unique.return_value = current + + async def update(where: Mapping[str, str], data: Mapping[str, object]) -> LiteLLM_ClaudeCodePluginTable: + return current.model_copy(update=dict(data)) + + publication_tx.litellm_claudecodeplugintable.update.side_effect = update + result = await publish_plugin( + cast(PluginPublicationTransaction, publication_tx), + PluginPublication("skill", None, "updated", { + "_litellm_review": {"approved_by": "forged", "security_scan_passed": True}, + "source": stored_manifest["source"] if same_source else {"source": "github", "repo": "example/skills"}, + }, "author", review=review), + PluginPublicationBase(current.id, current.updated_at), + ) + assert isinstance(result, PublishedPlugin) + manifest = json.loads(publication_tx.litellm_claudecodeplugintable.update.call_args.kwargs["data"]["manifest_json"]) + assert manifest["_litellm_published_by"] == (review.approved_by if review else "author") + expected = ( + {"approved_by": review.approved_by, "security_scan_passed": review.security_scan_passed} if review + else stored_review if same_source else None + ) + assert manifest.get("_litellm_review") == expected + assert (None if result.review is None else { + "approved_by": result.review.approved_by, "security_scan_passed": result.review.security_scan_passed, + }) == expected + assert json.loads(publication_tx.litellm_claudecodepluginversion.create.call_args.kwargs["data"]["manifest_json"]).get("_litellm_review") == expected diff --git a/tests/unit/repositories/test_skill_submission_repository.py b/tests/unit/repositories/test_skill_submission_repository.py new file mode 100644 index 00000000000..7187cdbe514 --- /dev/null +++ b/tests/unit/repositories/test_skill_submission_repository.py @@ -0,0 +1,31 @@ +import json + +import pytest + +from litellm.repositories.skill_submission_repository import scan_passed_for_archive + + +@pytest.mark.parametrize("status,digest,report,expected", [ + ("completed", "current", {"complete": True, "risk_score": 0, "findings": []}, True), + ("running", "current", {"complete": True, "risk_score": 0, "findings": []}, False), + ("failed", "current", {"complete": True, "risk_score": 0, "findings": []}, False), + ("completed", "previous", {"complete": True, "risk_score": 0, "findings": []}, False), + ("completed", "current", {"complete": False, "risk_score": 0, "findings": []}, False), + ("completed", "current", None, False), + ("completed", "current", {"complete": True, "risk_score": 90, "findings": []}, False), + ("completed", "current", {"complete": True, "risk_score": 0, "findings": [{"rule_id": "rule", "severity": "LOW", "message": "Issue", "file": "SKILL.md"}]}, False), + ("completed", "current", {"complete": True, "risk_score": 0}, False), +]) +def test_scan_pass_requires_clean_complete_report_for_published_archive( + status: str, digest: str, report: dict[str, object] | None, expected: bool, +) -> None: + assert scan_passed_for_archive(json.dumps({ + "id": "scan", "requested_by": "admin", "started_at": "2026-01-01T00:00:00Z", + "expires_at": "2026-01-01T00:03:00Z", "status": status, "archive_sha256": digest, + "report": None if report is None else {"scanner_version": "test", "scan_mode": "static-only", **report}, + }), "current") is expected + + +@pytest.mark.parametrize("record", [None, "", "{}", "null", "not-json"]) +def test_missing_or_invalid_scan_is_not_reported_as_passed(record: str | None) -> None: + assert scan_passed_for_archive(record, "current") is False diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 037004ab999..a39bc09cee6 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -2722,7 +2722,7 @@ async def test_stream_auto_execute_bills_a_follow_up_cut_after_streamed_text_wit @pytest.mark.asyncio -async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row( +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_one_failure_row( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -2743,10 +2743,22 @@ async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_r if not ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked): ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) await logging_obj.async_failure_handler(blocked, "blocked by output-policy") - await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + await _settle(lambda: len(recorder.failed) > 0) - _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) - assert [row["response_cost"] for row in recorder.failed] == [0.0] + assert recorder.logged == () + assert len(recorder.failed) == 1 + row: Final = recorder.failed[0] + charged_once: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN + assert row["status"] == "failure" + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + client: Final = _final_response(events) + assert "message" in _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (71 + 149, 19 + 5) + assert client.usage.cost == pytest.approx(charged_once) @pytest.mark.asyncio diff --git a/tests/unit/router_rust/test_walk.py b/tests/unit/router_rust/test_walk.py new file mode 100644 index 00000000000..2f2febea3eb --- /dev/null +++ b/tests/unit/router_rust/test_walk.py @@ -0,0 +1,49 @@ +import pytest + +from litellm.router_rust.invoke.walk import WalkMetadata + + +def _no_base_model(_model_id: str) -> str | None: + return None + + +@pytest.mark.parametrize("field", ["metadata", "litellm_metadata"]) +def test_hop_kwargs_keep_the_requests_key_labels_under_the_walk_stamps(field: str) -> None: + """The fallback loggers (prometheus' fallback counters) read the caller's key and team off + the hop's metadata, so a hop keeps the request's bucket and lays its own stamps over it.""" + request = { + "model": "primary", + field: { + "user_api_key_hash": "hashed-key", + "user_api_key_alias": "key-alias", + "user_api_key_team_id": "team-id", + "user_api_key_team_alias": "team-alias", + "attempted_fallbacks": "stale", + }, + } + walk = WalkMetadata(request, field, _no_base_model) + walk.hop({"parent": 0, "scope": 1, "from": "primary", "to": "fallback", "depth": 1, "max_fallbacks": 5}) + + kwargs = walk.hop_kwargs({"scope": 1, "to": "fallback"}) + + assert kwargs["model"] == "fallback" + metadata = kwargs[field] + assert isinstance(metadata, dict) + assert metadata["user_api_key_hash"] == "hashed-key" + assert metadata["user_api_key_alias"] == "key-alias" + assert metadata["user_api_key_team_id"] == "team-id" + assert metadata["user_api_key_team_alias"] == "team-alias" + assert metadata["model_group"] == "fallback" + assert metadata["original_model_group"] == "primary" + assert metadata["attempted_fallbacks"] == 1 + assert request[field]["attempted_fallbacks"] == "stale" + assert "model_group" not in request[field] + + +def test_hop_kwargs_without_request_metadata_hold_the_walk_stamps() -> None: + walk = WalkMetadata({"model": "primary"}, "metadata", _no_base_model) + walk.hop({"parent": 0, "scope": 1, "from": "primary", "to": "fallback", "depth": 1, "max_fallbacks": 5}) + + metadata = walk.hop_kwargs({"scope": 1, "to": "fallback"})["metadata"] + + assert metadata == {"original_model_group": "primary", "model_group": "fallback", "attempted_fallbacks": 1} diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 9487e4d6321..5169642bc43 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1480,3 +1480,58 @@ async def test_a_pin_written_without_tools_is_read_when_tools_is_an_empty_list() ) assert filtered == [deployments[0]] + + +class _LoggedMessagesProbe(CustomLogger): + def __init__(self) -> None: + self.logged_messages: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_messages.append((kwargs["standard_logging_object"]["messages"], kwargs["messages"])) + + +@pytest.mark.asyncio +async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_redacted_messages( + monkeypatch, local_model_cost_map +): + """ + `turn_off_message_logging` redacted the shared success-event payload before the check read it, so the + pin was keyed off `redacted-by-litellm`, never written, and a cacheable conversation bounced across the + group. The pin must land whatever the flag says, while a real logger alongside still sees only the + redacted messages. + """ + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + probe = _LoggedMessagesProbe() + monkeypatch.setattr(litellm, "callbacks", [probe]) + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-1", "dep-2") + ], + optional_pre_call_checks=["prompt_caching"], + ) + messages = _messages(word_count=5000) + + first = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok") + served_by = first._hidden_params["model_id"] + + affinity_key = PromptCachingCache.get_prompt_caching_cache_key(messages, None) + pin = await _eventually(lambda: router.prompt_caching_pin_cache.get_cache(key=affinity_key)) + assert pin == {"model_id": served_by}, pin + + subsequent = [ + (await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok"))._hidden_params[ + "model_id" + ] + for _ in range(4) + ] + assert subsequent == [served_by] * 4 + + await _eventually(lambda: len(probe.logged_messages) == 5) + redacted = [{"role": "user", "content": "redacted-by-litellm"}] + assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 8ef683c0aee..6a35c8a9602 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -36,6 +36,7 @@ from litellm.router_utils.cooldown_handlers import ( async_get_cooldown_deployments, cast_exception_status_to_int, get_cooldown_deployments, + is_response_lookup_not_found, mark_advisor_orchestration_failure, set_cooldown_deployments, should_cooldown_based_on_allowed_fails_policy, @@ -3565,3 +3566,25 @@ def test_the_refusal_names_each_cooled_deployment_some_cooldown_state_routes_the with pytest.raises(RouterRateLimitError) as refusal: _select(_tenant_router(EVERY_TENANT, **BOTH_ON), "async", "tenants", copy.deepcopy(request_kwargs)) assert {deployment["model_id"] for deployment in refusal.value.cooldown_deployments or ()} == reachable + + +@pytest.mark.parametrize( + ("status_code", "call_type", "expected"), + [ + (404, "aget_responses", True), + (404, "adelete_responses", True), + (404, "acancel_responses", True), + (404, "alist_input_items", True), + (404, "get_responses", True), + (404, "aresponses", False), + (404, "acompact_responses", False), + (404, "acompletion", False), + (500, "aget_responses", False), + (429, "alist_input_items", False), + ], +) +def test_is_response_lookup_not_found(status_code: int, call_type: str, expected: bool) -> None: + error: Final = litellm.NotFoundError(message="gone", llm_provider="openai", model="gpt-4o-mini") + error.status_code = status_code + + assert is_response_lookup_not_found(error, {"call_type": call_type}) is expected diff --git a/tests/unit/test_router/core/test_construction_precall_characterization.py b/tests/unit/test_router/core/test_construction_precall_characterization.py index 2985819cb5b..f58d9a35110 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization.py @@ -13,6 +13,11 @@ from collections.abc import Callable, Generator, Mapping from typing import Final import pytest +import redis +import redis.asyncio +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeConnection +from redis.asyncio import BlockingConnectionPool import litellm from litellm._logging import verbose_router_logger @@ -23,6 +28,7 @@ from litellm.router_api import ChangeSet, router_init_args from litellm.router_core import RouterCore from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck +from litellm.scheduler import FlowItem from litellm.types.router import ( AlertingConfig, AllowedFailsPolicy, @@ -620,6 +626,41 @@ def test_applying_a_redis_cache_attaches_it_to_the_routing_state_and_the_pins(ma assert core.registry().plugins("prompt_caching_pins").redis_cache is redis_cache +@pytest.fixture +def shared_redis(monkeypatch: pytest.MonkeyPatch) -> FakeServer: + server: Final = FakeServer() + + class Pool(BlockingConnectionPool): + def __init__(self, **kwargs: object) -> None: # kwargs-ok: forwarded to the redis pool + super().__init__(connection_class=FakeConnection, server=server, **kwargs) + + monkeypatch.setattr(redis, "Redis", FakeRedis) + monkeypatch.setattr(redis.asyncio, "BlockingConnectionPool", Pool) + return server + + +async def test_two_workers_on_one_redis_share_the_priority_queue(make_core: MakeCore, shared_redis: FakeServer) -> None: + first: Final = _one(make_core, redis_host="scheduler-shared", redis_port=6399) + second: Final = _one(make_core, redis_host="scheduler-shared", redis_port=6399) + await first.scheduler.add_request(FlowItem(priority=0, request_id="urgent", model_name="g")) # pyright: ignore[reportAttributeAccessIssue] + waiter: Final = FlowItem(priority=5, request_id="waiter", model_name="g") + await second.scheduler.add_request(waiter) # pyright: ignore[reportAttributeAccessIssue] + + assert not await second.scheduler.poll(waiter, []) # pyright: ignore[reportAttributeAccessIssue] + await first.scheduler.remove_request("urgent", "g") # pyright: ignore[reportAttributeAccessIssue] + assert await second.scheduler.poll(waiter, []) # pyright: ignore[reportAttributeAccessIssue] + + +async def test_without_redis_each_worker_keeps_its_own_queue(make_core: MakeCore) -> None: + first: Final = _one(make_core) + second: Final = _one(make_core) + await first.scheduler.add_request(FlowItem(priority=0, request_id="urgent", model_name="g")) # pyright: ignore[reportAttributeAccessIssue] + waiter: Final = FlowItem(priority=5, request_id="waiter", model_name="g") + await second.scheduler.add_request(waiter) # pyright: ignore[reportAttributeAccessIssue] + + assert await second.scheduler.poll(waiter, []) # pyright: ignore[reportAttributeAccessIssue] + + def test_cache_responses_installs_a_local_response_cache_on_the_host(make_core: MakeCore, host: ScriptedHost) -> None: core: Final = _one(make_core, cache_responses=True) assert isinstance(host.cache, litellm.Cache) diff --git a/tests/unit/test_router/test_router_provider_endpoints.py b/tests/unit/test_router/test_router_provider_endpoints.py index a78228b8658..baf3e94f1f0 100644 --- a/tests/unit/test_router/test_router_provider_endpoints.py +++ b/tests/unit/test_router/test_router_provider_endpoints.py @@ -18,6 +18,7 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import UserAPIKeyAuth +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.router_callbacks.track_deployment_metrics import ( get_deployment_successes_for_current_minute, ) @@ -921,6 +922,43 @@ async def test_video_lookup_404_never_cools_down_the_deployment_that_answered(re assert await router.cooldown_cache.async_get_active_cooldowns(["dep-a", "dep-b"], parent_otel_span=None) == [] +@pytest.mark.parametrize( + ("endpoint", "method", "path"), + [ + ("aget_responses", "GET", ""), + ("adelete_responses", "DELETE", ""), + ("acancel_responses", "POST", "/cancel"), + ("alist_input_items", "GET", "/input_items"), + ], +) +@pytest.mark.asyncio +async def test_response_lookup_404_never_cools_down_the_deployment_that_answered( + respx_mock: respx.MockRouter, endpoint: str, method: str, path: str +) -> None: + respx_mock.route(method=method, url=f"https://api.openai.com/v1/responses/resp_gone{path}").respond( + 404, json={"error": {"code": "NotFound", "message": "Response not found"}} + ) + router: Final = Router( + model_list=[ + { + "model_name": "gpt", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": api_key}, + "model_info": {"id": deployment_id}, + } + for api_key, deployment_id in (("key-a", "dep-a"), ("key-b", "dep-b")) + ], + num_retries=0, + allowed_fails=1, + ) + response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "dep-b", "resp_gone") + + for _ in range(3): + with pytest.raises(litellm.NotFoundError): + await getattr(router, endpoint)(response_id=response_id) + + assert await router.cooldown_cache.async_get_active_cooldowns(["dep-a", "dep-b"], parent_otel_span=None) == [] + + @pytest.mark.asyncio async def test_chat_404_still_cools_down_the_deployment(respx_mock: respx.MockRouter) -> None: respx_mock.post("https://api.openai.com/v1/chat/completions").respond( diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 924c3738638..c41c4db2c02 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -104,6 +104,7 @@ from litellm.utils import ( _snapshot_exception_for_hook, async_post_call_failure_deployment_hook, async_post_call_success_deployment_hook, + async_pre_call_deployment_hook, calculate_max_parallel_requests, client, get_non_default_completion_params, @@ -11626,3 +11627,25 @@ async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronous assert GLOBAL_LOGGING_WORKER.pending_count > 0, "Tracked enqueue must finish before worker draining starts" await GLOBAL_LOGGING_WORKER.flush() assert logging_obj.model_call_details["has_logged_async_success"] is True + + +class _RecordingPreCallDeploymentHook(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen_call_types: tuple[CallTypes | None, ...] = () + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + self.seen_call_types = (*self.seen_call_types, call_type) + return None + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_hands_responses_compact_its_call_type(monkeypatch: pytest.MonkeyPatch) -> None: + hook: Final = _RecordingPreCallDeploymentHook() + monkeypatch.setattr(litellm, "callbacks", [hook]) + + await async_pre_call_deployment_hook({"model": "gpt-5.4-mini"}, "acompact_responses") + + assert hook.seen_call_types == (CallTypes.acompact_responses,) diff --git a/tests/unit/types/test_utils.py b/tests/unit/types/test_utils.py index 32479c62372..bdbcbb698d4 100644 --- a/tests/unit/types/test_utils.py +++ b/tests/unit/types/test_utils.py @@ -1,7 +1,12 @@ import pytest from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_excluded_services_value -from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VALUES, CaptureMessageContent, captures_span_content +from litellm.types.utils import ( + CAPTURE_MESSAGE_CONTENT_VALUES, + CallTypes, + CaptureMessageContent, + captures_span_content, +) @pytest.mark.parametrize("value", ["", " ", "auth", "none,redis", "redis,", ",postgres"]) @@ -45,3 +50,8 @@ def test_capture_message_content_constructor_accepts_any_string() -> None: ) def test_captures_span_content(mode: str | None, captures: bool) -> None: assert captures_span_content(mode) is captures + + +@pytest.mark.parametrize("call_type", ["compact_responses", "acompact_responses"]) +def test_responses_compact_call_types_round_trip_their_public_names(call_type: str) -> None: + assert CallTypes(call_type).value == call_type diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts index 652b2c0b480..e44fca35d88 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts @@ -15,11 +15,12 @@ import { export interface FieldConfig { name: string; - label: string; + labelKey: ParseKeys; type: "text" | "textarea" | "url" | "switch" | "list" | "select"; required?: boolean; - tooltip?: string; + tooltipKey?: ParseKeys; placeholder?: string; + placeholderKey?: ParseKeys; defaultValue?: any; rows?: number; validation?: any[]; @@ -50,42 +51,41 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "name", - label: "Display Name", + labelKey: "agents.formConfig.fields.name.label", type: "text", required: true, - placeholder: "e.g., Customer Support Agent", + placeholderKey: "agents.formConfig.fields.name.placeholder", }, { name: "description", - label: "Description", + labelKey: "agents.formConfig.fields.description.label", type: "textarea", required: false, - placeholder: "Describe what this agent does...", + placeholderKey: "agents.formConfig.fields.description.placeholder", rows: 3, }, { name: "url", - label: "URL", + labelKey: "agents.formConfig.fields.url.label", type: "url", required: false, placeholder: "http://localhost:9999/", - tooltip: "Base URL where the agent is hosted (optional)", + tooltipKey: "agents.formConfig.fields.url.tooltip", }, { name: "version", - label: "Version", + labelKey: "agents.formConfig.fields.version.label", type: "text", placeholder: "1.0.0", defaultValue: "1.0.0", }, { name: "protocolVersion", - label: "Protocol Version", + labelKey: "agents.formConfig.fields.protocolVersion.label", type: "select", options: ["1.0", "0.3"], defaultValue: "1.0", - tooltip: - "The A2A protocol version LiteLLM serves to clients for this agent. LiteLLM converts the upstream agent's responses to this version, so clients always see the version you pick here regardless of the original agent's version.", + tooltipKey: "agents.formConfig.fields.protocolVersion.tooltip", helpTextKey: "agents.formConfig.basic.protocolVersionHelp", }, ], @@ -96,7 +96,7 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "skills", - label: "Skills", + labelKey: "agents.formConfig.fields.skills.label", type: "list", defaultValue: [], }, @@ -108,18 +108,18 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "streaming", - label: "Streaming", + labelKey: "agents.formConfig.fields.streaming.label", type: "switch", defaultValue: false, }, { name: "pushNotifications", - label: "Push Notifications", + labelKey: "agents.formConfig.fields.pushNotifications.label", type: "switch", }, { name: "stateTransitionHistory", - label: "State Transition History", + labelKey: "agents.formConfig.fields.stateTransitionHistory.label", type: "switch", }, ], @@ -130,19 +130,19 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "iconUrl", - label: "Icon URL", + labelKey: "agents.formConfig.fields.iconUrl.label", type: "url", placeholder: "https://example.com/icon.png", }, { name: "documentationUrl", - label: "Documentation URL", + labelKey: "agents.formConfig.fields.documentationUrl.label", type: "url", placeholder: "https://docs.example.com", }, { name: "supportsAuthenticatedExtendedCard", - label: "Supports Authenticated Extended Card", + labelKey: "agents.formConfig.fields.supportsAuthenticatedExtendedCard.label", type: "switch", }, ], @@ -153,12 +153,12 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "model", - label: "Model (Optional)", + labelKey: "agents.formConfig.fields.model.label", type: "text", }, { name: "make_public", - label: "Make Public", + labelKey: "agents.formConfig.fields.makePublic.label", type: "switch", }, ], @@ -169,24 +169,24 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "cost_per_query", - label: "Cost Per Query ($)", + labelKey: "agents.formConfig.fields.costPerQuery.label", type: "text", placeholder: "0.0", - tooltip: "Fixed cost per query", + tooltipKey: "agents.formConfig.fields.costPerQuery.tooltip", }, { name: "input_cost_per_token", - label: "Input Cost Per Token ($)", + labelKey: "agents.formConfig.fields.inputCostPerToken.label", type: "text", placeholder: "0.000001", - tooltip: "Cost per input token", + tooltipKey: "agents.formConfig.fields.inputCostPerToken.tooltip", }, { name: "output_cost_per_token", - label: "Output Cost Per Token ($)", + labelKey: "agents.formConfig.fields.outputCostPerToken.label", type: "text", placeholder: "0.000002", - tooltip: "Cost per output token", + tooltipKey: "agents.formConfig.fields.outputCostPerToken.tooltip", }, ], }, @@ -196,10 +196,10 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "enable_tracing", - label: "Enable Tracing", + labelKey: "agents.formConfig.fields.enableTracing.label", type: "switch", defaultValue: false, - tooltip: "Enable request tracing for this agent", + tooltipKey: "agents.formConfig.fields.enableTracing.tooltip", }, ], }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx index 41e25d14054..2f69336e6c6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx @@ -256,23 +256,24 @@ const AgentFormFields: React.FC = ({ panels, showAgentName {({ value, onChange, ref, ...control }) => { const text = typeof value === "string" ? value : ""; + const placeholder = field.placeholderKey ? t(field.placeholderKey) : field.placeholder; if (field.type === "textarea") { return (