Merge remote-tracking branch 'origin/main' into litellm_responses_calltypes
Some checks failed
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled

This commit is contained in:
mateo-berri 2026-10-28 19:20:18 -07:00
commit 7b55d9056a
1115 changed files with 62089 additions and 13328 deletions

View file

@ -1,6 +1,8 @@
---
name: database-access
description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope
metadata:
internal: true
---
# Database access

1
.github/CODEOWNERS vendored
View file

@ -2,3 +2,4 @@
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri
/helm/ @yassin-berriai
/terraform/ @yassin-berriai
/.github/workflows/zizmor.yml @yuneng-berri

View file

@ -1,7 +1,7 @@
name: Cost map guard
on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed
pull_request_target:
on:
pull_request:
branches:
- main
- "litellm_**"
@ -19,6 +19,7 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ github.event.pull_request.base.sha }}
persist-credentials: false
- name: Fetch the pull request head and its merge base
id: revisions

View file

@ -19,7 +19,7 @@ jobs:
name: interactions-spec-canary
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 15
timeout-minutes: 30
permissions:
contents: read
outputs:
@ -39,8 +39,11 @@ jobs:
with:
version: "0.10.9"
- name: Set up the Rust bridge
uses: ./.github/actions/rust-bridge
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
- name: Run the compliance tests against Google's live spec
id: compliance

View file

@ -1,13 +1,9 @@
name: Spend logs guard
on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's diff is read as data and never executed
pull_request_target:
on:
pull_request:
branches:
- main
pull_request_review:
types:
- submitted
- dismissed
permissions:
contents: read
@ -23,7 +19,7 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ github.event.pull_request.base.ref }}
ref: ${{ github.event.pull_request.base.sha }}
persist-credentials: false
- name: Fetch the pull request head, its merge base and its reviews
id: revisions

View file

@ -561,6 +561,9 @@ jobs:
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py
- name: router_surface_check
run: uv run --no-sync python ./tests/code_coverage_tests/router_surface_check.py
- name: test_chat_completion_imports
run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py

View file

@ -132,7 +132,7 @@ jobs:
windows-sdk:
name: windows-sdk
runs-on: windows-2025
timeout-minutes: 10
timeout-minutes: 30
env:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
defaults:
@ -154,8 +154,11 @@ jobs:
with:
version: "0.10.9"
- name: Set up the Rust bridge
uses: ./.github/actions/rust-bridge
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --no-install-workspace --group dev --python 3.11
run: .github/scripts/uv_sync_with_retries.sh --frozen --group dev --python 3.11
- name: Run the Windows tests
run: uv run --no-sync python -m pytest --tb=short -v tests/windows_tests

View file

@ -4,6 +4,7 @@ on:
push:
paths:
- "litellm-rust/**"
- "model_prices_and_context_window.json"
- "litellm/rust_bridge/**"
- "scripts/generate_trace_types.py"
- "scripts/trace_codegen/**"
@ -38,6 +39,7 @@ on:
- "litellm_**"
paths:
- "litellm-rust/**"
- "model_prices_and_context_window.json"
- "litellm/rust_bridge/**"
- "scripts/generate_trace_types.py"
- "scripts/trace_codegen/**"

View file

@ -42,7 +42,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/**', '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') }}
run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT"
- name: Set up Python
@ -137,6 +137,7 @@ jobs:
include:
- shard: core-utils
test-path: |-
tests/unit/accounting
tests/unit/batch_completion
tests/unit/decisions
tests/unit/extraction
@ -244,6 +245,7 @@ jobs:
tests/unit/llms/clarifai
tests/unit/llms/claude_code
tests/unit/llms/cloudflare
tests/unit/llms/codestral
tests/unit/llms/codex
tests/unit/llms/cohere
tests/unit/llms/cometapi
@ -1306,6 +1308,8 @@ jobs:
done
openai-sdk-compat:
name: OpenAI ${{ matrix.openai-version }} compatibility
needs: rust-bridge
if: ${{ !cancelled() }}
permissions:
contents: read
pull-requests: read
@ -1342,10 +1346,18 @@ jobs:
with:
version: "0.10.9"
- name: Set up the Rust bridge
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 5
uses: ./.github/actions/rust-bridge
with:
artifact: ${{ needs.rust-bridge.outputs.artifact }}
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace
cp rust-bridge-dist/litellm/rust_bridge/_native.* litellm/rust_bridge/
uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z
uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__'

View file

@ -27,10 +27,15 @@ jobs:
with:
persist-credentials: false
- name: Install zizmor
run: |
curl -fsSL --retry 3 -o "$RUNNER_TEMP/zizmor.tar.gz" \
https://github.com/zizmorcore/zizmor/releases/download/v1.31.1/zizmor-x86_64-unknown-linux-gnu.tar.gz
echo "d185f726bdaf92df78273217dfdd24ebcd6b0ef48cf3ce776ff5762b6e138c99 $RUNNER_TEMP/zizmor.tar.gz" | sha256sum -c -
tar xzf "$RUNNER_TEMP/zizmor.tar.gz" -C "$RUNNER_TEMP" zizmor
- name: Run zizmor
uses: zizmorcore/zizmor-action@5f14fd08f7cf1cb1609c1e344975f152c7ee938d # v0.5.6
with:
version: "1.24.1"
min-severity: medium
advanced-security: false
annotations: true
env:
GH_TOKEN: ${{ github.token }}
run: |
"$RUNNER_TEMP/zizmor" --no-ignores --no-config --min-severity=medium --format=github .

View file

@ -427,6 +427,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [Watsonx Text (`watsonx_text`)](https://docs.litellm.ai/docs/providers/watsonx) | ✅ | ✅ | ✅ | | | | | | | |
| [xAI (`xai`)](https://docs.litellm.ai/docs/providers/xai) | ✅ | ✅ | ✅ | | | | | | | |
| [Xinference (`xinference`)](https://docs.litellm.ai/docs/providers/xinference) | | | | ✅ | | | | | | |
| [ZeroGPU (`zerogpu`)](https://docs.litellm.ai/docs/providers/zerogpu) | ✅ | ✅ | ✅ | | | | | | | |
[**Read the Docs**](https://docs.litellm.ai/docs/)
@ -660,6 +661,8 @@ This covers:
We welcome contributions to LiteLLM! Whether you're fixing bugs, adding features, or improving documentation, we appreciate your help.
If you want to work closely with our team, check out [LiteLLM OSS Champions](https://www.litellm.ai/champions), our program for trusted OSS contributors with their own Slack channel and a voice in the roadmap
## Quick Start for Contributors
This requires uv to be installed.

View file

@ -2,9 +2,12 @@
needs an approving review from the spend-log owner before it can merge.
CODEOWNERS can only match file paths, and every schema change edits the same schema.prisma copies,
so a path rule would put every migration on the owner. This gate reads the diff instead: when no
added line in the watched files mentions the table it passes at once, otherwise it passes only once
the owner's latest review on the pull request is an approval. The workflow re-runs on every review.
so a path rule would put every migration on the owner. This gate reads the diff instead: a schema.prisma
hunk counts when it adds or removes a line inside the `model LiteLLM_SpendLogs { ... }` block, a migration
hunk counts when an added line names the table. Without such a hunk it passes at once, otherwise it passes
only once the owner's latest review on the pull request is an approval that still covers every spend-log line
in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never saw).
Reviews do not trigger the workflow, so the check is re-run by hand after approving.
"""
from __future__ import annotations
@ -13,8 +16,10 @@ import argparse
import re
import subprocess
import sys
from collections.abc import Sequence
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
from pydantic import BaseModel, ConfigDict, TypeAdapter
@ -25,7 +30,10 @@ WATCHED: Final = (
"litellm-proxy-extras/litellm_proxy_extras/schema.prisma",
)
MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/"
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs", re.IGNORECASE)
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs(?![A-Za-z0-9])", re.IGNORECASE)
DECIDING_STATES: Final = frozenset({"APPROVED", "CHANGES_REQUESTED", "DISMISSED"})
MODEL_START: Final = re.compile(r"^model\s+LiteLLM_SpendLogs\s*\{")
HUNK: Final = re.compile(r"^@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@[^\n]*\n", re.MULTILINE)
class Reviewer(BaseModel):
@ -37,31 +45,77 @@ class Review(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
state: str
user: Reviewer
commit_id: str | None = None
_REVIEWS: Final = TypeAdapter(list[Review])
@dataclass(frozen=True, slots=True)
class Block:
"""Lines of `model LiteLLM_SpendLogs { ... }` in one schema.prisma copy at the merge base and at the head."""
base: range
head: range
@dataclass(frozen=True, slots=True)
class Hunk:
removed: tuple[tuple[int, str], ...]
added: tuple[tuple[int, str], ...]
def watched(path: str) -> bool:
return path in WATCHED or (path.startswith(MIGRATIONS_PREFIX) and path.endswith("/migration.sql"))
def _added_lines(section: str) -> tuple[str, ...]:
body: Final = section.split("\n@@", 1)
return tuple(line[1:] for line in body[1].splitlines() if line.startswith("+")) if len(body) == 2 else ()
def model_block(schema: str) -> range:
"""1-based line range of `model LiteLLM_SpendLogs { ... }` including both braces, empty when the schema has none."""
lines: Final = schema.splitlines()
starts: Final = tuple(number for number, line in enumerate(lines, 1) if MODEL_START.match(line))
if not starts:
return range(0)
ends: Final = tuple(number for number, line in enumerate(lines, 1) if number > starts[0] and line.startswith("}"))
return range(starts[0], (ends[0] if ends else len(lines)) + 1)
def spend_log_hits(diff: str) -> tuple[str, ...]:
"""`path: added line` for every added line in a watched file that names the table."""
def _hunks(section: str) -> tuple[Hunk, ...]:
headers: Final = tuple(HUNK.finditer(section))
ends: Final = tuple(header.start() for header in headers[1:]) + (len(section),)
return tuple(_hunk(header, section[header.end() : end]) for header, end in zip(headers, ends))
def _hunk(header: re.Match[str], body: str) -> Hunk:
lines: Final = body.splitlines()
return Hunk(
removed=tuple(enumerate((line[1:] for line in lines if line.startswith("-")), int(header.group(1)))),
added=tuple(enumerate((line[1:] for line in lines if line.startswith("+")), int(header.group(2)))),
)
def spend_log_hits(diff: str, blocks: Mapping[str, Block]) -> tuple[str, ...]:
"""`path: +line` or `path: -line` for every changed line that touches the table.
A schema.prisma line counts when its number falls inside that copy's `blocks` entry (head for added lines,
merge base for removed ones), a migration line when it is added and names the table."""
sections: Final = re.split(r"^diff --git [^\n]*\n", diff, flags=re.MULTILINE)
return tuple(f"{path}: {line.strip()}" for path, line in _watched_added_lines(sections) if TABLE.search(line))
def _watched_added_lines(sections: Sequence[str]) -> tuple[tuple[str, str], ...]:
pairs: Final = tuple((_path_of(section), section) for section in sections)
paths: Final = tuple((_path_of(section), section) for section in sections)
return tuple(
(path, line) for path, section in pairs if watched(path) for line in _added_lines(section)
) # comprehension-ok: path then its lines
hit for path, section in paths if watched(path) for hit in _section_hits(path, _hunks(section), blocks)
) # comprehension-ok: path then its hits
def _section_hits(path: str, hunks: Sequence[Hunk], blocks: Mapping[str, Block]) -> tuple[str, ...]:
if path not in blocks:
return tuple(f"{path}: +{line.strip()}" for hunk in hunks for _, line in hunk.added if TABLE.search(line))
block: Final = blocks[path]
removed: Final = tuple(
f"{path}: -{line.strip()}" for hunk in hunks for number, line in hunk.removed if number in block.base
) # comprehension-ok: hunk then its lines
added: Final = tuple(
f"{path}: +{line.strip()}" for hunk in hunks for number, line in hunk.added if number in block.head
) # comprehension-ok: hunk then its lines
return removed + added
def _path_of(section: str) -> str:
@ -70,11 +124,33 @@ def _path_of(section: str) -> str:
def approved_by(reviews: Sequence[Review], login: str) -> bool:
"""True when `login`'s most recent review on the pull request is an approval."""
states: Final = tuple(review.state for review in reviews if review.user.login == login)
"""True when `login`'s most recent review that changes GitHub's approval state is an approval.
COMMENTED and PENDING reviews leave an earlier approval in force on GitHub, so they are skipped here too.
"""
own: Final = tuple(review.state for review in reviews if review.user.login == login)
states: Final = tuple(state for state in own if state in DECIDING_STATES)
return bool(states) and states[-1] == "APPROVED"
def latest_approval(reviews: Sequence[Review], login: str) -> Review | None:
"""`login`'s most recent APPROVED review on the pull request, the one GitHub keeps in force after pushes."""
own: Final = tuple(review for review in reviews if review.user.login == login and review.state == "APPROVED")
return own[-1] if own else None
def approval_covers(
approval: Review, head: str, hits: Sequence[str], hits_at: Callable[[str], Sequence[str] | None]
) -> bool:
"""True when the approval is on the head, or the head adds no spend-log line the approved commit did not."""
if approval.commit_id is None:
return False
if approval.commit_id == head:
return True
reviewed: Final = hits_at(approval.commit_id)
return reviewed is not None and set(hits) <= set(reviewed)
def added_diff(base: str, head: str) -> str:
result: Final = subprocess.run(
("git", "diff", "--unified=0", "--no-color", f"{base}..{head}", "--", *WATCHED, MIGRATIONS_PREFIX),
@ -85,6 +161,27 @@ def added_diff(base: str, head: str) -> str:
return result.stdout
def file_at(revision: str, path: str) -> str:
"""Contents of `path` at `revision`, empty when the file does not exist there."""
result: Final = subprocess.run(("git", "show", f"{revision}:{path}"), check=False, capture_output=True, text=True)
return result.stdout if result.returncode == 0 else ""
def schema_blocks(base: str, head: str) -> Mapping[str, Block]:
return MappingProxyType(
{path: Block(base=model_block(file_at(base, path)), head=model_block(file_at(head, path))) for path in WATCHED}
)
def fetched(commit: str) -> bool:
result: Final = subprocess.run(("git", "fetch", "--no-tags", "--depth=1", "origin", commit), check=False)
return result.returncode == 0
def reviewed_hits(base: str, commit: str) -> tuple[str, ...] | None:
return spend_log_hits(added_diff(base, commit), schema_blocks(base, commit)) if fetched(commit) else None
class Arguments(BaseModel):
model_config = ConfigDict(frozen=True)
base: str
@ -100,18 +197,28 @@ def main(argv: Sequence[str]) -> int:
parser.add_argument("--approver", required=True, help="GitHub login whose approval unlocks the merge")
parser.add_argument("--reviews", required=True, help="JSON file with the pull request's reviews")
args: Final = Arguments.model_validate(vars(parser.parse_args(argv)))
hits: Final = spend_log_hits(added_diff(args.base, args.head))
hits: Final = spend_log_hits(added_diff(args.base, args.head), schema_blocks(args.base, args.head))
if not hits:
sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; nothing to review.\n")
sys.stdout.write(
"No change inside model LiteLLM_SpendLogs in schema.prisma and no added migration line names it; "
"nothing to review.\n"
)
return 0
if approved_by(_REVIEWS.validate_json(args.reviews.read_bytes()), args.approver):
reviews: Final = _REVIEWS.validate_json(args.reviews.read_bytes())
approval: Final = latest_approval(reviews, args.approver) if approved_by(reviews, args.approver) else None
if approval is not None and approval_covers(
approval, args.head, hits, lambda commit: reviewed_hits(args.base, commit)
):
sys.stdout.write(f"LiteLLM_SpendLogs change approved by {args.approver}: {', '.join(hits)}\n")
return 0
listed: Final = "".join(f" {hit}\n" for hit in hits)
sys.stdout.write(
reason: Final = (
f"This pull request changes LiteLLM_SpendLogs and needs an approving review from @{args.approver}:\n"
f"{listed}The check re-runs when a review is submitted.\n"
if approval is None
else f"LiteLLM_SpendLogs lines changed after @{args.approver} approved {approval.commit_id}, "
f"a new approving review is needed:\n"
)
sys.stdout.write(f"{reason}{listed}After approving, re-run this job.\n")
return 1

View file

@ -40,6 +40,8 @@ env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" }
Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway.
The `lite` CLI can set up this provider each time it starts Codex, with no `config.toml` edit. Export `LITELLM_PROXY_URL` and `LITELLM_PROXY_API_KEY`, then run `lite codex --subscription`. Run `lite config set codex_subscription true` to make it the default.
For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway).
## Native compatibility and Pi

View file

@ -32,6 +32,7 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.base_llm.guardrail_translation.utils import effective_skip_tool_message_for_guardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text
from litellm.types.guardrails import GuardrailEventHooks
@ -753,7 +754,12 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
_masked_entity_count.set({})
# Covers multimodal list content + Responses-API input.
inspected = walk_user_text(data, self.redact_text, self.redact_tool_payload) + self._redact_prompt(data)
inspected = walk_user_text(
data,
self.redact_text,
self.redact_tool_payload,
skip_tool_message=effective_skip_tool_message_for_guardrail(self),
) + self._redact_prompt(data)
if inspected == 0:
# Image-only, empty-text, and unsupported payloads inspected

View file

@ -5,9 +5,9 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
import asyncio
import json
import re
from collections.abc import Mapping, Sequence
from collections.abc import Callable, Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
@ -28,6 +28,7 @@ from litellm.constants import (
from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
from litellm.utils import get_utc_datetime
if TYPE_CHECKING:
from prisma import models as prisma_models
@ -42,11 +43,14 @@ if TYPE_CHECKING:
def _error_names_resource(error: Exception, resource_id: str) -> bool:
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
tokens: Final[list[str]] = re.findall(r"""[^\s"'`<>(]+""", str(error))
return any(token.rstrip(".,:;)") == resource_id for token in tokens)
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
@ -91,6 +95,10 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
return TeamRepository(prisma_client).table
def _unique_jobs(jobs: Sequence[ManagedObjectRow]) -> tuple[ManagedObjectRow, ...]:
return tuple({job.id: job for job in jobs}.values())
class CheckBatchCost:
def __init__(
self,
@ -98,6 +106,7 @@ class CheckBatchCost:
prisma_client: "PrismaClient",
llm_router: "Router",
track_unmanaged_batch_cost: bool = False,
clock: Callable[[], datetime] = get_utc_datetime,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -106,13 +115,14 @@ class CheckBatchCost:
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
self._clock: Final = clock
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost")
# Cached after the first poll cycle. Once we know the column is absent we skip
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
self.batch_processed_support_confirmed: bool = False
self._poll_cursor: tuple[datetime, str] | None = None
self._poll_before: datetime = datetime.now(timezone.utc)
self._poll_before: datetime = clock()
self._arrival_cursor: tuple[datetime, str] | None = None
self._arrival_before: datetime = self._poll_before
self._poll_arrivals: bool = False
@ -147,27 +157,6 @@ class CheckBatchCost:
return
self.batch_processed_support_confirmed = True
async def _cleanup_stale_managed_objects(self) -> None:
"""
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
in non-terminal states as 'stale_expired'. These will never complete and
should not be polled.
"""
cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
result: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
"created_at": {"lt": cutoff},
},
data={"status": "stale_expired"},
)
if result > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: marked {result} stale managed objects "
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await self._find_jobs(
@ -187,20 +176,43 @@ class CheckBatchCost:
)
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
arrivals: Final = self._poll_arrivals
page: Final = await self._find_job_page(where, arrivals=arrivals)
self._poll_arrivals = not arrivals if page else arrivals
if page:
return page
return await self._find_job_page(where, arrivals=not arrivals)
lead_arrivals: Final = self._poll_arrivals
self._poll_arrivals = not lead_arrivals
lead_share: Final = MAX_OBJECTS_PER_POLL_CYCLE - MAX_OBJECTS_PER_POLL_CYCLE // 2
lead: Final = await self._find_job_page(where, arrivals=lead_arrivals, take=lead_share)
trail: Final = await self._find_job_page(
where, arrivals=not lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead)
)
spill: Final = (
await self._find_job_page(
where, arrivals=lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead) - len(trail)
)
if len(lead) == lead_share
else ()
)
return _unique_jobs((*lead, *trail, *spill))
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
async def _find_job_page(
self, where: Mapping[str, object], *, arrivals: bool, take: int
) -> Sequence[ManagedObjectRow]:
if take == 0:
return ()
resumed_mid_lane: Final = (self._arrival_cursor if arrivals else self._poll_cursor) is not None
page: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take)
if len(page) == take or not resumed_mid_lane:
return page
wrapped: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take - len(page))
return _unique_jobs((*page, *wrapped))
async def _find_job_page_from_cursor(
self, where: Mapping[str, object], *, arrivals: bool, take: int
) -> Sequence[ManagedObjectRow]:
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
if cursor is None:
if arrivals:
self._arrival_before = datetime.now(timezone.utc)
self._arrival_before = self._clock()
else:
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
self._poll_before = self._clock() - RECENT_BATCH_POLL_WINDOW
before: Final = self._arrival_before if arrivals else self._poll_before
page: Final = await _managed_object_table(self.prisma_client).find_many(
where={
@ -220,10 +232,10 @@ class CheckBatchCost:
else {}
),
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
take=take,
order=[{"created_at": "asc"}, {"id": "asc"}],
)
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == take else None
if arrivals:
self._arrival_cursor = next_cursor
else:
@ -231,11 +243,47 @@ class CheckBatchCost:
return page
def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
return {
"id": job.id,
"updated_at": job.updated_at,
**({"batch_processed": False} if self._has_batch_processed_column else {}),
}
if self._has_batch_processed_column:
return {"id": job.id, "batch_processed": False}
return {"id": job.id, "updated_at": job.updated_at}
def _is_past_give_up_age(self, job: ManagedObjectRow) -> bool:
return job.created_at < self._clock() - GIVE_UP_POLLING_AFTER
async def _expire_if_old_and_unfinished(self, job: ManagedObjectRow) -> None:
if job.status in TERMINAL_MANAGED_OBJECT_STATUSES or not self._is_past_give_up_age(job):
return
await self._expire_unfinished(job)
@staticmethod
def _is_completed_with_output_id_pending(response: "LiteLLMBatch") -> bool:
from litellm.proxy.openai_files_endpoints.common_utils import (
_completed_batch_safe_to_retire,
)
return (
response.status in ("completed", "complete")
and response.output_file_id is None
and not _completed_batch_safe_to_retire(response)
)
async def _expire_unfinished(self, job: ManagedObjectRow) -> None:
try:
expired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
**self._unprocessed_snapshot(job),
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
},
data={"status": "stale_expired"},
)
except Exception as db_err:
verbose_proxy_logger.error(f"CheckBatchCost: failed to mark stale job {job.id} as stale_expired: {db_err}")
return
if expired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} is older than {GIVE_UP_POLLING_AFTER.days} days "
"and still not finished, so it is marked stale_expired and no longer polled"
)
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
"""
@ -356,15 +404,21 @@ class CheckBatchCost:
def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool:
"""A 404 naming the output file means there is nothing to fetch on this or any
later poll: providers like Vertex AI advertise an output path for every batch,
including terminal ones that never wrote it. Any other failure may be
transient, so it keeps retrying until the staleness sweep bounds it."""
including terminal ones that never wrote it. GCS names a gs:// object without its scheme,
and S3 names an s3:// object by its key alone. Any other failure keeps retrying."""
import openai
from litellm.exceptions import NotFoundError
from litellm.llms.base_llm.chat.transformation import BaseLLMException
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
is_not_found: Final = isinstance(error, (NotFoundError, openai.NotFoundError)) or (
isinstance(error, BaseLLMException) and error.status_code == 404
)
object_path: Final = output_file_id.split("://", 1)[-1]
names: Final = (output_file_id, object_path, object_path.split("/", 1)[-1])
return is_not_found and any(_error_names_resource(error, name) for name in names)
async def _finalize_unbilled_terminal_job(
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
@ -1094,13 +1148,6 @@ class CheckBatchCost:
processed_models: List[Tuple[Optional[str], Optional[str]]] = []
try:
await self._cleanup_stale_managed_objects()
except Exception as cleanup_err:
verbose_proxy_logger.warning(
f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}"
)
# Look for all batches that have not yet been processed by CheckBatchCost.
# self._has_batch_processed_column is cached after the first probe so that
# older schemas don't pay a guaranteed-failing primary query + warning on
@ -1142,6 +1189,8 @@ class CheckBatchCost:
if routing is None:
if self._has_unified_id_without_model(job):
await self._retire_job(job, "unified object id has no model id")
else:
await self._expire_if_old_and_unfinished(job)
continue
model_id, batch_id = routing
@ -1166,6 +1215,14 @@ class CheckBatchCost:
prom_logger.record_check_batch_cost_error("provider_retrieval_error")
if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
else:
await self._expire_if_old_and_unfinished(job)
continue
if self._is_completed_with_output_id_pending(response):
verbose_proxy_logger.info(
f"CheckBatchCost: batch {batch_id} is completed but its output file id "
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
)
continue
## RETRIEVE THE BATCH JOB OUTPUT FILE
@ -1217,17 +1274,9 @@ class CheckBatchCost:
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
from litellm.proxy.openai_files_endpoints.common_utils import (
_completed_batch_safe_to_retire,
)
if response.status in ("completed", "complete") and not _completed_batch_safe_to_retire(response):
verbose_proxy_logger.info(
f"CheckBatchCost: batch {batch_id} is completed but its output file id "
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
)
continue
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
else:
await self._expire_if_old_and_unfinished(job)
# Record polling run metrics (always, even if nothing was processed)
if prom_logger:

View file

@ -6,6 +6,7 @@ import base64
import json
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from datetime import datetime, timedelta
from itertools import chain
from types import MappingProxyType
from typing import (
@ -27,7 +28,7 @@ import httpx
from fastapi import HTTPException
from openai import APIConnectionError, AsyncOpenAI
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, Unpack
import litellm
@ -39,6 +40,7 @@ from litellm.constants import (
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_FILE_LIST_LIMIT,
PENDING_RESPONSE_STATUSES,
)
@ -51,7 +53,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_metadata,
)
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
@ -113,6 +115,7 @@ from litellm.types.utils import (
LLMResponseTypes,
SpecialEnums,
)
from litellm.utils import get_utc_datetime
class _ManagedFileRetrieve(Protocol):
@ -306,6 +309,7 @@ def _is_xai_batch_output(
class _ManagedFileRow(Protocol):
unified_file_id: str
file_object: OpenAIFileObject
model_mappings: Mapping[str, str]
flat_model_file_ids: Sequence[str]
storage_backend: Optional[str]
storage_url: Optional[str]
@ -359,6 +363,8 @@ class _ManagedResourceDatabase(Protocol):
@property
def litellm_managedfiletable(self) -> _ManagedFileTableActions: ...
async def query_raw(self, query: str, *args: object) -> object: ...
@property
def litellm_managedobjecttable(self) -> _ManagedObjectTableActions: ...
@ -382,10 +388,83 @@ class _RouterFileCallKwargs(TypedDict, total=False):
custom_llm_provider: ReadOnly[str | None]
def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
return prisma_client.db
def _managed_file_table(prisma_client: _ManagedResourcePrismaClient) -> _ManagedFileTableActions:
return prisma_client.db.litellm_managedfiletable
class _ListedFileRow(Protocol):
@property
def unified_file_id(self) -> str: ...
@property
def file_object(self) -> object: ...
@property
def model_mappings(self) -> Mapping[str, object]: ...
@property
def storage_backend(self) -> str | None: ...
class _ScopedFileRow(BaseModel):
model_config = ConfigDict(frozen=True)
unified_file_id: str
file_object: JsonValue = None
model_mappings: Mapping[str, JsonValue]
storage_backend: str | None = None
_SCOPED_FILE_ROWS: Final = TypeAdapter(tuple[_ScopedFileRow, ...])
def _owner_condition(owner_filter: Mapping[str, object]) -> tuple[str, tuple[str, ...]]:
match owner_filter:
case {"OR": [{"created_by": str() as user_id}, {"team_id": str() as team_id}]}:
return '("created_by" = $1 OR "team_id" = $2)', (user_id, team_id)
case {"created_by": str() as user_id}:
return '"created_by" = $1', (user_id,)
case {"team_id": str() as team_id}:
return '"team_id" = $1', (team_id,)
case _:
return ("FALSE", ()) if owner_filter else ("TRUE", ())
async def _scoped_file_rows(
database: _ManagedResourceDatabase,
owner_filter: Mapping[str, object],
model_mapping_keys: frozenset[str],
cursor_id: str | None,
direction: Literal["asc", "desc"],
take: int,
) -> tuple[_ScopedFileRow, ...]:
owner_condition, owner_params = _owner_condition(owner_filter)
keys_param: Final = len(owner_params) + 1
comparison: Final = "<" if direction == "desc" else ">"
cursor_condition: Final = (
f'AND ("created_at", "unified_file_id") {comparison} (SELECT "created_at", "unified_file_id" '
f'FROM "LiteLLM_ManagedFileTable" WHERE "unified_file_id" = ${keys_param + 1}) '
if cursor_id
else ""
)
order: Final = direction.upper()
rows: Final = await database.query_raw(
'SELECT "unified_file_id", "file_object", "model_mappings", "storage_backend" FROM "LiteLLM_ManagedFileTable" '
f'WHERE {owner_condition} AND ("model_mappings" ?| ${keys_param}::text[] OR "storage_backend" IS NOT NULL) '
f'{cursor_condition}ORDER BY "created_at" {order}, "unified_file_id" {order} '
f"LIMIT ${keys_param + (2 if cursor_id else 1)}",
*owner_params,
sorted(model_mapping_keys),
*((cursor_id,) if cursor_id else ()),
take,
)
return _SCOPED_FILE_ROWS.validate_python(rows)
def _iter_provider_file_id_pairs(
rows: Sequence[_ManagedFileRow],
requested_provider_file_ids: frozenset[str],
@ -507,10 +586,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prisma_client: PrismaClient,
*,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
clock: Callable[[], datetime] = get_utc_datetime,
):
self.internal_usage_cache = internal_usage_cache
self.prisma_client = prisma_client
self._sleep = sleep
self._clock = clock
@staticmethod
def _get_prometheus_logger():
@ -2260,14 +2341,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict: UserAPIKeyAuth,
limit: Optional[int] = None,
after: Optional[str] = None,
model_scope: Optional[FileListModelScope] = None,
before: Optional[str] = None,
**data: Dict,
) -> FileListPage:
"""List the managed files the caller owns, newest first.
Pagination is keyset based on ``unified_file_id`` so a key that owns
every file on the proxy still reads one bounded page at a time.
``purpose`` is applied after parsing, because the managed file table
keeps it inside the ``file_object`` blob instead of a column, and rows
``after`` pages toward older files and ``before`` toward newer ones,
each page still ordered newest first.
``model_scope`` is applied in the query, so only rows whose
``model_mappings`` name one of its deployment ids or model names, and
storage-backed rows, leave the database; a storage-backed row then
stays only when its model name is one the router sends to the listed
provider. ``purpose`` is applied after parsing, because the managed
file table keeps it inside the ``file_object`` blob instead of a
column. Rows
whose blob will not parse drop out there too, so a chunk of rows can
yield fewer matches than the page holds. Successive chunks are read
until the page is full or the caller's rows run out, which keeps
@ -2276,9 +2366,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
scan has to continue past it, the chunk widens to
``FILE_LIST_CONTINUATION_CHUNK_SIZE``, so the walk costs one query per
that many rows instead of one per page. That bound is per query, not
per request: the work is still linear in the rows the caller owns, and
a filter matching nothing reads every one of them, with no index
covering either the owner filter or the sort.
per request: a ``purpose`` filter matching nothing still reads every
row that passes the query.
"""
validate_file_list_limit(limit)
validate_file_list_purpose(purpose)
@ -2287,36 +2376,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if owner_filter is None:
return FileListPage.model_validate(build_list_page([]))
if after:
cursor_row = await _managed_file_table(self.prisma_client).find_first(
where={**owner_filter, "unified_file_id": after}
if after and before:
raise ProxyException(
message="Pass either 'after' (or 'after_id') or 'before_id' when listing files, not both.",
type="invalid_request_error",
param="before_id",
code=400,
openai_code="invalid_value",
)
cursor_param: Final = "before_id" if before else "after"
cursor: Final = before or after
database: Final = _managed_database(self.prisma_client)
if cursor:
cursor_row = await database.litellm_managedfiletable.find_first(
where={**owner_filter, "unified_file_id": cursor}
)
if cursor_row is None:
raise ProxyException(
message=f"Invalid 'after' cursor: no file found with id '{after}'.",
message=f"Invalid '{cursor_param}' cursor: no file found with id '{cursor}'.",
type="invalid_request_error",
param="after",
param=cursor_param,
code=400,
openai_code="invalid_value",
)
if model_scope is not None and not model_scope.model_mapping_keys:
return FileListPage.model_validate(build_list_page([]))
page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT)
matches: Final[List[OpenAIFileObject]] = []
cursor_id = after
cursor_id = cursor
chunk_size = page_size + 1
direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc"
while len(matches) <= page_size:
cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {}
chunk = await _managed_file_table(self.prisma_client).find_many(
where=owner_filter,
take=chunk_size,
order=[{"created_at": "desc"}, {"unified_file_id": "desc"}],
**cursor_args,
chunk: Sequence[_ListedFileRow] = (
await _scoped_file_rows(
database, owner_filter, model_scope.model_mapping_keys, cursor_id, direction, chunk_size
)
if model_scope is not None
else await database.litellm_managedfiletable.find_many(
where=owner_filter,
take=chunk_size,
order=[{"created_at": direction}, {"unified_file_id": direction}],
**cursor_args,
)
)
matches.extend(
_public_file_object(parsed_file_object, row.unified_file_id)
for row in chunk
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
if (
model_scope is None or model_scope.holds(row.model_mappings, stored=row.storage_backend is not None)
)
and (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
and (purpose is None or parsed_file_object.purpose == purpose)
)
if len(chunk) < chunk_size:
@ -2324,7 +2437,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
cursor_id = chunk[-1].unified_file_id
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
return FileListPage.model_validate(build_list_page(matches[:page_size], has_more=len(matches) > page_size))
page: Final = matches[:page_size]
return FileListPage.model_validate(
build_list_page(page[::-1] if before else page, has_more=len(matches) > page_size)
)
def _is_batch_polling_enabled(self) -> bool:
"""
@ -2387,12 +2503,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"not_in": ["failed", "expired", "cancelled"]},
"status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]},
},
take=MAX_MATCHES_TO_RETURN,
order={"created_at": "desc"},
)
input_file_cutoff: Final = self._clock() - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
referencing_batches: Final[list[dict[str, object]]] = []
for batch in batches:
try:
@ -2405,7 +2522,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Extract file IDs from batch
# Batches typically reference the unified file ID in input_file_id
# Output and error files are generated by the provider
input_file_id = batch_data.get("input_file_id")
input_file_id = batch_data.get("input_file_id") if batch.created_at >= input_file_cutoff else None
output_file_id = batch_data.get("output_file_id")
error_file_id = batch_data.get("error_file_id")

View file

@ -15,7 +15,11 @@ from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import TypeAdapter
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
describe_missing_quota,
missing_project_model_quota,
)
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -27,6 +31,15 @@ from litellm.proxy.management.teams.dependencies import get_team_access
from litellm.proxy.management.users.service import org_admin_org_ids
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
from litellm.proxy.management_helpers.model_quota_policy import (
ENFORCE_PROJECT_MODEL_QUOTA_SETTING,
REQUIRE_BOTH_LIMITS,
ModelQuotaPolicy,
RouteOf,
no_route,
parse_model_quota_setting,
router_model_routes,
)
from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
@ -276,26 +289,6 @@ def _check_team_project_limits(
)
def _project_models_missing_positive_quota(
models: list[str] | None,
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
) -> list[str]:
"""Return the models that lack a positive `rpm` AND `tpm` quota.
A valid quota is a positive integer; null, zero, and negative are rejected
because downstream rate limiters treat a non-positive limit as immediately
exhausted (every request blocked).
"""
def _is_positive(value: object) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and value > 0
rpm = rpm_limits or {}
tpm = tpm_limits or {}
return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))]
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
@ -322,64 +315,79 @@ def _raise_on_project_models_expanding_at_request_time(
raise HTTPException(
status_code=400,
detail={
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead."
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' is enabled. List concrete model names instead."
},
)
def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None:
try:
return parse_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING))
except ValidationError as e:
raise HTTPException(
status_code=500,
detail={"error": f"general_settings.{ENFORCE_PROJECT_MODEL_QUOTA_SETTING} is invalid: {e}"},
) from e
def _raise_on_missing_project_model_quota(
data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset()
data: NewProjectRequest | UpdateProjectRequest,
access_group_names: frozenset[str] = frozenset(),
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
route_of: RouteOf = no_route,
) -> None:
"""Require a positive `rpm`/`tpm` quota for every model on project CREATE.
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE.
`model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request
model's `set_model_info` validator, so they are read from there.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
_raise_on_project_models_expanding_at_request_time(data.models, access_group_names)
metadata = data.metadata or {}
missing = _project_models_missing_positive_quota(
data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit")
missing: Final = missing_project_model_quota(
data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, route_of
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
"error": f"models added to project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit."
},
)
def _raise_on_missing_project_model_quota_on_update(
data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset()
data: UpdateProjectRequest,
existing_project: object,
access_group_names: frozenset[str] = frozenset(),
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
route_of: RouteOf = no_route,
) -> None:
"""Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE.
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE.
`/project/update` replaces `models` and `metadata` when they are provided, so the
check runs on what the project WILL look like: a partial update that doesn't touch
models/quota keeps the existing values, while one that adds a model or clears a
model's quota must leave every resulting model with a positive limit.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
model's quota must leave every resulting model with the limits it needs.
"""
resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or [])
resulting_metadata = (
data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {})
)
_raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names)
missing = _project_models_missing_positive_quota(
resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit")
missing: Final = missing_project_model_quota(
resulting_models,
resulting_metadata.get("model_rpm_limit"),
resulting_metadata.get("model_tpm_limit"),
policy,
route_of,
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
"error": f"models would be left on the project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit."
},
)
@ -530,6 +538,7 @@ async def new_project(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -578,10 +587,6 @@ async def new_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model added to the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router))
# Check if user has permission to create projects for this team
# only team admins can create projects for their team
has_permission = await _check_user_permission_for_project(
@ -600,6 +605,15 @@ async def new_project(
},
)
quota_policy: Final = _project_model_quota_policy(general_settings_view())
if quota_policy is not None:
_raise_on_missing_project_model_quota(
data,
_router_access_group_names(llm_router),
quota_policy,
router_model_routes(llm_router, data.team_id),
)
# Generate project_id if not provided
if data.project_id is None:
data.project_id = str(uuid.uuid4())
@ -721,6 +735,7 @@ async def update_project(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -828,10 +843,14 @@ async def update_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model the update would leave on the project.
if general_settings.get("enforce_project_model_quota", False):
quota_policy: Final = _project_model_quota_policy(general_settings_view())
if quota_policy is not None:
_raise_on_missing_project_model_quota_on_update(
data, existing_project, _router_access_group_names(llm_router)
data,
existing_project,
_router_access_group_names(llm_router),
quota_policy,
router_model_routes(llm_router, target_team_id),
)
# Prepare update data

View file

@ -0,0 +1,49 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from litellm.proxy.management_helpers.model_quota_policy import (
ModelQuotaPolicy,
QuotaLimit,
RouteOf,
required_limits,
unmet_limits,
)
_NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
class MissingModelQuota:
model: str
limits: tuple[QuotaLimit, ...]
def _limits_set_for(
model: str, rpm_set: Mapping[str, object], tpm_set: Mapping[str, object]
) -> Mapping[QuotaLimit, object]:
return MappingProxyType({"rpm": rpm_set.get(model), "tpm": tpm_set.get(model)})
def missing_project_model_quota(
models: Sequence[str],
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
policy: ModelQuotaPolicy,
route_of: RouteOf,
) -> tuple[MissingModelQuota, ...]:
rpm_set: Final = rpm_limits or _NO_LIMITS_SET
tpm_set: Final = tpm_limits or _NO_LIMITS_SET
gaps: Final = tuple(
MissingModelQuota(
model,
unmet_limits(required_limits(policy, model, route_of(model)), _limits_set_for(model, rpm_set, tpm_set)),
)
for model in models
)
return tuple(gap for gap in gaps if gap.limits)
def describe_missing_quota(missing: Sequence[MissingModelQuota]) -> str:
return ", ".join(f"{gap.model} ({' and '.join(f'model_{limit}_limit' for limit in gap.limits)})" for gap in missing)

View file

@ -153,7 +153,7 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
python -c "import litellm, litellm.rust_bridge._native; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__"
USER nonroot

View file

@ -343,11 +343,16 @@ def btree_max_entry_bytes(block_size: int) -> int:
def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed":
"""The bytes a value takes in the index entry: its length, or its stored size when
Postgres compressed it, since `index_form_tuple` keeps the compressed bytes."""
from psycopg import sql
if column.type_name in _OCTET_LENGTH_TYPES:
return sql.SQL("octet_length({})").format(identifier)
return sql.SQL("octet_length({}::text)").format(identifier)
length: Final = (
sql.SQL("octet_length({})").format(identifier)
if column.type_name in _OCTET_LENGTH_TYPES
else sql.SQL("octet_length({}::text)").format(identifier)
)
return sql.SQL("least({}, pg_column_size({}))").format(length, identifier)
def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed":

View file

@ -0,0 +1,2 @@
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_ManagedFileTable_created_by_created_at_idx"
ON "LiteLLM_ManagedFileTable" ("created_by", "created_at" DESC);

View file

@ -0,0 +1,47 @@
SET LOCAL lock_timeout = '10s';
CREATE TYPE "AccountingEventKind" AS ENUM ('inbound_usage', 'outbound_usage', 'cost');
CREATE TYPE "AccountingDeliveryState" AS ENUM ('pending', 'applied', 'parked');
CREATE TYPE "AccountingDeliveryTarget" AS ENUM ('LiteLLM_VerificationToken', 'LiteLLM_UserTable', 'LiteLLM_TeamTable', 'LiteLLM_TeamMembership', 'LiteLLM_OrganizationTable', 'LiteLLM_OrganizationMembership', 'LiteLLM_EndUserTable', 'LiteLLM_AgentsTable', 'LiteLLM_TagTable', 'LiteLLM_ProjectTable', 'LiteLLM_ModelAccessGroupBudgetTable', 'LiteLLM_BudgetTable', 'LiteLLM_BudgetWindowSpend', 'LiteLLM_DailyUserSpend', 'LiteLLM_DailyTeamSpend', 'LiteLLM_DailyOrganizationSpend', 'LiteLLM_DailyEndUserSpend', 'LiteLLM_DailyAgentSpend', 'LiteLLM_DailyTagSpend', 'LiteLLM_DailyModelUsage', 'LiteLLM_DailyToolSpend', 'LiteLLM_AutoRouterDailySpend');
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingEvent" (
"event_id" TEXT NOT NULL,
"kind" "AccountingEventKind" NOT NULL,
"source_event_id" TEXT,
"schema_version" INTEGER NOT NULL,
"payload_hash" TEXT NOT NULL,
"occurred_at" TIMESTAMP(3) NOT NULL,
"recorded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"payload" JSONB NOT NULL,
CONSTRAINT "LiteLLM_AccountingEvent_schema_version_check" CHECK ("schema_version" > 0),
CONSTRAINT "LiteLLM_AccountingEvent_pkey" PRIMARY KEY ("event_id")
);
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingDelivery" (
"event_id" TEXT NOT NULL,
"target" "AccountingDeliveryTarget" NOT NULL,
"version" INTEGER NOT NULL,
"state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending',
"attempts" INTEGER NOT NULL DEFAULT 0,
"next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"applied_at" TIMESTAMP(3),
"last_error" TEXT,
CONSTRAINT "LiteLLM_AccountingDelivery_version_check" CHECK ("version" > 0),
CONSTRAINT "LiteLLM_AccountingDelivery_attempts_check" CHECK ("attempts" >= 0),
CONSTRAINT "LiteLLM_AccountingDelivery_pkey" PRIMARY KEY ("event_id","target","version")
);
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_recorded_at_idx" ON "LiteLLM_AccountingEvent"("recorded_at");
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_kind_source_event_id_idx" ON "LiteLLM_AccountingEvent"("kind", "source_event_id");
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingDelivery_target_state_next_attempt_at_idx" ON "LiteLLM_AccountingDelivery"("target", "state", "next_attempt_at");

View file

@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
@@index([unified_file_id])
@@index([flat_model_file_ids], type: Gin)
@@index([team_id, created_at(sort: Desc)])
@@index([created_by, created_at(sort: Desc)])
}
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the
@ -2060,3 +2061,73 @@ model LiteLLM_TelemetryReport {
@@index([window_end])
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
parked
}
enum AccountingDeliveryTarget {
LiteLLM_VerificationToken
LiteLLM_UserTable
LiteLLM_TeamTable
LiteLLM_TeamMembership
LiteLLM_OrganizationTable
LiteLLM_OrganizationMembership
LiteLLM_EndUserTable
LiteLLM_AgentsTable
LiteLLM_TagTable
LiteLLM_ProjectTable
LiteLLM_ModelAccessGroupBudgetTable
LiteLLM_BudgetTable
LiteLLM_BudgetWindowSpend
LiteLLM_DailyUserSpend
LiteLLM_DailyTeamSpend
LiteLLM_DailyOrganizationSpend
LiteLLM_DailyEndUserSpend
LiteLLM_DailyAgentSpend
LiteLLM_DailyTagSpend
LiteLLM_DailyModelUsage
LiteLLM_DailyToolSpend
LiteLLM_AutoRouterDailySpend
}
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {
event_id String
target AccountingDeliveryTarget
version Int
state AccountingDeliveryState @default(pending)
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
@@id([event_id, target, version])
@@index([target, state, next_attempt_at])
}

View file

@ -0,0 +1,38 @@
---
name: rust-error
description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings
metadata:
internal: true
---
# Rust errors
`crates/http/src/error.rs` is the reference implementation. Read the reference for the part you are touching before editing, each one carries the reasoning and worked examples behind its rules
- [shape](references/shape.md)
- One top-level `Error` per crate in `src/error.rs`. A public enum by default, an opaque struct over a private `Repr` enum when the variants must stay private. No `#[non_exhaustive]`
- A variant is a failure mode a caller handles differently. Provider names and guidance are data, never variants
- [messages and causes](references/messages-and-causes.md)
- Every message is written once, in `#[error(...)]`, and describes its own layer only. Call sites pass typed facts, never a string
- The cause is the `source()`, never part of the message. Keep it typed, never stringify or drop it
- `reqwest::Error` enters our types only through `litellm_http::Error`, which strips its URL
- Anything shown to someone else is rendered with `litellm_error::Report`, never `to_string()`
- [conversions](references/conversions.md)
- One `From` per source and target pair, only when the conversion needs no context. It classifies and sanitizes once
- Nest a lower crate's error in one variant. Translate when the meaning changes, never mirror its variants
- [classification](references/classification.md)
- Callers ask predicates named after their decision, they do not match nested variants. Predicates match exhaustively
- Classification states facts. Retry and fallback policy lives with the caller that knows idempotency
- [boundaries](references/boundaries.md)
- An error becomes text for someone else only in one mapping function per error type, in the gateway crates, the Python bridge and host-python. The inference gateway and the bridge decide by the same predicates
- A failure before the provider was called is a decline. Both ask `declined()` first and hand it back to the host unchanged
- [gateway](references/gateway.md)
- One `Error` with `IntoResponse` per gateway crate that serves a router. Every failure flows through the same `From` chain, is classified once in `status()` and rendered by surface and framing
- [pyo3](references/pyo3.md)
- Two `create_exception!` classes with the payload in `args`, `ValueError` for the caller's fault and `RuntimeError` otherwise, decoded in `failures.py`
- [testing](references/testing.md)
- Test predicates, chain shape and redaction on errors built the way production builds them
- [recipes](references/recipes.md)
- Common tasks walked through end to end
The workspace uses thiserror 2 for every error type, and no library crate uses `anyhow` or `Box<dyn Error>` in a public signature. The [derive docs](https://docs.rs/thiserror/latest/thiserror/) are the authority on macro behavior, and the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html) on what `Display` and `source()` mean

View file

@ -0,0 +1,45 @@
# Boundaries
A boundary is where an error becomes something another program reads: an HTTP response, a Python exception, a stored or logged message. Each boundary has one mapping function per error type it receives, and nothing outside those functions formats an error for someone else
| Boundary | Mapping | Reference |
|---|---|---|
| Inference HTTP response | `Error::status()` and its renderers in `crates/gateway-inference/src/error.rs` | [gateway](gateway.md) |
| Other gateway surfaces | the `IntoResponse` of each `crates/gateway-*` crate's `Error` | [gateway](gateway.md) |
| Python bridge, inference and OCR routes | `route_failure` in `crates/python-bridge/src/errors.rs` and in `crates/python-bridge/src/routes/ocr/errors.rs` | [pyo3](pyo3.md) |
| Python bridge, everything else | one `From<Error> for PyErr` or mapping function per crate error, all in `crates/python-bridge` | [pyo3](pyo3.md) |
| Python host runtime | `From<Error> for PyErr` in `crates/host-python/src/error.rs` | [pyo3](pyo3.md) |
## Asking the nested error
The inference gateway and the Python bridge's route mapping both receive a nested route error and decide by the same predicates in the same order, so a new failure lands correctly on both once its predicates answer right
| Predicate | Gateway | Bridge |
|---|---|---|
| `declined()` | 501 | `RustBridgeDeclined`, without the route's public mapping |
| `upstream_response()` | the provider's status and body | `RustUpstreamError` |
| `is_missing_credential()` | 401 | no class of its own, a `RuntimeError` because `is_request()` answers false |
| `is_request()` | 400 | `ValueError` |
| none of these | 500 | `RuntimeError` |
If you find yourself adding an arm for your variant at a boundary, check whether a predicate is missing first
The two boundaries are not identical. The gateway has variants of its own for body parsing, an unknown model and an unsupported path, and the OCR bridge has Python parity arms such as `FileNotFoundError` for a missing file. Those arms match the boundary's own variants or the top level of a public enum, never a nested error's variants. A missing API key is 401 at the gateway and `RuntimeError` at both bridges, pinned by tests to what Python's own path does
## Declines
A failure before anything reached the provider is a decline, and `declined()` on the route error names its reason. It means another host may serve the same request as it was received, so each boundary hands it to the host unchanged: the gateway answers 501 and the Python bridge raises `RustBridgeDeclined`, never a public exception or a `ValueError`. See [gateway](gateway.md) and [pyo3](pyo3.md) for each host's side
Anything after the send is terminal, never a decline. `declined()` answers `None` for a response fault and for a post-call hook failure even when the wrapped error was a decline, and every route that transforms the provider's answer collapses its errors with `Error::after_send`, which keeps a provider answer and a response fault as they are and nests anything else as the source of an invalid response, so a decline can never describe a call that was billed. Ask `declined()`, do not match the variant behind it, because that after-send cut lives in the predicate
## What reaches the client
A boundary renders anything that is not a provider answer with `litellm_error::Report`, so the reader sees every layer of the chain once. That is why messages are written for a reader, each layer leaves its cause to `source()`, and reqwest URLs are stripped on entry. A boundary that must not show the chain, such as the bridge's TLS settings error that names the setting but not the file, renders fixed text of its own and says so in its test
Provider bodies are passed through intact for gateway and Python error mapping. `litellm_http::Error::upstream(status, body)` retains the complete body, but its `Display`, `Debug` and `Report` diagnostics show only the first 256 characters and mark the cut. Do not log the raw body or `RustUpstreamError` args. OCR's `Provider` error also keeps the whole body for Python's `get_error_class`, bounded by the connection's response size limit
Never format `Debug` of an error into a response or a Python exception. A type whose sources may hold secrets derives `veil::Redact`, so its `Debug` is safe to log
## Logging
Propagate with context, and let the terminal handler own the one log or trace event for a failed call. Do not log at every `?`, and do not log an error you also return. Diagnostic tracing excludes raw error strings and bodies, see [`rust-tracing`](../../rust-tracing/SKILL.md)

View file

@ -0,0 +1,66 @@
# Classification
## Ask, don't match
Downstream crates ask a lower error questions instead of matching its nested variants. `litellm_http::Error` answers these
| Predicate | Caller decides |
|---|---|
| `dispatched()` | whether retrying elsewhere could repeat a billed call |
| `is_request()` | the caller's input is at fault, so 400 and no reroute |
| `is_response()` | the provider answered with something we cannot use |
| `upstream_response()` | the provider's own status and body to pass through |
| `status()` | an HTTP status to report |
| `is_connect()` | the connection was never made, for a host's own connect-error mapping |
| `tls_source()` | which TLS setting to blame in a config error |
| `is_computed_header()` | the caller forwarded a header the signer owns, which `litellm_inference::Error` turns into a decline |
Upper types add the questions their callers ask. `litellm_inference::Error` and the OCR error answer `is_missing_credential()`, and `litellm_llms::Error` and `litellm_inference::Error` answer `declined()`, which the boundaries ask first instead of reaching into `Auth` or `Declined`, see [boundaries](boundaries.md)
```rust
// reaches through two crates' variants and breaks when either changes
Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => StatusCode::UNAUTHORIZED,
// asks the question it cares about
Self::Route(error) if error.is_missing_credential() => StatusCode::UNAUTHORIZED,
```
A predicate is an exhaustive `match` with no `_` arm, never `matches!`, so a new variant does not compile until it has answered every question on purpose
```rust
pub fn is_request(&self) -> bool {
match &*self.0 {
Repr::Request(_) | Repr::BlockedUrl | Repr::DownloadTooLarge => true,
Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => false,
// ... every other variant, listed by name
}
}
```
Add a predicate when a real caller needs a new distinction, and name it after the fact the caller needs, in the caller's words, not after the variant's spelling. An upper type's predicates call through to the nested error, as in `Self::Http(error) => error.is_request()` and `Self::Auth(error) => error.is_missing_credential()`
A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `Unsupported` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and the old `Unsupported` route variant
## Facts, not policy
A predicate reports what happened. The decision to retry, reroute or fall back depends on things the error cannot know, such as whether the request is idempotent, whether its body can be replayed, the deadline and the attempt budget. That decision lives with the caller. `dispatched()` exists for the router's reroute decision, so every transport failure answers it even before that caller is wired
A timeout is the main trap. A request that timed out may already have reached the provider and been billed, so `Timeout` answers `dispatched() == true` and is not safe to retry as if it were a connect failure. Only a failure that provably happened before any byte reached the provider answers `false`: connecting, building or signing the request, and media downloads, which fetch the caller's document before the provider call is made
```rust
pub fn dispatched(&self) -> bool {
match &*self.0 {
Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => true,
Repr::Connect(_) | Repr::RequestBody(_) | Repr::Download(_) | Repr::BlockedUrl => false,
// ... every other variant, listed by name
}
}
```
When the transport cannot tell whether a request reached the provider, it must answer `true`. Wrongly answering `false` risks a duplicate billed call, wrongly answering `true` only costs a reroute
Keep one predicate per question. `status()` currently mixes the provider's status, the media download's status and a synthesized 408 for timeouts, so a caller that wants "did the provider answer" must use `upstream_response()` instead. Do not add more meanings to an existing predicate, add a new one
## Equality
`Arc<Repr>` errors compare with `Arc::ptr_eq`, so a clone equals its original and two separately built errors never compare equal. That exists only so a wrapping enum such as `litellm_inference::Error` can keep its `PartialEq` derive. Do not give a new type pointer equality for any other reason, and do not rely on it in tests. Assert predicates instead, see [testing](testing.md)

View file

@ -0,0 +1,92 @@
# Conversions
## One From per source and target pair
A `From<Foreign> for Ours` exists when the conversion is obvious, infallible and needs no context from the call site. It does all classification and sanitizing for that pair, so every call site is `?` or `.map_err(Ours::from)`
```rust
impl From<reqwest::Error> for Error {
fn from(error: reqwest::Error) -> Self {
let error = error.without_url();
if error.is_timeout() {
Repr::Timeout(error).into()
} else if error.is_connect() || error.is_builder() {
Repr::Connect(error).into()
} else {
Repr::Network(error).into()
}
}
}
```
A call site never re-derives that classification
```rust
// every route used to repeat this, and each drifted a little
.map_err(|err| {
if err.is_connect() || err.is_builder() {
Error::Transport(Connect(err.to_string()))
} else {
Error::Transport(Network(err.to_string()))
}
})?
// now
.map_err(litellm_http::Error::from)?
```
The same foreign type may convert into several of our types, one `From` each. What is not allowed is two places converting the same pair differently
## When the call site knows something the From cannot
One `io::Error` can mean "reading the CA bundle failed" or "resolving the media host failed". A global `From<io::Error>` cannot know which, so there is none. Name the operation at the call site instead
```rust
.map_err(Repr::Resolve)?
.map_err(|source| Repr::Read { path, tls_source, source })?
```
Do not drop meaningful context just to use `?`
## Reclassify once, in the shared From
When a lower error needs a different meaning for every caller, do it in the upper type's `From`, not per route
```rust
impl From<litellm_http::Error> for Error {
fn from(error: litellm_http::Error) -> Self {
if error.is_computed_header() {
return Self::Declined(Decline::new("request forwards a header AWS SigV4 computes"));
}
Self::Http(error)
}
}
```
All four inference routes classify a header the signer computes the same way, because none of them does it themselves
## Nest, translate or add context, never mirror
| The upper layer | Do | Example |
|---|---|---|
| has nothing to add | nest in one transparent variant and delegate predicates | `Self::Http(error) => error.is_request()` |
| changes what the failure means | translate to an upper variant, keep the cause attached when it helps | computed header -> `Declined` |
| knows something the lower one does not | a variant with the extra fields plus the cause as its source | `Read { path, tls_source, source }` |
Never re-declare a lower crate's variants in the upper type with an upper-layer name and a 1:1 mapping function. OCR once did that for every media download failure, `MediaError::BlockedUrl => Error::BlockedDocumentUrl` and so on, which doubled every change and let the two lists drift.
## Third-party and host faults
Converting a fault from a crate we do not own into a string variant is still flattening. Prefer a variant that carries the fault, and when an existing `ErrorDetail` constructor fits, use it
```rust
// flattens
.map_err(|error| Error::InvalidRequest(error.to_string().into()))
// keeps the cause typed
.map_err(|error| Error::InvalidRequest(ErrorDetail::invalid("header", error)))
```
## Through a foreign error type
A trait we implement for another crate sometimes forces its error type on us, such as `io::Error` from a body stream or `sqlx::Error` from a driver. Wrap ours as that type's source (`io::Error::other(error)`, `azure_core::Error::new(kind, error)`) and recover it with `downcast_ref` on the other side, never `with_message(error.to_string())` or `from_msg`. The error must come back out as the same value it went in as, so one failure classifies the same way on every path

View file

@ -0,0 +1,56 @@
# Gateway responses
The gateway is where a Rust error becomes an HTTP response. This page is about the path an error takes to get there and the one place it is classified. The gateway crates' own conventions, such as which envelope each API surface uses, which extractors exist and how the router is composed, live in the [`rust-gateway`](../../rust-gateway/SKILL.md) skill
## The flow
```text
reqwest / provider failure
-> litellm_http::Error From<reqwest::Error>, strips the URL, classifies once
-> litellm_inference::Error::Http(_) From, one transparent variant, predicates delegate
-> gateway_inference::Error::Route(_) From, one variant per nested crate
-> handler returns Result<_, Error> every `?` ends here
-> IntoResponse status() picks the code, a renderer builds the body
```
Each hop is one `From` that nests the lower error in one variant, see [conversions](conversions.md). No hop re-derives what the lower crate already decided, and nothing between the provider and `IntoResponse` formats an error for the client. That is what makes "one `Error` with `IntoResponse` per gateway crate that serves a router" work: a handler can `?` any failure from any depth and the response comes out right
## Classify once
`status()` is pure and runs once per error. Variants the gateway declares itself, such as a bad body, an unknown model or an oversized body, are matched at the top of its own enum. A nested route error goes to `route_status` and an OCR error to `ocr_status`, which ask the predicates [boundaries](boundaries.md) lists
| Answer | Status |
|---|---|
| `declined()` | 501, for a route error, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider |
| `upstream_response()` | the provider's own status, 502 when it is not a valid code |
| `is_missing_credential()` | 401 |
| `is_request()` | 400, for a route error |
| `http_status_code()` | the transport's status, for an OCR error, such as 408 for a timeout |
| otherwise | 500 |
Do not reach into a nested enum's variants from the gateway. When a new nested failure needs a different status, add a predicate on the nested error, see [classification](classification.md)
## Render by surface and framing
Rendering is separate from classification. The surface picks the envelope and the moment picks the framing, a JSON body before the headers went out and an `event: error` SSE frame after. Both framings of one surface render the same status and `Report` message, so a streaming failure and a JSON failure of the same error agree on everything but framing
| Surface | JSON | SSE |
|---|---|---|
| Anthropic, `/v1/messages` | `body(request_id)` | `sse_frame()` |
| OpenAI, chat, responses and the rest | `openai_response()` | `responses_sse_frame()`, the Responses API `error` event |
The message of anything that is not a provider answer is the `litellm_error::Report` of the error, and in the OpenAI and Anthropic JSON envelopes the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Responses `error` event has the fixed type `error` and the status as its `code`. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else. The OpenAI envelope always wraps, because that is what the Python proxy returns there
## Every way in uses the same Error
A request can fail before the handler runs, in the handler, or after the headers went out. All three end in the same `Error`, so they get the same classification and envelope
- An extractor sets `type Rejection = Error`, so a bad body is the surface's envelope rather than axum's plain text
- A path the gateway knows but does not serve returns `Error::Unsupported(path)` from one handler, so the client gets a 501 in the surface's envelope rather than a bare 404
- A stream that fails after its headers went out hands the surface's SSE frame of `Error::from(error)` to `litellm_host_http::Sse::new`. `litellm_host_http::Error<E>` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the `From` maps the second to `Internal`
The axum hooks behind each line are in [`rust-gateway`](../../rust-gateway/SKILL.md)
## Testing
`status()` and `body()` are pure, so unit tests call them on errors built the way production builds them, one `#[case]` per fault, as `status_follows_who_is_at_fault`, `an_anthropic_provider_body_passes_through` and `other_failures_get_the_anthropic_envelope_with_their_message` in `crates/gateway-inference/src/error.rs` do. Assert the status and the fields the surface promises, not the whole body, except where passing the body through untouched is the contract. Driving a whole router belongs to the gateway crate's own tests, see [`rust-gateway`](../../rust-gateway/references/testing.md)

View file

@ -0,0 +1,112 @@
# Messages and causes
## Messages belong to the variant
Every message template is written once, in the variant's `#[error(...)]`. A caller passes the typed facts the template needs and nothing else
```rust
// none of these may appear at a call site
Error::Network("WebSocket handshake returned a wrong Sec-WebSocket-Accept".into())
Error::Request(format!("bad header {name}"))
Error::Signature(error.to_string())
// the call site names what happened, the variant owns the words
Repr::Handshake(HandshakeFault::Accept).into()
Error::signature(error)
Error::NonAsciiHeader { name }
```
When a check fails for a distinct reason, that reason is a variant or a field, never a string inside a generic variant. A `Message(String)` variant, or a `From<String>` / `From<&str>` into an error, is the hole every stringly call site goes through, so neither exists
Format arguments can call functions and reach into fields, so a hand-written `Display` is rarely needed
```rust
#[error("could not read {}", path.display())]
#[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))]
```
Keep messages short and lowercase, with no trailing period, unless Python parity pins the exact text. A reporter appends `: {cause}` after them. Never format `Debug` (`{:?}`) of an error, or of a type with no `Display`, into a message
## The cause is the source, never the message
This is the rule from the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html): an error's `Display` describes that layer only, and the error underneath is returned by `source()`. thiserror wires `source()` from a field marked `#[source]`, a `#[from]` field, or a field named `source`
```rust
// prints the cause twice under any reporter
#[error("could not read {}: {source}", path.display())]
Read { path: PathBuf, source: io::Error },
#[error("invalid JSON: {0}")]
Json(#[from] serde_json::Error),
// this layer's words, the cause is the source
#[error("could not read {}", path.display())]
Read { path: PathBuf, source: io::Error },
#[error("invalid JSON")]
Json(#[from] serde_json::Error),
```
| Field | thiserror does | Use for |
|---|---|---|
| `source: E` or `#[source] cause: E` | exposes it as `source()` | a cause with context fields next to it |
| `#[from] E` | implies `#[source]` and adds `From` | context-free wrapping, see [conversions](conversions.md) |
| `#[error(transparent)]` | forwards both `Display` and `source()` to the inner value | nesting another of our errors that adds no text |
| no attribute, any other name | nothing, the field is plain data | facts the message renders, such as a path or a status |
`#[error(transparent)]` is an invisible layer, not an extra chain node, so a transparent wrapper's `source()` is the inner value's `source()`
A field that holds an error but is neither rendered nor a source hides that error from every reader. Every error-typed field is a source. A fieldless reason enum, such as `HandshakeFault` or `PagesFault`, is data and not a cause, so the variant renders it in its message
`source()` is one chain, so an error with several causes, such as `CredentialChain` trying each credential in turn, makes the most relevant one its source and renders the others in its message, each with `Report`
## Keep causes typed
Store the lower-level error itself, never its string, and never drop it with `map_err(|_| ...)`
| Flattening or dropping loses | Because |
|---|---|
| the chain | a `String` has no `source()` |
| downcasting | the type is gone |
| redaction | `reqwest::Error`'s `Display` includes the request URL, which can carry credentials in the path or query |
If the cause is not `Clone` and the error must be, hold it as `#[source] litellm_error::Source`, a cloneable `Arc<dyn Error + Send + Sync>` that is transparent to `Display` and `source()`, or make the crate's error opaque. `litellm_error::find::<T>(&error)` walks the chain and looks through `Source`, so it is how anyone reaches a typed cause. A third-party `Box<dyn Error + Send + Sync>` converts with `Source::from`, and an `anyhow::Error` a dependency returns goes through `Box` first, `Source::from(Box::<dyn Error + Send + Sync>::from(error))`, and a dependency's own boxed alias such as `tokenizers::Error` is held as it is. A source boxed to keep the variant small is a `Box<E>` node of the chain, so it is reached as `find::<Box<E>>`. Never put `#[error(transparent)]` on a `Source`, the erased error would then be no node of the chain and `find` could not reach it. Do not stringify or drop a cause to get `Clone` or `PartialEq`
```rust
// loses the type and leaks the URL
.map_err(|error| Error::Network(error.to_string()))
// loses everything
.map_err(|_| Error::Transport)
// keeps the cause, the From impl strips the URL
.map_err(litellm_http::Error::from)
```
A cause that must not be shown at all, such as a secret manager's response, is still a source. The type derives [`veil::Redact`](https://docs.rs/veil) with `#[redact]` on the field, so `Debug` hides it, and the boundary decides how much of the chain to show, see [boundaries](boundaries.md)
## Strip reqwest URLs on entry
Every `reqwest::Error` enters our types through `From<reqwest::Error> for litellm_http::Error`, which calls `without_url()`. A crate that talks HTTP holds `litellm_http::Error`, never `reqwest::Error`. When a dependency's own error wraps a `reqwest::Error`, as vaultrs and rustify do, one `From` takes the reqwest error out and converts it, as `secrets-hashicorp` does. That way so every later `Display`, `Debug` or chain walk is safe from that URL
`without_url()` only removes the URL reqwest attached. It does not scrub headers, response bodies, provider messages, file paths or anything a variant stores next to it. Those need their own policy, see [boundaries](boundaries.md)
## Rendering
Anything that shows an error to someone else renders the whole chain with `litellm_error::Report`, which writes each layer once joined by `: `, the single-line form of the unstable [`std::error::Report`](https://doc.rust-lang.org/std/error/struct.Report.html). `to_string()` of an error is only its top layer and is never what a client, a Python exception or a log line shows
```rust
let message = litellm_error::Report::new(&error).to_string();
// "could not reach the provider: error sending request: client error (Connect): tcp connect error: Connection refused"
```
## #[from] limits
`#[from]` only works on a variant whose only fields are the source and an optional backtrace. A variant that needs context such as a path, a phase or a `TlsSource` takes its cause through an explicit `map_err` or a named constructor
```rust
std::fs::read(&path).map_err(|source| Repr::Read {
path: path.clone(),
tls_source: TlsSource::CaBundle,
source,
})?;
```
`#[backtrace]` and automatic `provide()` need a nightly compiler and are not used here

View file

@ -0,0 +1,76 @@
# Python bridge
`crates/python-bridge` turns a Rust error into a Python exception through one mapping function per error type. A `#[pyfunction]` returns `Err(PyErr)` and pyo3 raises it, so those functions are the whole boundary
| Error | Mapping |
|---|---|
| `litellm_inference::Error` from a route | `errors::route_failure`, with `route_error_to_pyerr` as its shorthand for routes that set no marker |
| the OCR error | `route_failure` and `to_pyerr` in `routes/ocr/errors.rs` |
| token counter, cache, spend, secrets, argument checks | one `From<Error> for PyErr` or mapping function next to the code that registers the functions |
A route host never builds a `PyErr` from a route error itself. It calls the mapping, and the mapping is the only place that knows the exception classes and their attributes
## Exception classes
| Class | Meaning | Payload |
|---|---|---|
| `RustUpstreamError` | the provider call was issued and the provider answered with a failure | `args == (status, body)` and a `headers` attribute of `(name, value)` pairs, always set |
| `RustBridgeDeclined` | the route stopped before calling the provider, so the host may run its own path | `args == (message,)` |
| a standard class from `pyo3::exceptions` | whichever Python class means the same thing, such as `ValueError` or `TypeError` for the caller's input, `FileNotFoundError` or `OSError` for a file, `KeyError` or `OverflowError` where Python raises them, and `RuntimeError` for anything else | the rendered chain, or a parity text where Python pins one |
The two native classes come from `create_exception!` in `crates/python-bridge/src/errors.rs` and are exported with `#[pymodule_export]`, which is how the [pyo3 guide](https://github.com/PyO3/pyo3/blob/main/guide/src/exception.md) makes an exception importable. Their docstrings are the contract Python reads. `errors::upstream_error(py, status, body, headers)` is the one constructor of `RustUpstreamError`, so the payload cannot differ between routes
Python decodes a `RustUpstreamError` in `litellm/rust_bridge/failures.py`, which validates `args` and `headers` with a `TypeAdapter` and rebuilds the provider's `httpx.Response`. `litellm/rust_bridge/runtime.py` catches `RustBridgeDeclined` and reads its reason. `_raise_upstream` in `runtime.py` still reads `args` a second time by hand for the `attempt` path, which is the one decode to fold into `failures.py`. A new fact goes into `args` or `headers`, and the decoder learns to read it in the same change
Two routes also flag a request fault with a boolean attribute, read by their Python route module: `ocr_request_format_error` in `host/routes/ocr.py` and `messages_request_error` in `host/routes/messages.py`. OCR additionally sets `status_code` and `message` for Python parity, and `native_settings_error`, which only the OCR host reads back to tell a settings failure from the caller's exception. Do not add another attribute, nothing on the Python side types them
## Why args, not a typed subclass
The bridge builds with pyo3's `abi3-py310` feature. Under abi3, `#[pyclass(extends = PyException)]` is only allowed on Python 3.12 and newer, so a class with typed fields is not available. Even where it is allowed, pyo3 passes an empty tuple to the native base's `tp_new`, so `.args` and `str()` would need manual overrides. The payload therefore travels in `args`, which a `create_exception!` class fills from `new_err`
## Mapping order
A core decline is returned first, as `RustBridgeDeclined` with the decline's reason as its only arg, without calling the route's public mapping. Then a Python exception raised inside the call, such as a secret manager callback's, is returned as it is. Then the provider's answer passes through as `RustUpstreamError`. On `/v1/messages` a request the route rejected becomes a `ValueError` flagged with `messages_request_error`. Everything else is split by `is_request()` into `ValueError` or `RuntimeError`, with the message rendered by `litellm_error::Report`. A missing API key is a `RuntimeError` because `is_request()` answers false for it, and `missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors` pins that
OCR puts Python parity arms in front of that order, each matching a top-level variant of its own public enum
| Arm | Exception |
|---|---|
| the provider answered | `RustUpstreamError` with the provider's headers |
| `RequestFormat` | `ValueError` flagged with `ocr_request_format_error` |
| `FileRead` whose cause is `NotFound` | `FileNotFoundError` naming the path |
| other `FileRead` | `OSError` with the io message |
Both routes ask `litellm_auth::Error::is_request()` through their own `is_request()`, so a missing key or API base is a `RuntimeError` on both, and a missing setting the caller can pass, such as a region, is a `ValueError`
Use the standard class whose Python meaning matches. `PyFileNotFoundError` for a missing file is right, a `RuntimeError` whose text says the file is missing is not. `import_exception!` exists for an exception defined in Python, none is used today because `failures.py` maps to LiteLLM's public exceptions through `litellm.exception_type`
## Declining
`RustBridgeDeclined` means no provider call was made and the host may run its own path for the same call. The tokenizer, embeddings and token counter routes raise it from their own pre-checks, and `route_failure` raises it for a route error whose `declined()` answers, so a decline from deep inside a transformation reaches Python the same way. Never raise it for a failure after the request went out, the host would repeat a billed call. `a_core_decline_hands_the_call_back_to_python_before_the_public_mapping` and `route_failures_are_terminal` in `errors.rs` pin both sides
Route hosts pass a terminal failure through the Python route module's `map_failure`, which runs `litellm.exception_type` and returns the public exception a Python caller expects. A decline must skip that step. Once wrapped it looks like a provider failure to the dispatcher and the host no longer falls back, so `route_failure` returns it before it calls `public`
## Rules
- Never format `Debug` of an error into a `PyErr`
- Never build a `PyErr` from a string at a call site. A failed argument check is a variant of a typed error, as `ProjectionError` in `coercion.rs` is, mapped once
- A Python exception raised inside the call stays that exception. Return it as is, never wrap it in a Rust message or stringify it into a new `ValueError`. The one policy exception is a caller-supplied callable, whose failure `callable::wrap_failure` reports as a `RuntimeError` with the original as `__cause__`
- A `panic!`, `unwrap` or `expect` inside a `#[pyfunction]` or `#[getter]` reaches Python as `pyo3_runtime.PanicException`, which is not an `Exception` subclass. Model the failure as a value instead
- When the mapping itself calls Python, as `setattr` does, propagate its failure with `?` rather than discarding it with `.ok()`
## Testing
Build the Rust error the way production builds it, map it, and assert the contract Python sees. `Python::initialize()` once, then `Python::attach`
```rust
Python::initialize();
Python::attach(|py| {
let failure = route_error_to_pyerr(litellm_http::Error::upstream(429, "slow down".into()).into());
assert!(failure.is_instance_of::<RustUpstreamError>(py));
let args: (u16, String) = failure.value(py).getattr("args").unwrap().extract().unwrap();
assert_eq!(args, (429, "slow down".into()));
});
```
One `#[case]` per fault class, asserting `is_instance_of` for the class and `extract` for the payload. The Python side of the contract is tested in `tests/unit/rust_bridge`, through `failures.py`

View file

@ -0,0 +1,64 @@
# Recipes
## A new failure in litellm_http
Say media downloads must reject a `Content-Type` outside an allow-list
1. Add a `Repr` variant with typed facts and the message, `#[error("media type {0:?} is not allowed")] DisallowedMediaType(String)`
2. Return it where the check fails, `return Err(Repr::DisallowedMediaType(content_type.to_owned()).into())`
3. Add an arm to every predicate's exhaustive `match`, the compiler lists them, the way its neighbor `DownloadTooLarge` answers. A media download happens before the provider call, so it is not in `dispatched()`. The caller chose the URL, so it is in `is_request()`
4. Leave both boundaries alone. With those answers the gateway returns 400 and Python raises `ValueError`
5. Test it with a mock server answering `Content-Type: application/x-msdownload` and assert `error.is_request()`
## A provider-specific failure
A second OCR provider needs a missing-key error whose text names its environment variable and how to set it
1. Look for an existing provider-neutral variant whose template fits. `litellm_auth::Error::MissingApiKey { provider, environment_variable }` renders `Missing {provider} API Key - Set ...`
2. If none fits, find what varies between providers with the same failure, here the environment variable and the guidance after the dash, and add one neutral variant with those as fields, `#[error("Missing {environment_variable} - {guidance}")] MissingCredential { environment_variable: &'static str, guidance: &'static str }`
3. Pass the provider's values from its transformation. The next provider passes its own and adds no variant, and `is_missing_credential()` lists `MissingCredential { .. }` once
## Wrapping a new lower crate in litellm_inference
Say routes start calling `litellm_cache` and its errors must surface
1. Add one variant, `#[error(transparent)] Cache(#[from] litellm_cache::Error)`
2. Extend every predicate's `match` with an arm that calls through to the cache error's own predicates, adding a predicate there if it has none
3. Do not copy the cache error's variants into `litellm_inference::Error`
## A foreign error that needs context
Reading a file fails with `io::Error`, and the message must name the file
1. A variant with the context and the cause, `Read { path: PathBuf, source: io::Error }` with `#[error("could not read {}", path.display())]`
2. Build it with `map_err(|source| Repr::Read { path: path.clone(), source })`, not a `From<io::Error>`
3. If the error must be `Clone`, the crate's error is opaque or the field is `litellm_error::Source`
## Migrating a stringly error
Call sites look like `Error::Network(error.to_string())`, `Error::Request(format!(...))` or `.map_err(|_| Error::Unavailable)`
1. List every call site of the string variant and group them by what actually failed
2. Each group becomes a variant holding the typed cause or facts. Groups no caller tells apart share one
3. Where a foreign error was stringified with classification logic around it, move that logic into the `From` and turn the call sites into `?`
4. Delete the string variant. If it cannot be deleted yet, nothing new may construct it
5. Re-run the boundary tests, since messages a client sees may change
## Moving a crate to source chains
A variant reads `#[error("invalid JSON: {0}")] Json(#[from] serde_json::Error)` or stores `cause` and renders `{cause}`
1. Drop the cause from the template, `#[error("invalid JSON")]`
2. Make the cause the source, `#[from]`, `#[source]` or a field named `source`
3. Find what renders this error for someone else and make it use `litellm_error::Report`
4. Update tests that asserted the cause in `to_string()` to assert it in the report
## Making a public enum opaque
Downstream crates match nested variants, or a variant needs a non-`Clone` cause
1. Rename the enum to `pub(crate) enum Repr` and add `pub struct Error(Arc<Repr>)` with `#[error(transparent)]` and `From<Repr>`
2. For each downstream `match`, find the question it asks and add a predicate. Replace the `match` with the predicate
3. For each downstream construction, add a constructor taking typed data, or move the failure into this crate
4. Swap test assertions to predicates and add the `#[cfg(test)] repr()` accessor for in-crate tests
5. Implement `PartialEq` with `Arc::ptr_eq` only if a wrapping enum's derive needs it

View file

@ -0,0 +1,105 @@
# Error shape
## One top-level Error per crate
A crate's errors live in `src/error.rs` and are re-exported from `lib.rs` as `Error`. A caller of any public function in the crate should only ever see `Result<_, crate::Error>`, so one `?` and one set of predicates cover the whole crate
This deliberately departs from the advice to put an error type next to each unit of fallibility ([Sabrina Jewson, "Modular Errors"](https://sabrinajewson.org/blog/errors)). Our callers do not branch on which function failed, they ask the predicates in [classification](classification.md), and the boundaries map one type per crate. Splitting by function would multiply those predicates and mappings without giving any caller a decision it lacks. Layering, the other half of that advice, is kept: each crate nests the one below it, see [conversions](conversions.md)
That one type takes one of two forms
| Form | Use when | Example |
|---|---|---|
| Public enum | Callers outside the crate match variants, and every variant holds our own data or another crate's `Error` that is already `Clone` and comparable | `litellm_inference::Error`, `litellm_auth::Error` |
| Opaque struct over a private `Repr` | The variants hold third-party causes, or the variants must be free to change without breaking callers | `litellm_http::Error` |
Both forms are still one top-level `Error`. The opaque form just moves the enum behind `Repr`, where it is `pub(crate)`
```rust
// public enum
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error(transparent)]
Http(litellm_http::Error),
}
// opaque struct
#[derive(Clone, Debug, thiserror::Error)]
#[error(transparent)]
pub struct Error(Arc<Repr>);
#[derive(Debug, thiserror::Error)]
pub(crate) enum Repr {
#[error("could not reach the provider")]
Connect(#[source] reqwest::Error),
#[error("upstream request failed with status {status}: {body}")]
Status { status: u16, body: String },
}
```
Start with the public enum. Move to opaque when a variant needs a non-`Clone` third-party cause, or when downstream crates keep matching variants in ways that duplicate classification. `reqwest::Error`, `std::io::Error` and `serde_json::Error` are built the opaque way
In the opaque form other crates cannot name a variant, so they cannot build an error with their own message. Public constructors exist only for errors another crate genuinely produces, take typed data and never a message, as in `Error::upstream(status, body)`, `Error::computed_header(name)`, `Error::unsignable_body()` and `Error::signature(cause)`
`Arc` makes cloning a refcount bump, so `Repr` never has to be `Clone`. `Box<Repr>` is enough when the goal is only a small `Result`. `Arc` is for when callers clone
Every crate here is internal to the workspace, so do not mark errors or variants `#[non_exhaustive]`. It only forces a wildcard arm on the other workspace crates and turns off the exhaustiveness check that makes a new variant visit every `match`. When the variants must be free to change, use the opaque form
## What else may sit next to Error
These are not a second top-level error and need no folding in
- Leaf data types whose fields are typed and whose message is in `#[error(...)]`, such as `litellm_http::request::HeaderError`. A wrapping crate nests it as one variant
- Plain data used by predicates, such as `litellm_http::TlsSource`
- A private fault enum one variant carries, such as `PemFault` inside `Repr::InvalidPem`
## Variants are failure modes
A public variant is something a caller handles differently: a phase, a status code, retry versus no retry, fault on the request versus the environment, or a message Python parity pins exactly. Failures no caller tells apart share one variant and differ in their data
```rust
// one variant per provider, so every new provider adds a variant and a match arm
#[error("Missing Azure API Base - set AZURE_API_BASE")]
MissingAzureApiBase,
#[error("Missing Vertex API Base - set VERTEXAI_API_BASE")]
MissingVertexApiBase,
// one failure mode, the provider and its guidance are data
#[error("Missing {provider} API Base - {guidance}")]
MissingApiBase { provider: &'static str, guidance: &'static str },
```
An exact parity message does not justify a provider-specific variant when caller-supplied data can produce the same text
Inside a private `Repr` a variant costs no API, so one per distinct reason is fine when each one carries a different message. When several only differ in a word and share every predicate answer, prefer one variant with a fieldless reason enum, so predicates list one arm instead of four, as `Repr::Handshake(HandshakeFault)` does
```rust
#[error("WebSocket handshake returned {0}")]
Handshake(HandshakeFault),
#[derive(Debug, thiserror::Error)]
pub(crate) enum HandshakeFault {
#[error("a wrong Sec-WebSocket-Accept")]
Accept,
#[error("an invalid Upgrade header")]
Upgrade,
#[error("an invalid Connection header")]
Connection,
#[error("an invalid subprotocol")]
Subprotocol,
}
```
## When to split
Split into a separate type only when different functions fail in disjoint ways and different callers see them, so no caller matches variants its function can never return. A type per function overdoes it
Name a split type after what went wrong, not after the function that returns it. `ConnectError` reads better than `OpenSessionError`, and a unit struct is fine for a single failure mode
A split type still converts into the crate's `Error` when a caller further up needs one `Result` type
`litellm-llms` has one recorded exception, the OCR error in `src/base_llm/ocr/error.rs`, kept apart until OCR folds into `litellm_llms::Error`, see `crates/inference-ocr/AGENTS.md`. Do not add another
A utility crate whose modules are unrelated, such as `core-utils`, has no crate-wide failure, so each module that can fail owns one error next to its code instead of a `src/error.rs`. The rest of this skill applies to each of them

View file

@ -0,0 +1,37 @@
# Testing errors
Assert behavior through predicates and accessors, on errors produced the way production produces them
```rust
// guesses the message and the variant, and cannot be written once the type is opaque
assert_eq!(error, Error::Transport(TransportError::Http { status: 429, body: "slow down".into() }));
// asserts what callers rely on
assert_eq!(error.upstream_response(), Some((429, "slow down")));
assert!(error.dispatched());
assert!(!error.is_request());
```
## Producing the error
In order of preference
1. A real failing call. Bind a `TcpListener`, drop it, then connect to its address for a connect failure. Accept and never answer for a timeout
2. A real cause, such as `tungstenite::Error::ConnectionClosed` or a `serde_json::from_str("{")` error, passed through our `From`
3. A public constructor, such as `Error::upstream(500, "boom".into())`
Inside the owning crate a test may match the private representation through the `#[cfg(test)] fn repr(&self) -> &Repr` accessor in `crates/http/src/error.rs`, as the `media.rs` and `tls.rs` tests do. Anywhere else, reach the cause the test expects with `litellm_error::find::<T>(&error)`. A transparent layer is no node of the chain, so `find` reaches foreign causes underneath it but never the nested error of ours itself. Ask that error's predicates through the outer type instead, which is what callers do too
## What to cover
| Property | Test shape |
|---|---|
| classification | one `#[case]` per input class, asserting the predicate |
| redaction | build from a URL holding `api_key=secret`, assert the message lacks `secret` and the path |
| cause kept | `litellm_error::find::<T>(&error)` returns the expected cause |
| rendered once | the cause's text appears exactly once in `litellm_error::Report::new(&error).to_string()` and not in `to_string()` |
| boundary mapping | the gateway status and body, or the Python exception type and its `status` attribute |
Assert exact message text only where it is a contract, such as a message Python parity pins. Everywhere else assert that the message contains the facts it must carry. Matching a fieldless variant with `matches!` is fine, full `assert_eq!` on an error that holds a cause is not, since causes compare by pointer
A test that returns `Result` reports only `Debug` of the error on failure. Prefer `expect` with a short reason, or `unwrap_err` followed by assertions on the predicates

View file

@ -0,0 +1,22 @@
---
name: rust-gateway
description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test
metadata:
internal: true
---
# 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
- [surfaces](references/surfaces.md)
- Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose
- The error type string is derived from the status and provider bodies pass through bounded
- [requests](references/requests.md)
- Extractors reject with the crate's `Error`, handlers take the body as a `Result` when the error needs the request id
- Every alias of a path is registered together, a known but unserved path goes to one handler
- Streams relay through `litellm_host_http::Sse`, with the error frame rendered by the crate's `Error`
- [testing](references/testing.md)
- Build the router through `tests/support` and drive it with `oneshot`, assert the fields the surface promises
How an error is classified before it reaches a response is the [`rust-error`](../rust-error/SKILL.md) skill, in particular its [gateway](../rust-error/references/gateway.md) page. Axum's own conventions are in its [docs](https://docs.rs/axum/latest/axum/)

View file

@ -0,0 +1,34 @@
# Requests
## Extractors reject with the crate's Error
A custom extractor sets `type Rejection = Error`, so a bad body gets the surface's envelope rather than axum's plain-text rejection. Never take axum's own `Json`, `Query` or `Path` in a handler, their rejection is plain text. The rejection is the variant's source, never its string. `JsonObject` and `InferenceBody` in `crates/gateway-inference/src/request.rs` are the pattern
```rust
impl<S: Send + Sync> FromRequest<S> for JsonObject {
type Rejection = Error;
async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
let body = Bytes::from_request(request, state).await.map_err(|error| {
if error.status() == StatusCode::PAYLOAD_TOO_LARGE {
Error::BodyTooLarge
} else {
BodyError::Read(error).into()
}
})?;
Ok(Self(object(&body)?))
}
}
```
When the handler needs something from the request in the error body, such as the request id on `/v1/messages`, take the extractor as `Result<JsonObject, Error>` so the handler maps the rejection itself, as `messages::create` does. `Infallible` is the right rejection for an extractor that only reads an optional header, as `RequestId` does
## Paths
Every alias the Python proxy registers for a route is registered together in `crates/gateway-inference/src/lib.rs`, such as the bare path, its `/v1` form and the `/engines/{model}` and `/openai/deployments/{model}` forms, otherwise one alias silently behaves differently from the others. A path the gateway knows but does not serve, such as `/v1/completions`, is routed to `request::unsupported`, which returns `Error::Unsupported(path)` so the client gets a 501 in the surface's envelope rather than a bare 404
`crates/gateway` merges the mounted routers and adds the shared inbound middleware as `route_layer`s, so authentication runs only for matched routes. Endpoint paths and parsing stay in the mounted crate, see the two crates' `AGENTS.md`
## Streaming
Once the response headers are sent the status cannot change. A streaming handler builds `litellm_host_http::Sse::new(Json, |error| Bytes::from(Error::from(error).sse_frame()))`, where the closure renders a mid-stream failure as an `event: error` frame with the same status and message a JSON failure on that surface carries, `sse_frame()` for Anthropic and `responses_sse_frame()` for the Responses API's `error` event. `litellm_host_http::Error<E>` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the crate's `From` maps the second to `Internal`

View file

@ -0,0 +1,20 @@
# Envelopes follow the API surface
Each surface mirrors what the Python proxy returns there, so the shapes differ on purpose
| Surface | Envelope | Where |
|---|---|---|
| `/v1/messages` | Anthropic, `{"type": "error", "error": {"type", "message"}, "request_id"}` | `Error::body()` and `sse_frame()` in `crates/gateway-inference/src/error.rs` |
| OpenAI-shaped routes | `{"error": {"message", "type", "param", "code"}}` | `Error::openai_response()` in the same file |
| `/v1/responses` stream | the Responses `error` event, `{"type": "error", "code", "message", "param"}` | `Error::responses_sse_frame()` in the same file |
| MCP | `{"detail": message}` | `crates/gateway-mcp/src/error.rs` |
| UI | `{"error": {"message"}}`, plus `Retry-After` on 429 | `crates/gateway-ui/src/error.rs` |
| auth rejections | plain text, the top-level message only | `crates/gateway-auth/src/error.rs` |
Every surface takes its status from one `status()` on its crate's `Error` and renders the message of anything that is not a provider answer with `litellm_error::Report`. In `gateway-inference` the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else, and the OpenAI envelope always wraps, as the Python proxy does. A provider body is bounded by `litellm_http::Error::upstream` when the error is built
Two surfaces show less than the chain on purpose, and a test in each pins it. An auth rejection goes to a caller who has not authenticated, so it never shows what a secret manager or key store said. An MCP upstream or authorization failure shows only its own layer, because rmcp renders the transport error, URL included, into its own `Display`, and an authorizer's cause is the same backend detail auth rejections hide. The inference surfaces show only the top layer of an auth rejection for the same reason
`/v1/messages` adds the request id to the envelope when the provider did not set one, which is why its handler takes the body as `Result<JsonObject, Error>` and renders the rejection itself, see [requests](requests.md)
A new surface gets its own renderer next to the others, named after the shape it produces, and its status still comes from `Error::status()`. Picking the status is the [`rust-error`](../../rust-error/references/gateway.md) skill's concern, not the surface's

View file

@ -0,0 +1,9 @@
# Testing
Test HTTP contracts in the crate that owns the endpoint: status codes, forwarded headers, error envelopes and streaming behavior. Inference and provider behavior is tested in the `inference-<fmt>` and `llms` crates, and mounting and server lifecycle in `crates/gateway`
`crates/gateway-inference/tests/support/mod.rs` builds a `Router` with injected deployments, pools and settings through `app()` and its variants, and `post()` drives it with `tower::ServiceExt::oneshot`. `json()` reads the response body back. Use these instead of building a router or a client by hand in a test
Assert the status and the fields the surface promises, not the whole body. A test that pins a whole envelope breaks on every unrelated field, one that checks `status`, `error.type` and the presence of `request_id` catches the regressions that matter
`status()` and `body()` of the error type are pure, so a classification case belongs in the error module's unit tests, see the [`rust-error`](../../rust-error/references/gateway.md) skill. Router tests cover what only the full path shows, such as an extractor rejection arriving in the right envelope or an SSE error frame after the headers went out

View file

@ -1,6 +1,8 @@
---
name: rust-string-enums
description: Define or refactor Rust string-valued enums and their Serde adapters in litellm-rust, using Strum and serde_with while preserving parsing, wire values, and schemas
metadata:
internal: true
---
# Rust string enums

View file

@ -1,6 +1,8 @@
---
name: rust-tracing
description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery
metadata:
internal: true
---
# Rust tracing

View file

@ -4,6 +4,8 @@ For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.a
For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md)
For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md)
For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches
Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all`
@ -24,13 +26,3 @@ Use [`#[rstest]`](https://docs.rs/rstest/latest/rstest/attr.rstest.html) for new
- Never loop over inputs (`for`, `.iter().for_each`, `.all`) inside a test body. Give each input its own `#[case::name(...)]`, or use `#[values(...)]` for a cross product
- Exception: a test pinning a Rust table against a repo-owned data file (for example `include_str!` of a JSON config) may iterate that file's entries
## Error definitions
- A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs`
- Put message templates in the variant's `#[error(...)]` declaration. Callers pass only the small typed arguments needed to fill them, never `Error::Variant(format!(...))` or a preformatted message. Keep the smallest set of neutral variants that callers need to distinguish. Different wording or providers do not justify new variants
- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each. A failure mode is something a caller handles differently (phase, status code, retry, a message Python parity pins exactly). Failures no caller tells apart share one variant and differ only in its message
- Keep shared error enums minimal and provider-neutral. Provider names, credential types, configuration fields, and setup guidance belong in caller-supplied data, not dedicated variants or hardcoded shared messages. Reuse a variant for the same failure mode across providers, such as `MissingApiBase { provider: "Azure", guidance: "..." }`. An exact parity message does not justify a provider-specific variant when caller-supplied context can preserve it
- Wrap a lower-level error as a variant with `#[from]` or `#[source]` instead of flattening it to a string
- Exception: split into separate types when different functions fail in disjoint ways, especially when different callers see them. A shared enum would force every caller to match variants its function can never return
- Name a split type after what went wrong (a unit struct is fine for a single failure mode), not after the function that returns it

221
litellm-rust/Cargo.lock generated
View file

@ -2342,7 +2342,7 @@ dependencies = [
"hmac 0.13.0",
"http 1.4.2",
"jiff",
"reqwest 0.13.5",
"reqwest",
"rustc_version",
"rustls 0.23.42",
"rustls-pki-types",
@ -2400,7 +2400,7 @@ dependencies = [
"pin-project",
"prost",
"prost-types",
"reqwest 0.13.5",
"reqwest",
"rustc_version",
"serde",
"serde_json",
@ -2896,7 +2896,6 @@ dependencies = [
"tokio",
"tokio-rustls 0.26.4",
"tower-service",
"webpki-roots",
]
[[package]]
@ -3151,6 +3150,7 @@ dependencies = [
"defmt",
"jiff-core",
"jiff-static",
"jiff-tzdb",
"log",
"portable-atomic",
"portable-atomic-util",
@ -3179,6 +3179,12 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "jiff-tzdb"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fa8377070c6bae868759445e5a77f66d84f0b72f3a054bfb00e6d038b8282da7"
[[package]]
name = "jni"
version = "0.22.4"
@ -3386,10 +3392,13 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"bytes",
"futures-util",
"litellm-auth-types",
"litellm-error",
"litellm-http",
"moka",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -3406,6 +3415,7 @@ dependencies = [
"azure_core",
"azure_identity",
"litellm-auth-types",
"litellm-error",
"moka",
"rstest",
"serde_json",
@ -3424,6 +3434,7 @@ dependencies = [
"google-cloud-auth",
"http 1.4.2",
"litellm-auth-types",
"litellm-error",
"moka",
"rstest",
"serde",
@ -3438,6 +3449,7 @@ dependencies = [
name = "litellm-auth-types"
version = "0.1.0"
dependencies = [
"litellm-error",
"rstest",
"serde",
"serde_json",
@ -3451,6 +3463,7 @@ dependencies = [
name = "litellm-cache"
version = "0.1.0"
dependencies = [
"litellm-error",
"rstest",
"serde",
"serde_json",
@ -3473,9 +3486,10 @@ dependencies = [
"litellm-cache-response",
"litellm-cache-testing",
"litellm-http",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
"url",
"wiremock",
@ -3494,6 +3508,7 @@ dependencies = [
"serde-pickle",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
]
@ -3506,11 +3521,13 @@ dependencies = [
"litellm-auth-types",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"percent-encoding",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
"wiremock",
]
@ -3533,9 +3550,10 @@ dependencies = [
"futures-util",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"qdrant-client",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -3553,11 +3571,13 @@ version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"r2d2",
"redis",
"redis-test",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
@ -3610,8 +3630,9 @@ dependencies = [
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde_json",
"tokio",
@ -3632,11 +3653,13 @@ dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-testing",
"litellm-error",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"uuid",
]
@ -3646,6 +3669,7 @@ name = "litellm-callbacks-legacy-python"
version = "0.1.0"
dependencies = [
"litellm-auth",
"litellm-error",
"litellm-host",
"litellm-host-python",
"proptest",
@ -3653,6 +3677,7 @@ dependencies = [
"rstest",
"serde_json",
"strum",
"thiserror 2.0.19",
]
[[package]]
@ -3668,6 +3693,7 @@ name = "litellm-config"
version = "0.1.0"
dependencies = [
"litellm-auth-types",
"litellm-error",
"litellm-router-types",
"rstest",
"serde",
@ -3685,6 +3711,7 @@ name = "litellm-core-utils"
version = "0.1.0"
dependencies = [
"fancy-regex 0.19.2",
"litellm-error",
"litellm-llms-types",
"litellm-tracing",
"rstest",
@ -3709,9 +3736,13 @@ dependencies = [
name = "litellm-cost"
version = "0.1.0"
dependencies = [
"criterion",
"proptest",
"jiff",
"rstest",
"schemars 1.2.2",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
]
[[package]]
@ -3741,6 +3772,14 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-error"
version = "0.1.0"
dependencies = [
"rstest",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-framer"
version = "0.1.0"
@ -3768,6 +3807,7 @@ dependencies = [
"litellm-auth-types",
"litellm-circuit-breaker",
"litellm-config",
"litellm-error",
"litellm-gateway-auth",
"litellm-gateway-inference",
"litellm-gateway-mcp",
@ -3801,6 +3841,7 @@ dependencies = [
"futures-util",
"litellm-auth-types",
"litellm-config",
"litellm-error",
"litellm-secrets",
"rstest",
"serde",
@ -3824,6 +3865,7 @@ dependencies = [
"litellm-auth",
"litellm-cache-memory",
"litellm-cache-response",
"litellm-error",
"litellm-gateway-auth",
"litellm-host",
"litellm-host-http",
@ -3843,6 +3885,7 @@ dependencies = [
"serde_json",
"thiserror 2.0.19",
"tokio",
"tokio-tungstenite",
"tower",
"wiremock",
]
@ -3853,6 +3896,7 @@ version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-auth-types",
"litellm-error",
"litellm-gateway-auth",
"rand 0.8.7",
"rstest",
@ -3870,6 +3914,7 @@ dependencies = [
"http 1.4.2",
"litellm-auth-types",
"litellm-config",
"litellm-error",
"litellm-secrets",
"moka",
"rmcp",
@ -3884,6 +3929,7 @@ dependencies = [
"tower",
"url",
"uuid",
"veil",
]
[[package]]
@ -3893,9 +3939,11 @@ dependencies = [
"axum",
"axum-login",
"base64 0.22.1",
"futures-util",
"governor",
"jsonwebtoken",
"litellm-auth-types",
"litellm-error",
"litellm-gateway-auth",
"rand 0.8.7",
"rstest",
@ -3920,6 +3968,7 @@ dependencies = [
"litellm-coroutine",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
@ -3957,6 +4006,7 @@ dependencies = [
"bytes",
"futures-util",
"litellm-coroutine",
"litellm-error",
"litellm-host",
"parking_lot",
"pyo3",
@ -3977,10 +4027,9 @@ dependencies = [
"http 1.4.2",
"hyper-util",
"litellm-core-utils",
"litellm-error",
"rcgen",
"reqwest 0.12.28",
"reqwest 0.13.5",
"rmcp",
"reqwest",
"rstest",
"rustls 0.23.42",
"rustls-native-certs",
@ -3990,8 +4039,10 @@ dependencies = [
"tempfile",
"thiserror 2.0.19",
"tokio",
"tokio-rustls 0.26.4",
"tokio-tungstenite",
"tracing",
"url",
"veil",
"webpki-roots",
]
@ -4007,13 +4058,14 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-http",
"litellm-llms",
"litellm-secrets",
"litellm-tracing",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -4034,6 +4086,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4061,6 +4114,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-host-native",
@ -4072,7 +4126,7 @@ dependencies = [
"litellm-router-types",
"litellm-secrets",
"litellm-tracing",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -4092,6 +4146,7 @@ dependencies = [
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4109,6 +4164,7 @@ dependencies = [
"time",
"tokio",
"tracing",
"url",
"wiremock",
]
@ -4124,6 +4180,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4133,7 +4190,7 @@ dependencies = [
"litellm-llms-types",
"litellm-secrets",
"litellm-tracing",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde_json",
"tokio",
@ -4161,6 +4218,7 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-inference",
"litellm-inference-testing",
@ -4189,6 +4247,7 @@ dependencies = [
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-http",
@ -4196,7 +4255,7 @@ dependencies = [
"litellm-python-compat",
"litellm-router-types",
"litellm-secrets",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -4228,16 +4287,35 @@ name = "litellm-model-catalog"
version = "0.1.0"
dependencies = [
"indexmap 2.14.0",
"jiff",
"jsonschema",
"litellm-llms-types",
"litellm-cost",
"litellm-error",
"macro_rules_attribute",
"rstest",
"schemars 1.2.2",
"serde",
"serde_json",
"serde_with",
"strum",
"thiserror 2.0.19",
"time",
]
[[package]]
name = "litellm-model-catalog-refresh"
version = "0.1.0"
dependencies = [
"litellm-error",
"litellm-http",
"litellm-model-catalog",
"reqwest",
"rstest",
"thiserror 2.0.19",
"tokio",
"url",
]
[[package]]
name = "litellm-python-bridge"
version = "0.1.0"
@ -4255,6 +4333,7 @@ dependencies = [
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-python",
"litellm-http",
@ -4278,13 +4357,14 @@ dependencies = [
"pyo3",
"pyo3-async-runtimes",
"redis",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
"serde_with",
"sqlx",
"strum",
"tempfile",
"thiserror 2.0.19",
"tokio",
"tokio-tungstenite",
@ -4335,12 +4415,17 @@ name = "litellm-router"
version = "0.1.0"
dependencies = [
"litellm-config",
"litellm-core-utils",
"litellm-inference",
"litellm-inference-messages",
"litellm-router-types",
"rstest",
]
[[package]]
name = "litellm-router-core"
version = "0.1.0"
[[package]]
name = "litellm-router-types"
version = "0.1.0"
@ -4362,6 +4447,7 @@ dependencies = [
"google-cloud-auth",
"google-cloud-kms-v1",
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-python-compat",
"litellm-secrets-aws",
@ -4371,7 +4457,7 @@ dependencies = [
"litellm-secrets-hashicorp",
"litellm-secrets-types",
"moka",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -4413,13 +4499,14 @@ dependencies = [
"litellm-http",
"litellm-secrets-types",
"percent-encoding",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4436,13 +4523,14 @@ dependencies = [
"moka",
"percent-encoding",
"rcgen",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4463,12 +4551,13 @@ dependencies = [
"litellm-secrets-types",
"moka",
"percent-encoding",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4478,7 +4567,10 @@ name = "litellm-secrets-hashicorp"
version = "0.1.0"
dependencies = [
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-secrets-types",
"reqwest",
"rstest",
"rustify",
"rustify_derive",
@ -4512,6 +4604,7 @@ name = "litellm-spend-clickhouse"
version = "0.1.0"
dependencies = [
"flate2",
"litellm-error",
"litellm-http",
"litellm-storage-clickhouse",
"rstest",
@ -4532,6 +4625,7 @@ name = "litellm-storage-clickhouse"
version = "0.1.0"
dependencies = [
"flate2",
"litellm-error",
"litellm-http",
"rstest",
"serde",
@ -4551,7 +4645,8 @@ version = "0.1.0"
dependencies = [
"flate2",
"futures-util",
"reqwest 0.12.28",
"litellm-http",
"reqwest",
"rstest",
"semver",
"serde",
@ -4574,6 +4669,7 @@ dependencies = [
"criterion",
"indexmap 2.14.0",
"itoa",
"litellm-error",
"litellm-token-counter-fast",
"litellm-token-counter-huggingface",
"litellm-token-counter-tiktoken",
@ -4590,6 +4686,7 @@ name = "litellm-token-counter-fast"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-error",
"rand 0.8.7",
"rstest",
"rustc-hash",
@ -4615,6 +4712,7 @@ name = "litellm-token-counter-tiktoken"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-error",
"once_cell",
"rstest",
"rustc-hash",
@ -5936,49 +6034,6 @@ version = "1.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2"
[[package]]
name = "reqwest"
version = "0.12.28"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-core",
"futures-util",
"h2 0.4.15",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.10.1",
"hyper-rustls 0.27.9",
"hyper-util",
"js-sys",
"log",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"quinn",
"rustls 0.23.42",
"rustls-pki-types",
"serde",
"serde_json",
"serde_urlencoded",
"sync_wrapper",
"tokio",
"tokio-rustls 0.26.4",
"tokio-util",
"tower",
"tower-http",
"tower-service",
"url",
"wasm-bindgen",
"wasm-bindgen-futures",
"wasm-streams 0.4.2",
"web-sys",
"webpki-roots",
]
[[package]]
name = "reqwest"
version = "0.13.5"
@ -6000,6 +6055,7 @@ dependencies = [
"js-sys",
"log",
"mime",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"quinn",
@ -6019,7 +6075,7 @@ dependencies = [
"url",
"wasm-bindgen",
"wasm-bindgen-futures",
"wasm-streams 0.5.0",
"wasm-streams",
"web-sys",
]
@ -6066,7 +6122,7 @@ dependencies = [
"pin-project-lite",
"process-wrap",
"rand 0.10.2",
"reqwest 0.13.5",
"reqwest",
"schemars 1.2.2",
"serde",
"serde_json",
@ -6189,7 +6245,7 @@ dependencies = [
"async-trait",
"bytes",
"http 1.4.2",
"reqwest 0.13.5",
"reqwest",
"rustify_derive",
"serde",
"serde_json",
@ -7221,7 +7277,7 @@ dependencies = [
"memchr",
"parse-display",
"pin-project-lite",
"reqwest 0.13.5",
"reqwest",
"serde",
"serde_json",
"serde_with",
@ -7470,11 +7526,7 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9"
dependencies = [
"futures-util",
"log",
"rustls 0.23.42",
"rustls-native-certs",
"rustls-pki-types",
"tokio",
"tokio-rustls 0.26.4",
"tungstenite",
]
@ -7819,8 +7871,6 @@ dependencies = [
"httparse",
"log",
"rand 0.8.7",
"rustls 0.23.42",
"rustls-pki-types",
"sha1 0.10.7",
"thiserror 1.0.69",
"utf-8",
@ -7860,7 +7910,7 @@ dependencies = [
"futures",
"pin-project",
"rand 0.10.2",
"reqwest 0.13.5",
"reqwest",
"serde",
"serde_json",
"time",
@ -8074,7 +8124,7 @@ dependencies = [
"async-trait",
"derive_builder",
"http 1.4.2",
"reqwest 0.13.5",
"reqwest",
"rustify",
"rustify_derive",
"serde",
@ -8221,19 +8271,6 @@ dependencies = [
"unicode-ident",
]
[[package]]
name = "wasm-streams"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
dependencies = [
"futures-util",
"js-sys",
"wasm-bindgen",
"wasm-bindgen-futures",
"web-sys",
]
[[package]]
name = "wasm-streams"
version = "0.5.0"

View file

@ -10,6 +10,7 @@ repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
litellm-config = { path = "crates/config" }
litellm-error = { path = "crates/error" }
litellm-router = { path = "crates/router" }
litellm-router-types = { path = "crates/router-types" }
litellm-tracing = { path = "crates/tracing" }
@ -72,6 +73,9 @@ litellm-python-compat = { path = "crates/python-compat" }
litellm-circuit-breaker = { path = "crates/circuit-breaker" }
litellm-redis-client = { path = "crates/redis-client" }
litellm-redis-testing = { path = "crates/redis-testing" }
litellm-cost = { path = "crates/cost" }
litellm-model-catalog = { path = "crates/model-catalog" }
litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" }
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
tracing = "0.1"
@ -90,7 +94,7 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
rand = "0.8"
macro_rules_attribute = "0.2.3"
schemars = "1"
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] }
redis = "1.8.0"
redis-test = "1.0.4"
testcontainers-modules = "0.15.0"
@ -100,14 +104,15 @@ rstest = "0.26.1"
wiremock = "0.6.5"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustls-native-certs = "0.8"
rustify = "=0.7.0"
rustify_derive = "=0.5.5"
vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] }
sha2 = "0.10"
sha1 = "0.10"
syn = { version = "2", default-features = false }
sqlx = { version = "0.9.0", default-features = false, features = ["json", "macros", "postgres", "runtime-tokio", "chrono", "tls-rustls-ring-native-roots"] }
chrono = { version = "0.4.45", default-features = false, features = ["std"] }
@ -118,7 +123,7 @@ thiserror = "2.0"
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
tiktoken-rs = "0.12.0"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] }
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
tokio-tungstenite = { version = "0.24", default-features = false, features = ["handshake"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
flate2 = "1"

View file

@ -10,8 +10,11 @@ disallowed-methods = [
{ path = "reqwest::Client::new", reason = "take litellm_http::Client from HttpClientPool" },
{ path = "reqwest::Client::builder", reason = "HttpClientConfig owns client construction" },
{ path = "reqwest::ClientBuilder::danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" },
{ path = "reqwest::ClientBuilder::tls_danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" },
{ path = "reqwest::ClientBuilder::tls_danger_accept_invalid_hostnames", reason = "set HttpClientConfig::verify instead" },
{ path = "reqwest::ClientBuilder::identity", reason = "set HttpClientConfig::client_certificate instead" },
{ path = "reqwest::ClientBuilder::use_preconfigured_tls", reason = "HttpClientConfig owns the TLS configuration" },
{ path = "reqwest::ClientBuilder::tls_backend_preconfigured", reason = "HttpClientConfig owns the TLS configuration" },
{ path = "sqlx::query", reason = "use sqlx::query! or query_file! so the SQL is checked against the migrated schema" },
{ path = "sqlx::query_as", reason = "use sqlx::query_as! or query_file_as! so the SQL is checked against the migrated schema" },
{ path = "sqlx::query_scalar", reason = "use sqlx::query_scalar! so the SQL is checked against the migrated schema" },

View file

@ -7,7 +7,9 @@ repository.workspace = true
[dependencies]
litellm-auth-types.workspace = true
litellm-error.workspace = true
litellm-http.workspace = true
reqwest.workspace = true
moka = { workspace = true, features = ["sync"] }
serde.workspace = true
@ -26,5 +28,6 @@ aws-smithy-runtime-api = "1.13.0"
[dev-dependencies]
rstest.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
reqwest.workspace = true
bytes.workspace = true
futures-util.workspace = true
tokio.workspace = true

View file

@ -1,4 +1,3 @@
use std::collections::BTreeMap;
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};
@ -12,16 +11,17 @@ use aws_sigv4::http_request::{
};
use aws_sigv4::sign::v4;
use aws_smithy_runtime_api::client::identity::Identity;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use litellm_auth_types::AwsParams;
use super::Error;
use super::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN,
AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES,
};
use crate::error::{CredentialFlow, Error, SigningStage};
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
@ -271,7 +271,7 @@ impl AwsAuthService {
provider
.provide_credentials()
.await
.map_err(|error| Error::AwsProfile(error.to_string()))
.map_err(Error::credentials(CredentialFlow::Profile))
}
AwsAuthFlow::AssumeRole { role, session_name } => {
if is_already_running_as_role(&role, &resolved).await? {
@ -287,7 +287,7 @@ impl AwsAuthService {
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::DefaultChain))?;
self.set_cached_credentials(
key,
credentials.clone(),
@ -327,7 +327,7 @@ impl AwsAuthService {
provider
.provide_credentials()
.await
.map_err(|error| Error::AwsAssumeRole(error.to_string()))
.map_err(Error::credentials(CredentialFlow::AssumeRole))
}
AwsAuthFlow::WebIdentity {
token,
@ -350,12 +350,14 @@ impl AwsAuthService {
.web_identity_token(token)
.send()
.await
.map_err(|error| Error::AwsWebIdentity(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::WebIdentity))?;
let credentials = response
.credentials()
.ok_or(Error::AwsMissingWebIdentityCredentials)?;
let expiration = SystemTime::try_from(*credentials.expiration())
.map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?;
.ok_or(Error::MissingWebIdentityCredentials)?;
let expiration =
SystemTime::try_from(*credentials.expiration()).map_err(|source| {
Error::WebIdentityExpiration(litellm_error::Source::new(source))
})?;
Ok(Credentials::new(
credentials.access_key_id(),
credentials.secret_access_key(),
@ -376,7 +378,7 @@ impl AwsAuthService {
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::DefaultChain))?;
self.set_cached_credentials(
key,
credentials.clone(),
@ -435,12 +437,12 @@ fn default_session_name() -> String {
/// client header cannot change the canonical request and invalidate the
/// signature. Signing everything instead makes the request 403 on a header the
/// caller supplied, on a deployment that works on the Python path.
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
pub fn aws_signature_headers(headers: &HeaderMap) -> HeaderMap {
headers
.iter()
.filter(|(name, _)| {
let name = name.to_ascii_lowercase();
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
let name = name.as_str();
AWS_SIGNED_HEADER_NAMES.contains(&name)
|| name.starts_with("x-amz-")
|| name.starts_with("x-amzn-")
})
@ -456,45 +458,62 @@ pub fn is_sigv4_computed_header(name: &str) -> bool {
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
}
pub fn sign_post(
#[derive(Clone, Copy, Debug)]
pub struct SigningScope<'a> {
pub region: &'a str,
pub service: &'a str,
pub credentials: &'a Credentials,
pub time: SystemTime,
}
pub fn sign_headers(
method: &str,
url: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
service: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> Result<BTreeMap<String, String>, Error> {
let identity: Identity = credentials.clone().into();
headers: &HeaderMap,
scope: SigningScope<'_>,
) -> Result<HeaderMap, Error> {
let identity: Identity = scope.credentials.clone().into();
let params = v4::SigningParams::builder()
.identity(&identity)
.region(region)
.name(service)
.time(signing_time)
.region(scope.region)
.name(scope.service)
.time(scope.time)
.settings(SigningSettings::default())
.build()
.map(SigningParams::from)
.map_err(|error| Error::AwsSigningParameters(error.to_string()))?;
.map_err(Error::signing(SigningStage::Parameters))?;
let header_refs = headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_str()));
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
.map_err(|error| Error::AwsSignableRequest(error.to_string()))?;
.map(|(name, value)| {
value
.to_str()
.map(|value| (name.as_str(), value))
.map_err(|_| Error::NonAsciiHeader {
name: name.to_string(),
})
})
.collect::<Result<Vec<_>, _>>()?;
let request = SignableRequest::new(
method,
url,
header_refs.into_iter(),
SignableBody::Bytes(body),
)
.map_err(Error::signing(SigningStage::Request))?;
let (instructions, _) = sign(request, &params)
.map_err(|error| Error::AwsSigning(error.to_string()))?
.map_err(Error::signing(SigningStage::Signature))?
.into_parts();
Ok(instructions
instructions
.headers()
.map(|(name, value)| {
let normalized_name = match name {
"authorization" => "Authorization",
"x-amz-date" => "X-Amz-Date",
"x-amz-security-token" => "X-Amz-Security-Token",
_ => name,
};
(normalized_name.to_string(), value.to_string())
let name = HeaderName::from_bytes(name.as_bytes())
.map_err(Error::signing(SigningStage::Output))?;
let value =
HeaderValue::from_str(value).map_err(Error::signing(SigningStage::Output))?;
Ok((name, value))
})
.collect())
.collect()
}
/// Model-id and region parsing shared by every Bedrock route.
@ -639,6 +658,7 @@ pub fn host_supplied_credentials(params: &AwsParams) -> Option<Credentials> {
mod tests {
use super::*;
use crate::constants::{AWS_REGION, BEDROCK_SERVICE};
use reqwest::header::CONTENT_TYPE;
fn no_env(_: &str) -> Option<String> {
None
@ -731,12 +751,12 @@ mod tests {
);
}
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
fn parity_inputs() -> (String, Vec<u8>, HeaderMap) {
(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
.to_string(),
br#"{"input":"hello"}"#.to_vec(),
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]),
HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]),
)
}
@ -961,9 +981,9 @@ mod tests {
// cannot change the canonical request. Signing it instead makes the
// request 403 the moment anything on the wire rewrites or drops it.
let (url, body, mut headers) = parity_inputs();
headers.insert("x-request-id".to_string(), "abc-123".to_string());
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
headers.insert("x-request-id", HeaderValue::from_static("abc-123"));
headers.insert("Accept-Encoding", HeaderValue::from_static("gzip"));
headers.insert("x-amzn-trace-id", HeaderValue::from_static("Root=1-abc"));
let signable = aws_signature_headers(&headers);
assert!(!signable.contains_key("x-request-id"));
@ -979,18 +999,21 @@ mod tests {
None,
"test",
);
let signed = sign_post(
let signed = sign_headers(
"POST",
&url,
&body,
&signable,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
SystemTime::UNIX_EPOCH,
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: SystemTime::UNIX_EPOCH,
},
)
.expect("signs");
let authorization = signed
.get("Authorization")
let authorization = signed["Authorization"]
.to_str()
.expect("carries an authorization header");
assert!(
!authorization.contains("x-request-id"),
@ -1012,29 +1035,24 @@ mod tests {
None,
"test",
);
let signed = sign_post(
let signed = sign_headers(
"POST",
&url,
&body,
&headers,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
},
)
.expect("golden signature");
assert_eq!(signed["X-Amz-Date"], "20240102T030405Z");
assert_eq!(signed["X-Amz-Security-Token"], "session-token");
assert_eq!(
signed.get("X-Amz-Date").map(String::as_str),
Some("20240102T030405Z")
);
assert_eq!(
signed.get("X-Amz-Security-Token").map(String::as_str),
Some("session-token")
);
assert_eq!(
signed.get("Authorization").map(String::as_str),
Some(
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
)
signed["Authorization"],
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
);
}
@ -1048,19 +1066,71 @@ mod tests {
None,
"test",
);
let signed = sign_post(
let signed = sign_headers(
"POST",
&url,
&body,
&headers,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
},
)
.expect("signature");
assert!(!signed.contains_key("X-Amz-Security-Token"));
}
#[test]
fn a_header_value_outside_visible_ascii_is_named_instead_of_signed() {
let (url, body, mut headers) = parity_inputs();
headers.insert("x-amz-meta-note", HeaderValue::from_str("café").unwrap());
let credentials = Credentials::new("AKIDEXAMPLE", "secret", None, None, "test");
let error = sign_headers(
"POST",
&url,
&body,
&headers,
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: UNIX_EPOCH,
},
)
.unwrap_err();
assert_eq!(
error,
Error::NonAsciiHeader {
name: "x-amz-meta-note".into()
}
);
}
#[test]
fn the_shared_auth_error_nests_the_aws_failure_and_its_cause() {
let cause = std::io::Error::new(std::io::ErrorKind::NotFound, "profile not found");
let error =
litellm_auth_types::Error::from(Error::credentials(CredentialFlow::Profile)(cause));
assert!(matches!(
litellm_error::find::<Error>(&error),
Some(Error::Credentials {
flow: CredentialFlow::Profile,
..
})
));
let report = litellm_error::Report::new(&error).to_string();
assert!(
report.contains("could not load AWS profile credentials"),
"{report}"
);
assert_eq!(report.matches("profile not found").count(), 1, "{report}");
}
#[ignore]
#[tokio::test]
async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box<dyn std::error::Error>> {
@ -1068,7 +1138,7 @@ mod tests {
let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?;
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec();
let headers =
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]);
HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]);
let credentials = AwsAuthService::default()
.resolve_credentials(
AwsAuthConfig {
@ -1080,30 +1150,32 @@ mod tests {
&no_env,
)
.await?;
let client = litellm_http::Client::plain_for_test();
let client = litellm_http::Client::for_test();
let mut failures = Vec::new();
for region in ["us-west-2", "us-east-1"] {
let url = format!(
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
);
let signed_headers = sign_post(
let signed_headers = sign_headers(
"POST",
&url,
&body,
&headers,
region,
BEDROCK_SERVICE,
&credentials,
SystemTime::now(),
SigningScope {
region,
service: BEDROCK_SERVICE,
credentials: &credentials,
time: SystemTime::now(),
},
)?;
let mut request = client.post(&url).body(body.clone());
for (name, value) in &headers {
request = request.header(name, value);
}
for (name, value) in signed_headers {
request = request.header(name, value);
}
let response = request.send().await?;
let response = client
.post(&url)
.headers(headers.clone())
.headers(signed_headers)
.body(body.clone())
.send()
.await?;
let status = response.status();
let response_body = response.text().await?;
let snippet: String = response_body.chars().take(240).collect();

View file

@ -1,46 +1,73 @@
use thiserror::Error as ThisError;
use litellm_error::Source;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("AWS profile credentials failed: {0}")]
AwsProfile(String),
#[error("AWS default credentials failed: {0}")]
AwsDefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AwsAssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
AwsWebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
AwsWebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
AwsSigningParameters(String),
#[error("AWS signable request failed: {0}")]
AwsSignableRequest(String),
#[error("AWS request signing failed: {0}")]
AwsSigning(String),
#[error("could not load AWS {flow} credentials")]
Credentials {
flow: CredentialFlow,
#[source]
source: Source,
},
#[error("AWS web identity response had no credentials")]
AwsMissingWebIdentityCredentials,
MissingWebIdentityCredentials,
#[error("AWS web identity response had an invalid expiration")]
WebIdentityExpiration(#[source] Source),
#[error("header {name} has a value that is not visible ASCII")]
NonAsciiHeader { name: String },
#[error("{stage}")]
Signing {
stage: SigningStage,
#[source]
source: Source,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum CredentialFlow {
#[error("profile")]
Profile,
#[error("default chain")]
DefaultChain,
#[error("assumed role")]
AssumeRole,
#[error("web identity")]
WebIdentity,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum SigningStage {
#[error("SigV4 signing parameters are invalid")]
Parameters,
#[error("SigV4 cannot sign this request")]
Request,
#[error("SigV4 could not compute the signature")]
Signature,
#[error("SigV4 produced an invalid header")]
Output,
}
impl Error {
pub(crate) fn credentials<E: std::error::Error + Send + Sync + 'static>(
flow: CredentialFlow,
) -> impl FnOnce(E) -> Self {
move |source| Self::Credentials {
flow,
source: Source::new(source),
}
}
pub(crate) fn signing<E: std::error::Error + Send + Sync + 'static>(
stage: SigningStage,
) -> impl FnOnce(E) -> Self {
move |source| Self::Signing {
stage,
source: Source::new(source),
}
}
}
impl From<Error> for litellm_auth_types::Error {
fn from(error: Error) -> Self {
Self::ProviderAuthentication(error.to_string())
}
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn converts_to_shared_auth_error_without_losing_context() {
let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into()));
assert_eq!(
error,
litellm_auth_types::Error::ProviderAuthentication(
"AWS profile credentials failed: profile not found".into()
)
);
Self::provider(error)
}
}

View file

@ -5,5 +5,5 @@ mod signer;
pub use aws::*;
pub use aws_credential_types::Credentials;
pub use error::Error;
pub use error::{CredentialFlow, Error, SigningStage};
pub use signer::SigV4Signer;

View file

@ -1,11 +1,13 @@
use std::{collections::BTreeMap, time::SystemTime};
use std::time::SystemTime;
use aws_credential_types::Credentials;
use reqwest::{Request, header::HeaderMap};
use crate::{
AwsAuthService, AwsCredentialSource, Error, aws_signature_headers, is_sigv4_computed_header,
sign_post,
AwsAuthService, AwsCredentialSource, Error, SigningScope, aws_signature_headers,
is_sigv4_computed_header, sign_headers,
};
use aws_credential_types::Credentials;
use litellm_http::outbound::{RequestSigner, UnsignedRequest};
use litellm_http::outbound::RequestSigner;
#[derive(Clone, Debug)]
pub struct SigV4Signer {
@ -45,41 +47,54 @@ impl SigV4Signer {
}
impl RequestSigner for SigV4Signer {
fn sign(
&self,
request: UnsignedRequest<'_>,
) -> Result<Vec<(String, String)>, litellm_http::Error> {
if let Some((name, _)) = request
.headers
.iter()
.find(|(name, _)| is_sigv4_computed_header(name))
fn sign(&self, request: &Request) -> Result<HeaderMap, litellm_http::Error> {
if let Some(name) = request
.headers()
.keys()
.find(|name| is_sigv4_computed_header(name.as_str()))
{
return Err(litellm_http::Error::ComputedHeader(name.clone()));
return Err(litellm_http::Error::computed_header(name.as_str()));
}
let headers: BTreeMap<String, String> = request.headers.iter().cloned().collect();
sign_post(
request.url,
request.body,
&aws_signature_headers(&headers),
&self.region,
self.service,
&self.credentials,
(self.clock)(),
let body = match request.body() {
None => &[],
Some(body) => body
.as_bytes()
.ok_or_else(litellm_http::Error::unsignable_body)?,
};
sign_headers(
request.method().as_str(),
request.url().as_str(),
body,
&aws_signature_headers(request.headers()),
SigningScope {
region: &self.region,
service: self.service,
credentials: &self.credentials,
time: (self.clock)(),
},
)
.map(|signature| signature.into_iter().collect())
.map_err(|error| litellm_http::Error::Signature(error.to_string()))
.map_err(|error| match error {
Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error),
Error::Credentials { .. }
| Error::MissingWebIdentityCredentials
| Error::WebIdentityExpiration(_)
| Error::Signing { .. } => litellm_http::Error::signature(error),
})
}
}
#[cfg(test)]
mod tests {
use reqwest::{Body, Method, Url};
use serde_json::{Value, json as json_value};
use std::time::{Duration, UNIX_EPOCH};
use litellm_http::outbound::OutboundRequest;
use serde_json::{Value, json};
use litellm_http::outbound::{json, signed};
use super::*;
const URL: &str = "https://textract.us-east-1.amazonaws.com/";
fn fixed_clock() -> SystemTime {
UNIX_EPOCH + Duration::from_secs(1_700_000_000)
}
@ -93,82 +108,121 @@ mod tests {
.with_clock(fixed_clock)
}
fn target() -> (String, String) {
("X-Amz-Target".into(), "Textract.DetectDocumentText".into())
}
fn authorization(body: &Value, service: &'static str) -> String {
OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
body,
None,
signed(
json(URL, &[target()], body, None).unwrap(),
&signer(service),
)
.unwrap()
.header("Authorization")
.unwrap()
.to_string()
.headers()["authorization"]
.to_str()
.unwrap()
.to_string()
}
#[test]
fn the_signature_verifies_against_the_bytes_that_are_sent() {
let sent = OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
&json!({"Document": {"Bytes": "aGk="}}),
None,
let sent = signed(
json(
URL,
&[target()],
&json_value!({"Document": {"Bytes": "aGk="}}),
None,
)
.unwrap(),
&signer("textract"),
)
.unwrap();
let unsigned: BTreeMap<String, String> = sent
let unsigned: HeaderMap = sent
.headers()
.iter()
.filter(|(name, _)| !is_sigv4_computed_header(name))
.cloned()
.filter(|(name, _)| !is_sigv4_computed_header(name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect();
let recomputed = sign_post(
sent.url(),
sent.body(),
let recomputed = sign_headers(
"POST",
sent.url().as_str(),
sent.body().unwrap().as_bytes().unwrap(),
&aws_signature_headers(&unsigned),
"us-east-1",
"textract",
&Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
fixed_clock(),
SigningScope {
region: "us-east-1",
service: "textract",
credentials: &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
time: fixed_clock(),
},
)
.unwrap();
assert_eq!(
sent.header("Authorization"),
Some(recomputed["Authorization"].as_str())
);
assert_eq!(sent.headers()["authorization"], recomputed["authorization"]);
}
#[test]
fn the_signature_depends_on_the_body_and_the_service() {
let original = authorization(&json!({"text": "card 4111"}), "textract");
let original = authorization(&json_value!({"text": "card 4111"}), "textract");
assert_ne!(
original,
authorization(&json!({"text": "card [REDACTED]"}), "textract")
authorization(&json_value!({"text": "card [REDACTED]"}), "textract")
);
assert_ne!(
original,
authorization(&json!({"text": "card 4111"}), "bedrock")
authorization(&json_value!({"text": "card 4111"}), "bedrock")
);
assert!(original.contains("/us-east-1/textract/aws4_request"));
}
#[test]
fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() {
let error = OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("authorization".into(), "Bearer caller".into())],
&json!({}),
None,
let error = signed(
json(
URL,
&[("authorization".into(), "Bearer caller".into())],
&json_value!({}),
None,
)
.unwrap(),
&signer("textract"),
)
.unwrap_err();
assert_eq!(
error,
litellm_http::Error::ComputedHeader("authorization".into())
);
assert!(error.is_computed_header());
assert!(error.to_string().contains("authorization"), "{error}");
}
#[test]
fn a_forwarded_header_that_cannot_be_signed_is_the_callers_fault() {
let error = signed(
json(
URL,
&[("x-amz-meta-note".into(), "café".into())],
&json_value!({}),
None,
)
.unwrap(),
&signer("textract"),
)
.unwrap_err();
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains("x-amz-meta-note"), "{report}");
assert!(error.is_request());
assert!(!error.dispatched());
}
#[test]
fn a_streaming_body_cannot_be_signed() {
let mut request = Request::new(Method::POST, Url::parse(URL).unwrap());
*request.body_mut() = Some(Body::wrap_stream(futures_util::stream::once(async {
Ok::<_, std::io::Error>(bytes::Bytes::from_static(b"{}"))
})));
assert!(matches!(
signed(request, &signer("textract")),
Err(error) if !error.dispatched() && !error.is_request()
));
}
}

View file

@ -18,5 +18,6 @@ azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
tokio.workspace = true

View file

@ -12,7 +12,7 @@ use azure_identity::{
};
use sha2::{Digest, Sha256};
use litellm_auth_types::Error;
use litellm_auth_types::{Error, ErrorDetail};
use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced};
use super::credential_provider_cache::{
@ -133,7 +133,9 @@ impl NativeAzureTokenAcquirer {
let token = credential
.get_token(&[scope.as_str()], None)
.await
.map_err(|error| Error::CredentialAcquisition(error.to_string().into()))?;
.map_err(|error| {
Error::CredentialAcquisition(ErrorDetail::failed("Azure token request", error))
})?;
let expires_on = u64::try_from(token.expires_on.unix_timestamp())
.ok()
.map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds));
@ -250,11 +252,8 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
let Some(authority) = authority else {
return Ok(());
};
let url = url::Url::parse(authority.value()).map_err(|_| {
Error::InvalidConfiguration(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into(),
)
let url = url::Url::parse(authority.value()).map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed("Azure authority parsing", error))
})?;
if url.scheme() != "https"
|| url.host_str().is_none()
@ -264,10 +263,10 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
|| url.fragment().is_some()
|| !matches!(url.path(), "" | "/")
{
return Err(Error::InvalidConfiguration(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into(),
));
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
field: "Azure authority".into(),
expected: "an HTTPS origin without credentials, query, or fragment",
}));
}
Ok(())
}
@ -377,7 +376,7 @@ fn trusted_source(sources: &[InputSource]) -> InputSource {
fn mixed_sources<T>() -> Result<T, Error> {
Err(Error::InvalidConfiguration(
"request-controlled Azure auth inputs cannot be combined with host credentials".into(),
ErrorDetail::MixedInputSources { provider: "Azure" },
))
}
@ -444,7 +443,7 @@ fn build_credential(
.map(|credential| credential as Arc<dyn TokenCredential>),
}
.map_err(|error| {
Error::InvalidConfiguration(litellm_auth_types::ErrorDetail::failed(
Error::InvalidConfiguration(ErrorDetail::failed(
"Azure credential initialization",
error,
))
@ -662,12 +661,11 @@ mod tests {
))
.unwrap_err();
assert_eq!(
error,
litellm_auth_types::Error::InvalidConfiguration(
assert!(
litellm_error::Report::new(&error).to_string().contains(
"request-controlled Azure auth inputs cannot be combined with host credentials"
.into()
)
),
"{error:?}"
);
}
@ -695,12 +693,11 @@ mod tests {
authority,
))
.unwrap_err();
assert_eq!(
error,
litellm_auth_types::Error::InvalidConfiguration(
assert!(
litellm_error::Report::new(&error).to_string().contains(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into()
)
),
"{error:?}"
);
}
}

View file

@ -1,8 +1,8 @@
use litellm_auth_types::Error;
use litellm_auth_types::{
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
SecretValue, Sourced, TokenProviderHandle,
};
use litellm_auth_types::{Error, ErrorDetail};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
@ -91,9 +91,9 @@ impl AzureAuthService {
AzureCredentialPlan::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(Error::EmptyCallerCredential(
"Azure AD token provider returned an empty token",
));
return Err(Error::EmptyCallerCredential {
caller: "Azure AD token provider",
});
}
Ok(Some(Sourced::new(credential, InputSource::Deployment)))
}
@ -106,11 +106,9 @@ impl AzureAuthService {
} => {
let assertion = resolve_reference(inputs, env_lookup, reference.value())
.await?
.ok_or_else(|| {
Error::CredentialAcquisition(
"Azure OIDC reference did not resolve to a value".into(),
)
})?;
.ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved {
subject: "Azure OIDC reference",
}))?;
let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion {
tenant_id,
client_id,
@ -142,7 +140,13 @@ impl AzureAuthService {
Err(error) => failures.push(error),
}
}
Err(Error::CredentialChain(failures))
match failures.pop() {
Some(last) => Err(Error::CredentialChain {
last: Box::new(last),
earlier: failures,
}),
None => Ok(None),
}
}
AzureCredentialPlan::Missing => Ok(None),
}
@ -173,7 +177,12 @@ pub(crate) fn select_auth_plan(
.map(|selector| Sourced::new(selector, value.source()))
})
.transpose()
.map_err(|_| Error::InvalidConfiguration("invalid Azure credential selector".into()))?;
.map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed(
"Azure credential selector parsing",
error,
))
})?;
let federated_token_file = configured_string(
&inputs.federated_token_file,
AZURE_FEDERATED_TOKEN_FILE_ENV,
@ -263,9 +272,12 @@ fn select_native_plan(
let selection_source = selected.source();
match selected.into_value() {
AzureCredentialType::ClientSecretCredential => Err(Error::InvalidConfiguration(
"ClientSecretCredential requires tenant_id, client_id, and client_secret".into(),
)),
AzureCredentialType::ClientSecretCredential => {
Err(Error::InvalidConfiguration(ErrorDetail::MissingSettings {
credential: "ClientSecretCredential",
settings: "tenant_id, client_id, and client_secret",
}))
}
AzureCredentialType::WorkloadIdentityCredential => {
Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
workload_request(tenant_id, client_id, federated_token_file, scope, authority)?,
@ -349,22 +361,21 @@ fn workload_request(
authority: Option<Sourced<String>>,
) -> Result<NativeAzureRequest, Error> {
Ok(NativeAzureRequest::WorkloadIdentity {
tenant_id: tenant_id.ok_or_else(|| {
Error::InvalidConfiguration("WorkloadIdentityCredential requires tenant_id".into())
})?,
client_id: client_id.ok_or_else(|| {
Error::InvalidConfiguration("WorkloadIdentityCredential requires client_id".into())
})?,
token_file_path: token_file_path.ok_or_else(|| {
Error::InvalidConfiguration(
"WorkloadIdentityCredential requires azure_federated_token_file".into(),
)
})?,
tenant_id: tenant_id.ok_or(workload_requires("tenant_id"))?,
client_id: client_id.ok_or(workload_requires("client_id"))?,
token_file_path: token_file_path.ok_or(workload_requires("azure_federated_token_file"))?,
scope,
authority,
})
}
fn workload_requires(settings: &'static str) -> Error {
Error::InvalidConfiguration(ErrorDetail::MissingSettings {
credential: "WorkloadIdentityCredential",
settings,
})
}
fn configured_string(
configured: &ConfigValue<String>,
environment_name: &str,
@ -410,11 +421,10 @@ async fn resolve_reference(
.map_or(CredentialLookup::Missing, CredentialLookup::Found),
CredentialRef::None => return Ok(None),
CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => {
let resolver = inputs.credential_resolver.as_ref().ok_or_else(|| {
Error::InvalidConfiguration(
"credential reference requires a host credential resolver".into(),
)
})?;
let resolver = inputs
.credential_resolver
.as_ref()
.ok_or(Error::InvalidConfiguration(ErrorDetail::MissingResolver))?;
resolver.resolve(reference).await?
}
};
@ -433,7 +443,7 @@ fn oidc_reference(
let value = token.value().expose();
if token.source() == InputSource::Request && value.starts_with("oidc/") {
return Err(Error::InvalidConfiguration(
"request-controlled Azure credential references are not allowed".into(),
ErrorDetail::RequestControlledReference { provider: "Azure" },
));
}
if let Some(name) = value.strip_prefix("oidc/env/") {
@ -457,19 +467,17 @@ fn oidc_reference(
}
if value.starts_with("oidc/") {
return Err(Error::InvalidConfiguration(
"unsupported OIDC reference".into(),
ErrorDetail::UnsupportedReference,
));
}
Ok(None)
}
fn non_empty_reference(value: &str, kind: &str) -> Result<String, Error> {
fn non_empty_reference(value: &str, kind: &'static str) -> Result<String, Error> {
if value.is_empty() {
return Err(Error::InvalidConfiguration(
litellm_auth_types::ErrorDetail::Empty {
subject: kind.into(),
},
));
return Err(Error::InvalidConfiguration(ErrorDetail::Empty {
subject: kind,
}));
}
Ok(value.to_string())
}
@ -518,7 +526,9 @@ mod tests {
expires_on: None,
})
} else {
Err(Error::CredentialAcquisition(kind.into()))
Err(Error::CredentialAcquisition(
litellm_auth_types::ErrorDetail::Unresolved { subject: kind },
))
}
})
}
@ -624,7 +634,11 @@ mod tests {
)))
.expect_err("unsupported backend must fail validation");
assert!(error.to_string().contains("unsupported OIDC reference"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("unsupported OIDC reference")
);
}
#[rstest::rstest]
@ -649,11 +663,11 @@ mod tests {
})
.unwrap_err();
assert_eq!(
error,
Error::InvalidConfiguration(
"request-controlled Azure credential references are not allowed".into()
)
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("request-controlled Azure credential references are not allowed"),
"{error:?}"
);
}
@ -715,7 +729,16 @@ mod tests {
.await
.unwrap_err();
assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
assert!(matches!(&error, Error::CredentialChain { earlier, .. } if earlier.len() == 1));
let report = litellm_error::Report::new(&error).to_string();
assert!(
report.contains("managed-identity did not resolve"),
"{report}"
);
assert!(
report.contains("developer-tools did not resolve"),
"{report}"
);
}
#[derive(Debug)]
@ -763,7 +786,9 @@ mod tests {
assert_eq!(
error,
Error::EmptyCallerCredential("Azure AD token provider returned an empty token")
Error::EmptyCallerCredential {
caller: "Azure AD token provider"
}
);
}
}

View file

@ -24,5 +24,6 @@ google-cloud-auth = { workspace = true, optional = true }
http = { workspace = true, optional = true }
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
tempfile.workspace = true

View file

@ -73,13 +73,17 @@ impl VertexAuth {
.or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV));
let project_id = get_vertex_ai_project(config, env_lookup);
if !has_authorization && static_token.is_none() {
let access = self.get_access_token(config, env_lookup).await?;
return Ok(VertexEnvironment {
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
project_id: project_id.unwrap_or(access.project_id),
});
}
let static_token = match (has_authorization, static_token) {
(true, _) => None,
(false, Some(token)) => Some(token),
(false, None) => {
let access = self.get_access_token(config, env_lookup).await?;
return Ok(VertexEnvironment {
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
project_id: project_id.unwrap_or(access.project_id),
});
}
};
let project_id = match project_id {
Some(project_id) => project_id,
@ -90,14 +94,9 @@ impl VertexAuth {
.await?
}
};
let headers = if has_authorization {
headers
} else {
apply_credential(
headers,
static_token.as_deref().expect("static token was checked"),
CredentialPlacement::Bearer,
)?
let headers = match static_token {
Some(token) => apply_credential(headers, &token, CredentialPlacement::Bearer)?,
None => headers,
};
Ok(VertexEnvironment {
headers,
@ -234,14 +233,13 @@ enum CredentialCacheKey {
Adc,
}
const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
Set it to a readable file path, or to the credentials JSON itself.";
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
if is_base64_encoded_json(configured) {
return Err(Error::InvalidConfiguration(ErrorDetail::Message(
BASE64_CREDENTIALS.into(),
)));
return Err(Error::InvalidConfiguration(
ErrorDetail::Base64EncodedJson {
field: "vertex_credentials",
},
));
}
if names_a_credentials_file(configured) {
CustomServiceAccount::from_file(configured)
@ -299,7 +297,8 @@ mod tests {
else {
panic!("{configured:?} should not load");
};
assert!(error.to_string().contains(expected), "{error}");
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains(expected), "{report}");
}
#[rstest::rstest]
@ -321,7 +320,7 @@ mod tests {
.await
{
Ok(_) => panic!("{configured:?} should not load"),
Err(error) => error.to_string(),
Err(error) => litellm_error::Report::new(&error).to_string(),
}
}

View file

@ -55,7 +55,7 @@ impl CredentialsProvider for GoogleCredentials {
entity_tag: EntityTag::new(),
data,
})
.map_err(|_| CredentialsError::from_msg(false, "Google authentication failed"))
.map_err(|error| CredentialsError::from_source(false, error))
}
async fn universe_domain(&self) -> Option<String> {
@ -102,5 +102,9 @@ mod tests {
);
let error = credentials.headers(Extensions::new()).await.unwrap_err();
assert!(!format!("{error:?}").contains("private"));
assert_eq!(
std::error::Error::source(&error).and_then(|cause| cause.downcast_ref::<Error>()),
Some(&Error::InvalidHeader)
);
}
}

View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-error.workspace = true
serde.workspace = true
serde_json.workspace = true
subtle.workspace = true

View file

@ -86,9 +86,9 @@ impl CredentialPlan {
Self::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(Error::EmptyCallerCredential(
"credential caller returned an empty credential",
));
return Err(Error::EmptyCallerCredential {
caller: "credential caller",
});
}
Ok(CredentialPlanResolution::Resolved(credential))
}
@ -104,8 +104,8 @@ mod tests {
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
CredentialRef, CredentialResolver, CredentialResolverHandle,
};
use crate::Error;
use crate::SecretValue;
use crate::{Error, ErrorDetail};
#[derive(Debug)]
struct HostResolver;
@ -150,9 +150,10 @@ mod tests {
impl CredentialResolver for FailingResolver {
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async {
Err(Error::CredentialAcquisition(
"host credential lookup failed".into(),
))
Err(Error::CredentialAcquisition(ErrorDetail::failed(
"host credential lookup",
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
)))
})
}
}
@ -168,9 +169,14 @@ mod tests {
.await
.expect_err("acquisition errors cannot become fallback");
assert_eq!(
error,
Error::CredentialAcquisition("host credential lookup failed".into())
assert!(
matches!(error, Error::CredentialAcquisition(_)),
"{error:?}"
);
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("host credential lookup failed")
);
}
}

View file

@ -1,17 +1,23 @@
use thiserror::Error as ThisError;
use litellm_error::{Report, Source};
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
use crate::{ParamKind, ParamSpec};
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("invalid authentication configuration: {0}")]
#[error("invalid authentication configuration")]
InvalidConfiguration(#[source] ErrorDetail),
#[error("credential acquisition failed: {0}")]
#[error("credential acquisition failed")]
CredentialAcquisition(#[source] ErrorDetail),
#[error("credential caller failed: {0}")]
EmptyCallerCredential(&'static str),
#[error("{0}")]
ProviderAuthentication(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
CredentialChain(Vec<Error>),
#[error("{caller} returned an empty credential")]
EmptyCallerCredential { caller: &'static str },
#[error("provider authentication failed")]
Provider(#[source] Source),
#[error("no credential in the chain succeeded{}", earlier_attempts(earlier))]
CredentialChain {
#[source]
last: Box<Error>,
earlier: Vec<Error>,
},
#[error(
"Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable"
)]
@ -22,7 +28,7 @@ pub enum Error {
#[error("Missing {provider} {} - {}", .spec.setting, .spec.guidance())]
MissingParam {
provider: &'static str,
spec: &'static crate::ParamSpec,
spec: &'static ParamSpec,
},
#[error("Missing {provider} API Base - {guidance}")]
MissingApiBase {
@ -33,78 +39,282 @@ pub enum Error {
InvalidHeader,
}
fn earlier_attempts(earlier: &[Error]) -> String {
if earlier.is_empty() {
return String::new();
}
let attempts = earlier
.iter()
.map(|failure| Report::new(failure).to_string())
.collect::<Vec<_>>()
.join("; ");
format!(" (earlier attempts: {attempts})")
}
impl Error {
pub fn provider(cause: impl std::error::Error + Send + Sync + 'static) -> Self {
Self::Provider(Source::new(cause))
}
pub fn is_missing_credential(&self) -> bool {
match self {
Self::MissingApiKey { .. } => true,
Self::MissingParam { spec, .. } => match spec.kind {
ParamKind::Credential => true,
ParamKind::Setting => false,
},
Self::InvalidConfiguration(_)
| Self::CredentialAcquisition(_)
| Self::EmptyCallerCredential { .. }
| Self::Provider(_)
| Self::CredentialChain { .. }
| Self::MissingApiBase { .. }
| Self::InvalidHeader => false,
}
}
pub fn is_request(&self) -> bool {
match self {
Self::InvalidConfiguration(_)
| Self::EmptyCallerCredential { .. }
| Self::InvalidHeader => true,
Self::CredentialAcquisition(_) | Self::Provider(_) | Self::CredentialChain { .. } => {
false
}
Self::MissingParam { spec, .. } => match spec.kind {
ParamKind::Setting => true,
ParamKind::Credential => false,
},
Self::MissingApiKey { .. } | Self::MissingApiBase { .. } => false,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum ErrorDetail {
#[error("{0}")]
Message(String),
#[error("{field} must be {expected}")]
InvalidType {
field: String,
expected: &'static str,
},
#[error("{subject} cannot be empty")]
Empty { subject: String },
Empty { subject: &'static str },
#[error("credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("{operation} failed: {source}")]
#[error("credential plan is not allowed by the provider auth policy")]
PlanNotAllowed,
#[error("{credential} requires {settings}")]
MissingSettings {
credential: &'static str,
settings: &'static str,
},
#[error("credentials cannot be sent to a request-controlled {provider} endpoint")]
RequestControlledEndpoint { provider: &'static str },
#[error("request-controlled {provider} credential references are not allowed")]
RequestControlledReference { provider: &'static str },
#[error("request-controlled {provider} auth inputs cannot be combined with host credentials")]
MixedInputSources { provider: &'static str },
#[error("unsupported OIDC reference")]
UnsupportedReference,
#[error("credential reference requires a host credential resolver")]
MissingResolver,
#[error("{subject} did not resolve to a value")]
Unresolved { subject: &'static str },
#[error(
"The `{field}` value is base64-encoded JSON, not the JSON itself. Set it to a readable file path, or to the credentials JSON itself."
)]
Base64EncodedJson { field: &'static str },
#[error("{operation} failed")]
Failed {
operation: &'static str,
#[source]
source: ErrorSource,
cause: Source,
},
}
impl ErrorDetail {
pub fn failed(
operation: &'static str,
source: impl std::error::Error + Send + Sync + 'static,
cause: impl std::error::Error + Send + Sync + 'static,
) -> Self {
Self::Failed {
operation,
source: ErrorSource::new(source),
cause: Source::new(cause),
}
}
}
impl From<String> for ErrorDetail {
fn from(message: String) -> Self {
Self::Message(message)
#[cfg(test)]
mod tests {
use std::io;
use litellm_error::{Report, Source};
use rstest::rstest;
use super::{Error, ErrorDetail};
use crate::{AwsParams, VertexParams};
#[rstest]
#[case::api_key(
Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" },
"Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable"
)]
#[case::another_api_key(
Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" },
"Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable"
)]
#[case::api_base(
Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" },
"Missing Example API Base - Pass api_base"
)]
#[case::another_api_base(
Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" },
"Missing Custom API Base - Set CUSTOM_ENDPOINT"
)]
fn missing_inputs_keep_the_python_text(#[case] error: Error, #[case] expected: &str) {
assert_eq!(error.to_string(), expected);
}
#[rstest]
#[case::configuration(
Error::InvalidConfiguration(ErrorDetail::Empty { subject: "credential selector" }),
"invalid authentication configuration: credential selector cannot be empty"
)]
#[case::acquisition(
Error::CredentialAcquisition(ErrorDetail::Unresolved { subject: "OIDC reference" }),
"credential acquisition failed: OIDC reference did not resolve to a value"
)]
#[case::empty_caller(
Error::EmptyCallerCredential { caller: "token provider" },
"token provider returned an empty credential"
)]
fn the_report_names_every_layer_in_order(#[case] error: Error, #[case] expected: &str) {
assert!(
Report::new(&error).to_string().contains(expected),
"{}",
Report::new(&error)
);
}
#[rstest]
#[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed(
"test credential",
denied()
)))]
#[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed(
"test credential",
denied()
)))]
fn a_failure_keeps_its_cause_as_the_source_and_renders_it_once(#[case] error: Error) {
let cause =
std::iter::successors(std::error::Error::source(&error), |cause| cause.source())
.find_map(|cause| cause.downcast_ref::<Source>()?.downcast_ref::<io::Error>())
.expect("the io cause is in the source chain");
assert_eq!(cause.kind(), io::ErrorKind::PermissionDenied);
let report = Report::new(&error).to_string();
assert_eq!(report.matches("access denied").count(), 1, "{report}");
}
#[rstest]
#[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed(
"test credential",
denied()
)))]
#[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed(
"test credential",
denied()
)))]
fn a_contextual_failure_leaves_its_cause_out_of_its_own_message(#[case] error: Error) {
assert!(!error.to_string().contains("access denied"), "{error}");
assert!(
Report::new(&error)
.to_string()
.contains("test credential failed: access denied")
);
}
#[test]
fn a_provider_failure_keeps_the_provider_error_as_its_source() {
let error = Error::provider(denied());
assert_eq!(
litellm_error::find::<io::Error>(&error).map(io::Error::kind),
Some(io::ErrorKind::PermissionDenied)
);
assert!(!error.to_string().contains("access denied"), "{error}");
let report = Report::new(&error).to_string();
assert_eq!(report.matches("access denied").count(), 1, "{report}");
}
#[test]
fn a_chain_keeps_its_last_failure_as_the_source_and_renders_every_attempt_once() {
let error = Error::CredentialChain {
last: Box::new(Error::provider(denied())),
earlier: vec![
Error::CredentialAcquisition(ErrorDetail::Unresolved {
subject: "OIDC reference",
}),
Error::EmptyCallerCredential {
caller: "token provider",
},
],
};
assert_eq!(
litellm_error::find::<io::Error>(&error).map(io::Error::kind),
Some(io::ErrorKind::PermissionDenied)
);
let report = Report::new(&error).to_string();
assert_eq!(report.matches("access denied").count(), 1, "{report}");
assert_eq!(
report.matches("OIDC reference did not resolve").count(),
1,
"{report}"
);
assert_eq!(
report
.matches("token provider returned an empty credential")
.count(),
1,
"{report}"
);
}
fn denied() -> io::Error {
io::Error::new(io::ErrorKind::PermissionDenied, "access denied")
}
#[rstest]
#[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, true)]
#[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), true)]
#[case::aws_secret_access_key(AwsParams::SECRET_ACCESS_KEY.missing("AWS"), true)]
#[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), false)]
#[case::aws_region(AwsParams::REGION.missing("AWS"), false)]
#[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)]
#[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), false)]
fn only_an_absent_credential_is_a_missing_credential(
#[case] error: Error,
#[case] expected: bool,
) {
assert_eq!(error.is_missing_credential(), expected);
}
#[rstest]
#[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), true)]
#[case::aws_region(AwsParams::REGION.missing("AWS"), true)]
#[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)]
#[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), false)]
#[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, false)]
#[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), true)]
#[case::invalid_header(Error::InvalidHeader, true)]
#[case::identity_provider_outage(
Error::CredentialAcquisition(ErrorDetail::MissingResolver),
false
)]
#[case::provider_sdk_failure(Error::provider(denied()), false)]
fn a_setting_the_caller_can_pass_is_theirs_to_fix_but_credentials_and_api_bases_are_not(
#[case] error: Error,
#[case] expected: bool,
) {
assert_eq!(error.is_request(), expected);
}
}
impl From<&str> for ErrorDetail {
fn from(message: &str) -> Self {
Self::Message(message.into())
}
}
#[derive(Clone, Debug)]
pub struct ErrorSource(std::sync::Arc<dyn std::error::Error + Send + Sync>);
impl std::ops::Deref for ErrorSource {
type Target = dyn std::error::Error + Send + Sync;
fn deref(&self) -> &Self::Target {
self.0.as_ref()
}
}
impl std::fmt::Display for ErrorSource {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(&self.0, formatter)
}
}
impl ErrorSource {
pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self {
Self(std::sync::Arc::new(error))
}
}
impl PartialEq for ErrorSource {
fn eq(&self, other: &Self) -> bool {
std::sync::Arc::ptr_eq(&self.0, &other.0)
}
}
impl Eq for ErrorSource {}

View file

@ -21,9 +21,9 @@ pub fn apply_credential(
placement: CredentialPlacement,
) -> Result<Vec<(String, String)>, Error> {
if credential.trim().is_empty() {
return Err(Error::InvalidConfiguration(
"credential cannot be empty".into(),
));
return Err(Error::InvalidConfiguration(crate::ErrorDetail::Empty {
subject: "credential",
}));
}
if headers
.iter()
@ -71,6 +71,10 @@ mod tests {
)
.expect_err("provider policy must handle existing credentials");
assert!(error.to_string().contains("already exists"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("already exists")
);
}
}

View file

@ -51,9 +51,9 @@ pub use credential::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
};
pub use error::{Error, ErrorDetail, ErrorSource};
pub use error::{Error, ErrorDetail};
pub use http::CredentialPlacement;
pub use params::{AwsParams, ParamSpec, VertexParams};
pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams};
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
pub use secret::SecretValue;
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};

View file

@ -4,15 +4,17 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use veil::Redact;
use super::ParamSpec;
use super::{ParamKind, ParamSpec};
const fn spec(
setting: &'static str,
kind: ParamKind,
wire: &'static [&'static str],
env: &'static [&'static str],
) -> ParamSpec {
ParamSpec {
setting,
kind,
wire,
module_global: None,
env,
@ -53,40 +55,67 @@ pub struct AwsParams {
impl AwsParams {
pub const ACCESS_KEY_ID: ParamSpec = spec(
"access key id",
ParamKind::Credential,
&["aws_access_key_id"],
&["AWS_ACCESS_KEY_ID"],
);
pub const SECRET_ACCESS_KEY: ParamSpec = spec(
"secret access key",
ParamKind::Credential,
&["aws_secret_access_key"],
&["AWS_SECRET_ACCESS_KEY"],
);
pub const SESSION_TOKEN: ParamSpec = spec(
"session token",
ParamKind::Credential,
&["aws_session_token"],
&["AWS_SESSION_TOKEN"],
);
pub const REGION: ParamSpec = spec(
"region",
ParamKind::Setting,
&["aws_region_name"],
&["AWS_REGION_NAME", "AWS_REGION"],
);
pub const SESSION_NAME: ParamSpec =
spec("session name", &["aws_session_name"], &["AWS_SESSION_NAME"]);
pub const PROFILE_NAME: ParamSpec =
spec("profile name", &["aws_profile_name"], &["AWS_PROFILE_NAME"]);
pub const ROLE_NAME: ParamSpec = spec("role name", &["aws_role_name"], &["AWS_ROLE_NAME"]);
pub const SESSION_NAME: ParamSpec = spec(
"session name",
ParamKind::Setting,
&["aws_session_name"],
&["AWS_SESSION_NAME"],
);
pub const PROFILE_NAME: ParamSpec = spec(
"profile name",
ParamKind::Setting,
&["aws_profile_name"],
&["AWS_PROFILE_NAME"],
);
pub const ROLE_NAME: ParamSpec = spec(
"role name",
ParamKind::Setting,
&["aws_role_name"],
&["AWS_ROLE_NAME"],
);
pub const WEB_IDENTITY_TOKEN: ParamSpec = spec(
"web identity token",
ParamKind::Credential,
&["aws_web_identity_token"],
&["AWS_WEB_IDENTITY_TOKEN"],
);
pub const STS_ENDPOINT: ParamSpec =
spec("STS endpoint", &["aws_sts_endpoint"], &["AWS_STS_ENDPOINT"]);
pub const EXTERNAL_ID: ParamSpec =
spec("external id", &["aws_external_id"], &["AWS_EXTERNAL_ID"]);
pub const STS_ENDPOINT: ParamSpec = spec(
"STS endpoint",
ParamKind::Setting,
&["aws_sts_endpoint"],
&["AWS_STS_ENDPOINT"],
);
pub const EXTERNAL_ID: ParamSpec = spec(
"external id",
ParamKind::Setting,
&["aws_external_id"],
&["AWS_EXTERNAL_ID"],
);
pub const BEDROCK_RUNTIME_ENDPOINT: ParamSpec = spec(
"Bedrock runtime endpoint",
ParamKind::Setting,
&["aws_bedrock_runtime_endpoint"],
&["AWS_BEDROCK_RUNTIME_ENDPOINT"],
);

View file

@ -10,11 +10,18 @@ pub use vertex::VertexParams;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ParamSpec {
pub setting: &'static str,
pub kind: ParamKind,
pub wire: &'static [&'static str],
pub module_global: Option<&'static str>,
pub env: &'static [&'static str],
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ParamKind {
Credential,
Setting,
}
impl ParamSpec {
/// The first non-blank value: each wire name through `params`, then each environment
/// name through `env`.

View file

@ -2,7 +2,7 @@ use serde::{Deserialize, Deserializer, Serialize, de::Unexpected};
use serde_json::Value;
use veil::Redact;
use super::ParamSpec;
use super::{ParamKind, ParamSpec};
/// The Vertex AI fields of Python's `GenericLiteLLMParams`, both the current names and the
/// `vertex_ai_*` spellings `VertexBase.safe_get_vertex_ai_*` still read.
@ -27,18 +27,21 @@ pub struct VertexParams {
impl VertexParams {
pub const CREDENTIALS: ParamSpec = ParamSpec {
setting: "credentials",
kind: ParamKind::Credential,
wire: &["vertex_credentials", "vertex_ai_credentials"],
module_global: None,
env: &["VERTEXAI_CREDENTIALS"],
};
pub const PROJECT: ParamSpec = ParamSpec {
setting: "project",
kind: ParamKind::Setting,
wire: &["vertex_project", "vertex_ai_project"],
module_global: Some("vertex_project"),
env: &["VERTEXAI_PROJECT"],
};
pub const LOCATION: ParamSpec = ParamSpec {
setting: "location",
kind: ParamKind::Setting,
wire: &["vertex_location", "vertex_ai_location"],
module_global: Some("vertex_location"),
env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"],

View file

@ -1,5 +1,5 @@
use crate::http::apply_credential;
use crate::{CredentialPlacement, Error, ResolvedCredential};
use crate::{CredentialPlacement, Error, ErrorDetail, ResolvedCredential};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlanKind {
@ -31,11 +31,18 @@ pub struct ProviderAuthPolicy {
impl ProviderAuthPolicy {
pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, _)| {
self.accepted_existing_headers
.iter()
.any(|accepted| name.eq_ignore_ascii_case(accepted))
})
self.existing_credential_header(headers).is_some()
}
fn existing_credential_header(&self, headers: &[(String, String)]) -> Option<&'static str> {
self.accepted_existing_headers
.iter()
.copied()
.find(|accepted| {
headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case(accepted))
})
}
pub fn apply(
@ -44,11 +51,11 @@ impl ProviderAuthPolicy {
kind: CredentialPlanKind,
credential: &ResolvedCredential,
) -> Result<Vec<(String, String)>, Error> {
if self.has_existing_credential(&headers) {
if let Some(existing) = self.existing_credential_header(&headers) {
return match self.existing_header_behavior {
ExistingHeaderBehavior::Preserve => Ok(headers),
ExistingHeaderBehavior::Reject => Err(Error::InvalidConfiguration(
"credential header already exists".into(),
ErrorDetail::DuplicateHeader(existing),
)),
};
}
@ -56,11 +63,7 @@ impl ProviderAuthPolicy {
.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or_else(|| {
Error::InvalidConfiguration(
"credential plan is not allowed by the provider auth policy".into(),
)
})?;
.ok_or(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed))?;
apply_credential(headers, credential.secret().expose(), rule.placement)
}
}
@ -108,6 +111,10 @@ mod tests {
)
.unwrap_err();
assert!(error.to_string().contains("not allowed"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("not allowed")
);
}
}

View file

@ -1,74 +0,0 @@
use litellm_auth_types::Error;
use rstest::rstest;
#[rstest]
#[case::api_key(
Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" },
"Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable"
)]
#[case::another_api_key(
Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" },
"Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable"
)]
#[case::api_base(
Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" },
"Missing Example API Base - Pass api_base"
)]
#[case::another_api_base(
Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" },
"Missing Custom API Base - Set CUSTOM_ENDPOINT"
)]
#[case::configuration(
Error::InvalidConfiguration("credential selector is invalid".into()),
"invalid authentication configuration: credential selector is invalid"
)]
#[case::acquisition(
Error::CredentialAcquisition("token expired".into()),
"credential acquisition failed: token expired"
)]
#[case::caller(
Error::EmptyCallerCredential("empty token"),
"credential caller failed: empty token"
)]
#[case::provider(
Error::ProviderAuthentication("provider rejected credentials".into()),
"provider rejected credentials"
)]
#[case::chain(
Error::CredentialChain(vec![
Error::CredentialAcquisition("token expired".into()),
Error::EmptyCallerCredential("empty token"),
]),
"credential acquisition failed: credential acquisition failed: token expired; credential caller failed: empty token"
)]
fn display_preserves_failure_phase_and_caller_context(
#[case] error: Error,
#[case] expected: &str,
) {
assert_eq!(error.to_string(), expected);
}
#[rstest]
#[case::configuration(true)]
#[case::acquisition(false)]
fn contextual_failures_keep_the_original_source(#[case] configuration: bool) {
use litellm_auth_types::ErrorDetail;
let detail = ErrorDetail::failed(
"test credential",
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
);
let error = if configuration {
Error::InvalidConfiguration(detail)
} else {
Error::CredentialAcquisition(detail)
};
let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| {
error.source()
})
.find_map(|error| error.downcast_ref::<std::io::Error>())
.expect("the original credential error remains available");
assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied);
assert!(error.to_string().contains("test credential failed:"));
assert!(error.to_string().ends_with(&source.to_string()));
}

View file

@ -1,8 +1,9 @@
use litellm_auth_types::{AwsParams, ParamSpec};
use litellm_auth_types::{AwsParams, ParamKind, ParamSpec};
use rstest::rstest;
const LOCATION: ParamSpec = ParamSpec {
setting: "location",
kind: ParamKind::Setting,
wire: &["vertex_location", "vertex_ai_location"],
module_global: Some("vertex_location"),
env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"],

View file

@ -1,5 +1,4 @@
use std::{
error::Error as StdError,
future::{Future, poll_fn},
sync::{
Arc,
@ -53,7 +52,7 @@ async fn callbacks_acquire_fresh_credentials_on_demand(#[case] access_token: boo
#[rstest]
#[tokio::test]
async fn callback_errors_preserve_the_original_source() {
async fn callback_errors_keep_their_typed_cause() {
let provider = TokenProviderHandle::from_callback(|| async {
Err(Error::CredentialAcquisition(ErrorDetail::failed(
"caller credential",
@ -62,13 +61,15 @@ async fn callback_errors_preserve_the_original_source() {
});
let error = provider.acquire().await.unwrap_err();
assert!(matches!(error, Error::CredentialAcquisition(_)));
let source = std::iter::successors(Some(&error as &(dyn StdError + 'static)), |error| {
(*error).source()
})
.find_map(|error| error.downcast_ref::<std::io::Error>())
.unwrap();
assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied);
let Error::CredentialAcquisition(ErrorDetail::Failed { cause, .. }) = &error else {
panic!("the callback's error passes through unchanged: {error:?}");
};
assert_eq!(
cause
.downcast_ref::<std::io::Error>()
.map(std::io::Error::kind),
Some(std::io::ErrorKind::PermissionDenied)
);
}
struct Release(Arc<AtomicBool>);

View file

@ -16,6 +16,7 @@ azure_core = "1.1.0"
azure_storage_blob = "1.1.0"
futures-util.workspace = true
reqwest.workspace = true
thiserror.workspace = true
tokio.workspace = true
url.workspace = true

View file

@ -18,6 +18,10 @@ use url::Url;
use crate::{credential::AzureBlobCredential, transport::ReqwestTransport};
#[derive(Debug, thiserror::Error)]
#[error("account URL cannot hold a container path")]
struct NotABaseUrl;
pub struct AzureBlobCache<C> {
container: BlobContainerClient,
codec: C,
@ -57,12 +61,12 @@ impl<C: CacheCodec> AzureBlobCache<C> {
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?;
let parsed = Url::parse(account_url).map_err(Error::unavailable)?;
let account_url = parsed.as_str().trim_end_matches('/').to_string();
let container_url = {
let mut url = parsed;
url.path_segments_mut()
.map_err(|()| Error::Unavailable)?
.map_err(|()| Error::unavailable(NotABaseUrl))?
.pop_if_empty()
.push(container);
url
@ -75,7 +79,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
..BlobContainerClientOptions::default()
}),
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let cache = Self {
container: client,
codec,
@ -101,7 +105,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => {
Ok(())
}
Err(_) => Err(Error::Unavailable),
Err(error) => Err(Error::unavailable(error)),
}
}
@ -116,7 +120,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
{
Ok(_) => Ok(()),
Err(error) if !overwrite && is_already_present(&error) => Ok(()),
Err(_) => Err(Error::Unavailable),
Err(error) => Err(Error::unavailable(error)),
}
}
@ -126,13 +130,9 @@ impl<C: CacheCodec> AzureBlobCache<C> {
Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => {
return Ok(None);
}
Err(_) => return Err(Error::Unavailable),
Err(error) => return Err(Error::unavailable(error)),
};
let bytes = response
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?;
let bytes = response.body.collect().await.map_err(Error::unavailable)?;
self.codec.decode(&bytes).map(Some)
}
@ -140,16 +140,16 @@ impl<C: CacheCodec> AzureBlobCache<C> {
let mut pages = self
.container
.list_blobs(None)
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.into_pages();
while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? {
let page = page.into_model().map_err(|_| Error::Unavailable)?;
while let Some(page) = pages.try_next().await.map_err(Error::unavailable)? {
let page = page.into_model().map_err(Error::unavailable)?;
for name in page.blob_items.into_iter().filter_map(|item| item.name) {
self.container
.blob_client(&name)
.delete(None)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
}
Ok(())

View file

@ -17,6 +17,10 @@ const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
type EnvLookup = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
#[derive(Debug, thiserror::Error)]
#[error("no Azure credential is available for blob storage")]
struct NoCredential;
pub struct AzureBlobCredential {
service: AzureAuthService,
env_lookup: EnvLookup,
@ -62,15 +66,8 @@ impl TokenCredential for AzureBlobCredential {
&lookup,
)
.await
.map_err(|error| {
azure_core::Error::with_message(ErrorKind::Credential, error.to_string())
})?
.ok_or_else(|| {
azure_core::Error::with_message(
ErrorKind::Credential,
"no Azure credential is available for blob storage",
)
})?;
.map_err(|error| azure_core::Error::new(ErrorKind::Credential, error))?
.ok_or_else(|| azure_core::Error::new(ErrorKind::Credential, NoCredential))?;
let (token, expires_on) = match credential.into_value() {
ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on),
ResolvedCredential::Static(token) => (token, None),

View file

@ -24,6 +24,7 @@ impl HttpClient for ReqwestTransport {
Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)),
};
let response = outgoing.send().await.map_err(|error| {
let error = litellm_http::Error::from(error);
let kind = if error.is_connect() {
ErrorKind::Connection
} else {
@ -41,9 +42,9 @@ impl HttpClient for ReqwestTransport {
);
}
}
let body = response
.bytes_stream()
.map_err(|error| azure_core::Error::new(ErrorKind::Io, error));
let body = response.bytes_stream().map_err(|error| {
azure_core::Error::new(ErrorKind::Io, litellm_http::Error::from(error))
});
Ok(AsyncRawResponse::new(status, headers, Box::pin(body)))
}
}

View file

@ -128,7 +128,7 @@ fn connect_surfaces_service_failures() {
ResponseCacheCodec,
runtime.handle().clone(),
));
assert!(matches!(result, Err(Error::Unavailable)));
assert!(matches!(result, Err(Error::Unavailable(_))));
}
#[rstest]
@ -263,7 +263,7 @@ fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[ca
fixture.service.seed_blob(key, bytes);
assert!(matches!(
fixture.cache.get_cache(key, &no_ttl()),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(_))
));
}
@ -378,15 +378,15 @@ fn service_failures_map_to_unavailable(fixture: Fixture) {
fixture.service.set_failing(true);
assert!(matches!(
fixture.cache.get_cache("key", &no_ttl()),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.cache.flush_cache(),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.runtime.block_on(
@ -394,7 +394,7 @@ fn service_failures_map_to_unavailable(fixture: Fixture) {
.cache
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
}

View file

@ -31,7 +31,7 @@ async fn connect(server: &MockServer) -> AzureBlobCache<JsonCodec<Value>> {
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(ReqwestTransport(
litellm_http::Client::plain_for_test(),
litellm_http::Client::for_test(),
)))),
..ClientOptions::default()
},

View file

@ -12,6 +12,7 @@ rand.workspace = true
rusqlite = { version = "0.39", features = ["bundled"] }
serde-pickle = "1.2"
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]

View file

@ -64,7 +64,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DiskCache<S, D, A> {
{
tokio::task::spawn_blocking(move || operation(&store))
.await
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
}
}
@ -162,7 +162,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D
.map(|key| match self.get_cache(key, context) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
@ -186,7 +186,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D
Some(value) => match self.decode_stored(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
},
})
@ -272,7 +272,7 @@ fn increment<A: ValueAdapter, D: DiskStore>(
Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64())))
};
store.update(key, unix_now(), &mut apply)?;
result.ok_or(Error::InvalidEntry)
result.ok_or(Error::InvalidEntry(None))
}
fn unix_now() -> f64 {

View file

@ -3,8 +3,8 @@ use py_literal::Value;
use serde_json::{Map, Number};
pub(crate) fn from_pickle(bytes: &[u8]) -> Result<Value, Error> {
let value = serde_pickle::value_from_slice(bytes, Default::default())
.map_err(|_| Error::InvalidEntry)?;
let value =
serde_pickle::value_from_slice(bytes, Default::default()).map_err(Error::invalid_entry)?;
from_pickle_value(value)
}
@ -70,7 +70,7 @@ fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result<Value, Err
}
fn integer(value: String) -> Result<Value, Error> {
value.parse().map_err(|_| Error::InvalidEntry)
value.parse().map_err(Error::invalid_entry)
}
pub(crate) fn from_json(value: serde_json::Value) -> Value {
@ -101,7 +101,7 @@ pub(crate) fn from_json(value: serde_json::Value) -> Value {
pub(crate) fn from_json_text(value: &str) -> Result<Value, Error> {
serde_json::from_str(value)
.map(from_json)
.map_err(|_| Error::InvalidEntry)
.map_err(Error::invalid_entry)
}
pub(crate) fn is_truthy(value: &Value) -> bool {
@ -131,7 +131,7 @@ pub(crate) fn to_f64(value: &Value) -> Option<f64> {
}
pub(crate) fn to_json(value: &Value) -> Result<Vec<u8>, Error> {
serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry)
serde_json::to_vec(&to_json_value(value)?).map_err(Error::invalid_entry)
}
fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
@ -142,12 +142,12 @@ fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
value
.to_string()
.parse::<Number>()
.map_err(|_| Error::InvalidEntry)?,
.map_err(Error::invalid_entry)?,
),
Value::Float(value) => {
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?)
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry(None))?)
}
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry),
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry(None)),
Value::String(value) => serde_json::Value::String(value.clone()),
Value::Tuple(values) | Value::List(values) | Value::Set(values) => {
serde_json::Value::Array(
@ -162,7 +162,7 @@ fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
.iter()
.map(|(key, value)| {
let Value::String(key) = key else {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
};
Ok((key.clone(), to_json_value(value)?))
})

View file

@ -48,13 +48,13 @@ struct Row {
impl DiskcacheSqliteStore {
pub fn open(directory: impl AsRef<Path>) -> Result<Self, Error> {
let directory = directory.as_ref().to_path_buf();
fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?;
let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?;
fs::create_dir_all(&directory).map_err(Error::unavailable)?;
let directory = std::path::absolute(&directory).map_err(Error::unavailable)?;
let database = directory.join("cache.db");
let connection = Connection::open(database).map_err(|_| Error::Unavailable)?;
let connection = Connection::open(database).map_err(Error::unavailable)?;
connection
.busy_timeout(std::time::Duration::from_secs(60))
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let mut settings = read_settings(&connection)?;
for (key, value) in default_settings() {
@ -74,7 +74,7 @@ impl DiskcacheSqliteStore {
value
)",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
for (key, value) in &settings {
if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") {
connection
@ -82,7 +82,7 @@ impl DiskcacheSqliteStore {
"INSERT OR REPLACE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
}
for (key, value) in [
@ -96,7 +96,7 @@ impl DiskcacheSqliteStore {
"INSERT OR IGNORE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
connection
.execute_batch(
@ -117,7 +117,7 @@ impl DiskcacheSqliteStore {
CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw);
CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let eviction_policy = setting_string(&settings, "eviction_policy")
.unwrap_or_else(|| "least-recently-stored".to_string());
@ -128,23 +128,27 @@ impl DiskcacheSqliteStore {
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
"least-recently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
"least-frequently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
_ => {
return Err(Error::unavailable(SettingFault::EvictionPolicy(
eviction_policy,
)));
}
_ => return Err(Error::Unavailable),
}
connection
.execute_batch(
@ -170,12 +174,12 @@ impl DiskcacheSqliteStore {
UPDATE Settings SET value = value - OLD.size
WHERE key = \"size\"; END;",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let min_file_size = setting_i64(&settings, "disk_min_file_size")
.unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE)
.try_into()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT);
let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT);
let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0;
@ -207,7 +211,7 @@ impl DiskcacheSqliteStore {
|row| row.get::<_, Option<String>>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.flatten()
{
cleanup.push(old_filename);
@ -221,7 +225,7 @@ impl DiskcacheSqliteStore {
|row| row.get::<_, i64>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if let Some(rowid) = rowid {
connection
.execute(
@ -230,7 +234,7 @@ impl DiskcacheSqliteStore {
WHERE rowid = ?",
params![now, expire_time, now, size, mode, filename, value, rowid],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
} else {
connection
.execute(
@ -240,7 +244,7 @@ impl DiskcacheSqliteStore {
) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)",
params![key, now, expire_time, now, size, mode, filename, value],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
cleanup.extend(self.cull(connection, now)?);
Ok(cleanup)
@ -257,13 +261,13 @@ impl DiskcacheSqliteStore {
WHERE expire_time IS NOT NULL AND expire_time < ?
ORDER BY expire_time LIMIT ?",
)
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.query_map(params![now, self.cull_limit], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
for (_, filename) in &expired {
if let Some(filename) = filename {
cleanup.push(filename.clone());
@ -272,7 +276,7 @@ impl DiskcacheSqliteStore {
for (rowid, _) in &expired {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit);
if remaining <= 0 || self.volume(connection)? < self.size_limit {
@ -283,19 +287,23 @@ impl DiskcacheSqliteStore {
"least-recently-stored" => "store_time",
"least-recently-used" => "access_time",
"least-frequently-used" => "access_count",
_ => return Err(Error::Unavailable),
policy => {
return Err(Error::unavailable(SettingFault::EvictionPolicy(
policy.to_owned(),
)));
}
};
let rows = connection
.prepare(&format!(
"SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?"
))
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.query_map(params![remaining], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
for (_, filename) in &rows {
if let Some(filename) = filename {
cleanup.push(filename.clone());
@ -304,7 +312,7 @@ impl DiskcacheSqliteStore {
for (rowid, _) in rows {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
Ok(cleanup)
}
@ -312,15 +320,15 @@ impl DiskcacheSqliteStore {
fn volume(&self, connection: &Connection) -> Result<i64, Error> {
let page_count: i64 = connection
.query_row("PRAGMA page_count", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let page_size: i64 = connection
.query_row("PRAGMA page_size", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let size: i64 = connection
.query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| {
row.get(0)
})
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(page_count.saturating_mul(page_size).saturating_add(size))
}
}
@ -331,13 +339,13 @@ impl DiskStore for DiskcacheSqliteStore {
}
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)";
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if !self.statistics && !has_get_update(&self.eviction_policy) {
return row
.map(|row| fetch_row(&self.directory, row))
@ -348,7 +356,7 @@ impl DiskStore for DiskcacheSqliteStore {
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(row) = row else {
if self.statistics {
connection
@ -356,7 +364,7 @@ impl DiskStore for DiskcacheSqliteStore {
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
return Ok(None);
};
@ -369,14 +377,14 @@ impl DiskStore for DiskcacheSqliteStore {
"UPDATE Settings SET value = value + 1 WHERE key = 'hits'",
[],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
} else if !hit && self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
if has_get_update(&self.eviction_policy) && hit {
let update = match self.eviction_policy.as_str() {
@ -384,16 +392,20 @@ impl DiskStore for DiskcacheSqliteStore {
"least-frequently-used" => {
"UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?"
}
_ => return Err(Error::Unavailable),
policy => {
return Err(Error::unavailable(SettingFault::EvictionPolicy(
policy.to_owned(),
)));
}
};
if self.eviction_policy == "least-recently-used" {
connection
.execute(update, params![now, rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
} else {
connection
.execute(update, params![rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
}
value
@ -409,7 +421,7 @@ impl DiskStore for DiskcacheSqliteStore {
) -> Result<(), Error> {
let columns = store_value(&self.directory, self.min_file_size, value)?;
let new_filename = columns.filename.clone();
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let result = transactional(&connection, |connection| {
self.set_locked(connection, key, columns, expire_time, now)
});
@ -428,7 +440,7 @@ impl DiskStore for DiskcacheSqliteStore {
}
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let selected = transactional(&connection, |connection| {
let row = connection
.query_row(
@ -439,13 +451,13 @@ impl DiskStore for DiskcacheSqliteStore {
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(row) = row else {
return Ok(None);
};
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(Some(row))
})?;
let Some(row) = selected else {
@ -460,7 +472,7 @@ impl DiskStore for DiskcacheSqliteStore {
}
fn clear(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let mut last_rowid = 0_i64;
loop {
let batch = transactional(&connection, |connection| {
@ -469,13 +481,13 @@ impl DiskStore for DiskcacheSqliteStore {
"SELECT rowid, filename FROM Cache
WHERE rowid > ? ORDER BY rowid LIMIT 100",
)
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.query_map(params![last_rowid], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if rows.is_empty() {
return Ok(rows);
}
@ -486,7 +498,7 @@ impl DiskStore for DiskcacheSqliteStore {
.join(",");
connection
.execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), [])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(rows)
})?;
if batch.is_empty() {
@ -509,7 +521,7 @@ impl DiskStore for DiskcacheSqliteStore {
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let mut created_filename = None;
let result = transactional(&connection, |connection| {
let current = connection
@ -521,7 +533,7 @@ impl DiskStore for DiskcacheSqliteStore {
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.map(|row| fetch_row(&self.directory, row))
.transpose()?
.flatten();
@ -582,21 +594,35 @@ fn read_settings(connection: &Connection) -> Result<HashMap<String, Value>, Erro
};
statement
.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.collect::<Result<HashMap<_, _>, _>>()
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
#[derive(Debug, thiserror::Error)]
enum SettingFault {
#[error("unknown eviction policy {0:?}")]
EvictionPolicy(String),
#[error("setting {0:?} is not an SQLite pragma")]
NotPragma(String),
#[error("setting {0:?} has an unsupported value type")]
PragmaValue(String),
}
fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> {
let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?;
let pragma = key
.strip_prefix("sqlite_")
.ok_or_else(|| Error::unavailable(SettingFault::NotPragma(key.to_owned())))?;
match value {
Value::Integer(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
.map_err(Error::unavailable),
Value::Text(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
_ => Err(Error::Unavailable),
.map_err(Error::unavailable),
_ => Err(Error::unavailable(SettingFault::PragmaValue(
key.to_owned(),
))),
}
}
@ -634,7 +660,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
Value::Text(value) => Ok(Some(StoredValue::Text(value))),
Value::Integer(value) => Ok(Some(StoredValue::Integer(value))),
Value::Real(value) => Ok(Some(StoredValue::Float(value))),
Value::Null => Err(Error::InvalidEntry),
Value::Null => Err(Error::InvalidEntry(None)),
},
MODE_BINARY | MODE_PICKLE => {
let bytes = match row.value {
@ -645,7 +671,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
};
value
}
_ => return Err(Error::InvalidEntry),
_ => return Err(Error::InvalidEntry(None)),
};
Ok(Some(if row.mode == MODE_BINARY {
StoredValue::Bytes(bytes)
@ -663,24 +689,24 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
}
Value::Blob(value) => value,
Value::Text(value) => value.into_bytes(),
_ => return Err(Error::InvalidEntry),
_ => return Err(Error::InvalidEntry(None)),
};
Ok(Some(StoredValue::Text(
String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?,
String::from_utf8(bytes).map_err(Error::invalid_entry)?,
)))
}
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
fn read_file(directory: &Path, filename: Option<&str>) -> Result<Option<Vec<u8>>, Error> {
let Some(filename) = filename else {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
};
match fs::read(directory.join(filename)) {
Ok(value) => Ok(Some(value)),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(_) => Err(Error::Unavailable),
Err(error) => Err(Error::unavailable(error)),
}
}
@ -712,7 +738,7 @@ fn store_value(
let bytes = value.into_bytes();
let filename = write_file(directory, &bytes)?;
Ok(StoredColumns {
size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?,
size: i64::try_from(bytes.len()).map_err(Error::unavailable)?,
mode: MODE_TEXT,
filename: Some(filename),
value: None,
@ -727,7 +753,7 @@ fn store_value(
StoredValue::Bytes(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
size: i64::try_from(value.len()).map_err(Error::unavailable)?,
mode: MODE_BINARY,
filename: Some(filename),
value: None,
@ -742,7 +768,7 @@ fn store_value(
StoredValue::Pickle(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
size: i64::try_from(value.len()).map_err(Error::unavailable)?,
mode: MODE_PICKLE,
filename: Some(filename),
value: None,
@ -761,14 +787,14 @@ fn write_file(directory: &Path, bytes: &[u8]) -> Result<String, Error> {
let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]);
let path = directory.join(&filename);
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?;
fs::create_dir_all(parent).map_err(Error::unavailable)?;
}
let mut file = OpenOptions::new()
.write(true)
.create_new(true)
.open(path)
.map_err(|_| Error::Unavailable)?;
file.write_all(bytes).map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
file.write_all(bytes).map_err(Error::unavailable)?;
Ok(filename)
}
@ -789,12 +815,12 @@ fn transactional<T>(
) -> Result<T, Error> {
connection
.execute_batch("BEGIN IMMEDIATE")
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
match operation(connection) {
Ok(value) => {
connection
.execute_batch("COMMIT")
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(value)
}
Err(error) => {

View file

@ -84,7 +84,7 @@ fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation
match expected {
ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected),
ReadExpectation::Miss => assert_eq!(result.unwrap(), None),
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))),
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry(_)))),
}
}

View file

@ -13,9 +13,11 @@ litellm-auth-types.workspace = true
litellm-cache.workspace = true
percent-encoding.workspace = true
reqwest.workspace = true
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-cache-testing.workspace = true
rstest.workspace = true

View file

@ -25,6 +25,14 @@ pub fn key_prefix(gcs_path: Option<&str>) -> String {
}
}
#[derive(Debug, thiserror::Error)]
#[error("GCS answered {0}")]
struct UnexpectedStatus(reqwest::StatusCode);
#[derive(Debug, thiserror::Error)]
#[error("GCS cache worker thread panicked")]
struct WorkerPanicked;
#[derive(Clone, Debug)]
pub struct GcsConfig {
pub bucket_name: String,
@ -114,9 +122,9 @@ impl<S: CacheCodec> GcsCache<S> {
.body(payload)
.send()
.await
.map_err(|_| Error::Unavailable)?;
.map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?;
if !response.status().is_success() {
return Err(Error::Unavailable);
return Err(Error::unavailable(UnexpectedStatus(response.status())));
}
Ok(())
}
@ -135,18 +143,18 @@ impl<S: CacheCodec> GcsCache<S> {
.header(reqwest::header::CONTENT_TYPE, "application/json")
.send()
.await
.map_err(|_| Error::Unavailable)?;
.map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?;
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Ok(None);
}
if !response.status().is_success() {
return Err(Error::Unavailable);
return Err(Error::unavailable(UnexpectedStatus(response.status())));
}
let body = response.bytes().await.map_err(|_| Error::Unavailable)?;
self.codec
.decode(&body)
.map(Some)
.map_err(|_| Error::InvalidEntry)
let body = response
.bytes()
.await
.map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?;
self.codec.decode(&body).map(Some)
}
fn run_sync<T, F>(future: F) -> Result<T, Error>
@ -158,7 +166,7 @@ impl<S: CacheCodec> GcsCache<S> {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
.and_then(|runtime| runtime.block_on(future))
};
match tokio::runtime::Handle::try_current() {
@ -169,7 +177,7 @@ impl<S: CacheCodec> GcsCache<S> {
scope
.spawn(|| run(future))
.join()
.map_err(|_| Error::Unavailable)
.map_err(|_| Error::unavailable(WorkerPanicked))
.and_then(|result| result)
}),
Err(_) => run(future),
@ -242,7 +250,7 @@ impl<S: CacheCodec> BatchCache for GcsCache<S> {
match self.async_get_cache(&key, &context).await {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}

View file

@ -30,7 +30,7 @@ impl TokenSource for GcpTokenSource {
self.auth
.access_token(&self.config, &|name| std::env::var(name).ok())
.await
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
}
}

View file

@ -55,19 +55,27 @@ async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer
Ok(Some(json!({"value": "entry"})))
)]
#[case::missing("missing", ResponseTemplate::new(404), Ok(None))]
#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))]
#[case::unauthorized("unauthorized", ResponseTemplate::new(401), Err(Error::Unavailable))]
#[case::server_error(
"server-error",
ResponseTemplate::new(500),
Err("cache is unavailable: GCS answered 500 Internal Server Error")
)]
#[case::unauthorized(
"unauthorized",
ResponseTemplate::new(401),
Err("cache is unavailable: GCS answered 401 Unauthorized")
)]
#[case::invalid(
"invalid",
ResponseTemplate::new(200).set_body_string("not json"),
Err(Error::InvalidEntry)
Err("invalid cache entry: ")
)]
#[tokio::test]
async fn get_maps_statuses_and_decode_failures(
#[future(awt)] server: MockServer,
#[case] key: &str,
#[case] response: ResponseTemplate,
#[case] expected: Result<Option<Value>, Error>,
#[case] expected: Result<Option<Value>, &str>,
) {
Mock::given(method("GET"))
.and(path(format!("/storage/v1/b/bucket/o/{key}")))
@ -76,8 +84,21 @@ async fn get_maps_statuses_and_decode_failures(
.mount(&server)
.await;
let cache = support::cache(&server, None);
assert_eq!(cache.get_cache(key, &context()), expected);
assert_eq!(cache.async_get_cache(key, &context()).await, expected);
let outcome = |result: Result<Option<Value>, Error>| {
result.map_err(|error| litellm_error::Report::new(&error).to_string())
};
let sync = outcome(cache.get_cache(key, &context()));
let asynchronous = outcome(cache.async_get_cache(key, &context()).await);
match expected {
Ok(value) => {
assert_eq!(sync, Ok(value.clone()));
assert_eq!(asynchronous, Ok(value));
}
Err(prefix) => {
assert!(sync.unwrap_err().starts_with(prefix));
assert!(asynchronous.unwrap_err().starts_with(prefix));
}
}
}
#[rstest]
@ -97,7 +118,7 @@ async fn cache_exposes_its_configuration(#[future(awt)] server: MockServer) {
path_service_account: Some("/secrets/sa.json".into()),
..support::config(&server, Some("folder"))
},
litellm_http::Client::plain_for_test(),
litellm_http::Client::for_test(),
litellm_cache::JsonCodec::<Value>::new(),
);
assert_eq!(cache.bucket_name(), "bucket");
@ -257,7 +278,7 @@ struct FailingTokenSource;
impl TokenSource for FailingTokenSource {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async { Err(Error::Unavailable) })
Box::pin(async { Err(Error::Unavailable(None)) })
}
}
@ -267,14 +288,14 @@ async fn token_source_failure_skips_http(#[future(awt)] server: MockServer) {
let cache = support::cache_with_token(&server, None, Arc::new(FailingTokenSource));
assert_eq!(
cache.get_cache("key", &context()).unwrap_err(),
Error::Unavailable
Error::Unavailable(None)
);
assert_eq!(
cache
.async_set_cache("key", json!(1), context())
.await
.unwrap_err(),
Error::Unavailable
Error::Unavailable(None)
);
assert_eq!(server.received_requests().await.unwrap().len(), 0);
}

View file

@ -29,7 +29,7 @@ pub fn cache_with_token(
) -> JsonGcsCache {
GcsCache::with_token_source(
config(server, gcs_path),
litellm_http::Client::plain_for_test(),
litellm_http::Client::for_test(),
JsonCodec::new(),
token,
)

View file

@ -2,7 +2,7 @@ use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap, HashSet},
hash::Hash,
sync::{Arc, Mutex},
sync::{Arc, Mutex, MutexGuard},
time::{Duration, SystemTime, UNIX_EPOCH},
};
@ -84,6 +84,10 @@ impl<V: Clone> InMemoryCache<V> {
}
}
fn lock(&self) -> Result<MutexGuard<'_, CacheState<V>>, Error> {
Ok(self.state.lock()?)
}
pub fn set_cache(
&self,
key: impl Into<String>,
@ -94,13 +98,13 @@ impl<V: Clone> InMemoryCache<V> {
return Ok(CacheWrite::Disabled);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
self.store(&mut state, key.into(), value, ttl, now)
}
pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
if state
.expirations
.get(key)
@ -124,7 +128,7 @@ impl<V: Clone> InMemoryCache<V> {
/// fits.
pub fn evict_cache(&self) -> Result<(), Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, None);
Ok(())
}
@ -132,7 +136,7 @@ impl<V: Clone> InMemoryCache<V> {
/// `evict_element_if_expired`: `true` when `key` had expired and was removed.
pub fn evict_element_if_expired(&self, key: &str) -> Result<bool, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
let expired = state
.expirations
.get(key)
@ -153,12 +157,7 @@ impl<V: Clone> InMemoryCache<V> {
/// The number of stored entries, expired ones included until they are evicted.
pub fn len(&self) -> Result<usize, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.values
.len())
Ok(self.lock()?.values.len())
}
pub fn is_empty(&self) -> Result<bool, Error> {
@ -167,12 +166,7 @@ impl<V: Clone> InMemoryCache<V> {
/// Entries in the expiration heap, stale ones included; bounded by eviction.
pub fn expiration_heap_len(&self) -> Result<usize, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expiration_heap
.len())
Ok(self.lock()?.expiration_heap.len())
}
pub fn max_size_in_memory(&self) -> usize {
@ -184,13 +178,7 @@ impl<V: Clone> InMemoryCache<V> {
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expirations
.get(key)
.copied())
Ok(self.lock()?.expirations.get(key).copied())
}
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
@ -198,7 +186,7 @@ impl<V: Clone> InMemoryCache<V> {
}
pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result<Vec<String>, Error> {
let state = self.state.lock().map_err(|_| Error::Unavailable)?;
let state = self.lock()?;
let mut expirations = state
.expirations
.iter()
@ -213,13 +201,13 @@ impl<V: Clone> InMemoryCache<V> {
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
Self::remove(&mut state, key);
Ok(())
}
pub fn flush_cache(&self) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
state.values.clear();
state.expirations.clear();
state.expiration_heap.clear();
@ -316,7 +304,7 @@ where
return Ok(candidate);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
let existing = state
.values
@ -351,7 +339,7 @@ impl CounterCache for InMemoryCache<f64> {
return Ok(amount);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
let value = Self::live(&mut state, key, now).unwrap_or_default() + amount;
self.store(&mut state, key.into(), value, self.get_ttl(&context), now)?;
Ok(value)
@ -424,7 +412,7 @@ where
return Ok(values);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
let mut stored = Self::live(&mut state, key, now).unwrap_or_default();
stored.extend(values.iter().cloned());
self.store(&mut state, key.into(), stored, ttl, now)?;

View file

@ -52,7 +52,7 @@ fn measured(capacity: usize) -> InMemoryCache<String> {
Some(4),
Some(Arc::new(|value: &String| {
if value.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
Ok(value.len())
})),
@ -275,7 +275,7 @@ fn evict_element_if_expired_reports_removal(clock: Clock) {
#[case::fits("ok", Ok(true))]
#[case::at_limit("four", Ok(true))]
#[case::too_large("oversized", Ok(false))]
#[case::measure_error("", Err(Error::InvalidEntry))]
#[case::measure_error("", Err(Error::InvalidEntry(None)))]
fn check_value_size_applies_the_entry_limit(
#[case] value: &str,
#[case] expected: Result<bool, Error>,
@ -310,7 +310,7 @@ fn disabled_size_limited_and_validated_writes_are_observable() {
assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into()));
assert_eq!(
cache.set_cache("invalid", String::new(), None),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
assert_eq!(cache.get_cache("invalid").unwrap(), None);
cache.delete_cache("small").unwrap();

View file

@ -18,6 +18,7 @@ tokio.workspace = true
uuid.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
futures-executor = "0.3"
litellm-cache-testing.workspace = true

View file

@ -36,7 +36,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
let exists = client
.collection_exists(config.collection_name.clone())
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if !exists {
client
.create_collection(
@ -48,7 +48,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
.quantization_config(quantization(&config.quantization)),
)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
let _ = client
.create_field_index(CreateFieldIndexCollectionBuilder::new(
@ -99,13 +99,13 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let response =
String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?;
String::from_utf8(self.codec.encode(&value)?).map_err(Error::invalid_entry)?;
let payload = Payload::try_from(json!({
"litellm_cache_key": key,
"text": prompt,
"response": response,
}))
.map_err(|_| Error::InvalidEntry)?;
.map_err(Error::invalid_entry)?;
self.client
.upsert_points(
UpsertPointsBuilder::new(
@ -119,7 +119,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
.wait(true),
)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(())
}
@ -152,7 +152,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
),
)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(point) = result.result.into_iter().next() else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
@ -170,7 +170,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
let response = payload
.get("response")
.and_then(Value::as_str)
.ok_or(Error::InvalidEntry)?;
.ok_or(Error::InvalidEntry(None))?;
Ok(SemanticLookup {
value: Some(self.codec.decode(response.as_bytes())?),
similarity: Some(similarity),

View file

@ -55,10 +55,9 @@ impl Embedder for OpenAiEmbedder {
}
.send()
.await
.map_err(|_| Error::Unavailable)?
.error_for_status()
.map_err(|_| Error::Unavailable)?;
let body: Value = response.json().await.map_err(|_| Error::Unavailable)?;
.and_then(reqwest::Response::error_for_status)
.map_err(http_error)?;
let body: Value = response.json().await.map_err(http_error)?;
body.get("data")
.and_then(Value::as_array)
.and_then(|data| data.first())
@ -70,6 +69,14 @@ impl Embedder for OpenAiEmbedder {
.map(|value| value.as_f64().map(|value| value as f32))
.collect::<Option<Vec<_>>>()
})
.ok_or(Error::Unavailable)
.ok_or(Error::unavailable(MissingEmbedding))
}
}
fn http_error(error: reqwest::Error) -> Error {
Error::unavailable(litellm_http::Error::from(error))
}
#[derive(Debug, thiserror::Error)]
#[error("embeddings response has no embedding")]
struct MissingEmbedding;

View file

@ -108,7 +108,7 @@ fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
async fn posts_embeddings_request_and_parses_vector() {
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
let embedder = OpenAiEmbedder::new(
litellm_http::Client::plain_for_test(),
litellm_http::Client::for_test(),
config(
format!("{}/", server.base_url()),
Some(Duration::from_secs(1)),
@ -134,13 +134,19 @@ async fn posts_embeddings_request_and_parses_vector() {
}
#[rstest]
#[case::error_status("500 Internal Server Error", "{}", 0, None, Err(Error::Unavailable))]
#[case::error_status(
"500 Internal Server Error",
"{}",
0,
None,
Err(Error::Unavailable(None))
)]
#[case::timed_out(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
500,
Some(Duration::from_millis(200)),
Err(Error::Unavailable)
Err(Error::Unavailable(None))
)]
#[case::within_timeout(
"200 OK",
@ -149,7 +155,7 @@ async fn posts_embeddings_request_and_parses_vector() {
Some(Duration::from_secs(1)),
Ok(vec![0.1, 0.2])
)]
#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable))]
#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable(None)))]
#[tokio::test]
async fn status_timeout_and_body_errors_are_unavailable(
#[case] status: &str,
@ -161,16 +167,23 @@ async fn status_timeout_and_body_errors_are_unavailable(
let server =
TestHttpServer::response_after(status, body, Duration::from_millis(delay_ms)).await;
let embedder = OpenAiEmbedder::new(
litellm_http::Client::plain_for_test(),
litellm_http::Client::for_test(),
config(server.base_url(), timeout),
);
assert_eq!(embedder.async_embed("hello", None).await, expected);
let result = embedder.async_embed("hello", None).await;
if let Err(error) = &result {
let report = litellm_error::Report::new(error).to_string();
assert!(!report.contains(&server.base_url()), "{report}");
}
let kind =
|result: Result<Vec<f32>, Error>| result.map_err(|error| std::mem::discriminant(&error));
assert_eq!(kind(result), kind(expected));
}
#[rstest]
fn sync_embedding_is_unsupported() {
let embedder = OpenAiEmbedder::new(
litellm_http::Client::plain_for_test(),
litellm_http::Client::for_test(),
config("http://127.0.0.1:9".to_owned(), None),
);
assert_eq!(

View file

@ -53,7 +53,10 @@ impl Embedder for FixedEmbedder {
.lock()
.unwrap()
.push((input.to_owned(), metadata.cloned()));
self.vectors.get(input).cloned().ok_or(Error::Unavailable)
self.vectors
.get(input)
.cloned()
.ok_or(Error::Unavailable(None))
}
}
@ -396,9 +399,9 @@ async fn lookup_reports_python_semantic_similarity(
#[rstest]
#[case::codec_decodes_the_payload(Some(json!("{\"a\":1}")), Ok(Some(json!({"a": 1}))))]
#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry))]
#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry))]
#[case::missing_response(None, Err(Error::InvalidEntry))]
#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry(None)))]
#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry(None)))]
#[case::missing_response(None, Err(Error::InvalidEntry(None)))]
#[tokio::test(flavor = "multi_thread")]
async fn stored_responses_go_through_the_codec(
#[case] response: Option<JsonValue>,
@ -419,9 +422,12 @@ async fn stored_responses_go_through_the_codec(
.into(),
});
let kind = |result: Result<Option<JsonValue>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(
cache.async_get_cache("key", &context("hello")).await,
expected
kind(cache.async_get_cache("key", &context("hello")).await),
kind(expected)
);
server.stop();
}
@ -434,7 +440,7 @@ async fn embedding_failures_propagate() {
assert_eq!(
cache.async_get_cache("key", &context("unknown")).await,
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
server.stop();
}
@ -493,10 +499,10 @@ async fn stopped_qdrant_server_maps_to_unavailable() {
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
server.stop();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(
assert!(matches!(
cache.async_get_cache("key", &context("hello")).await,
Err(Error::Unavailable)
);
Err(Error::Unavailable(Some(_)))
));
}
/// `_payload_matches_cache_key` compares `str(cached_key) == str(key)`, so a point whose stored

View file

@ -63,13 +63,13 @@ impl Inner {
.arg(CACHE_KEY_FIELD)
.arg(tag)
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if let Some(ttl) = ttl {
redis::cmd("EXPIRE")
.arg(&hash_key)
.arg(ttl_seconds(ttl))
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
Ok(())
}
@ -111,7 +111,7 @@ impl Inner {
.arg("vector")
.arg(vector_buffer(vector))
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(fields) = first_document(&result) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};

View file

@ -40,10 +40,10 @@ impl Index {
Some(false) => self.isolated(connection, dims)?,
None => match create_index(connection, &self.name, dims) {
Ok(()) => self.name.clone(),
Err(_) => match index_compatible(connection, &self.name, dims)? {
Err(error) => match index_compatible(connection, &self.name, dims)? {
Some(true) => self.name.clone(),
Some(false) => self.isolated(connection, dims)?,
None => return Err(Error::Unavailable),
None => return Err(error),
},
},
};
@ -59,7 +59,7 @@ impl Index {
redis::cmd("FT.DROPINDEX")
.arg(&name)
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
create_index(connection, &name, dims)?;
Ok(name)
}
@ -109,7 +109,7 @@ fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) ->
.arg("SEPARATOR")
.arg(",")
.query::<()>(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
fn index_compatible(
@ -123,14 +123,18 @@ fn index_compatible(
{
Ok(info) => info,
Err(error) if unknown_index(&error) => return Ok(None),
Err(_) => return Err(Error::Unavailable),
Err(error) => return Err(Error::unavailable(error)),
};
Ok(Some(schema_compatible(&info, dims)))
}
fn unknown_index(error: &redis::RedisError) -> bool {
let message = error.to_string().to_lowercase();
message.contains("unknown") && message.contains("index")
error
.code()
.is_some_and(|code| code.eq_ignore_ascii_case("unknown"))
&& error
.detail()
.is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index"))
}
struct Attribute {

View file

@ -79,7 +79,7 @@ fn s(value: &str) -> redis::Value {
}
fn unknown_index_error() -> redis::RedisError {
redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name"))
support::server_error("Unknown index name")
}
fn info_missing(index: &str) -> MockCmd {
@ -320,7 +320,7 @@ fn store_writes_the_redisvl_hash(
#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))]
#[case::above_distance_threshold("key1", "0.5", None, Ok(None))]
#[case::other_cache_key("other", "0.05", None, Ok(None))]
#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry))]
#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))]
fn lookup_applies_threshold_scope_and_codec(
#[case] stored_tag: &str,
#[case] distance: &str,
@ -343,7 +343,10 @@ fn lookup_applies_threshold_scope_and_codec(
FakeEmbedder::new(&[]),
);
assert_eq!(cache.get_cache("key1", &context), expected);
let kind = |result: Result<Option<Value>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected));
}
#[rstest]
@ -451,10 +454,7 @@ fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticC
info_missing(INDEX),
MockCmd::new(
create_index_command(INDEX, 3),
Err::<&str, _>(redis::RedisError::from((
redis::ErrorKind::Extension,
"Index already exists",
))),
Err::<&str, _>(support::server_error("Index already exists")),
),
info(INDEX, compatible_info(3)),
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
@ -673,10 +673,10 @@ fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value)
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert_eq!(
worker_b.get_cache("key1", &context()).unwrap_err(),
Error::Unavailable
);
assert!(matches!(
worker_b.get_cache("key1", &context()),
Err(Error::Unavailable(Some(_)))
));
}
#[fixture]
@ -732,10 +732,10 @@ fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option<Strin
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert_eq!(
worker_b.get_cache("key1", &context()).unwrap_err(),
Error::Unavailable
);
assert!(matches!(
worker_b.get_cache("key1", &context()),
Err(Error::Unavailable(Some(_)))
));
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
for index in [&base, &isolated] {

View file

@ -74,7 +74,7 @@ impl FakeSearch {
"FT.CREATE" => {
let name = text(1);
if state.indexes.contains_key(&name) {
return Err(error("Index already exists"));
return Err(server_error("Index already exists"));
}
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
let prefix = args[position("PREFIX").unwrap() + 2].clone();
@ -94,7 +94,7 @@ impl FakeSearch {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| error("Unknown index name"))?;
.ok_or_else(|| server_error("Unknown index name"))?;
Ok(index_info(index))
}
"FT.DROPINDEX" => {
@ -118,7 +118,7 @@ impl FakeSearch {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| error("no such index"))?;
.ok_or_else(|| server_error("no such index"))?;
let query = text(2);
let tag = query_tag(&query);
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
@ -155,7 +155,7 @@ impl FakeSearch {
]))
}
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
_ => Err(error("unsupported command")),
_ => Err(server_error("unsupported command")),
}
}
}
@ -192,8 +192,11 @@ impl redis::ConnectionLike for FakeSearch {
}
}
fn error(message: &'static str) -> redis::RedisError {
redis::RedisError::from((redis::ErrorKind::Extension, message))
pub fn server_error(reply: &str) -> redis::RedisError {
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
.expect("error reply")
.extract_error()
.expect_err("server error")
}
fn bulk(bytes: &[u8]) -> redis::Value {

View file

@ -7,8 +7,10 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-error.workspace = true
redis = { workspace = true, features = ["cluster", "tls-rustls"] }
r2d2 = "0.8.10"
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]

View file

@ -141,7 +141,7 @@ where
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
@ -152,7 +152,7 @@ where
match self.decode_response(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}

View file

@ -17,7 +17,7 @@ fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
@ -37,13 +37,13 @@ where
{
let payload = codec.encode(&candidate)?;
if payload.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
for _ in 0..CLAIM_ATTEMPTS {
let current = stored_bytes(
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)?,
.map_err(Error::unavailable)?,
)?
.filter(|bytes| !bytes.is_empty());
let existing = current
@ -63,14 +63,18 @@ where
.arg(write)
.arg(u8::from(refresh))
.query::<bool>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if applied {
return Ok(existing.unwrap_or(candidate));
}
}
Err(Error::Unavailable)
Err(Error::unavailable(ClaimContended))
}
#[derive(Debug, thiserror::Error)]
#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")]
struct ClaimContended;
impl<S, C> ClaimCache for RedisCache<S, C>
where
S: CacheCodec + Clone + 'static,

View file

@ -36,19 +36,19 @@ where
) -> Result<T, Error> {
match self {
Self::Pool(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let mut pooled = pool.get().map_err(Error::unavailable)?;
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
result
}
Self::Cluster(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let mut pooled = pool.get().map_err(Error::unavailable)?;
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
let mut connection = connection.lock()?;
operation(&mut ConnectionRef::Node(&mut *connection))
}
}
@ -61,7 +61,7 @@ where
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
}
pub fn fixed(connection: C) -> Self {
@ -95,7 +95,7 @@ fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(manager)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
fn close_idle<M, T>(pool: &r2d2::Pool<M>)
@ -123,7 +123,7 @@ impl ConnectionManager {
fn open(url: &str) -> Result<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
}
@ -151,14 +151,26 @@ impl r2d2::ManageConnection for ConnectionManager {
}
}
#[derive(Debug, thiserror::Error)]
enum TopologyFault {
#[error("Redis Cluster needs at least one startup node")]
NoStartupNodes,
#[error("Redis Cluster nodes need a TCP address")]
NonTcpAddress,
#[error("Redis Cluster answered a fan-out command without a per-node map")]
FanOutReply,
#[error("scan target does not match the connection topology")]
ScanTarget,
}
pub(crate) struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
if startup_nodes.is_empty() {
return Err(Error::Unavailable);
return Err(Error::unavailable(TopologyFault::NoStartupNodes));
}
let info = url.into_connection_info().map_err(|_| Error::Unavailable)?;
let info = url.into_connection_info().map_err(Error::unavailable)?;
let nodes = startup_nodes
.iter()
.map(|node| node_info(&info, node))
@ -168,7 +180,7 @@ impl ClusterConnectionManager {
.response_timeout(REDIS_TIMEOUT)
.build()
.map(Self)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
}
@ -185,7 +197,7 @@ fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo,
insecure: *insecure,
tls_params: tls_params.clone(),
},
_ => return Err(Error::Unavailable),
_ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)),
};
Ok(info.clone().set_addr(addr))
}
@ -298,7 +310,7 @@ impl ConnectionRef<'_> {
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
}
}
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
pub(crate) fn scan(
@ -311,7 +323,7 @@ impl ConnectionRef<'_> {
Self::Node(connection) => {
let page = scan_command(0, pattern, count)
.query::<ScanPage>(*connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
vec![(None, page)]
}
Self::Cluster(connection) => connection
@ -322,7 +334,7 @@ impl ConnectionRef<'_> {
Some(ResponsePolicy::Special),
)),
)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
.and_then(primary_pages)?
.into_iter()
.map(|(node, page)| (Some(node), page))
@ -362,7 +374,7 @@ impl ConnectionRef<'_> {
pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
match self {
Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable),
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
Self::Cluster(connection) => {
let value = connection
.route_command(
@ -372,18 +384,16 @@ impl ConnectionRef<'_> {
Some(ResponsePolicy::Special),
)),
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
return Err(Error::unavailable(TopologyFault::FanOutReply));
};
let mut replies = entries
.into_iter()
.map(|(node, reply)| {
Ok((
redis::from_redis_value::<String>(node)
.map_err(|_| Error::Unavailable)?,
redis::from_redis_value::<String>(reply)
.map_err(|_| Error::Unavailable)?,
redis::from_redis_value::<String>(node).map_err(Error::unavailable)?,
redis::from_redis_value::<String>(reply).map_err(Error::unavailable)?,
))
})
.collect::<Result<Vec<(String, String)>, Error>>()?;
@ -400,7 +410,7 @@ impl ConnectionRef<'_> {
pub(crate) fn flushall(&mut self) -> Result<(), Error> {
let command = redis::cmd("FLUSHALL");
match self {
Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable),
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
Self::Cluster(connection) => connection
.route_command(
&command,
@ -410,7 +420,7 @@ impl ConnectionRef<'_> {
)),
)
.map(|_| ())
.map_err(|_| Error::Unavailable),
.map_err(Error::unavailable),
}
}
@ -424,7 +434,7 @@ impl ConnectionRef<'_> {
let command = scan_command(cursor, pattern, count);
match (self, node) {
(Self::Node(connection), None) => {
command.query(*connection).map_err(|_| Error::Unavailable)
command.query(*connection).map_err(Error::unavailable)
}
(Self::Cluster(connection), Some(node)) => connection
.route_command(
@ -434,9 +444,9 @@ impl ConnectionRef<'_> {
port: node.port(),
}),
)
.map_err(|_| Error::Unavailable)
.and_then(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable)),
_ => Err(Error::Unavailable),
.map_err(Error::unavailable)
.and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)),
_ => Err(Error::unavailable(TopologyFault::ScanTarget)),
}
}
}
@ -445,14 +455,14 @@ type ScanPage = (u64, Vec<String>);
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
return Err(Error::unavailable(TopologyFault::FanOutReply));
};
entries
.into_iter()
.map(|(node, page)| {
let node = redis::from_redis_value::<String>(node).map_err(|_| Error::Unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(|_| Error::Unavailable)?;
let page = redis::from_redis_value::<ScanPage>(page).map_err(|_| Error::Unavailable)?;
let node = redis::from_redis_value::<String>(node).map_err(Error::unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?;
let page = redis::from_redis_value::<ScanPage>(page).map_err(Error::unavailable)?;
Ok((node, page))
})
.collect()

View file

@ -100,7 +100,7 @@ fn increment(
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable);
.map_err(Error::unavailable);
}
connection
.query_pipeline(
@ -140,7 +140,7 @@ where
}
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
redis::from_redis_value(value).map_err(|_| Error::InvalidEntry)
redis::from_redis_value(value).map_err(Error::invalid_entry)
}
impl<S, C> BoundedCounterCache for RedisCache<S, C>
@ -182,7 +182,7 @@ where
.arg(value)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await
}
@ -201,5 +201,5 @@ fn increment_with_floor(
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}

View file

@ -16,7 +16,7 @@ where
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await?;
Ok(u64::try_from(ttl).ok().map(Duration::from_secs))
@ -36,7 +36,7 @@ where
.arg(key)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await
}

View file

@ -11,11 +11,11 @@ where
C: redis::ConnectionLike + Send + 'static,
{
fn sync_ping(&self) -> Result<bool, Error> {
self.execute(|connection| connection.ping().map_err(|_| Error::Unavailable))
self.execute(|connection| connection.ping().map_err(Error::unavailable))
}
async fn ping(&self) -> Result<bool, Error> {
self.run(|connection| connection.ping().map_err(|_| Error::Unavailable))
self.run(|connection| connection.ping().map_err(Error::unavailable))
.await
}
}
@ -33,9 +33,10 @@ where
RedisTopology::Cluster { .. } => "Redis Cluster",
};
let ping = self
.run(|connection| Ok(connection.ping().map_err(|error| error.to_string())))
.run(|connection| Ok(connection.ping().map_err(Error::unavailable)))
.await
.unwrap_or_else(|error| Err(error.to_string()));
.and_then(|ping| ping)
.map_err(|error| litellm_error::Report::new(&error).to_string());
Ok(match ping {
Ok(true) => CacheConnectionResult {
status: CacheConnectionStatus::Success,

Some files were not shown because too many files have changed in this diff Show more