diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 8ac157c94f5..fb29a92e306 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -182,6 +182,7 @@ jobs: tests/unit/google_genai tests/unit/router_strategy tests/unit/router_utils + tests/unit/redis tests/unit/proxy/common_utils/test_cache_aware_routing.py tests/unit/enterprise/enterprise_callbacks/send_emails tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py @@ -1034,9 +1035,9 @@ jobs: run: | diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))" @@ -1293,7 +1294,7 @@ jobs: - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard - name: Cache Prisma binaries if: steps.changes.outputs.decision != 'skip' @@ -1496,9 +1497,9 @@ jobs: RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }} run: | if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils else - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so fi diff --git a/AGENTS.md b/AGENTS.md index 01127303da3..d8de447be90 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -52,7 +52,6 @@ If you ever write any human-facing text (pull requests, issues, commit messages, - don't use "—". Don't use ";" or ":" unless it would be nonsensical not to. Instead, reach for ",", ".", conjunction words, etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them also feels AI-y. A word cap does not penalize you for adding more sentences - don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. - unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose -- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a ".". Exception is that docs (i.e., litellm-docs) should have trailing "." - don't use →. Instead, prefer not to use arrows, and if need be, use -> instead - use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure diff --git a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md index e1f606a2eb7..538a5f8a404 100644 --- a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -46,7 +46,43 @@ For a managed workspace rollout, use OpenAI's [gateway configuration guide](http ## Native compatibility and Pi -For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name. +For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials. A ChatGPT model keeps its upstream name and needs no deployment of its own. + +### Mix ChatGPT and LiteLLM models + +On the native route, Codex's model picker shows OpenAI's catalog first and then the other chat and responses models the virtual key can use. Pick one of those and Codex's requests go through LiteLLM's `/v1/responses` with the key's usual routing, budgets, guardrails, and logging. The proxy removes the ChatGPT bearer and account header from that request, so the other provider and your callbacks never receive them. A name is routed this way when another provider serves it as an exact `model_name`, through a router `model_group_alias`, or as a team's public model name. A ChatGPT model, a name that only a wildcard deployment matches, or a name OpenAI's catalog lists for the caller's ChatGPT account keeps the passthrough, so a model OpenAI adds later still uses the subscription even when a deployment of another provider has the same name. The proxy learns the account's list from the catalog Codex fetches, or asks OpenAI with the caller's login. When it cannot get the list, for example because OpenAI is down, the request times out, or the client is not Codex, every request keeps the passthrough for the next minute, as it did before LiteLLM models were added to the picker. When the key can use many models, the picker lists the ones that fit in Codex's catalog size limit, and the proxy logs a warning with how many it left out. + +Codex's picker sends a `none` reasoning effort for a LiteLLM model whose catalog entry lists no effort levels. With `litellm_settings: drop_params: true`, the proxy drops it for a model that does not support `none`, and the provider applies its own default. A model whose cost map entry or `model_info` sets `supports_none_reasoning_effort: true` keeps it. Without `drop_params`, the effort reaches the provider unchanged. + +Codex also sends its ChatGPT app tools with every turn, which on a LiteLLM model adds a few hundred thousand prompt tokens to each turn. Set `apps = false` under `[features]` in Codex's `config.toml` to turn them off. + +Only a key that may call `/v1/responses` gets these models. A key whose `allowed_routes` covers only the native routes sees OpenAI's catalog unchanged, and every request it sends uses the subscription. + +Codex sends earlier reasoning items back with each turn. The proxy tags the reasoning that a LiteLLM model returns on this route, OpenAI and Azure deployments included. When a session switches to a ChatGPT model, the proxy drops those items before forwarding the request to OpenAI, which would reject them. Messages, tool calls, and ChatGPT's own reasoning are kept, and a LiteLLM model still gets its own reasoning back. + +When the catalog adds LiteLLM models for a key, its `ETag` is a digest of OpenAI's ETag and those models, and that key's responses carry the same value in `X-Models-Etag`. Codex keeps its cached catalog while the two match and fetches it again when OpenAI's catalog or the key's models change. A key whose catalog adds no LiteLLM models keeps OpenAI's ETags. + +### Restrict native models with an access group + +A key with a model allowlist needs a rule for ChatGPT names that OpenAI adds over time. Put the subscription wildcards in an access group under `chatgpt/` names: + +```yaml +model_list: + - model_name: chatgpt/gpt-* + litellm_params: + model: chatgpt/gpt-* + model_info: + access_groups: ["chatgpt-native"] + - model_name: chatgpt/codex-* + litellm_params: + model: chatgpt/codex-* + model_info: + access_groups: ["chatgpt-native"] +``` + +Then add `chatgpt-native` to the key's models, next to any other model names it should reach. On the native route, a key that holds only the `chatgpt/` form of a name is checked as that form, so a request for `gpt-5.6-sol` is allowed when the key may call `chatgpt/gpt-5.6-sol`. The key can call every ChatGPT model that matches the wildcards, including ones the proxy configuration does not name. + +On `/v1` these wildcards match only a request that asks for a `chatgpt/` name, so a request for a plain `gpt-` or `codex-` name still goes where it went before, such as a `*` catch-all deployment. A `/v1` request for a `chatgpt/` name without a ChatGPT login falls back to the proxy's own saved login if it has one, so keep that login unset as described above. ### Configure Pi @@ -71,7 +107,7 @@ Keep Pi's subscription login as the provider credential. The header value names Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person. -The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions. +The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key needs the `chatgpt-native` access group above or that exact name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions. The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate. diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index a31e926256a..4cfc6d4e081 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -8,7 +8,6 @@ import re from collections.abc import Callable, Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta -from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast from pydantic import TypeAdapter, ValidationError @@ -766,7 +765,6 @@ class CheckBatchCost: credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {} _file_content = await afile_content( file_id=cast(str, raw_output_file_id), - _litellm_internal_model_credentials=MappingProxyType(dict(credentials)), **cast(FileContentCallOptions, credentials), ) @@ -867,10 +865,7 @@ class CheckBatchCost: error_file_failed_requests: Final = await count_error_file_failed_requests( response, custom_llm_provider=batch_file_provider, - litellm_params={ - **credentials, - "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), - }, + litellm_params=credentials, ) batch_result: Final = ( output_file_result @@ -1029,11 +1024,7 @@ class CheckBatchCost: custom_llm_provider=deployment_info.litellm_params.custom_llm_provider, ) cost_provider: Final = cast(BatchFileProvider, llm_provider) - credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {} - file_access: Final = { - **credentials, - "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), - } + file_access: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {} output_content: Final = ( await fetch_batch_managed_file_content( response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index a4a40a7a3d7..1c433d04819 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -182,7 +182,7 @@ class ManagedObjectCreatorAttribution: team_id: Final = getattr(job, "team_id", None) rows: Final = _CreatorRows( key=await self._live_key_row(api_key, job.created_by), - team=await self._live_team_row(team_id), + team=await self._live_team_row(team_id, job.created_at), user=await self._user_row(object_id, _creator_user_id(job.created_by)), ) return ResponseCreator( @@ -269,10 +269,15 @@ class ManagedObjectCreatorAttribution: verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}") return None - async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: + async def _live_team_row( + self, team_id: str | None, created_before: datetime | None = None + ) -> prisma_models.LiteLLM_TeamTable | None: if not team_id: return None - return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + team: Final = await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + if team is not None and created_before is not None and team.created_at > created_before: + return None + return team def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 91c5e2c5f39..626a3b0bb83 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -78,6 +78,7 @@ from litellm.proxy.litellm_pre_call_utils import ( from litellm.proxy.openai_files_endpoints.common_utils import ( BATCH_CREATE_HIDDEN_PARAM, FILE_LIST_CONTINUATION_CHUNK_SIZE, + TeamProviderDeployment, _is_base64_encoded_unified_file_id, apply_unified_file_ids, decode_model_from_file_id, @@ -127,7 +128,6 @@ class _ManagedFileRetrieve(Protocol): self, *, file_id: str, - _litellm_internal_model_credentials: Mapping[str, object] | None = None, **kwargs: Unpack[FileRetrieveCallOptions], ) -> OpenAIFileObject: ... @@ -400,7 +400,7 @@ class _RouterFileCallKwargs(TypedDict, total=False): _MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str]) _FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"}) -_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id", "_litellm_internal_model_credentials"}) +_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id"}) @dataclass(frozen=True, slots=True) @@ -459,15 +459,7 @@ async def _delete_provider_file( provider_file_id: str, request_data: Mapping[str, object], ) -> None: - credentials: Final = llm_router.registry().credentials(model_id, form="with_provider") - delete_data: Final = { - **{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS}, - **( - {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))} - if credentials is not None - else {} - ), - } + delete_data: Final = {k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS} await llm_router.afile_delete( model=model_id, file_id=provider_file_id, @@ -1160,7 +1152,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): *, provider_file_id: str, call_options: Mapping[str, object], - internal_model_credentials: Mapping[str, object] | None = None, ) -> OpenAIFileObject: """Retrieve provider file details with transient retries and SDK retries disabled.""" retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve) @@ -1172,14 +1163,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if attempt > 0: await self._sleep(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS[attempt - 1]) try: - if internal_model_credentials is None: - return await retrieve_file( - file_id=provider_file_id, - **retrieve_options, - ) return await retrieve_file( file_id=provider_file_id, - _litellm_internal_model_credentials=MappingProxyType(dict(internal_model_credentials)), **retrieve_options, ) except Exception as error: @@ -1221,7 +1206,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self._afile_retrieve_with_retries( provider_file_id=provider_file_id, call_options=router_credentials, - internal_model_credentials=router_credentials, ), timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, ) @@ -1443,6 +1427,43 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): user_api_key_dict=owner, ) + async def register_provider_file( + self, + *, + provider_file_id: str, + content_type: str, + deployment: TeamProviderDeployment, + owner: UserAPIKeyAuth, + llm_router: Router | None, + ) -> str | None: + unified_file_id: Final = self.get_unified_output_file_id( + provider_file_id, deployment.deployment_id, deployment.public_model_name, file_type=content_type + ) + provider_file_object, _ = await self._fetch_provider_file_object( + unified_file_id=unified_file_id, + provider_file_id=provider_file_id, + model_id=deployment.deployment_id, + model_name=None, + llm_router=llm_router, + ) + if provider_file_object is None: + return None + try: + await self.store_unified_file_id( + file_id=unified_file_id, + file_object=provider_file_object, + litellm_parent_otel_span=None, + model_mappings={deployment.deployment_id: provider_file_id}, + user_api_key_dict=owner, + ) + except Exception as error: # noqa: BLE001 # the upload succeeded, so the caller keeps the provider id + verbose_logger.warning( + f"Failed to register provider_file_id={sanitize_for_log(provider_file_id)} as a managed file: " + f"{type(error).__name__} {sanitize_for_log(error)}" + ) + return None + return unified_file_id + async def list_user_batches( self, user_api_key_dict: UserAPIKeyAuth, @@ -2270,10 +2291,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") - def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str: + def get_unified_output_file_id( + self, output_file_id: str, model_id: str, model_name: Optional[str], file_type: str = "application/json" + ) -> str: deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id) unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - "application/json", + file_type, str(deterministic_uuid), model_name or "", output_file_id, @@ -2866,15 +2889,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): exception_dict = {} for model_id, provider_file_id in specific_model_file_id_mapping.items(): try: - # Cloud-storage providers (e.g. Bedrock S3) validate file ids - # against the deployment's configured bucket, which they only - # trust from this immutable server-side snapshot, never from - # request params. - credentials = llm_router.registry().credentials(model_id, form="with_provider") - if credentials is not None: - data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials))) - else: - data.pop("_litellm_internal_model_credentials", None) return cast( HttpxBinaryResponseContent, await llm_router.afile_content( diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 7db2de1f874..64fe03769b8 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.81" +version = "0.1.83" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.81" +version = "0.1.83" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql new file mode 100644 index 00000000000..4965ed3cfc4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql @@ -0,0 +1,9 @@ +SET LOCAL lock_timeout = '10s'; + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BackgroundResponseSettlement" ( + "response_id" TEXT NOT NULL, + "settled_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_BackgroundResponseSettlement_pkey" PRIMARY KEY ("response_id") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql new file mode 100644 index 00000000000..6b77194e3b4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql @@ -0,0 +1,7 @@ +SET LOCAL lock_timeout = '10s'; + +-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_tpm_limit" JSONB; + +-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_rpm_limit" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 818f2dc4858..37783275d9b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -19,6 +19,8 @@ model LiteLLM_BudgetTable { rpm_limit BigInt? tpd_limit BigInt? model_max_budget Json? + model_tpm_limit Json? + model_rpm_limit Json? budget_duration String? budget_reset_at DateTime? allowed_models String[] @default([]) // per-member model scope; empty = inherit team models @@ -2053,6 +2055,11 @@ model LiteLLM_BackgroundInteractionSettlement { @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } +model LiteLLM_BackgroundResponseSettlement { + response_id String @id + settled_at DateTime @default(now()) +} + model LiteLLM_Lens { reviews LiteLLM_LensReview[] id String @id diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index a12e1e27cfc..1c0c708ae7a 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.112" +version = "0.4.114" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -30,7 +30,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.112" +version = "0.4.114" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md b/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md new file mode 100644 index 00000000000..5ead8f287ab --- /dev/null +++ b/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md @@ -0,0 +1,53 @@ +--- +name: rust-python-boundary +description: Decide where Python-shaped input becomes typed Rust in litellm-rust. Use when reading config.yaml values, os.environ/ references, REDIS_* or other environment layering, get_secret fallbacks, or kwargs and objects handed over from Python, when coercing text to a bool, int or float the way Python does, or when adding a dependency on litellm-python-compat +metadata: + internal: true +--- + +# Python boundary + +Python-shaped input becomes typed Rust in exactly two crates. Every crate past them handles typed values, so no crate re-guesses how Python would have read a value + +## The two boundaries + +- `crates/config` reads config.yaml, which is Python-shaped: `os.environ/` references, scalars that Helm or the environment turned into text, raw kwargs maps. It also owns the environment layering and `get_secret` fallbacks Python applies while it loads that config. It has no Python runtime, so it reproduces Python value semantics with `litellm-python-compat` +- `crates/python-bridge` receives live Python objects, such as `RedisCache.redis_kwargs`. It reads them through PyO3 with the real protocol (`is_truthy()`, `str()`, `extract`), so it never imitates Python +- Both produce the same typed input, defined in the data crate the consumer already depends on (`redis-types` for Redis), so a downstream crate cannot tell which host it runs under + +## Python value semantics + +Reading a value the way Python itself reads it. This lives only at the two boundaries + +- Truthiness and `x or fallback`, `str(x)`, `int(text)` and `float(text)` +- Each surface's true and false text: `_str_to_bool`, pydantic's lax bool, `get_secret`'s `true` and `false` +- `os.environ/` references and `REDIS_` style environment layering + +A boundary knows which surface a value came from. `general_settings.coordination_redis` is a pydantic model and `cache_params` goes through `_str_to_bool`, so the boundary applies that surface's rule exactly instead of refusing the text two surfaces disagree on + +Each typed field keeps where its value came from (the key, `REDIS_` or the fallback name), so an error or warning downstream still names it + +## What stays downstream + +A crate that stands in for a Python library keeps that library's semantics, because they are its compatibility contract and not value coercion. `redis-resolve` stands in for redis-py, so it keeps + +- Which target wins, and the defaults redis-py applies +- How `parse_url` reads a url, including `int()` on `?db=` and redis-py's `to_bool` on the TLS query parameters +- Which kwargs redis-py accepts, and whether the Rust client reads, warns about or refuses each one + +Its input is typed. A rule it applies reads typed fields, never a `serde_json::Value` it has to coerce first. A value it only carries back to Python untouched, like `passthrough`, may stay untyped + +## litellm-python-compat + +The rule is about where config and kwargs get coerced, not about who may depend on the crate + +- Its data formats (`pickle`, `json.dumps`, `repr`) are fine anywhere Rust reads or writes state Python also reads or writes, such as cache values in a shared Redis +- Its coercions (`truthy`, `number`) belong at the two boundaries, and in a crate that stands in for a Python library when that library calls them itself + +## Existing code + +Some code still coerces Python-shaped values past the boundary, for example the `REDIS_*` merge in `gateway/src/redis/layers.rs` and the flag and number readers in `redis-resolve/src/options.rs`. Do not add more. When you change one, move the coercion to its boundary instead of extending it in place + +## Tests + +Pin a boundary's coercion with named cases next to the Python function it mirrors, in the Rust test and in its Python twin, like the Redis contract cases in `redis-resolve/tests/interpret.rs` and `tests/unit/test_redis.py`. A downstream crate's tests build its typed input directly and never feed it raw text to coerce diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 5eaf876e3b2..5fcf05f61f7 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -8,6 +8,10 @@ For string-valued enums and their Serde conversions, follow [.agents/skills/rust For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md) +For where config.yaml values, environment layering and kwargs from Python become typed Rust, follow [.agents/skills/rust-python-boundary/SKILL.md](.agents/skills/rust-python-boundary/SKILL.md) + +The process-wide rustls `CryptoProvider` (ring) is installed once by the host. rust-proxy does it at the top of `main` through `litellm_gateway::install_crypto_provider`, and python-bridge must do the same in its module init before it builds any rustls config. Library crates never install it. The AWS SDK links aws-lc-rs next to ring, so without an installed default the first rustls config a library builds, such as redis-rs TLS, panics + 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` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 107b55e89cb..ed658e074ee 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -8,6 +8,17 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures 0.2.17", +] + [[package]] name = "ahash" version = "0.8.12" @@ -883,6 +894,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" dependencies = [ "fastrand", + "tokio", ] [[package]] @@ -982,6 +994,15 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "block-padding" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93" +dependencies = [ + "generic-array", +] + [[package]] name = "bollard" version = "0.20.2" @@ -1134,6 +1155,15 @@ dependencies = [ "rustversion", ] +[[package]] +name = "cbc" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6" +dependencies = [ + "cipher", +] + [[package]] name = "cc" version = "1.3.0" @@ -1210,6 +1240,16 @@ dependencies = [ "half", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common 0.1.7", + "inout", +] + [[package]] name = "clap" version = "4.6.7" @@ -1774,6 +1814,15 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "des" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e" +dependencies = [ + "cipher", +] + [[package]] name = "digest" version = "0.10.7" @@ -3105,6 +3154,16 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "block-padding", + "generic-array", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -3599,19 +3658,15 @@ dependencies = [ name = "litellm-cache-s3" version = "0.1.0" dependencies = [ - "aws-credential-types", "aws-sdk-s3", - "aws-smithy-runtime-api", "aws-smithy-types", - "aws-types", "futures-util", - "http 1.4.2", "litellm-auth-aws", "litellm-cache", "litellm-cache-testing", "litellm-error", "litellm-http", - "reqwest", + "litellm-s3", "rstest", "serde_json", "tokio", @@ -3663,8 +3718,7 @@ dependencies = [ "strum", "tempfile", "thiserror 2.0.19", - "tracing", - "tracing-subscriber", + "veil", ] [[package]] @@ -3767,6 +3821,7 @@ dependencies = [ "http-body-util", "litellm-circuit-breaker", "litellm-config", + "litellm-core-utils", "litellm-error", "litellm-gateway-auth", "litellm-gateway-forward", @@ -3777,11 +3832,15 @@ dependencies = [ "litellm-inference", "litellm-llms", "litellm-redis-client", + "litellm-redis-connect", + "litellm-redis-resolve", "litellm-redis-testing", + "litellm-redis-types", "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", + "rustls 0.23.42", "serde", "serde_json", "tempfile", @@ -3792,6 +3851,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "veil", ] [[package]] @@ -4283,12 +4343,14 @@ dependencies = [ "jsonschema", "litellm-cost", "litellm-error", + "litellm-tracing", "macro_rules_attribute", "rstest", "schemars 1.2.2", "serde", "serde_json", "serde_with", + "sha1 0.10.7", "strum", "thiserror 2.0.19", "time", @@ -4298,9 +4360,13 @@ dependencies = [ name = "litellm-model-catalog-refresh" version = "0.1.0" dependencies = [ + "aws-credential-types", + "backon", + "bytes", "litellm-error", "litellm-http", - "litellm-model-catalog", + "litellm-s3", + "litellm-tracing", "reqwest", "rstest", "thiserror 2.0.19", @@ -4318,11 +4384,13 @@ dependencies = [ "criterion", "fancy-regex 0.19.2", "futures-util", + "indexmap 2.14.0", "litellm-auth", "litellm-auth-aws", "litellm-cache", "litellm-cache-response", "litellm-callbacks-legacy-python", + "litellm-circuit-breaker", "litellm-core-utils", "litellm-error", "litellm-host", @@ -4337,6 +4405,13 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-model-catalog-refresh", + "litellm-python-compat", + "litellm-redis-client", + "litellm-redis-connect", + "litellm-redis-resolve", + "litellm-redis-store", + "litellm-redis-types", "litellm-router-core", "litellm-router-types", "litellm-secret-value", @@ -4352,6 +4427,7 @@ dependencies = [ "pyo3-async-runtimes", "reqwest", "rstest", + "rustls 0.23.42", "serde", "serde_json", "serde_with", @@ -4387,23 +4463,110 @@ version = "0.1.0" dependencies = [ "futures-util", "litellm-circuit-breaker", + "litellm-python-compat", "litellm-redis-testing", + "litellm-redis-types", "litellm-secret-value", + "pem", + "pkcs8", "redis", "rstest", + "rustls 0.23.42", + "rustls-native-certs", + "socket2 0.6.5", "thiserror 2.0.19", "tokio", + "tracing", +] + +[[package]] +name = "litellm-redis-connect" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-auth-aws", + "litellm-auth-azure", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-circuit-breaker", + "litellm-core-utils", + "litellm-error", + "litellm-http", + "litellm-redis-client", + "litellm-redis-resolve", + "litellm-redis-testing", + "litellm-redis-types", + "litellm-secret-value", + "rstest", + "rustls 0.23.42", + "serde", + "serde_json", + "strum", + "thiserror 2.0.19", + "time", + "tokio", + "tracing", + "url", + "veil", +] + +[[package]] +name = "litellm-redis-resolve" +version = "0.1.0" +dependencies = [ + "litellm-python-compat", + "litellm-redis-types", + "litellm-secret-value", + "rstest", + "serde", + "serde_json", + "strum", + "thiserror 2.0.19", + "tracing", + "url", +] + +[[package]] +name = "litellm-redis-store" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-cache", + "litellm-cache-memory", + "litellm-circuit-breaker", + "litellm-python-compat", + "litellm-redis-client", + "litellm-redis-testing", + "litellm-secret-value", + "litellm-tracing", + "rstest", + "strum", + "thiserror 2.0.19", + "tokio", + "tracing", ] [[package]] name = "litellm-redis-testing" version = "0.1.0" dependencies = [ + "pkcs8", + "rcgen", + "tempfile", "testcontainers-modules", "thiserror 2.0.19", "tokio", ] +[[package]] +name = "litellm-redis-types" +version = "0.1.0" +dependencies = [ + "litellm-secret-value", + "rstest", + "url", +] + [[package]] name = "litellm-router" version = "0.1.0" @@ -4450,6 +4613,23 @@ dependencies = [ "serde_with", ] +[[package]] +name = "litellm-s3" +version = "0.1.0" +dependencies = [ + "aws-config", + "aws-credential-types", + "aws-sdk-s3", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "http 1.4.2", + "litellm-http", + "reqwest", + "rstest", + "thiserror 2.0.19", +] + [[package]] name = "litellm-secret-value" version = "0.1.0" @@ -5286,6 +5466,16 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" +[[package]] +name = "pbkdf2" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" +dependencies = [ + "digest 0.10.7", + "hmac 0.12.1", +] + [[package]] name = "pem" version = "4.0.0" @@ -5396,6 +5586,22 @@ dependencies = [ "spki", ] +[[package]] +name = "pkcs5" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" +dependencies = [ + "aes", + "cbc", + "der", + "des", + "pbkdf2", + "scrypt", + "sha2 0.10.9", + "spki", +] + [[package]] name = "pkcs8" version = "0.10.2" @@ -5403,6 +5609,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" dependencies = [ "der", + "pkcs5", + "rand_core 0.6.4", "spki", ] @@ -5590,6 +5798,7 @@ version = "0.29.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b" dependencies = [ + "indexmap 2.14.0", "libc", "once_cell", "portable-atomic", @@ -6411,6 +6620,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "salsa20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" +dependencies = [ + "cipher", +] + [[package]] name = "same-file" version = "1.0.6" @@ -6474,6 +6692,17 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scrypt" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0516a385866c09368f0b5bcd1caff3366aace790fcd46e2bb032697bb172fd1f" +dependencies = [ + "pbkdf2", + "salsa20", + "sha2 0.10.9", +] + [[package]] name = "sct" version = "0.7.1" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 7121df06bc0..cc72cd01dc0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -60,6 +60,7 @@ litellm-cache = { path = "crates/cache" } litellm-cache-azure-blob = { path = "crates/cache-azure-blob" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-s3 = { path = "crates/cache-s3" } +litellm-s3 = { path = "crates/s3" } litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-disk = { path = "crates/cache-disk" } litellm-cache-response = { path = "crates/cache-response" } @@ -73,7 +74,11 @@ litellm-host-python = { path = "crates/host-python" } litellm-python-compat = { path = "crates/python-compat" } litellm-circuit-breaker = { path = "crates/circuit-breaker" } litellm-redis-client = { path = "crates/redis-client" } +litellm-redis-connect = { path = "crates/redis-connect" } +litellm-redis-store = { path = "crates/redis-store" } litellm-redis-testing = { path = "crates/redis-testing" } +litellm-redis-resolve = { path = "crates/redis-resolve" } +litellm-redis-types = { path = "crates/redis-types" } litellm-cost = { path = "crates/cost" } litellm-model-catalog = { path = "crates/model-catalog" } litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" } @@ -86,6 +91,8 @@ axum = { version = "0.8.9", default-features = false, features = ["http1", "toki axum-login = "0.18.0" tower-sessions = { version = "0.14.0", features = ["memory-store"] } bytes = "1" +backon = { version = "1.6.0", default-features = false, features = ["std", "tokio-sleep"] } +aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rt-tokio"] } http = "1" google-cloud-auth = { version = "1.16.0", default-features = false } jsonwebtoken = { version = "11.1.0", default-features = false } diff --git a/litellm-rust/crates/auth-aws/src/lib.rs b/litellm-rust/crates/auth-aws/src/lib.rs index bd0f866371f..c8c5b5e8653 100644 --- a/litellm-rust/crates/auth-aws/src/lib.rs +++ b/litellm-rust/crates/auth-aws/src/lib.rs @@ -1,9 +1,11 @@ mod aws; pub mod constants; mod error; +mod provider; mod signer; pub use aws::*; pub use aws_credential_types::Credentials; pub use error::{CredentialFlow, Error, SigningStage}; +pub use provider::AwsCredentialsProvider; pub use signer::SigV4Signer; diff --git a/litellm-rust/crates/auth-aws/src/provider.rs b/litellm-rust/crates/auth-aws/src/provider.rs new file mode 100644 index 00000000000..6cb2325fc11 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/provider.rs @@ -0,0 +1,135 @@ +use std::sync::Arc; + +use aws_credential_types::Credentials; +use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; + +use crate::{AwsAuthConfig, AwsAuthService}; + +type Environment = Arc Option + Send + Sync>; + +#[derive(Clone)] +pub struct AwsCredentialsProvider { + auth: AwsAuthService, + config: Box, + environment: Environment, +} + +impl AwsCredentialsProvider { + /// Explicit keys in `config` as given, otherwise LiteLLM's `AWS_*` settings read through `environment` + pub fn configured( + auth: AwsAuthService, + config: AwsAuthConfig, + environment: impl Fn(&str) -> Option + Send + Sync + 'static, + ) -> Self { + Self { + auth, + config: Box::new(config), + environment: Arc::new(environment), + } + } + + async fn credentials(&self) -> Result { + match (&self.config.access_key_id, &self.config.secret_access_key) { + (Some(access_key_id), Some(secret_access_key)) => Ok(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + self.config + .session_token + .as_ref() + .map(|token| token.expose().to_owned()), + None, + "litellm-configured", + )), + _ => { + self.auth + .resolve_credentials( + AwsAuthConfig::clone(&self.config), + self.environment.as_ref(), + ) + .await + } + } + } +} + +impl ProvideCredentials for AwsCredentialsProvider { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::new(async { + self.credentials() + .await + .map_err(CredentialsError::provider_error) + }) + } +} + +impl std::fmt::Debug for AwsCredentialsProvider { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("AwsCredentialsProvider") + .finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn explicit(session_token: Option<&str>) -> AwsAuthConfig { + AwsAuthConfig { + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), + session_token: session_token.map(Into::into), + region_name: Some("us-east-1".to_string()), + ..Default::default() + } + } + + fn ambient_token(name: &str) -> Option { + (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()) + } + + fn environment_keys(name: &str) -> Option { + match name { + "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()), + "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()), + "AWS_SESSION_TOKEN" => Some("env-token".to_string()), + _ => None, + } + } + + #[rstest] + #[case::explicit_keys_ignore_an_ambient_session_token( + explicit(None), ambient_token, ("key", "secret", None) + )] + #[case::explicit_keys_keep_their_session_token( + explicit(Some("t")), ambient_token, ("key", "secret", Some("t")) + )] + #[case::environment_keys_resolve_with_their_session_token( + AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token")) + )] + #[tokio::test] + async fn configured_credentials_read_the_environment_they_are_given( + #[case] config: AwsAuthConfig, + #[case] environment: fn(&str) -> Option, + #[case] expected: (&str, &str, Option<&str>), + ) { + let credentials = + AwsCredentialsProvider::configured(AwsAuthService::default(), config, environment) + .provide_credentials() + .await + .unwrap(); + assert_eq!( + ( + credentials.access_key_id(), + credentials.secret_access_key(), + credentials.session_token(), + ), + expected + ); + } +} diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 73fbf14b07a..809ffe3044e 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -2,7 +2,10 @@ use std::{ cmp::Reverse, collections::{BinaryHeap, HashMap, HashSet}, hash::Hash, - sync::{Arc, Mutex, MutexGuard}, + sync::{ + Arc, Mutex, MutexGuard, + atomic::{AtomicUsize, Ordering}, + }, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -31,7 +34,7 @@ struct CacheState { pub struct InMemoryCache { state: Mutex>, - max_size_in_memory: usize, + max_size_in_memory: AtomicUsize, default_ttl: Duration, max_entry_bytes: Option, measure_value: Option>, @@ -74,7 +77,9 @@ impl InMemoryCache { expirations: HashMap::new(), expiration_heap: BinaryHeap::new(), }), - max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + max_size_in_memory: AtomicUsize::new( + max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + ), default_ttl: default_ttl .filter(|ttl| !ttl.is_zero()) .unwrap_or(DEFAULT_TTL), @@ -94,7 +99,7 @@ impl InMemoryCache { value: V, ttl: Option, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(CacheWrite::Disabled); } let now = (self.now)(); @@ -129,7 +134,7 @@ impl InMemoryCache { pub fn evict_cache(&self) -> Result<(), Error> { let now = (self.now)(); let mut state = self.lock()?; - Self::evict(&mut state, self.max_size_in_memory, now, None); + Self::evict(&mut state, self.max_size_in_memory(), now, None); Ok(()) } @@ -170,7 +175,16 @@ impl InMemoryCache { } pub fn max_size_in_memory(&self) -> usize { - self.max_size_in_memory + self.max_size_in_memory.load(Ordering::Relaxed) + } + + /// `DualCache.update_in_memory_max_size`: `None` restores the default. Entries already held stay + /// until the next write evicts down to the new size. + pub fn update_max_size_in_memory(&self, max_size_in_memory: Option) { + self.max_size_in_memory.store( + max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + Ordering::Relaxed, + ); } pub fn max_entry_bytes(&self) -> Option { @@ -200,6 +214,23 @@ impl InMemoryCache { .collect()) } + /// Every live entry, for callers that inspect the whole tier. + pub fn entries(&self) -> Result, Error> { + let now = (self.now)(); + let state = self.lock()?; + Ok(state + .values + .iter() + .filter(|(key, _)| { + state + .expirations + .get(*key) + .is_none_or(|expiration| *expiration >= now) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect()) + } + pub fn delete_cache(&self, key: &str) -> Result<(), Error> { let mut state = self.lock()?; Self::remove(&mut state, key); @@ -276,7 +307,7 @@ impl InMemoryCache { ttl: Option, now: Duration, ) -> Result { - Self::evict(state, self.max_size_in_memory, now, None); + Self::evict(state, self.max_size_in_memory(), now, None); if !self.check_value_size(&value)? { return Ok(CacheWrite::TooLarge); } @@ -300,12 +331,12 @@ where eligible: &[V], context: ExactCacheContext, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(candidate); } let now = (self.now)(); let mut state = self.lock()?; - Self::evict(&mut state, self.max_size_in_memory, now, Some(key)); + Self::evict(&mut state, self.max_size_in_memory(), now, Some(key)); let existing = state .values .get(key) @@ -335,7 +366,7 @@ impl CounterCache for InMemoryCache { amount: f64, context: ExactCacheContext, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(amount); } let now = (self.now)(); @@ -408,7 +439,7 @@ where values: Vec, ttl: Option, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(values); } let now = (self.now)(); diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml index d51c282d4d7..b0382d62b44 100644 --- a/litellm-rust/crates/cache-s3/Cargo.toml +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -9,14 +9,10 @@ repository.workspace = true litellm-http.workspace = true litellm-cache.workspace = true litellm-auth-aws.workspace = true -aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] } -aws-credential-types = "1.3.0" -aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] } -aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] } -aws-types = "1.6.0" +litellm-s3.workspace = true +aws-sdk-s3.workspace = true +aws-smithy-types = "1.6.0" futures-util.workspace = true -http.workspace = true -reqwest.workspace = true tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs deleted file mode 100644 index 02707c24881..00000000000 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ /dev/null @@ -1,61 +0,0 @@ -use aws_credential_types::{ - Credentials as AwsCredentials, - provider::{ProvideCredentials, error::CredentialsError, future}, -}; -use litellm_auth_aws::{AwsAuthConfig, AwsAuthService}; - -#[derive(Clone)] -pub struct S3Credentials { - auth: AwsAuthService, - config: AwsAuthConfig, - env: fn(&str) -> Option, -} - -impl S3Credentials { - pub fn new(config: AwsAuthConfig) -> Self { - Self::with_env(config, |name| std::env::var(name).ok()) - } - - pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option) -> Self { - Self { - auth: AwsAuthService::default(), - config, - env, - } - } -} - -impl ProvideCredentials for S3Credentials { - fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> - where - Self: 'a, - { - future::ProvideCredentials::new(async { - if let (Some(access_key_id), Some(secret_access_key)) = ( - self.config.access_key_id.clone(), - self.config.secret_access_key.clone(), - ) { - return Ok(AwsCredentials::new( - access_key_id.expose(), - secret_access_key.expose(), - self.config - .session_token - .as_ref() - .map(|value| value.expose().to_owned()), - None, - "litellm-s3-cache", - )); - } - self.auth - .resolve_credentials(self.config.clone(), &self.env) - .await - .map_err(CredentialsError::provider_error) - }) - } -} - -impl std::fmt::Debug for S3Credentials { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("S3Credentials").finish_non_exhaustive() - } -} diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index 5369bca0f18..7d5f3f50418 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -5,20 +5,18 @@ use std::{ }; use aws_sdk_s3::{ - config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation}, + config::{RequestChecksumCalculation, ResponseChecksumValidation}, error::SdkError, primitives::ByteStream, }; use aws_smithy_types::{DateTime, date_time::Format}; use futures_util::future::try_join_all; -use litellm_auth_aws::AwsAuthConfig; +use litellm_auth_aws::AwsCredentialsProvider; use litellm_cache::{ BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache, }; use tokio::runtime::Handle; -use crate::{auth::S3Credentials, transport::ReqwestHttpClient}; - pub struct S3Endpoint { pub url: String, } @@ -28,7 +26,7 @@ pub struct S3CacheConfig { pub key_prefix: String, pub region: String, pub endpoint: Option, - pub auth: AwsAuthConfig, + pub credentials: AwsCredentialsProvider, } pub struct S3Cache { @@ -49,17 +47,14 @@ impl S3Cache { runtime: Handle, ) -> Self { let endpoint_url: Option = config.endpoint.map(|endpoint| endpoint.url); - let base = aws_sdk_s3::Config::builder() - .behavior_version(BehaviorVersion::latest()) - .region(Region::new(config.region.clone())) - .http_client(ReqwestHttpClient(http)) - .credentials_provider(S3Credentials::new(config.auth)) - .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) - .response_checksum_validation(ResponseChecksumValidation::WhenRequired); - let builder = match &endpoint_url { - Some(url) => base.endpoint_url(url).force_path_style(true), - None => base, - }; + let builder = litellm_s3::config( + config.region.clone(), + endpoint_url.as_deref(), + http, + config.credentials, + ) + .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) + .response_checksum_validation(ResponseChecksumValidation::WhenRequired); Self { client: aws_sdk_s3::Client::from_conf(builder.build()), codec, diff --git a/litellm-rust/crates/cache-s3/src/lib.rs b/litellm-rust/crates/cache-s3/src/lib.rs index 9be0210aad3..43006532441 100644 --- a/litellm-rust/crates/cache-s3/src/lib.rs +++ b/litellm-rust/crates/cache-s3/src/lib.rs @@ -1,6 +1,3 @@ -mod auth; mod cache; -mod transport; -pub use auth::S3Credentials; pub use cache::{S3Cache, S3CacheConfig, S3Endpoint}; diff --git a/litellm-rust/crates/cache-s3/tests/auth.rs b/litellm-rust/crates/cache-s3/tests/auth.rs deleted file mode 100644 index 5aa251ac749..00000000000 --- a/litellm-rust/crates/cache-s3/tests/auth.rs +++ /dev/null @@ -1,57 +0,0 @@ -use aws_credential_types::provider::ProvideCredentials; -use litellm_auth_aws::AwsAuthConfig; -use litellm_cache_s3::S3Credentials; -use rstest::rstest; - -fn explicit(session_token: Option<&str>) -> AwsAuthConfig { - AwsAuthConfig { - access_key_id: Some("key".into()), - secret_access_key: Some("secret".into()), - session_token: session_token.map(Into::into), - region_name: Some("us-east-1".to_string()), - ..Default::default() - } -} - -fn ambient_token(name: &str) -> Option { - (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()) -} - -fn environment_keys(name: &str) -> Option { - match name { - "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()), - "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()), - "AWS_SESSION_TOKEN" => Some("env-token".to_string()), - _ => None, - } -} - -#[rstest] -#[case::explicit_keys_ignore_an_ambient_session_token( - explicit(None), ambient_token, ("key", "secret", None) -)] -#[case::explicit_keys_keep_their_session_token( - explicit(Some("t")), ambient_token, ("key", "secret", Some("t")) -)] -#[case::environment_keys_resolve_with_their_session_token( - AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token")) -)] -#[tokio::test] -async fn credentials_resolve( - #[case] config: AwsAuthConfig, - #[case] env: fn(&str) -> Option, - #[case] expected: (&str, &str, Option<&str>), -) { - let credentials = S3Credentials::with_env(config, env) - .provide_credentials() - .await - .unwrap(); - assert_eq!( - ( - credentials.access_key_id(), - credentials.secret_access_key(), - credentials.session_token(), - ), - expected - ); -} diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs index 311d2402e72..7b62d0f7208 100644 --- a/litellm-rust/crates/cache-s3/tests/support/mod.rs +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -5,7 +5,7 @@ use std::{ sync::{Arc, Mutex}, }; -use litellm_auth_aws::AwsAuthConfig; +use litellm_auth_aws::{AwsAuthConfig, AwsAuthService, AwsCredentialsProvider}; use litellm_cache::JsonCodec; use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint}; use serde_json::Value; @@ -22,12 +22,16 @@ pub fn config(endpoint: &str) -> S3CacheConfig { endpoint: Some(S3Endpoint { url: endpoint.to_string(), }), - auth: AwsAuthConfig { - access_key_id: Some("key".into()), - secret_access_key: Some("secret".into()), - region_name: Some("us-east-1".to_string()), - ..Default::default() - }, + credentials: AwsCredentialsProvider::configured( + AwsAuthService::default(), + AwsAuthConfig { + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + |_| None, + ), } } diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index ab887bbc014..8a212223523 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -27,6 +27,10 @@ "api_key", "additional_args" ], + "messages_input": [ + "body", + "kwargs" + ], "post_call": [ "logger", "original_response", @@ -47,9 +51,18 @@ "defers_async_logging": [ "logger" ], + "claim_async_success_log": [ + "logger" + ], + "mark_async_success_enqueued": [ + "logger" + ], "defer_success": [ "logger", - "pending" + "pending", + "response", + "start", + "end" ], "sync_success_for_async_call": [ "logger", @@ -117,7 +130,8 @@ "chunks", "start", "end", - "first_chunk" + "first_chunk", + "stream_teardown" ], "stream_failure": [ "logger", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..f68541c5c7b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -14,7 +14,7 @@ use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::{PyDateTime, PyDict, PyList}, + types::{PyDateTime, PyDict, PyList, PyString}, }; use serde_json::Value; @@ -22,7 +22,7 @@ use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, finalize, is_internal_call, - python::Streaming, + python::{Logging, Streaming}, setup, }; @@ -137,6 +137,40 @@ impl LegacyLogging { Some(params) } + // Compared as Python compares them, so a caller's `0` is kept over a typed float's `0.0`. + fn caller_value<'py>( + &self, + py: Python<'py>, + name: &str, + sent: &Value, + ) -> PyResult>> { + match self.call.lookup(py, name)? { + Some(value) if from_py::(&value).is_ok() && value.eq(to_py(py, sent)?)? => { + Ok(Some(value)) + } + _ => Ok(None), + } + } + + // Python's Messages handler logs the caller's own optional params, not their typed copies. + fn as_called(&self, py: Python<'_>, context: &RequestContext) -> PyResult { + let mut context = context.clone(); + if self.operation != LoggingOperation::Messages { + return Ok(context); + } + for (name, sent) in context + .optional_params + .as_object_mut() + .into_iter() + .flatten() + { + if let Some(value) = self.caller_value(py, name, sent)? { + *sent = from_py(&value)?; + } + } + Ok(context) + } + /// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync /// client is untimed. fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> { @@ -158,13 +192,20 @@ impl LegacyLogging { } } - fn input_description(&self) -> &'static str { - match self.operation { + fn logged_input<'py>( + &self, + py: Python<'py>, + body: &Bound<'py, PyDict>, + ) -> PyResult> { + let description = match self.operation { LoggingOperation::Completion => "Chat completions", LoggingOperation::Responses => "Responses", - LoggingOperation::Messages => "Messages", + LoggingOperation::Messages => { + return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py))); + } LoggingOperation::Ocr => "OCR document processing", - } + }; + Ok(PyString::new(py, description).into_any()) } /// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key. @@ -251,30 +292,49 @@ impl LegacyLogging { if !self.asynchronous { return pending().sync(py); } - if !self.internal - && self + if self.internal { + return Ok(()); + } + let defers = logger.defers_async_logging(py); + if logger.claim_async_success_log(py)? { + let enqueues = self .call .kwargs() .bind(py) .get_item("fallbacks")? - .is_none_or(|value| value.is_none()) - { - if logger.defers_async_logging(py) { + .is_none_or(|value| value.is_none()); + if defers { let pending = Py::new( py, PendingLogging { pending: Some(pending()), + enqueues, }, )?; - logger.defer_success(py, pending.bind(py).as_any())?; - } else { + return logger.defer_success( + py, + pending.bind(py).as_any(), + &self.response, + &self.start, + &self.end, + ); + } + if enqueues { pending().asynchronous(py)?; } } + if defers { + return Ok(()); + } logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) } - fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> { + fn stream_success( + &self, + py: Python<'_>, + stream: &DeliveredStream, + teardown: bool, + ) -> PyResult<()> { let logger = self.logger()?; let billing = self.stream_billing().ok_or(HostError::MissingState)?; let billed = Streaming::Success.call( @@ -291,6 +351,7 @@ impl LegacyLogging { &self.start, &self.end, &stream.first_chunk, + teardown, ), ); match billed { @@ -442,6 +503,7 @@ impl LegacyLogging { wire: Box, context: &RequestContext, ) -> PyResult>> { + let context = &self.as_called(py, context)?; let logger = self.logger()?; logger.update_from_kwargs( py, @@ -454,12 +516,17 @@ impl LegacyLogging { .into_bound(py) .cast_into::()?; for (name, sent) in wire.body.as_object().into_iter().flatten() { - if let Some(value) = self.call.lookup(py, name)? - && from_py::(&value).is_ok_and(|caller| caller == *sent) - { + if let Some(value) = self.caller_value(py, name, sent)? { body.set_item(name, value)?; } } + if self.operation == LoggingOperation::Messages { + // Python's Messages handler merges its transformed request body into the call details. + logger + .object(py) + .getattr("model_call_details")? + .call_method1("update", (&body,))?; + } let headers = PyDict::new(py); for (name, value) in &wire.headers { headers.set_item(name, value)?; @@ -469,9 +536,10 @@ impl LegacyLogging { headers: headers.clone().unbind(), context: context.clone(), }); + let input = self.logged_input(py, &body)?; self.logger()?.pre_call( py, - self.input_description(), + &input, self.logged_api_key(context), &body, &headers, @@ -525,7 +593,11 @@ impl LegacyLogging { .getattr("litellm_params")? .cast_into::()? .copy()?; - params.set_item("custom_llm_provider", &facts.provider.provider)?; + if self.operation != LoggingOperation::Messages { + // Python's Messages handler hands the provider to the logger beside its params, never + // in them, and the proxy's logger carries over to the next fallback hop. + params.set_item("custom_llm_provider", &facts.provider.provider)?; + } crate::python::Logging::Update.call( py, ( @@ -548,8 +620,17 @@ impl LegacyLogging { ResultSource::Provider => None, ResultSource::Cache { key } => Some(key.clone()), }; - details.set_item("cache_hit", self.cache_key.is_some())?; + if self.cache_key.is_some() { + details.set_item("cache_hit", true)?; + } else { + details.call_method1("pop", ("cache_hit", py.None()))?; + } details.set_item("cache_key", self.cache_key.as_deref())?; + if let (LoggingOperation::Messages, ResultSource::Provider, Some(request)) = + (self.operation, &facts.source, &self.request) + { + details.call_method1("update", (request.body.bind(py),))?; + } Ok(HookStep::Ready(())) } @@ -592,7 +673,7 @@ impl LegacyLogging { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); match &self.stream { - Some(stream) => self.stream_success(py, stream)?, + Some(stream) => self.stream_success(py, stream, false)?, None => self.dispatch_success(py)?, } Ok(HookStep::Ready(())) @@ -662,8 +743,28 @@ impl LegacyLogging { } } +impl LegacyLogging { + // A caller that goes away mid-stream cancels the call before it ends: Python's Messages + // stream iterator still bills the chunks it read, straight away rather than deferred. + fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> { + if !self.asynchronous || self.end.is_some() { + return Ok(()); + } + let Some(stream) = self.stream.take() else { + return Ok(()); + }; + self.end = Some(datetime(py, epoch_seconds())?); + self.stream_success(py, &stream, true) + } +} + impl PythonOwned for LegacyLogging { fn close(&mut self, py: Python<'_>) { + if self.logger.is_some() + && let Err(error) = self.bill_abandoned_stream(py) + { + error.write_unraisable(py, None); + } if let Some(logger) = self.logger.take() && let Err(error) = logger.restore_context(py) { @@ -1114,10 +1215,10 @@ logger._provider_reported_cost = None }; let previous = litellm_host::interceptors::ExecutionFacts { provider: facts.provider.clone(), - cost: litellm_host::interceptors::Cost::Reported { - amount: serde_json::Number::from_f64(9.0).unwrap(), + cost: litellm_host::interceptors::Cost::CacheHit, + source: litellm_host::interceptors::ResultSource::Cache { + key: "earlier-hit".into(), }, - source: litellm_host::interceptors::ResultSource::Provider, }; logging.result_ready(py, &previous).unwrap(); local(&locals, "logger") @@ -1135,13 +1236,59 @@ logger._provider_reported_cost = None c" assert logger.model_call_details.get('response_cost') is None, logger.model_call_details assert logger._provider_reported_cost == reported, logger._provider_reported_cost -assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details +assert logger.model_call_details.get('cache_hit') is (True if cache_hit else None), logger.model_call_details assert logger.update['custom_llm_provider'] == 'edenai', logger.update ", ); }); } + #[rstest] + #[case::messages(crate::LoggingOperation::Messages, false)] + #[case::ocr(crate::LoggingOperation::Ocr, true)] + fn result_ready_puts_the_provider_in_the_logged_params_except_for_messages( + #[case] operation: crate::LoggingOperation, + #[case] in_params: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + run( + py, + &locals, + c" +logger.litellm_params = {} +logger.optional_params = {} +logger.model_call_details = {} +logger._provider_reported_cost = None +", + ); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + operation, + ..legacy_call(py, &locals, false) + }; + let facts = litellm_host::interceptors::ExecutionFacts { + provider: litellm_host::interceptors::ProviderIdentity { + model: "native-test".into(), + provider: "anthropic".into(), + }, + cost: litellm_host::interceptors::Cost::Deferred, + source: litellm_host::interceptors::ResultSource::Provider, + }; + logging.result_ready(py, &facts).unwrap(); + locals.set_item("in_params", in_params).unwrap(); + run( + py, + &locals, + c" +assert logger.update['custom_llm_provider'] == 'anthropic', logger.update +assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params, logger.update +", + ); + }); + } + const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; @@ -1276,7 +1423,10 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update #[rstest] #[case::completion(crate::LoggingOperation::Completion, "Chat completions")] #[case::responses(crate::LoggingOperation::Responses, "Responses")] - #[case::messages(crate::LoggingOperation::Messages, "Messages")] + #[case::messages( + crate::LoggingOperation::Messages, + r#"[{'role': 'user', 'content': '{"pages": [1]}'}]"# + )] #[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")] fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases( #[case] operation: crate::LoggingOperation, @@ -1317,7 +1467,7 @@ prepared = {'pages': replacement} &locals, c" assert logger.pre['complete_input_dict']['pages'] is replacement -assert logger.pre_input == description +assert str(logger.pre_input) == description assert original == [0] ", ); @@ -1873,10 +2023,15 @@ mod terminal_tests { #[case::async_listened( true, c"", - &["async_success_handler", "enqueued", "sync_success_for_async_call"] + &["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call"] )] - #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] + #[case::async_deferred(true, c"logger._defer_async_logging = True", &[])] #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + #[case::async_deferred_with_fallbacks( + true, + c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}", + &[] + )] fn success_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, @@ -1902,12 +2057,88 @@ mod terminal_tests { &locals, c" assert all(value is response for name, value in logger.calls if name.endswith('_handler')) -assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) +assert hasattr(logger, 'enqueue_deferred_logging') == getattr(logger, '_defer_async_logging', False) +assert getattr(logger, 'deferred_success_result', (response,))[0] is response ", ); }); } + #[rstest] + #[case::accepted(c"", &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"])] + #[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + fn a_deferred_success_runs_sync_callbacks_only_once_released( + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"response = object()\nlogger._defer_async_logging = True", + ); + run(py, &locals, script); + let mut logging = logged(py, &locals, true); + succeed(py, &locals, &mut logging); + run( + py, + &locals, + c" +assert logger.names() == [], logger.calls +logger.enqueue_deferred_logging() +logger.enqueue_deferred_logging() +", + ); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + + #[rstest] + #[case::deferred( + c"logger._defer_async_logging = True", + c"logger.enqueue_deferred_logging()", + &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"] + )] + #[case::deferred_and_blocked(c"logger._defer_async_logging = True", c"", &[])] + #[case::immediate( + c"", + c"", + &["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call", "sync_success_for_async_call"] + )] + fn a_fallback_attempt_on_the_same_logger_leaves_the_first_success_claimed( + #[case] script: &CStr, + #[case] release: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"refused = object()\nresponse = refused"); + run(py, &locals, script); + let mut first = logged(py, &locals, true); + succeed(py, &locals, &mut first); + run(py, &locals, c"response = object()"); + let mut fallback = logged(py, &locals, true); + succeed(py, &locals, &mut fallback); + run(py, &locals, release); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c"assert all(value is refused for name, value in logger.calls if name == 'async_success_handler')", + ); + }); + } + #[rstest] fn dropped_observations_preserve_deferred_success_and_response_identity() { Python::initialize(); @@ -1933,10 +2164,10 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy )); assert_eq!(sender.dropped_events(), 1); run(py, &locals, c" -assert logger.names() == ['sync_success_for_async_call'], logger.calls -logger._native_pending_logging.release(True) -logger._native_pending_logging.release(True) -assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls +assert logger.names() == [], logger.calls +logger.enqueue_deferred_logging() +logger.enqueue_deferred_logging() +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls assert logger.calls[0][1] is response assert logger.calls[1][1] is response "); @@ -1981,6 +2212,45 @@ assert chunks[1] is last }); } + #[rstest] + #[case::cancelled_mid_stream(None, &["stamp_stream_head", "stream_opened", "stream_success", "restore"])] + #[case::after_success(Some(true), &["stamp_stream_head", "stream_opened", "stream_success", "restore"])] + #[case::after_failure(Some(false), &["stamp_stream_head", "stream_opened", "stream_failure", "restore"])] + fn closing_an_unfinished_stream_bills_its_delivered_chunks_once( + #[case] ended: Option, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"first = b'first'\nresponse = None\nfailure = ValueError('provider')\nhead = {}", + ); + let mut logging = LegacyLogging { + operation: crate::LoggingOperation::Messages, + ..logged(py, &locals, true) + }; + logging + .on_stream_open(py, &local(&locals, "head").unbind()) + .unwrap(); + logging + .on_stream_chunk(py, &local(&locals, "first").unbind()) + .unwrap(); + match ended { + Some(true) => drop(succeed(py, &locals, &mut logging)), + Some(false) => drop(fail(py, &locals, &mut logging)), + None => {} + } + litellm_host_python::PythonOwned::close(&mut logging, py); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + #[rstest] #[case::synchronous(false, &["failure_handler"])] #[case::asynchronous(true, &[])] @@ -2009,7 +2279,7 @@ assert chunks[1] is last } #[rstest::rstest] - fn internal_async_calls_skip_the_async_success_fan_out() { + fn internal_async_calls_skip_success_logging() { Python::initialize(); Python::attach(|py| { let locals = namespace(py, c"response = object()"); @@ -2018,11 +2288,7 @@ assert chunks[1] is last ..logged(py, &locals, true) }; succeed(py, &locals, &mut logging); - run( - py, - &locals, - c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", - ); + run(py, &locals, c"assert logger.calls == [], logger.calls"); }); } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index 1951c9ebfb4..093e86c43dc 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -26,7 +26,7 @@ pub trait LegacyCallbacks { fn pre_call( &self, py: Python<'_>, - input: &str, + input: &Bound<'_, PyAny>, api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, @@ -59,7 +59,16 @@ pub trait LegacyCallbacks { fn defers_async_logging(&self, py: Python<'_>) -> bool; - fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>; + fn claim_async_success_log(&self, py: Python<'_>) -> PyResult; + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()>; fn sync_success_for_async_call( &self, @@ -148,7 +157,7 @@ impl LegacyCallbacks for PythonLogger { fn pre_call( &self, py: Python<'_>, - input: &str, + input: &Bound<'_, PyAny>, api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, @@ -203,8 +212,21 @@ impl LegacyCallbacks for PythonLogger { .unwrap_or(false) } - fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> { - Logging::DeferSuccess.call(py, (self.object(py), pending))?; + fn claim_async_success_log(&self, py: Python<'_>) -> PyResult { + Logging::ClaimAsyncSuccess + .call(py, (self.object(py),))? + .extract() + } + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + Logging::DeferSuccess.call(py, (self.object(py), pending, response, start, end))?; Ok(()) } @@ -252,6 +274,7 @@ impl LegacyCallbacks for PythonLogger { ) -> PyResult<()> { let coroutine = Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?; + Logging::MarkAsyncSuccessEnqueued.call(py, (self.object(py),))?; let enqueue = Logging::Enqueue.call(py, (&coroutine,)); if enqueue.is_err() && let Err(error) = coroutine.call_method0("close") diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index 292648ba69c..ea72191ddb2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -22,21 +22,35 @@ impl PendingSuccess { self.logger .enqueue_success(py, &self.response, &self.start, &self.end) } + + pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .sync_success_for_async_call(py, &self.response, &self.start, &self.end) + } } #[pyclass] pub(crate) struct PendingLogging { pub(crate) pending: Option, + pub(crate) enqueues: bool, } #[pymethods] impl PendingLogging { - fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { - let pending = slf.borrow_mut().pending.take(); - if let Some(pending) = pending - && success - { - match pending.asynchronous(py) { + fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> { + let (pending, enqueues) = { + let mut slot = slf.borrow_mut(); + (slot.pending.take(), slot.enqueues) + }; + if let Some(pending) = pending { + let released = pending.sync_for_async_call(py).and_then(|()| { + if enqueues { + pending.asynchronous(py) + } else { + Ok(()) + } + }); + match released { Err(error) if error.is_instance_of::(py) => { error.write_unraisable(py, Some(pending.logger.object(py))); } @@ -87,6 +101,7 @@ mod tests { start: py.None(), end: Some(py.None()), }), + enqueues: true, }, ) .unwrap(); @@ -108,7 +123,7 @@ observed = [] def on_enqueue(coroutine): observed.append(marker.get()) - pending.release(True) + pending() logger.on_enqueue = on_enqueue ", @@ -118,28 +133,12 @@ logger.on_enqueue = on_enqueue &locals, c" marker.set('release') -pending.release(True) -pending.release(True) +pending() +pending() assert observed == ['release'], observed -assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls assert logger.calls[0][1] is response -", - ); - }); - } - - #[test] - fn a_blocked_release_drops_the_success_for_good() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -pending.release(False) -pending.release(True) -assert logger.calls == [], logger.calls +assert logger.calls[1][1] is response ", ); }); @@ -168,7 +167,7 @@ logger.on_enqueue = on_enqueue locals .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) .unwrap(); - let released = local(&locals, "pending").call_method1("release", (true,)); + let released = local(&locals, "pending").call0(); match released { Ok(_) => assert!(!propagates), Err(error) => { @@ -181,8 +180,8 @@ logger.on_enqueue = on_enqueue py, &locals, c" -pending.release(True) -assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls +pending() +assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued', 'closed'], logger.calls assert unraisable_from(logger) == ([] if propagates else [failure]) ", ); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6d910d35f53..5b6faa3dbc2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -42,6 +42,8 @@ pub(crate) enum Logging { Update, #[strum(serialize = "pre_call")] PreCall, + #[strum(serialize = "messages_input")] + MessagesInput, #[strum(serialize = "post_call")] PostCall, #[strum(serialize = "record_response")] @@ -50,6 +52,8 @@ pub(crate) enum Logging { RecordApiWindow, #[strum(serialize = "defers_async_logging")] DefersAsync, + #[strum(serialize = "claim_async_success_log")] + ClaimAsyncSuccess, #[strum(serialize = "defer_success")] DeferSuccess, #[strum(serialize = "sync_success_for_async_call")] @@ -58,6 +62,8 @@ pub(crate) enum Logging { SubmitSuccess, #[strum(serialize = "async_success_handler")] AsyncSuccessHandler, + #[strum(serialize = "mark_async_success_enqueued")] + MarkAsyncSuccessEnqueued, #[strum(serialize = "enqueue_logging")] Enqueue, #[strum(serialize = "failure_handler")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 8acfa7a898b..ee16e0c0059 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -62,6 +62,7 @@ FAKES = { custom_llm_provider=provider, ), 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'messages_input': lambda body, kwargs: [{'role': 'user', 'content': json.dumps(body)}], 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( original_response, api_key, additional_args ), @@ -72,7 +73,11 @@ FAKES = { llm_api_duration_ms=(end - start) * 1000 ), 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), - 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'claim_async_success_log': lambda logger: logger.claim_async_success_log(), + 'defer_success': lambda logger, pending, response, start, end: ( + setattr(logger, 'deferred_success_result', (response, start, end)), + setattr(logger, 'enqueue_deferred_logging', pending), + ), 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( response, start, end ), @@ -81,6 +86,7 @@ FAKES = { )(error, ''.join(traceback.format_exception(error)), start, end), 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None), 'enqueue_logging': lambda coroutine: coroutine.enqueue(), 'restore_context': lambda logger: logger.record('restore', None), 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), @@ -92,7 +98,7 @@ FAKES = { 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), 'stream_opened': lambda logger, head: logger.record('stream_opened', head), 'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head), - 'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record( 'stream_success', list(chunks) ), 'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error), @@ -129,6 +135,7 @@ class StubCoroutine: class StubLogger: def __init__(self): self.calls = [] + self.model_call_details = {} self.hooks = {} self.on_enqueue = lambda coroutine: None @@ -156,6 +163,10 @@ class StubLogger: self.record('async_success_handler', response) return StubCoroutine(self) + def claim_async_success_log(self): + claimed, self.claimed = getattr(self, 'claimed', False), True + return not claimed + def handle_sync_success_callbacks_for_async_calls(self, response, start, end): self.record('sync_success_for_async_call', response) diff --git a/litellm-rust/crates/circuit-breaker/AGENTS.md b/litellm-rust/crates/circuit-breaker/AGENTS.md index 52ef0ca7435..852cd72664a 100644 --- a/litellm-rust/crates/circuit-breaker/AGENTS.md +++ b/litellm-rust/crates/circuit-breaker/AGENTS.md @@ -1,13 +1,14 @@ - `fsm` is a generic state machine runtime and knows nothing about breakers. `breaker` is a generic circuit breaker and knows nothing about Redis. Deciding which errors count as failures belongs to the consumer - `Machine::step` stays pure. `&self` is config and time comes in on the event, never read inside the machine. The only clock read is `Admission`'s drop, at the call-site boundary - Transition rules are one `match (state, event)` per machine, with no macro DSL and no `_` arm that would hide a new state from the compiler -- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes, matching Python bumping `generation` only when `_state` changes. Data inside a variant (failure counters, `since`) changes without a transition +- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes. Data inside a variant (failure counters, `since`) changes without a transition - The observer runs under the lock so transitions are seen in generation order. Keep it to logging or a metric - A disabled breaker is `Option::None` at the consumer, not a flag inside the machine. `Config` comes in already resolved, reading `REDIS_CIRCUIT_BREAKER_*` env is the consumer's job -## Parity with Python `RedisCircuitBreaker` (`litellm/caching/redis_cache.py`) -- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against the Python breaker. A behavior change here needs the same check, by driving the Python class with the same steps -- Divergences, each on purpose +## The only Redis breaker +- Python's `RedisCache` admits through this breaker via `NativeRedisBreaker` (`python-bridge/src/redis/breaker.rs`), sharing one per Redis server with the native client through `redis-connect`'s `RedisHosts`, so the Python `RedisCircuitBreaker` it replaced is gone +- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against that Python class and now pins its behavior. A behavior change here changes what every Redis caller in the proxy sees +- Divergences from the old Python class, each on purpose - monotonic `Instant` instead of `time.time()`, so a wall clock jump can't open or close the breaker - a probe `Admission` dropped without an outcome records a hard failure and reopens. Python's probe that dies on `CancelledError` (a `BaseException`) records nothing and leaves the breaker HALF_OPEN for good, and a dropped Rust future is common - `_failure_count` is dropped, it only fed a log line diff --git a/litellm-rust/crates/circuit-breaker/src/breaker.rs b/litellm-rust/crates/circuit-breaker/src/breaker.rs index c9509561b37..cf4961ef099 100644 --- a/litellm-rust/crates/circuit-breaker/src/breaker.rs +++ b/litellm-rust/crates/circuit-breaker/src/breaker.rs @@ -1,5 +1,5 @@ use std::{ - sync::Arc, + sync::{Arc, Weak}, time::{Duration, Instant}, }; @@ -128,6 +128,9 @@ impl Machine for Config { #[derive(Clone)] pub struct Breaker(Arc>); +#[derive(Clone)] +pub struct WeakBreaker(Weak>); + #[derive(Debug, PartialEq, Eq)] pub struct Rejected; @@ -165,6 +168,16 @@ impl Breaker { pub fn state(&self) -> State { self.0.state() } + + pub fn downgrade(&self) -> WeakBreaker { + WeakBreaker(Arc::downgrade(&self.0)) + } +} + +impl WeakBreaker { + pub fn upgrade(&self) -> Option { + self.0.upgrade().map(Breaker) + } } impl Admission { diff --git a/litellm-rust/crates/circuit-breaker/src/lib.rs b/litellm-rust/crates/circuit-breaker/src/lib.rs index e76228122d1..af8df68246e 100644 --- a/litellm-rust/crates/circuit-breaker/src/lib.rs +++ b/litellm-rust/crates/circuit-breaker/src/lib.rs @@ -1,5 +1,7 @@ mod breaker; mod fsm; -pub use breaker::{Admission, Breaker, Config, Event, Failure, Output, Rejected, State}; +pub use breaker::{ + Admission, Breaker, Config, Event, Failure, Output, Rejected, State, WeakBreaker, +}; pub use fsm::{Machine, Observer, Shared, Transition}; diff --git a/litellm-rust/crates/config/AGENTS.md b/litellm-rust/crates/config/AGENTS.md new file mode 100644 index 00000000000..9f6f1db9a67 --- /dev/null +++ b/litellm-rust/crates/config/AGENTS.md @@ -0,0 +1,3 @@ +## Redis settings + +- This crate only deserializes YAML. `RedisSettings` is the raw `coordination_redis` or `cache_params` map and `CacheSettings` tells a `type: redis` cache apart from the rest. Reading those maps the way `litellm/_redis.py` does belongs to `gateway/src/redis`, and turning merged kwargs into a `Resolution` belongs to `redis-resolve` diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 278aea6200e..a433aacc1c9 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -6,17 +6,16 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-secret-value.workspace = true litellm-router-types.workspace = true +litellm-secret-value.workspace = true serde.workspace = true -serde_json.workspace = true serde_yaml_ng = "0.10.0" strum.workspace = true thiserror.workspace = true -tracing.workspace = true +veil.workspace = true [dev-dependencies] litellm-error.workspace = true rstest.workspace = true +serde_json.workspace = true tempfile.workspace = true -tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs index e5219ac992c..300f15380bd 100644 --- a/litellm-rust/crates/config/src/error.rs +++ b/litellm-rust/crates/config/src/error.rs @@ -11,11 +11,4 @@ pub enum Error { path: PathBuf, source: serde_yaml_ng::Error, }, - #[error("Redis host does not yet support configured setting {0}")] - RedisSetting(String), - #[error("invalid Redis setting {setting}: {reason}")] - RedisInvalid { - setting: &'static str, - reason: &'static str, - }, } diff --git a/litellm-rust/crates/config/src/lib.rs b/litellm-rust/crates/config/src/lib.rs index b8af6f11d2e..46657aa91e7 100644 --- a/litellm-rust/crates/config/src/lib.rs +++ b/litellm-rust/crates/config/src/lib.rs @@ -14,10 +14,7 @@ pub use error::Error; pub use litellm_router_types::{LitellmParams, Spelled}; pub use mcp::{McpAuth, McpServer, McpTransport}; pub use model::Model; -pub use redis::{ - CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings, - ResolvedRedis, environment_has_target, -}; +pub use redis::{CacheSettings, RedisSettings}; pub use settings::{ ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings, TracingStoreSettings, diff --git a/litellm-rust/crates/config/src/redis.rs b/litellm-rust/crates/config/src/redis.rs index 7a91691c54b..d0ba107242c 100644 --- a/litellm-rust/crates/config/src/redis.rs +++ b/litellm-rust/crates/config/src/redis.rs @@ -1,195 +1,21 @@ -use std::{collections::BTreeMap, fmt, time::Duration}; +use std::collections::BTreeMap; -use litellm_router_types::Spelled; -use litellm_secret_value::SecretValue; use serde::{Deserialize, Deserializer, de::Error as _}; +use veil::Redact; -use crate::{AdditionalFields, Error, Object, Value}; +use crate::{Object, Value}; -pub type Lookup<'a> = &'a dyn Fn(&str) -> Option; - -const SUPPORTED: &[&str] = &[ - "url", - "host", - "port", - "username", - "password", - "ssl", - "db", - "socket_timeout", - "startup_nodes", - "namespace", -]; - -const IGNORED: &[&str] = &[ - "max_connections", - "single_connection_client", - "health_check_interval", - "socket_keepalive", - "socket_keepalive_options", - "decode_responses", - "encoding", - "encoding_errors", - "charset", - "errors", - "client_name", - "lib_name", - "lib_version", - "sentinel_password", - "gcp_ssl_ca_certs", - "azure_client_id", - "azure_tenant_id", - "azure_client_secret", - "aws_iam_user_name", - "aws_iam_cache_name", - "aws_iam_region", - "aws_iam_serverless", -]; - -const WARN_IGNORED: &[&str] = &[ - "max_connections", - "health_check_interval", - "socket_keepalive", - "socket_keepalive_options", -]; - -const REDIS_PY_KWARGS: &[&str] = &[ - "aws_iam_auth", - "aws_iam_cache_name", - "aws_iam_region", - "aws_iam_serverless", - "aws_iam_user_name", - "azure_client_id", - "azure_client_secret", - "azure_redis_ad_token", - "azure_tenant_id", - "cache", - "cache_config", - "charset", - "client_name", - "db", - "decode_responses", - "encoding", - "encoding_errors", - "errors", - "event_dispatcher", - "gcp_service_account", - "gcp_ssl_ca_certs", - "health_check_interval", - "host", - "lib_name", - "lib_version", - "max_connections", - "password", - "port", - "protocol", - "redis_connect_func", - "retry_on_error", - "retry_on_timeout", - "single_connection_client", - "socket_connect_timeout", - "socket_keepalive", - "socket_keepalive_options", - "socket_timeout", - "ssl", - "ssl_ca_certs", - "ssl_ca_data", - "ssl_ca_path", - "ssl_cert_reqs", - "ssl_certfile", - "ssl_check_hostname", - "ssl_ciphers", - "ssl_keyfile", - "ssl_min_version", - "ssl_ocsp_context", - "ssl_ocsp_expected_cert", - "ssl_password", - "ssl_validate_ocsp", - "ssl_validate_ocsp_stapled", - "unix_socket_path", - "url", - "username", -]; - -const UNSUPPORTED_EXTRA: &[&str] = &[ - "credential_provider", - "sentinel_nodes", - "retry", - "connection_pool", -]; - -const CACHE_ONLY: &[&str] = &[ - "type", - "mode", - "ttl", - "default_in_memory_ttl", - "default_in_redis_ttl", - "supported_call_types", - "similarity_threshold", - "max_messages", - "_backend", -]; - -#[derive(Clone, Default)] +/// A Redis section as `config.yaml` spells it, before any surface reads it. +#[derive(Clone, Default, Deserialize, Redact)] +#[serde(transparent)] pub struct RedisSettings { - pub url: Option, - pub host: Option>, - pub port: Option>, - pub username: Option, - pub password: Option, - pub ssl: Option>, - pub db: Option>, - pub socket_timeout: Option>, - pub startup_nodes: Option, - pub namespace: Option>, - pub additional_fields: AdditionalFields, - unparsed: DeferredInputs, + #[redact(with = "[REDACTED]")] + fields: BTreeMap, } -impl<'de> Deserialize<'de> for RedisSettings { - fn deserialize>(deserializer: D) -> Result { - let fields = BTreeMap::::deserialize(deserializer)?; - let scalar = |key: &str| fields.get(key).and_then(as_text); - Ok(Self { - url: scalar("url").map(SecretValue::new), - host: scalar("host").map(Spelled::Value), - port: spelled(scalar("port")), - username: scalar("username").map(SecretValue::new), - password: scalar("password").map(SecretValue::new), - ssl: spelled(scalar("ssl")), - db: spelled(scalar("db")), - socket_timeout: spelled(scalar("socket_timeout")), - startup_nodes: fields - .get("startup_nodes") - .filter(|value| !value.is_null()) - .cloned(), - namespace: scalar("namespace").map(Spelled::Value), - additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)), - unparsed: deferred_inputs(&fields), - }) - } -} - -impl fmt::Debug for RedisSettings { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("RedisSettings") - .field("url", &self.url) - .field("host", &self.host) - .field("port", &self.port) - .field("username", &self.username) - .field("password", &self.password) - .field("ssl", &self.ssl) - .field("db", &self.db) - .field("socket_timeout", &self.socket_timeout) - .field( - "startup_nodes", - &self.startup_nodes.as_ref().map(|_| ""), - ) - .field("namespace", &self.namespace) - .field("additional_fields", &self.additional_fields.keys()) - .field("unparsed", &self.unparsed.keys()) - .finish() +impl RedisSettings { + pub fn fields(&self) -> &BTreeMap { + &self.fields } } @@ -216,452 +42,22 @@ impl<'de> Deserialize<'de> for CacheSettings { } } -#[derive(Clone, PartialEq)] -pub struct ResolvedRedis { - pub endpoint: RedisEndpoint, - pub namespace: Option, - pub timeout: Option, -} +#[cfg(test)] +mod tests { + use super::RedisSettings; + use rstest::rstest; -impl fmt::Debug for ResolvedRedis { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("ResolvedRedis") - .field("endpoint", &self.endpoint) - .field("namespace", &self.namespace) - .field("timeout", &self.timeout) - .finish() + #[rstest] + #[case::scalar(serde_json::json!("credential-value"))] + #[case::object(serde_json::json!({"secret": "credential-value"}))] + #[case::array(serde_json::json!(["credential-value"]))] + fn credentials_are_redacted(#[case] password: serde_json::Value) { + let settings: RedisSettings = serde_json::from_value(serde_json::json!({ + "host": "redis", "password": password, + })) + .unwrap(); + let debug = format!("{settings:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("[REDACTED]")); } } - -#[derive(Clone, PartialEq)] -pub enum RedisEndpoint { - Url(SecretValue), - Host(RedisHost), - Cluster(RedisCluster), -} - -impl fmt::Debug for RedisEndpoint { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Url(_) => formatter.debug_tuple("Url").field(&"").finish(), - Self::Host(host) => host.fmt(formatter), - Self::Cluster(cluster) => cluster.fmt(formatter), - } - } -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisHost { - pub host: String, - pub port: u16, - pub username: Option, - pub password: Option, - pub tls: bool, - pub db: i64, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisCluster { - pub nodes: Vec, - pub username: Option, - pub password: Option, - pub tls: bool, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisNode { - pub host: String, - pub port: u16, -} - -#[derive(Clone)] -struct Entry { - value: Value, - source: String, -} - -type Kwargs = BTreeMap; - -impl RedisSettings { - pub fn resolve(&self, lookup: Lookup<'_>) -> Result { - self.resolve_with(lookup, false) - } - - pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result { - self.resolve_with(lookup, true) - } - - fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result { - let block = self.kwargs(cache); - let block: Kwargs = block - .into_iter() - .filter_map(|(key, entry)| { - let value = environment_reference(&entry.value, lookup)?; - Some((key, Entry { value, ..entry })) - }) - .collect(); - let environment = environment_kwargs(lookup); - let names_a_target = ["host", "startup_nodes", "sentinel_nodes"] - .iter() - .any(|key| block.contains_key(*key)); - let drop_environment_url = names_a_target && !block.contains_key("url"); - let merged: Kwargs = environment - .into_iter() - .filter(|(key, _)| !(key == "url" && drop_environment_url)) - .chain(block) - .collect(); - reject_unsupported(&merged, lookup)?; - warn_ignored(&merged); - - let startup_nodes = match merged.get("startup_nodes") { - Some(entry) => Some(cluster_nodes(&entry.value)?), - None => lookup("REDIS_CLUSTER_NODES") - .map(|nodes| cluster_nodes(&Value::String(nodes))) - .transpose()?, - }; - let username = secret_value(&merged, "username"); - let password = secret_value(&merged, "password"); - let tls = bool_value(&merged, "ssl"); - let endpoint = match (startup_nodes, text_value(&merged, "url")) { - (Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster { - nodes, - username, - password, - tls, - }), - (_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)), - _ => RedisEndpoint::Host(RedisHost { - host: text_value(&merged, "host").ok_or(Error::RedisInvalid { - setting: "host", - reason: "either host or url must be specified for redis", - })?, - port: integer(&merged, "port")?.unwrap_or(6379), - username, - password, - tls, - db: integer(&merged, "db")?.unwrap_or(0), - }), - }; - let timeout = merged - .get("socket_timeout") - .and_then(|entry| as_text(&entry.value)) - .map(|seconds| { - seconds - .parse::() - .ok() - .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) - .ok_or(Error::RedisInvalid { - setting: "socket_timeout", - reason: "must be a non-negative number of seconds", - }) - }) - .transpose()?; - Ok(ResolvedRedis { - endpoint, - namespace: text_value(&merged, "namespace"), - timeout, - }) - } - - fn kwargs(&self, cache: bool) -> Kwargs { - let known = [ - ( - "url", - self.url.as_ref().map(|value| value.expose().to_owned()), - ), - ("host", spelled_text(self.host.as_ref())), - ("port", spelled_text(self.port.as_ref())), - ( - "username", - self.username - .as_ref() - .map(|value| value.expose().to_owned()), - ), - ( - "password", - self.password - .as_ref() - .map(|value| value.expose().to_owned()), - ), - ("ssl", spelled_text(self.ssl.as_ref())), - ("db", spelled_text(self.db.as_ref())), - ("socket_timeout", spelled_text(self.socket_timeout.as_ref())), - ("namespace", spelled_text(self.namespace.as_ref())), - ]; - let known = known.into_iter().filter_map(|(key, value)| { - value.map(|value| { - ( - key.to_owned(), - Entry { - value: Value::String(value), - source: key.to_owned(), - }, - ) - }) - }); - let startup_nodes = self.startup_nodes.clone().map(|value| { - ( - "startup_nodes".to_owned(), - Entry { - value, - source: "startup_nodes".to_owned(), - }, - ) - }); - let extras = self.additional_fields.iter().filter_map(|(key, value)| { - if cache && CACHE_ONLY.contains(&key.as_str()) { - return None; - } - let key = match (cache, key.as_str()) { - (true, "redis_startup_nodes") => "startup_nodes".to_owned(), - _ => key.clone(), - }; - Some(( - key.clone(), - Entry { - value: value.clone(), - source: key, - }, - )) - }); - let unparsed = self.unparsed.iter().map(|(key, value)| { - ( - key.clone(), - Entry { - value: value.0.clone(), - source: key.clone(), - }, - ) - }); - extras - .chain(unparsed) - .chain(known) - .chain(startup_nodes) - .collect() - } -} - -pub fn environment_has_target(lookup: Lookup<'_>) -> bool { - [ - "REDIS_HOST", - "REDIS_URL", - "REDIS_CLUSTER_NODES", - "REDIS_SENTINEL_NODES", - ] - .iter() - .any(|name| lookup(name).is_some()) -} - -fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> { - let active_auth = [ - merged - .get("gcp_service_account") - .map(|entry| entry.source.as_str()) - .or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")), - text_value(merged, "azure_redis_ad_token") - .or_else(|| lookup("REDIS_AZURE_AD_TOKEN")) - .filter(|value| value.eq_ignore_ascii_case("true")) - .map(|_| { - merged - .get("azure_redis_ad_token") - .map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str()) - }), - bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()), - merged - .get("sentinel_nodes") - .map(|entry| entry.source.as_str()) - .or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")), - ]; - if let Some(setting) = active_auth.into_iter().flatten().next() { - return Err(Error::RedisSetting(setting.to_owned())); - } - match merged - .iter() - .find(|(key, _)| is_unsupported(key)) - .map(|(_, entry)| entry.source.clone()) - { - Some(setting) => Err(Error::RedisSetting(setting)), - None => Ok(()), - } -} - -fn warn_ignored(merged: &Kwargs) { - for setting in WARN_IGNORED { - if merged.contains_key(*setting) { - tracing::warn!(setting, "Redis setting is accepted but ignored"); - } - } -} - -fn is_unsupported(key: &str) -> bool { - let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key); - let handled = SUPPORTED.contains(&key) - || IGNORED.contains(&key) - || matches!( - key, - "gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth" - ); - known && !handled -} - -fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs { - REDIS_PY_KWARGS - .iter() - .filter_map(|key| { - let source = format!("REDIS_{}", key.to_ascii_uppercase()); - let value = lookup(&source)?; - Some(( - (*key).to_owned(), - Entry { - value: Value::String(value), - source, - }, - )) - }) - .collect() -} - -fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option { - match value { - Value::String(text) => match text.strip_prefix("os.environ/") { - Some(name) => lookup(name).map(Value::String), - None => Some(value.clone()), - }, - Value::Null => None, - _ => Some(value.clone()), - } -} - -fn as_text(value: &Value) -> Option { - match value { - Value::String(text) => Some(text.clone()), - Value::Number(number) => Some(number.to_string()), - Value::Bool(flag) => Some(flag.to_string()), - _ => None, - } -} - -fn text_value(kwargs: &Kwargs, key: &str) -> Option { - kwargs.get(key).and_then(|entry| as_text(&entry.value)) -} - -fn secret_value(kwargs: &Kwargs, key: &str) -> Option { - text_value(kwargs, key).map(SecretValue::new) -} - -fn bool_value(kwargs: &Kwargs, key: &str) -> bool { - match kwargs.get(key).map(|entry| &entry.value) { - Some(Value::Bool(flag)) => *flag, - Some(value) => as_text(value) - .is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")), - None => false, - } -} - -fn integer(kwargs: &Kwargs, key: &'static str) -> Result, Error> { - kwargs - .get(key) - .and_then(|entry| as_text(&entry.value)) - .map(|text| { - text.parse().map_err(|_| Error::RedisInvalid { - setting: key, - reason: "must be an integer", - }) - }) - .transpose() -} - -const INVALID_NODES: Error = Error::RedisInvalid { - setting: "startup_nodes", - reason: "must be a list of {host, port} objects", -}; - -fn cluster_nodes(value: &Value) -> Result, Error> { - match value { - Value::String(json) => { - let nodes: Vec> = - serde_json::from_str(json).map_err(|_| INVALID_NODES)?; - nodes - .iter() - .map(|node| { - cluster_node( - node.get("host").and_then(json_text), - node.get("port").and_then(json_text), - ) - }) - .collect() - } - Value::Sequence(nodes) => nodes - .iter() - .map(|node| { - cluster_node( - node.get("host").and_then(as_text), - node.get("port").and_then(as_text), - ) - }) - .collect(), - _ => Err(INVALID_NODES), - } -} - -fn cluster_node(host: Option, port: Option) -> Result { - match (host, port.and_then(|port| port.parse().ok())) { - (Some(host), Some(port)) => Ok(RedisNode { host, port }), - _ => Err(INVALID_NODES), - } -} - -fn json_text(value: &serde_json::Value) -> Option { - match value { - serde_json::Value::String(text) => Some(text.clone()), - serde_json::Value::Number(number) => Some(number.to_string()), - _ => None, - } -} - -fn spelled_text(value: Option<&Spelled>) -> Option { - value.map(|value| match value { - Spelled::Value(value) => value.to_string(), - Spelled::Text(value) => value.clone(), - }) -} - -fn spelled(text: Option) -> Option> { - text.map(|text| match text.parse() { - Ok(value) => Spelled::Value(value), - Err(_) => Spelled::Text(text), - }) -} - -fn cloned( - fields: &BTreeMap, - keep: impl Fn(&str, &Value) -> bool, -) -> AdditionalFields { - fields - .iter() - .filter(|(key, value)| keep(key, value)) - .map(|(key, value)| (key.clone(), value.clone())) - .collect() -} - -#[derive(Clone)] -struct DeferredInput(Value); - -type DeferredInputs = BTreeMap; - -fn deferred_inputs(fields: &BTreeMap) -> DeferredInputs { - fields - .iter() - .filter(|(key, value)| is_unparsed(key, value)) - .map(|(key, value)| (key.clone(), DeferredInput(value.clone()))) - .collect() -} - -fn is_unparsed(key: &str, value: &Value) -> bool { - SUPPORTED.contains(&key) - && match key { - "startup_nodes" => value.is_null(), - _ => as_text(value).is_none(), - } -} diff --git a/litellm-rust/crates/config/tests/config.rs b/litellm-rust/crates/config/tests/config.rs index 3c7d02203e5..fec042e41be 100644 --- a/litellm-rust/crates/config/tests/config.rs +++ b/litellm-rust/crates/config/tests/config.rs @@ -440,3 +440,29 @@ fn mcp_config_redacts_nested_credentials_and_preserves_policy_for_validation() { fn rejects_invalid_typed_mcp_settings(#[case] setting: &str) { assert!(Config::from_yaml(&format!("mcp_servers:\n docs:\n {setting}\n")).is_err()); } + +#[rstest] +fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() { + let config = Config::from_yaml( + "litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid", + ) + .unwrap(); + + assert!(matches!( + config.litellm_settings.cache_params, + Some(litellm_config::CacheSettings::Redis(_)) + )); +} + +#[rstest] +fn non_redis_cache_settings_remain_opaque() { + let config = Config::from_yaml( + "litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true", + ) + .unwrap(); + + assert!(matches!( + config.litellm_settings.cache_params, + Some(litellm_config::CacheSettings::Other(_)) + )); +} diff --git a/litellm-rust/crates/config/tests/redis.rs b/litellm-rust/crates/config/tests/redis.rs deleted file mode 100644 index ab409541c8e..00000000000 --- a/litellm-rust/crates/config/tests/redis.rs +++ /dev/null @@ -1,405 +0,0 @@ -use std::{ - collections::HashMap, - fmt::Write, - sync::{Arc, Mutex}, -}; - -use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis}; -use rstest::rstest; -use tracing::{Event, Subscriber, field::Visit}; -use tracing_subscriber::{Layer, layer::Context, prelude::*}; - -fn resolve( - yaml: &str, - environment: &[(&str, &str)], -) -> Result { - let environment: HashMap = environment - .iter() - .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) - .collect(); - let config = Config::from_yaml(yaml).unwrap(); - config - .general_settings - .coordination_redis - .as_ref() - .unwrap_or(&RedisSettings::default()) - .resolve(&|name| environment.get(name).cloned()) -} - -fn resolve_cache(yaml: &str) -> Result { - let config = Config::from_yaml(yaml).unwrap(); - match config.litellm_settings.cache_params.unwrap() { - CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None), - CacheSettings::Other(_) => panic!("expected a Redis cache"), - } -} - -fn cluster_ports(resolved: &ResolvedRedis) -> Vec { - match &resolved.endpoint { - RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(), - endpoint => panic!("expected a cluster, got {endpoint:?}"), - } -} - -#[rstest] -#[case::host( - "general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5", - &[], - "host:redis:6380:2:true:app:1500" -)] -#[case::environment_url( - "model_list: []", - &[("REDIS_URL", "redis://environment:6379")], - "url:redis://environment:6379" -)] -#[case::explicit_host_drops_environment_url( - "general_settings:\n coordination_redis:\n host: configured", - &[("REDIS_URL", "redis://environment:6379")], - "host:configured:6379:0:false::none" -)] -#[case::yaml_cluster( - "general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"", - &[], - "cluster:node-a:7000,node-b:7001" -)] -#[case::json_environment_cluster( - "model_list: []", - &[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)], - "cluster:node-a:7000,node-b:7001" -)] -#[case::missing_reference_keeps_environment_fallback( - "general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST", - &[("REDIS_HOST", "environment")], - "host:environment:6379:0:false::none" -)] -#[case::configured_values_beat_the_same_environment_values( - "general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3", - &[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")], - "host:configured:6380:3:false::none" -)] -#[case::environment_url_beats_the_environment_host( - "model_list: []", - &[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")], - "url:redis://environment:6379" -)] -#[case::startup_nodes_beat_a_url( - "general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000", - &[], - "cluster:node-a:7000" -)] -#[case::environment_references_resolve_the_host_and_port( - "general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT", - &[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")], - "host:referenced:6400:0:false::none" -)] -#[case::a_null_host_falls_back_to_the_environment( - "general_settings:\n coordination_redis:\n host: null", - &[("REDIS_HOST", "environment")], - "host:environment:6379:0:false::none" -)] -fn resolves_redis_through_the_parsed_config( - #[case] yaml: &str, - #[case] environment: &[(&str, &str)], - #[case] expected: &str, -) { - let resolved = resolve(yaml, environment).unwrap(); - let description = match resolved.endpoint { - RedisEndpoint::Url(url) => format!("url:{}", url.expose()), - RedisEndpoint::Host(host) => format!( - "host:{}:{}:{}:{}:{}:{}", - host.host, - host.port, - host.db, - host.tls, - resolved.namespace.unwrap_or_default(), - resolved - .timeout - .map(|timeout| timeout.as_millis().to_string()) - .unwrap_or_else(|| "none".to_owned()) - ), - RedisEndpoint::Cluster(cluster) => format!( - "cluster:{}", - cluster - .nodes - .iter() - .map(|node| format!("{}:{}", node.host, node.port)) - .collect::>() - .join(",") - ), - }; - - assert_eq!(description, expected); -} - -#[rstest] -#[case::custom_tls_file( - "general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem", - &[], - "ssl_ca_certs" -)] -#[case::gcp_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account", - &[], - "gcp_service_account" -)] -#[case::gcp_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")], - "REDIS_GCP_SERVICE_ACCOUNT" -)] -#[case::azure_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true", - &[], - "azure_redis_ad_token" -)] -#[case::azure_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")], - "REDIS_AZURE_AD_TOKEN" -)] -#[case::aws_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true", - &[], - "aws_iam_auth" -)] -#[case::aws_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")], - "REDIS_AWS_IAM_AUTH" -)] -#[case::sentinel_from_config( - "general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]", - &[], - "sentinel_nodes" -)] -#[case::sentinel_from_environment( - "model_list: []", - &[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")], - "REDIS_SENTINEL_NODES" -)] -fn rejected_settings_name_their_source( - #[case] yaml: &str, - #[case] environment: &[(&str, &str)], - #[case] expected: &str, -) { - let error = resolve(yaml, environment).unwrap_err(); - - assert!(matches!(error, Error::RedisSetting(setting) if setting == expected)); -} - -#[rstest] -fn an_explicit_host_keeps_environment_credentials() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: configured", - &[ - ("REDIS_PASSWORD", "environment-secret"), - ("REDIS_SSL", "True"), - ], - ) - .unwrap(); - - let RedisEndpoint::Host(host) = resolved.endpoint else { - panic!("expected a host endpoint"); - }; - assert_eq!( - host.password.as_ref().map(|password| password.expose()), - Some("environment-secret") - ); - assert!(host.tls); -} - -#[rstest] -#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])] -#[case::a_structured_host_with_an_environment_target( - "general_settings:\n coordination_redis:\n host: []", - &[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")] -)] -fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) { - let error = resolve(yaml, environment).unwrap_err(); - - assert!(matches!( - error, - Error::RedisInvalid { - setting: "host", - .. - } - )); -} - -#[rstest] -fn the_cache_alias_names_cluster_nodes() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000", - ) - .unwrap(); - - assert_eq!(cluster_ports(&resolved), vec![7000]); -} - -#[rstest] -fn canonical_startup_nodes_beat_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101", - ) - .unwrap(); - - assert_eq!(cluster_ports(&resolved), vec![7000]); -} - -#[rstest] -fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100", - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -#[rstest] -fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100", - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -#[rstest] -fn rejected_environment_settings_never_print_the_value() { - let error = resolve( - "model_list: []", - &[ - ("REDIS_HOST", "redis"), - ("REDIS_SSL_CERT_REQS", "private-option-value"), - ], - ) - .unwrap_err(); - let message = error.to_string(); - - assert!(message.contains("REDIS_SSL_CERT_REQS")); - assert!(!message.contains("private-option-value")); -} - -#[rstest] -#[case::negative("-1")] -#[case::not_a_number("slow")] -fn rejects_bad_socket_timeout(#[case] timeout: &str) { - let yaml = format!( - "general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}" - ); - let error = resolve(&yaml, &[]).unwrap_err(); - - assert!(matches!( - error, - Error::RedisInvalid { - setting: "socket_timeout", - .. - } - )); -} - -#[rstest] -fn defaults_socket_timeout_for_the_gateway() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: redis", - &[], - ) - .unwrap(); - - assert_eq!(resolved.timeout, None); -} - -#[rstest] -fn debug_redacts_redis_secrets_and_url() { - let config = Config::from_yaml( - "general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password", - ) - .unwrap(); - let settings = config.general_settings.coordination_redis.unwrap(); - let resolved = settings.resolve(&|_| None).unwrap(); - let debug = format!("{settings:?} {resolved:?}"); - - assert!(!debug.contains("redis://user:password")); - assert!(!debug.contains("credential-user")); - assert!(!debug.contains("credential-password")); - assert!(debug.contains("redacted") || debug.contains("REDACTED")); -} - -#[rstest] -fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() { - let config = Config::from_yaml( - "litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid", - ) - .unwrap(); - - assert!(matches!( - config.litellm_settings.cache_params, - Some(litellm_config::CacheSettings::Redis(_)) - )); -} - -#[rstest] -fn non_redis_cache_settings_remain_opaque() { - let config = Config::from_yaml( - "litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true", - ) - .unwrap(); - - assert!(matches!( - config.litellm_settings.cache_params, - Some(litellm_config::CacheSettings::Other(_)) - )); -} - -#[rstest] -fn legacy_unknown_settings_do_not_change_resolution() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary", - &[], - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -struct WarningCapture(Arc>>); - -impl Layer for WarningCapture { - fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) { - let mut fields = Fields::default(); - event.record(&mut fields); - self.0.lock().unwrap().push(fields.0); - } -} - -#[derive(Default)] -struct Fields(String); - -impl Visit for Fields { - fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { - write!(self.0, "{}={value:?} ", field.name()).unwrap(); - } - - fn record_str(&mut self, field: &tracing::field::Field, value: &str) { - write!(self.0, "{}={value} ", field.name()).unwrap(); - } -} - -#[rstest] -fn ignored_setting_warning_names_the_setting_without_its_value() { - let events = Arc::new(Mutex::new(Vec::new())); - let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone())); - tracing::subscriber::with_default(subscriber, || { - resolve( - "general_settings:\n coordination_redis:\n host: redis\n max_connections: 47", - &[], - ) - .unwrap(); - }); - let warnings = events.lock().unwrap().join(" "); - - assert!(warnings.contains("max_connections")); - assert!(!warnings.contains("47")); -} diff --git a/litellm-rust/crates/gateway/AGENTS.md b/litellm-rust/crates/gateway/AGENTS.md index a33bd19cc80..16714cd9a1f 100644 --- a/litellm-rust/crates/gateway/AGENTS.md +++ b/litellm-rust/crates/gateway/AGENTS.md @@ -5,8 +5,11 @@ - Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate ## Redis -- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored +- `build_redis` builds the Redis clients once at boot (`src/redis/mod.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored - Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands -- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped +- `src/redis/layers.rs` is the Rust copy of `_merge_redis_kwargs` in `litellm/_redis.py`, which only this host needs because the Python host merges in Python. It puts the configured kwargs over their `REDIS_*` values and fills the `get_secret` fallbacks from the `SETTINGS` rows in `redis-resolve`, under the same conditions Python reads them, then hands the result to `litellm_redis_resolve::interpret`. Values come from a `SecretSnapshot` filled once at boot, read like `get_secret` (secret manager, then environment, with `true` and `false` read as bools), and the only straight environment read is Azure's `REDIS_USERNAME` +- A surface is a `ConfiguredRedis` constructor: `coordination` for `general_settings.coordination_redis`, which rejects a malformed `port` like the proxy's `CoordinationRedisParams` pydantic model, `cache_params` for `litellm_settings.cache_params`, which drops the cache-only keys and reads the `redis_startup_nodes` alias, and `redis_cache_kwargs` for a `RedisCache.redis_kwargs` dict in the contract tests. A new surface is a new constructor, never a flag +- Nulls depend on the surface. `coordination_redis` drops every null like `model_dump(exclude_none=True)`, `cache_params` keeps a null as an explicit `None` that hides the `REDIS_*` value unless `Cache` or `RedisCache` names that parameter, and `RedisCache.redis_kwargs` keeps every null. A missing `os.environ/` reference is a null, as `load_config` resolves it +- A setting the Rust client cannot honor warns and is dropped, and only a security check it cannot enforce fails boot, per `redis-resolve/AGENTS.md` - Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python - `/health/readiness` never touches Redis, matching the proxy diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index d0cedbb3b36..a14df9b9c53 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -20,9 +20,14 @@ futures-util.workspace = true litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true litellm-config.workspace = true +litellm-core-utils.workspace = true litellm-error.workspace = true litellm-circuit-breaker.workspace = true litellm-redis-client.workspace = true +litellm-redis-connect.workspace = true +litellm-redis-resolve.workspace = true +litellm-redis-types.workspace = true +rustls.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets.workspace = true @@ -32,10 +37,12 @@ serde.workspace = true tracing.workspace = true tokio = { workspace = true, features = ["signal"] } uuid.workspace = true +veil.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] litellm-redis-testing.workspace = true +litellm-http = { workspace = true, features = ["test-support"] } base64.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index e2f646c5448..c267e79353d 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -12,13 +12,12 @@ pub enum Error { Mcp(#[from] litellm_gateway_mcp::ConnectError), #[error(transparent)] Config(#[from] litellm_config::Error), - #[error("invalid Redis setting {setting}: {reason}")] - RedisInvalid { - setting: &'static str, - reason: &'static str, - }, #[error(transparent)] - RedisConnect(#[from] litellm_redis_client::ConnectError), + Redis(#[from] litellm_redis_resolve::Error), + #[error("Redis setting {setting} is not plain data")] + RedisNotPlainData { setting: String }, + #[error(transparent)] + RedisConnect(#[from] litellm_redis_connect::Error), #[error("MCP requires working gateway authentication")] Auth(#[from] litellm_gateway_auth::Error), } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 12a9bf9110a..0a0191b0499 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -68,15 +68,29 @@ pub fn build_inference(config: &Config) -> Result, Error> { Ok(Arc::new(Gateway::new(resources, http, secrets, router)?)) } -pub async fn build_redis(config: &Config) -> Result { - let environment = secrets::environment_values(&config.environment_variables)?; - let lookup = |name: &str| { +pub fn install_crypto_provider() { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } +} + +pub async fn build_redis( + config: &Config, + pool: &HttpClientPool, + http: &litellm_http::HttpClientConfig, +) -> Result { + let environment = Arc::new(secrets::environment_values(&config.environment_variables)?); + let lookup: litellm_redis_connect::SharedEnvironment = Arc::new(move |name: &str| { environment .get(name) .map(|value| value.expose().to_owned()) .or_else(|| std::env::var(name).ok()) + }); + let resources = redis::Resources { + http: pool.client(http, ClientVariant::Provider)?, + lookup: lookup.clone(), }; - redis::connect(config, &lookup).await + redis::connect(config, &*lookup, &resources).await } pub async fn build_mcp( diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index ff528723095..dc77cfe28eb 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -41,6 +41,7 @@ impl Sink for StderrSink { #[tokio::main] async fn main() -> Result<(), Box> { + litellm_gateway::install_crypto_provider(); if let Some(forward) = litellm_gateway_forward::Settings::from_env()? { let log: LogSettings = settings::from_iter(std::env::vars_os())?; Logger::new(StderrSink { @@ -55,7 +56,8 @@ async fn main() -> Result<(), Box> { Logger::new(StderrSink { level }).install_global()?; let config = Config::load(settings.litellm_config)?; let inference = litellm_gateway::build_inference(&config)?; - let _redis = litellm_gateway::build_redis(&config).await?; + let _redis = + litellm_gateway::build_redis(&config, &inference.resources.pool, &inference.http).await?; let ui = match std::env::var_os("LITELLM_UI_PATH") { Some(directory) => { let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?; diff --git a/litellm-rust/crates/gateway/src/redis/configured.rs b/litellm-rust/crates/gateway/src/redis/configured.rs new file mode 100644 index 00000000000..93297d0e57f --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/configured.rs @@ -0,0 +1,144 @@ +use std::collections::BTreeMap; + +use litellm_config::{RedisSettings, Value}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_resolve::{Expected, Kwargs}; +use litellm_redis_types::{ProviderHandle, Resolution}; +use veil::Redact; + +use super::{layers, secrets::SecretSnapshot}; +use crate::Error; + +const CACHE_ONLY: &[&str] = &[ + "type", + "mode", + "ttl", + "default_in_memory_ttl", + "default_in_redis_ttl", + "supported_call_types", + "similarity_threshold", + "max_messages", + "_backend", +]; + +/// The redis-py kwargs one surface hands `_get_redis_client_logic`, before `REDIS_*` and secrets. +#[derive(Clone, Redact)] +pub(crate) struct ConfiguredRedis { + #[redact(with = "[REDACTED]")] + pub(super) fields: BTreeMap, + pub(super) nulls: Nulls, + pub(super) credential_provider: Option, +} + +/// Whether a null setting reaches redis-py as an explicit `None`, which outranks its `REDIS_*` value. +#[derive(Clone, Copy, Debug)] +pub(super) enum Nulls { + /// `coordination_redis` reaches redis-py through `model_dump(exclude_none=True)`. + Dropped, + /// `Cache` drops a null only for the parameters it names and passes any other key on. + KeptUnlessNamedByCache, + /// `RedisCache.redis_kwargs` already holds only the nulls redis-py receives. + #[cfg(test)] + Kept, +} + +impl ConfiguredRedis { + /// `general_settings.coordination_redis`, which the proxy validates as a pydantic model first. + /// An `os.environ/` port is only known after merging, where it is read like any other port. + pub(crate) fn coordination( + settings: &RedisSettings, + ) -> Result { + if let Some(port) = settings.fields().get("port").and_then(text) + && !port.starts_with("os.environ/") + && port.parse::().is_err() + { + return Err(litellm_redis_resolve::Error::Invalid { + setting: "port".to_owned(), + expected: Expected::Integer, + }); + } + Ok(Self { + fields: settings.fields().clone(), + nulls: Nulls::Dropped, + credential_provider: None, + }) + } + + /// `litellm_settings.cache_params` with `type: redis`, as `litellm.Cache` hands it to `RedisCache`. + pub(crate) fn cache_params(settings: &RedisSettings) -> Self { + let fields = settings.fields(); + let alias = fields + .get("redis_startup_nodes") + .filter(|_| !fields.contains_key("startup_nodes")) + .map(|nodes| ("startup_nodes".to_owned(), nodes.clone())); + let fields = fields + .iter() + .filter(|(key, _)| { + key.as_str() != "redis_startup_nodes" && !CACHE_ONLY.contains(&key.as_str()) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .chain(alias) + .collect(); + Self { + fields, + nulls: Nulls::KeptUnlessNamedByCache, + credential_provider: None, + } + } + + /// `RedisCache.redis_kwargs`, the dict every Python surface ends in. + #[cfg(test)] + pub(crate) fn redis_cache_kwargs(fields: BTreeMap) -> Self { + Self { + fields, + nulls: Nulls::Kept, + credential_provider: None, + } + } + + /// The handle the host holds for the `credential_provider` object in these settings. + #[cfg(test)] + pub(crate) fn with_credential_provider(self, handle: ProviderHandle) -> Self { + Self { + credential_provider: Some(handle), + ..self + } + } + + /// No configured settings, so only `REDIS_*` and the other environment fallbacks apply. + pub(crate) fn environment_only() -> Self { + Self { + fields: BTreeMap::new(), + nulls: Nulls::Dropped, + credential_provider: None, + } + } + + /// `_merge_redis_kwargs` in `litellm/_redis.py`: the kwargs over their `REDIS_*` values with every + /// `get_secret` fallback filled. + pub(crate) fn merged(&self, secrets: &SecretSnapshot) -> Result { + layers::merged(self, secrets) + } + + /// The environment is only for what Python reads straight from `os.environ`, `REDIS_USERNAME` for Azure. + pub(crate) fn resolve( + &self, + secrets: &SecretSnapshot, + environment: &dyn Lookup, + ) -> Result { + Ok(litellm_redis_resolve::interpret( + &self.merged(secrets)?, + self.credential_provider, + environment.get("REDIS_USERNAME").as_deref(), + )?) + } +} + +pub(super) fn text(value: &Value) -> Option { + match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} diff --git a/litellm-rust/crates/gateway/src/redis/layers.rs b/litellm-rust/crates/gateway/src/redis/layers.rs new file mode 100644 index 00000000000..1cf68b2cd4f --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/layers.rs @@ -0,0 +1,226 @@ +use litellm_config::Value; +use litellm_core_utils::settings::Lookup; +use litellm_redis_resolve::{Entry, Kwargs, SETTINGS, is_present, setting}; + +use super::configured::{ConfiguredRedis, Nulls}; +use crate::Error; + +const NAMED_BY_CACHE: &[&str] = &[ + "host", + "port", + "password", + "socket_timeout", + "startup_nodes", + "namespace", + "gcp_service_account", + "gcp_ssl_ca_certs", +]; + +/// How `_merge_redis_kwargs` reads a fallback: `get_secret` turns `true` and `false` into a bool, and +/// `get_secret_str` treats that bool as unset. +#[derive(Clone, Copy)] +enum Read { + Secret, + SecretStr, +} + +const TARGETS: &[(&str, Read)] = &[ + ("startup_nodes", Read::Secret), + ("sentinel_nodes", Read::Secret), + ("sentinel_password", Read::SecretStr), + ("service_name", Read::Secret), +]; + +const IDENTITIES: &[(&str, Read)] = &[ + ("gcp_service_account", Read::SecretStr), + ("gcp_ssl_ca_certs", Read::SecretStr), + ("azure_redis_ad_token", Read::Secret), +]; + +const AZURE: &[(&str, Read)] = &[ + ("azure_client_id", Read::SecretStr), + ("azure_tenant_id", Read::SecretStr), + ("azure_client_secret", Read::SecretStr), +]; + +const AWS: &[(&str, Read)] = &[("aws_iam_region", Read::SecretStr)]; + +/// The configured settings over their `REDIS_*` values, `{**environment_kwargs, **env_overrides}` +/// in `_merge_redis_kwargs`, with the `get_secret` fallbacks it reads filled in. A null left in the result is +/// an explicit `None` that hid its environment value. +pub(super) fn merged(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result { + let block: Kwargs = configured + .fields + .iter() + .map(|(key, value)| (key, environment_reference(value, lookup))) + .filter(|(key, value)| !value.is_null() || keeps_null(configured.nulls, key)) + .map(|(key, value)| { + let value = serde_json::to_value(&value).map_err(|_| Error::RedisNotPlainData { + setting: key.clone(), + })?; + let entry = Entry { + value, + source: key.clone(), + }; + Ok((key.clone(), entry)) + }) + .collect::>()?; + let names_a_target = ["host", "startup_nodes", "sentinel_nodes"] + .iter() + .any(|key| is_set(&block, key)); + let drop_environment_url = names_a_target && !is_set(&block, "url"); + let layered: Kwargs = environment(lookup) + .filter(|(key, _)| !(key == "url" && drop_environment_url)) + .chain(block) + .collect(); + let fallbacks = fallbacks(&layered, configured.credential_provider.is_some(), lookup); + Ok(layered.into_iter().chain(fallbacks).collect()) +} + +// The fallbacks are read under the same conditions as _merge_redis_kwargs, so a secret Python never asks +// for is never filled in +fn fallbacks(layered: &Kwargs, has_provider: bool, lookup: &dyn Lookup) -> Kwargs { + let fill = |kwargs: &[(&str, Read)]| -> Kwargs { + kwargs + .iter() + .filter_map(|(kwarg, read)| fallback(layered, kwarg, *read, lookup)) + .collect() + }; + let targets = fill(TARGETS); + if has_provider || is_set(layered, "credential_provider") { + return targets; + } + let identities = fill(IDENTITIES); + let value = |kwarg: &str| { + layered + .get(kwarg) + .filter(|entry| is_present(&entry.value)) + .or_else(|| identities.get(kwarg)) + .map(|entry| &entry.value) + }; + let gcp = value("gcp_service_account").is_some(); + let azure = value("azure_redis_ad_token") + .is_some_and(|token| text(token).is_some_and(|token| token.eq_ignore_ascii_case("true"))); + let aws = layered + .get("aws_iam_auth") + .is_some_and(|entry| coerces_to_true(&entry.value)); + let azure_settings = match azure && !gcp { + true => fill(AZURE), + false => Kwargs::new(), + }; + let aws_settings = match aws && !gcp && !azure { + true => fill(AWS), + false => Kwargs::new(), + }; + targets + .into_iter() + .chain(identities) + .chain(azure_settings) + .chain(aws_settings) + .collect() +} + +fn fallback( + layered: &Kwargs, + kwarg: &str, + read: Read, + lookup: &dyn Lookup, +) -> Option<(String, Entry)> { + if layered + .get(kwarg) + .is_some_and(|entry| is_present(&entry.value)) + { + return None; + } + let names = setting(kwarg).map_or(&[][..], |setting| setting.fallback); + names.iter().find_map(|name| { + let value = lookup + .get(name) + .filter(|value| !value.is_empty()) + .map(secret)?; + match (read, &value) { + (Read::SecretStr, serde_json::Value::Bool(_)) => None, + _ => { + let entry = Entry { + value, + source: (*name).to_owned(), + }; + Some((kwarg.to_owned(), entry)) + } + } + }) +} + +// get_secret reads an environment value of true or false, in any case and padding, as a bool +fn secret_flag(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "true" => Some(true), + "false" => Some(false), + _ => None, + } +} + +fn secret(value: String) -> serde_json::Value { + secret_flag(&value).map_or(serde_json::Value::String(value), serde_json::Value::Bool) +} + +// _coerces_to_true: text is read like _str_to_bool, anything else by its truthiness +fn coerces_to_true(value: &serde_json::Value) -> bool { + match value { + serde_json::Value::String(text) => { + matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes") + } + serde_json::Value::Null => false, + serde_json::Value::Bool(flag) => *flag, + serde_json::Value::Number(number) => number.as_f64().is_some_and(|number| number != 0.0), + serde_json::Value::Array(items) => !items.is_empty(), + serde_json::Value::Object(entries) => !entries.is_empty(), + } +} + +fn text(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::String(text) => Some(text.clone()), + serde_json::Value::Number(number) => Some(number.to_string()), + serde_json::Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} + +fn keeps_null(nulls: Nulls, key: &str) -> bool { + match nulls { + Nulls::Dropped => false, + Nulls::KeptUnlessNamedByCache => !NAMED_BY_CACHE.contains(&key), + #[cfg(test)] + Nulls::Kept => true, + } +} + +fn environment(lookup: &dyn Lookup) -> impl Iterator + '_ { + SETTINGS.iter().filter_map(|setting| { + let source = setting.environment_name()?; + let value = lookup.get(&source)?; + let entry = Entry { + value: secret(value), + source, + }; + Some((setting.kwarg.to_owned(), entry)) + }) +} + +// load_config resolves os.environ/ to None when unset, so a missing reference is a null like any other +fn environment_reference(value: &Value, lookup: &dyn Lookup) -> Value { + match value { + Value::String(text) => match text.strip_prefix("os.environ/") { + Some(name) => lookup.get(name).map_or(Value::Null, |value| { + secret_flag(&value).map_or(Value::String(value), Value::Bool) + }), + None => value.clone(), + }, + _ => value.clone(), + } +} + +fn is_set(kwargs: &Kwargs, key: &str) -> bool { + kwargs.get(key).is_some_and(|entry| !entry.value.is_null()) +} diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis/mod.rs similarity index 60% rename from litellm-rust/crates/gateway/src/redis.rs rename to litellm-rust/crates/gateway/src/redis/mod.rs index c40b289029b..91428a825f7 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis/mod.rs @@ -1,19 +1,31 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; -use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State}; -use litellm_config::{ - CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value, - environment_has_target, -}; -use litellm_redis_client::{ - ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient, - RedisConfig, -}; +use litellm_circuit_breaker::Config as BreakerConfig; +use litellm_config::{CacheSettings, Config, RedisSettings, Value}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{RedisClient, RedisEvents}; +use litellm_redis_connect::{HostServices, RedisHosts, SharedEnvironment, TracingEvents}; +use litellm_redis_resolve::Expected; +use litellm_redis_types::Resolution; use crate::Error; +mod configured; +mod layers; +mod secrets; +#[cfg(test)] +mod surface_tests; + +use configured::ConfiguredRedis; +use secrets::SecretSnapshot; + +pub(crate) struct Resources { + pub(crate) http: litellm_http::Client, + pub(crate) lookup: SharedEnvironment, +} + const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2); -type Params = ResolvedRedis; +type Params = Resolution; pub struct RedisClients { pub coordination: Option, @@ -34,12 +46,15 @@ enum Target { Environment(Params), } -pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result { +pub(crate) fn plan(config: &Config, lookup: &dyn Lookup) -> Result { let response_cache = redis_cache_params(config) - .map(|settings| settings.resolve_cache(lookup)) + .map(|settings| resolve(&ConfiguredRedis::cache_params(settings), lookup)) .transpose()?; let coordination = match &config.general_settings.coordination_redis { - Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)), + Some(settings) => Some(Target::Configured(resolve( + &ConfiguredRedis::coordination(settings)?, + lookup, + )?)), None => match (&response_cache, config.litellm_settings.cache == Some(true)) { (Some(cache), _) => Some(Target::Configured(cache.clone())), (None, true) => environment(lookup)?.map(Target::Configured), @@ -57,12 +72,30 @@ pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result { }) } -pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result { +pub(crate) async fn connect( + config: &Config, + lookup: &dyn Lookup, + resources: &Resources, +) -> Result { let plan = plan(config, lookup)?; let breaker = breaker_config(lookup)?; + let events: Arc = Arc::new(TracingEvents::new(lookup)); + let hosts = RedisHosts::default(); + let client = |params: Params| -> Result { + Ok(hosts.client( + params, + HostServices { + events: Some(events.clone()), + http: resources.http.clone(), + environment: resources.lookup.clone(), + credential_provider: None, + }, + breaker, + )?) + }; let coordination = match plan.coordination { - Some(Target::Configured(params)) => Some(client(params, breaker)?), - Some(Target::Environment(params)) => match client(params, breaker) { + Some(Target::Configured(params)) => Some(client(params)?), + Some(Target::Environment(params)) => match client(params) { Ok(client) => reachable(client).await, Err(_) => { tracing::warn!( @@ -75,13 +108,10 @@ pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result Some(coordination.clone()), - (None, Some(params)) => Some(client(params, breaker)?), + (None, Some(params)) => Some(client(params)?), (None, None) => None, }; - let response_cache = plan - .response_cache - .map(|params| client(params, breaker)) - .transpose()?; + let response_cache = plan.response_cache.map(client).transpose()?; Ok(RedisClients { coordination, transaction_buffer, @@ -98,17 +128,38 @@ fn redis_cache_params(config: &Config) -> Option<&RedisSettings> { } } -fn environment(lookup: Lookup<'_>) -> Result, Error> { +fn resolve(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result { + let secrets = SecretSnapshot::read(configured.secret_names(), lookup); + let resolution = configured.resolve(&secrets, lookup)?; + for ignored in &resolution.ignored { + tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored: {}", + ignored.setting, + ignored.consequence + ); + } + Ok(resolution) +} + +fn environment(lookup: &dyn Lookup) -> Result, Error> { match environment_has_target(lookup) { - true => RedisSettings::default() - .resolve(lookup) - .map(Some) - .map_err(Into::into), + true => Ok(Some(resolve(&ConfiguredRedis::environment_only(), lookup)?)), false => Ok(None), } } -fn optional_environment(lookup: Lookup<'_>) -> Option { +fn environment_has_target(lookup: &dyn Lookup) -> bool { + [ + "REDIS_HOST", + "REDIS_URL", + "REDIS_CLUSTER_NODES", + "REDIS_SENTINEL_NODES", + ] + .iter() + .any(|name| lookup.get(name).is_some()) +} + +fn optional_environment(lookup: &dyn Lookup) -> Option { match environment(lookup) { Ok(params) => params, Err(error) => { @@ -149,122 +200,101 @@ async fn reachable(client: RedisClient) -> Option { } } -fn client(params: Params, breaker: Option) -> Result { - tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured"); - let endpoint = match params.endpoint { - RedisEndpoint::Url(url) => Endpoint::Url(url), - RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint { - host: host.host, - port: host.port, - username: host.username, - password: host.password, - tls: host.tls, - db: host.db, - }), - RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { - nodes: cluster - .nodes - .into_iter() - .map(|node| ClusterNode { - host: node.host, - port: node.port, - }) - .collect(), - username: cluster.username, - password: cluster.password, - tls: cluster.tls, - }), - }; - Ok(RedisClient::new(RedisConfig { - endpoint, - namespace: params.namespace.map(Namespace::new).unwrap_or_default(), - timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT), - breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))), - })?) +fn parse_seconds(text: &str) -> Option { + text.parse::() + .ok() + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) } -fn log_transition(transition: &litellm_circuit_breaker::Transition) { - match transition.to { - State::Open { .. } => tracing::warn!( - generation = transition.generation, - "redis circuit breaker opened, fast-failing redis calls" - ), - State::HalfOpen => tracing::info!( - generation = transition.generation, - "redis circuit breaker half open, probing redis" - ), - State::Closed { .. } => tracing::info!( - generation = transition.generation, - "redis circuit breaker closed, redis recovered" - ), - } -} - -fn breaker_config(lookup: Lookup<'_>) -> Result, Error> { - let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED") +fn breaker_config(lookup: &dyn Lookup) -> Result, Error> { + let enabled = lookup + .get("REDIS_CIRCUIT_BREAKER_ENABLED") .is_none_or(|value| value.eq_ignore_ascii_case("true")); if !enabled { return Ok(None); } let defaults = BreakerConfig::default(); - let seconds = |name: &'static str, default: Duration| -> Result { - lookup(name) - .map(|value| { - value - .parse::() - .ok() - .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) - .ok_or(Error::RedisInvalid { - setting: name, - reason: "must be a non-negative number of seconds", - }) - }) - .unwrap_or(Ok(default)) - }; - let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD") + Ok(Some(BreakerConfig { + failure_threshold: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", + Expected::Integer, + |text| text.parse().ok(), + )? + .unwrap_or(defaults.failure_threshold), + recovery_timeout: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", + Expected::Seconds, + parse_seconds, + )? + .unwrap_or(defaults.recovery_timeout), + timeout_min_duration: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", + Expected::Seconds, + parse_seconds, + )? + .unwrap_or(defaults.timeout_min_duration), + })) +} + +fn environment_value( + lookup: &dyn Lookup, + name: &str, + expected: Expected, + parse: impl Fn(&str) -> Option, +) -> Result, litellm_redis_resolve::Error> { + lookup + .get(name) .map(|value| { - value.parse().map_err(|_| Error::RedisInvalid { - setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", - reason: "must be an integer", + parse(&value).ok_or_else(|| litellm_redis_resolve::Error::Invalid { + setting: name.to_owned(), + expected, }) }) - .unwrap_or(Ok(defaults.failure_threshold))?; - Ok(Some(BreakerConfig { - failure_threshold, - recovery_timeout: seconds( - "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", - defaults.recovery_timeout, - )?, - timeout_min_duration: seconds( - "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", - defaults.timeout_min_duration, - )?, - })) + .transpose() } #[cfg(test)] mod tests { + use litellm_core_utils::settings::Lookup; use std::{collections::HashMap, time::Duration}; - use litellm_config::{Config, RedisEndpoint, RedisHost}; + use litellm_config::Config; + use litellm_redis_client::{Standalone, StaticAuth, commands}; use litellm_redis_testing::RedisServer; + use litellm_redis_types::{Auth, HealthCheck}; use rstest::rstest; - use super::{Params, Plan, Target, connect, plan}; + use std::sync::Arc; + + use super::{Params, Plan, Resources, Target, plan}; use crate::Error; + async fn connect(config: &Config, lookup: &dyn Lookup) -> Result { + let resources = Resources { + http: litellm_http::Client::for_test(), + lookup: Arc::new(|_: &str| None), + }; + super::connect(config, lookup, &resources).await + } + fn host(name: &str, port: u16) -> Params { Params { - endpoint: RedisEndpoint::Host(RedisHost { + target: litellm_redis_client::Target::Standalone(Standalone { host: name.to_owned(), port, - username: None, - password: None, - tls: false, + tls: None, db: 0, }), + auth: Auth::Static(StaticAuth::default()), namespace: None, timeout: None, + connect_timeout: None, + health_check: HealthCheck::Unset, + keepalive: true, + ignored: Vec::new(), } } @@ -273,7 +303,7 @@ mod tests { .iter() .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) .collect(); - plan(&Config::from_yaml(yaml).unwrap(), &|name| { + plan(&Config::from_yaml(yaml).unwrap(), &|name: &str| { environment.get(name).cloned() }) } @@ -350,36 +380,44 @@ mod tests { port_of(&server) )) .unwrap(); - let clients = connect(&config, &|_| None).await.unwrap(); + let clients = connect(&config, &|_: &str| None).await.unwrap(); let coordination = clients.coordination.unwrap(); let buffer = clients.transaction_buffer.unwrap(); let cache = clients.response_cache.unwrap(); - let unnamespaced = connect(&coordination_config(&server), &|_| None) + let unnamespaced = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap() .coordination .unwrap(); coordination - .set("rate-limit", "1", Duration::from_secs(60)) + .send(commands::set("rate-limit", "1", Duration::from_secs(60))) .await .unwrap(); buffer - .set("spend", "2", Duration::from_secs(60)) + .send(commands::set("spend", "2", Duration::from_secs(60))) .await .unwrap(); assert_eq!( - cache.get::("rate-limit").await.unwrap().as_deref(), + cache + .send(commands::get::("rate-limit")) + .await + .unwrap() + .as_deref(), Some("1") ); assert_eq!( - cache.get::("spend").await.unwrap().as_deref(), + cache + .send(commands::get::("spend")) + .await + .unwrap() + .as_deref(), Some("2") ); assert_eq!( unnamespaced - .get::("responses:rate-limit") + .send(commands::get::("responses:rate-limit")) .await .unwrap() .as_deref(), @@ -387,7 +425,7 @@ mod tests { ); assert_eq!( unnamespaced - .get::("responses:spend") + .send(commands::get::("responses:spend")) .await .unwrap() .as_deref(), @@ -395,16 +433,21 @@ mod tests { ); assert!( unnamespaced - .get::("rate-limit") + .send(commands::get::("rate-limit")) + .await + .unwrap() + .is_none() + ); + assert!( + unnamespaced + .send(commands::get::("spend")) .await .unwrap() .is_none() ); - assert!(unnamespaced.get::("spend").await.unwrap().is_none()); } #[rstest] - #[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])] #[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])] fn redis_migration_regression_malformed_optional_environment_is_absent( #[case] environment: &[(&str, &str)], @@ -413,54 +456,45 @@ mod tests { } #[rstest] - fn redis_migration_regression_explicit_malformed_setting_still_fails() { + #[case::optional_environment("model_list: []", &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])] + #[case::transaction_buffer_environment( + "general_settings:\n use_redis_transaction_buffer: true", + &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")] + )] + #[case::response_cache( + "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: env\n port: not-a-port", + &[] + )] + fn a_malformed_port_outside_the_coordination_block_falls_back_to_6379_like_redis_py( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + ) { + let plan = resolved(yaml, environment).unwrap(); + + let target = match (plan.coordination, plan.transaction_buffer_fallback) { + (Some(Target::Configured(params) | Target::Environment(params)), _) + | (None, Some(params)) => params, + (None, None) => panic!("expected a Redis target"), + }; + assert_eq!(target.target, host("env", 6379).target); + } + + #[rstest] + #[case::port("port: not-a-port", "port")] + #[case::ssl("ssl: maybe", "ssl")] + fn a_malformed_coordination_block_fails_like_the_pydantic_model( + #[case] setting: &str, + #[case] expected: &str, + ) { let error = resolved( - "general_settings:\n coordination_redis:\n host: coord\n port: not-a-port", + &format!("general_settings:\n coordination_redis:\n host: coord\n {setting}"), &[], ) .unwrap_err(); - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); - } - - #[rstest] - fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() { - let error = resolved( - "general_settings:\n use_redis_transaction_buffer: true", - &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")], - ) - .unwrap_err(); - - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); - } - - #[rstest] - fn redis_migration_regression_invalid_response_cache_still_fails() { - let error = resolved( - "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port", - &[], - ) - .unwrap_err(); - - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); + assert!( + matches!(error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == expected) + ); } #[rstest] @@ -468,7 +502,7 @@ mod tests { async fn redis_migration_regression_invalid_optional_environment_url_is_absent() { let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| { + let clients = connect(&config, &|name: &str| { (name == "REDIS_URL").then(|| "not a redis URL".to_owned()) }) .await @@ -506,7 +540,7 @@ mod tests { let server = RedisServer::start().await.unwrap(); server.stop().await.unwrap(); - let clients = connect(&coordination_config(&server), &|_| None) + let clients = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap(); let coordination = clients.coordination.unwrap(); @@ -521,12 +555,12 @@ mod tests { #[tokio::test] async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() { let server = RedisServer::start().await.unwrap(); - let clients = connect(&coordination_config(&server), &|_| None) + let clients = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap(); let coordination = clients.coordination.unwrap(); coordination - .set("written", "1", Duration::from_secs(60)) + .send(commands::set("written", "1", Duration::from_secs(60))) .await .unwrap(); @@ -537,6 +571,30 @@ mod tests { assert!(buffer.ping().await.is_ok()); } + #[rstest] + #[tokio::test] + async fn an_environment_only_tls_redis_is_reached_with_its_ca() { + crate::install_crypto_provider(); + let server = litellm_redis_testing::RedisTlsServer::start(false) + .await + .unwrap(); + let ca = server.files().ca.display().to_string(); + let port = server.port().to_string(); + let config = Config::from_yaml("model_list: []").unwrap(); + + let clients = connect(&config, &|name: &str| match name { + "REDIS_HOST" => Some(server.host().to_owned()), + "REDIS_PORT" => Some(port.clone()), + "REDIS_SSL" => Some("true".to_owned()), + "REDIS_SSL_CA_CERTS" => Some(ca.clone()), + _ => None, + }) + .await + .unwrap(); + + assert!(clients.coordination.is_some()); + } + #[rstest] #[tokio::test] async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() { @@ -545,7 +603,7 @@ mod tests { server.stop().await.unwrap(); let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| match name { + let clients = connect(&config, &|name: &str| match name { "REDIS_HOST" => Some("127.0.0.1".to_owned()), "REDIS_PORT" => Some(port.clone()), _ => None, @@ -564,7 +622,7 @@ mod tests { let port = port_of(&server); let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| match name { + let clients = connect(&config, &|name: &str| match name { "REDIS_HOST" => Some("127.0.0.1".to_owned()), "REDIS_PORT" => Some(port.clone()), _ => None, diff --git a/litellm-rust/crates/gateway/src/redis/secrets.rs b/litellm-rust/crates/gateway/src/redis/secrets.rs new file mode 100644 index 00000000000..a969497f4c3 --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/secrets.rs @@ -0,0 +1,61 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; + +use litellm_config::Value; +use litellm_redis_resolve::SETTINGS; + +use super::configured::ConfiguredRedis; + +#[derive(Clone, Debug, Default)] +pub(crate) struct SecretSnapshot(BTreeMap>); + +impl SecretSnapshot { + pub(crate) fn read(names: BTreeSet, lookup: &dyn Lookup) -> Self { + Self( + names + .into_iter() + .map(|name| { + let value = lookup.get(&name).map(SecretValue::new); + (name, value) + }) + .collect(), + ) + } +} + +impl Lookup for SecretSnapshot { + fn get(&self, name: &str) -> Option { + match self.0.get(name) { + Some(value) => value.as_ref().map(|value| value.expose().to_owned()), + None => { + debug_assert!( + false, + "{name} is missing from ConfiguredRedis::secret_names" + ); + tracing::error!( + name, + "a Redis secret was read without being in the snapshot" + ); + None + } + } + } +} + +impl ConfiguredRedis { + pub(crate) fn secret_names(&self) -> BTreeSet { + let references = self.fields.values().filter_map(|value| match value { + Value::String(text) => text.strip_prefix("os.environ/").map(str::to_owned), + _ => None, + }); + let declared = SETTINGS.iter().flat_map(|setting| { + setting + .environment_name() + .into_iter() + .chain(setting.fallback.iter().map(|name| (*name).to_owned())) + }); + declared.chain(references).collect() + } +} diff --git a/litellm-rust/crates/gateway/src/redis/surface_tests.rs b/litellm-rust/crates/gateway/src/redis/surface_tests.rs new file mode 100644 index 00000000000..9ee881e2e04 --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/surface_tests.rs @@ -0,0 +1,964 @@ +use std::collections::{BTreeMap, HashMap}; + +use litellm_config::{CacheSettings, Config, Value}; +use litellm_redis_resolve::Expected; +use litellm_redis_types::{ + Auth, AwsIamAuth, AzureAdAuth, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions, +}; +use litellm_secret_value::SecretValue; +use rstest::rstest; +use serde_json::json; + +use super::{configured::ConfiguredRedis, secrets::SecretSnapshot}; +use crate::Error; + +fn resolution( + configured: &ConfiguredRedis, + environment: &HashMap, +) -> Result { + let lookup = |name: &str| environment.get(name).cloned(); + configured.resolve( + &SecretSnapshot::read(configured.secret_names(), &lookup), + &lookup, + ) +} + +const PROVIDER: ProviderHandle = ProviderHandle::new(7); + +fn redis_cache_kwargs(kwargs: serde_json::Value) -> ConfiguredRedis { + let has_provider = kwargs + .get("credential_provider") + .is_some_and(|provider| !provider.is_null()); + let configured = ConfiguredRedis::redis_cache_kwargs( + serde_json::from_value::>(kwargs).unwrap(), + ); + match has_provider { + true => configured.with_credential_provider(PROVIDER), + false => configured, + } +} + +fn coordination_resolution(yaml: &str, environment: &[(&str, &str)]) -> Result { + let environment: HashMap = environment + .iter() + .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) + .collect(); + let config = Config::from_yaml(yaml).unwrap(); + let configured = config + .general_settings + .coordination_redis + .unwrap_or_default(); + let configured = ConfiguredRedis::coordination(&configured)?; + resolution(&configured, &environment) +} + +fn resolve(yaml: &str, environment: &[(&str, &str)]) -> Result { + coordination_resolution(yaml, environment) +} + +fn resolve_cache(yaml: &str) -> Result { + resolve_cache_in(yaml, &[]) +} + +fn resolve_cache_in(yaml: &str, environment: &[(&str, &str)]) -> Result { + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + let config = Config::from_yaml(yaml).unwrap(); + let configured = match config.litellm_settings.cache_params.unwrap() { + CacheSettings::Redis(settings) => ConfiguredRedis::cache_params(&settings), + CacheSettings::Other(_) => panic!("expected a Redis cache"), + }; + configured.resolve( + &SecretSnapshot::read(configured.secret_names(), &lookup), + &lookup, + ) +} + +fn cluster_ports(resolved: &Resolution) -> Vec { + match &resolved.target { + Target::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(), + target => panic!("expected a cluster, got {target:?}"), + } +} + +#[rstest] +#[case::host( + "general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5", + &[], + "host:redis:6380:2:true:app:1500" +)] +#[case::environment_references_resolve_the_host_and_port( + "general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT", + &[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")], + "host:referenced:6400:0:false::none" +)] +#[case::a_null_host_falls_back_to_the_environment( + "general_settings:\n coordination_redis:\n host: null", + &[("REDIS_HOST", "environment")], + "host:environment:6379:0:false::none" +)] +fn resolves_redis_through_the_parsed_config( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: &str, +) { + let resolved = resolve(yaml, environment).unwrap(); + let description = match resolved.target { + Target::Unix(socket) => format!("unix:{}", socket.path.display()), + Target::Url(url) => format!("url:{}", url.url.expose()), + Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name), + Target::Standalone(host) => format!( + "host:{}:{}:{}:{}:{}:{}", + host.host, + host.port, + host.db, + host.tls.is_some(), + resolved.namespace.unwrap_or_default(), + resolved + .timeout + .map(|timeout| timeout.as_millis().to_string()) + .unwrap_or_else(|| "none".to_owned()) + ), + Target::Cluster(cluster) => format!( + "cluster:{}", + cluster + .nodes + .iter() + .map(|node| format!("{}:{}", node.host, node.port)) + .collect::>() + .join(",") + ), + }; + + assert_eq!(description, expected); +} + +#[rstest] +#[case::tls_ciphers( + "general_settings:\n coordination_redis:\n host: redis\n ssl_ciphers: HIGH", + &[], + "ssl_ciphers" +)] +#[case::ssl_ciphers_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_SSL_CIPHERS", "HIGH")], + "REDIS_SSL_CIPHERS" +)] +#[case::connection_pool( + "general_settings:\n coordination_redis:\n host: redis\n connection_pool: shared", + &[], + "connection_pool" +)] +fn unsupported_settings_resolve_and_are_named_by_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: &str, +) { + let resolution = coordination_resolution(yaml, environment).unwrap(); + + assert_eq!( + resolution.target, + Target::Standalone(litellm_redis_types::Standalone { + host: "redis".into(), + port: 6379, + db: 0, + tls: None, + }) + ); + let [ignored] = resolution.ignored.as_slice() else { + panic!("expected one ignored setting, got {:?}", resolution.ignored); + }; + assert_eq!(ignored.setting, expected); + assert!(!ignored.consequence.is_empty()); +} + +#[rstest] +#[case::kwarg( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_validate_ocsp: true", + &[], + "ssl_validate_ocsp" +)] +#[case::stapled_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_SSL", "true"), ("REDIS_SSL_VALIDATE_OCSP_STAPLED", "yes")], + "REDIS_SSL_VALIDATE_OCSP_STAPLED" +)] +#[case::rediss_url_query( + "general_settings:\n coordination_redis:\n url: rediss://cache:6380/0?ssl_validate_ocsp=true", + &[], + "url ?ssl_validate_ocsp" +)] +fn an_ocsp_check_fails_instead_of_connecting_without_it( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] setting: &str, +) { + let error = coordination_resolution(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Unenforceable { setting: named, .. }) if named == setting), + "{error:?}" + ); +} + +#[rstest] +#[case::kwarg(" host: redis\n ssl: true\n ssl_validate_ocsp: false")] +#[case::rediss_url_query(" url: rediss://cache:6380/0?ssl_validate_ocsp=false")] +fn an_ocsp_flag_set_to_false_asks_for_nothing(#[case] block: &str) { + let yaml = format!("general_settings:\n coordination_redis:\n{block}"); + + let resolution = coordination_resolution(&yaml, &[]).unwrap(); + + assert!(resolution.ignored.is_empty(), "{:?}", resolution.ignored); +} + +fn static_auth(username: Option<&str>, password: Option<&str>) -> Auth { + Auth::Static(StaticAuth { + username: username.map(litellm_secret_value::SecretValue::new), + password: password.map(litellm_secret_value::SecretValue::new), + }) +} + +#[rstest] +fn an_explicit_host_keeps_environment_credentials() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: configured", + &[ + ("REDIS_PASSWORD", "environment-secret"), + ("REDIS_SSL", "True"), + ], + ) + .unwrap(); + + assert_eq!(resolved.auth, static_auth(None, Some("environment-secret"))); + assert!(matches!(resolved.target, Target::Standalone(host) if host.tls.is_some())); +} + +fn azure(username: Option<&str>) -> Auth { + Auth::AzureAd(AzureAdAuth { + client_id: Some("client".to_owned()), + tenant_id: None, + client_secret: None, + username: username.map(str::to_owned), + }) +} + +fn aws(serverless: bool) -> Auth { + Auth::AwsIam(AwsIamAuth { + user_name: "iam-user".to_owned(), + cache_name: "cache-1".to_owned(), + region: "us-east-1".to_owned(), + serverless, + }) +} + +#[rstest] +#[case::static_credentials( + "general_settings:\n coordination_redis:\n host: redis\n username: user\n password: pass", + &[], + static_auth(Some("user"), Some("pass")) +)] +#[case::gcp_beats_azure_and_aws( + "general_settings:\n coordination_redis:\n host: redis\n password: pass\n azure_redis_ad_token: \"True\"\n aws_iam_auth: true", + &[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com")], + Auth::GcpIam { service_account: "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com".to_owned() } +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + "general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client\n aws_iam_auth: true", + &[("REDIS_USERNAME", "environment-user")], + azure(Some("environment-user")) +)] +#[case::azure_without_an_environment_username_sends_only_the_token( + "general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client", + &[], + azure(None) +)] +#[case::azure_needs_the_literal_true( + "general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: \"1\"", + &[], + static_auth(None, None) +)] +#[case::aws_with_a_region_from_the_environment( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n password: dropped\n aws_iam_auth: \"yes\"\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1", + &[("AWS_DEFAULT_REGION", "us-east-1")], + aws(false) +)] +#[case::aws_serverless_over_a_rediss_url( + "general_settings:\n coordination_redis:\n url: rediss://cache-1:6379\n aws_iam_auth: true\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1\n aws_iam_region: us-east-1\n aws_iam_serverless: true", + &[], + aws(true) +)] +fn auth_follows_python_precedence( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: Auth, +) { + assert_eq!(resolve(yaml, environment).unwrap().auth, expected); +} + +#[rstest] +#[case::coordination( + "general_settings:\n coordination_redis:\n host: redis\n credential_provider: host-object", + &[] +)] +#[case::without_a_target( + "general_settings:\n coordination_redis:\n credential_provider: host-object", + &[("REDIS_HOST", "redis")] +)] +fn a_credential_provider_without_a_host_handle_fails( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!( + error, + Error::Redis(litellm_redis_resolve::Error::ProviderWithoutHandle) + ), + "{error}" + ); +} + +#[rstest] +#[case::without_tls( + "general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r", + "requires TLS" +)] +#[case::a_plain_url_ignores_ssl( + "general_settings:\n coordination_redis:\n url: redis://cache:6379\n ssl: true\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r", + "requires TLS" +)] +#[case::missing_settings( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n aws_iam_auth: true\n aws_iam_cache_name: c", + "requires: aws_iam_user_name, aws_iam_region" +)] +fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) { + let error = resolve(yaml, &[]).unwrap_err(); + + assert!(error.to_string().contains(message), "{error}"); +} + +#[rstest] +#[case::defaults_verify_the_certificate_but_not_the_hostname( + " ssl: true", + &[], + Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false }) +)] +#[case::every_option( + " ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_password: key-secret\n ssl_cert_reqs: none\n ssl_check_hostname: true", + &[], + Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), key_password: Some(SecretValue::new("key-secret")), verify_certificate: false, check_hostname: true }) +)] +#[case::options_from_the_environment( + " ssl: true", + &[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional"), ("REDIS_SSL_PASSWORD", "env-key-secret")], + Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: Some(SecretValue::new("env-key-secret")), verify_certificate: true, check_hostname: false }) +)] +#[case::gcp_ca_certs_replace_the_configured_ones( + " ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem", + &[], + Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false }) +)] +#[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)] +fn tls_options_resolve_like_redis_py( + #[case] block: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: Option, +) { + let yaml = format!("general_settings:\n coordination_redis:\n host: redis\n{block}"); + + let resolved = resolve(&yaml, environment).unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(host) if host.tls == expected)); +} + +#[rstest] +fn an_unknown_cert_requirement_fails() { + let error = resolve( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_cert_reqs: strict", + &[], + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::CertRequirement }) if setting == "ssl_cert_reqs" + )); +} + +#[rstest] +#[case::from_config( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[sentinel-a, 26379], [sentinel-b, \"26380\"]]\n service_name: primary\n sentinel_password: sentinel-secret\n db: 2", + &[] +)] +#[case::from_the_environment( + "model_list: []", + &[ + ("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379], ["sentinel-b", "26380"]]"#), + ("REDIS_SERVICE_NAME", "primary"), + ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), + ("REDIS_DB", "2"), + ] +)] +fn sentinel_nodes_resolve_to_a_sentinel_endpoint( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], +) { + let Target::Sentinel(sentinel) = resolve(yaml, environment).unwrap().target else { + panic!("expected a sentinel endpoint"); + }; + + assert_eq!( + sentinel.monitors, + vec![ + Node { + host: "sentinel-a".to_owned(), + port: 26379 + }, + Node { + host: "sentinel-b".to_owned(), + port: 26380 + }, + ] + ); + assert_eq!(sentinel.service_name, "primary"); + assert_eq!( + sentinel.monitor_auth, + Some(StaticAuth { + username: None, + password: Some(litellm_secret_value::SecretValue::new("sentinel-secret")), + }) + ); + assert_eq!(sentinel.db, 2); +} + +#[rstest] +#[case::a_url_beats_sentinels( + "general_settings:\n coordination_redis:\n url: redis://single:6379\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + "url:redis://single:6379" +)] +fn sentinel_precedence(#[case] yaml: &str, #[case] expected: &str) { + let resolved = resolve(yaml, &[]).unwrap(); + let description = match resolved.target { + Target::Url(url) => format!("url:{}", url.url.expose()), + Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name), + other => panic!("unexpected {other:?}"), + }; + + assert_eq!(description, expected); +} + +#[rstest] +fn sentinels_without_a_service_name_fail() { + let error = resolve( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]", + &[], + ) + .unwrap_err(); + + assert!( + matches!( + error, + Error::Redis(litellm_redis_resolve::Error::NoSentinelService) + ), + "{error:?}" + ); +} + +#[rstest] +#[case::configured_host( + "general_settings:\n coordination_redis:\n host: ignored\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + &[], + "host" +)] +#[case::environment_port( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + &[("REDIS_PORT", "6380")], + "REDIS_PORT" +)] +fn a_host_or_port_cannot_reach_sentinel( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::SentinelWithHost { setting }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::not_pairs( + "general_settings:\n coordination_redis:\n sentinel_nodes: [{host: s, port: 26379}]\n service_name: primary", + &[], + "sentinel_nodes" +)] +#[case::from_the_environment( + "model_list: []", + &[("REDIS_SENTINEL_NODES", "not-json"), ("REDIS_SERVICE_NAME", "primary")], + "REDIS_SENTINEL_NODES" +)] +fn invalid_sentinel_nodes_name_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::SentinelNodes }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::from_config("general_settings:\n coordination_redis:\n startup_nodes: [{host: a}]", &[], "startup_nodes")] +#[case::from_the_environment("model_list: []", &[("REDIS_CLUSTER_NODES", "not-json")], "REDIS_CLUSTER_NODES")] +#[case::out_of_range_port_from_the_environment("model_list: []", &[("REDIS_HOST", "redis"), ("REDIS_PORT", "70000")], "REDIS_PORT")] +fn invalid_values_name_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::falls_back_to_the_socket_timeout(" socket_timeout: 1.5", Some(1500), None)] +#[case::explicit_connect_timeout( + " socket_timeout: 1.5\n socket_connect_timeout: 0.25", + Some(1500), + Some(250) +)] +fn connect_timeout_is_resolved_separately( + #[case] block: &str, + #[case] timeout_ms: Option, + #[case] connect_ms: Option, +) { + let resolved = resolve( + &format!("general_settings:\n coordination_redis:\n host: redis\n{block}"), + &[], + ) + .unwrap(); + + assert_eq!( + resolved.timeout.map(|timeout| timeout.as_millis()), + timeout_ms + ); + assert_eq!( + resolved.connect_timeout.map(|timeout| timeout.as_millis()), + connect_ms + ); +} + +#[rstest] +#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])] +#[case::a_structured_host_with_an_environment_target( + "general_settings:\n coordination_redis:\n host: []", + &[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")] +)] +fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(error, Error::Redis(litellm_redis_resolve::Error::NoTarget)), + "{error:?}" + ); +} + +#[rstest] +fn the_cache_alias_names_cluster_nodes() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000", + ) + .unwrap(); + + assert_eq!(cluster_ports(&resolved), vec![7000]); +} + +#[rstest] +fn canonical_startup_nodes_beat_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101", + ) + .unwrap(); + + assert_eq!(cluster_ports(&resolved), vec![7000]); +} + +#[rstest] +fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100", + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100", + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +fn unsupported_environment_settings_never_carry_the_value() { + let resolution = coordination_resolution( + "model_list: []", + &[ + ("REDIS_HOST", "redis"), + ("REDIS_SSL_CIPHERS", "private-option-value"), + ], + ) + .unwrap(); + let ignored = format!("{:?}", resolution.ignored); + + assert!(ignored.contains("REDIS_SSL_CIPHERS"), "{ignored}"); + assert!(!ignored.contains("private-option-value"), "{ignored}"); +} + +#[rstest] +#[case::negative("-1")] +#[case::infinite("inf")] +fn rejects_bad_socket_timeout(#[case] timeout: &str) { + let yaml = format!( + "general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}" + ); + let error = resolve(&yaml, &[]).unwrap_err(); + + assert!(matches!( + error, + Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::Seconds }) if setting == "socket_timeout" + )); +} + +#[rstest] +fn debug_redacts_redis_secrets_and_url() { + let config = Config::from_yaml( + "general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password\n sentinel_kwargs: {password: credential-passthrough}", + ) + .unwrap(); + let settings = config.general_settings.coordination_redis.unwrap(); + let configured = ConfiguredRedis::coordination(&settings).unwrap(); + let resolved = resolution(&configured, &HashMap::new()).unwrap(); + let debug = format!("{settings:?} {configured:?} {resolved:?}"); + + assert!(!debug.contains("redis://user:password")); + assert!(!debug.contains("credential-user")); + assert!(!debug.contains("credential-password")); + assert!(!debug.contains("credential-passthrough")); + assert!(debug.contains("redacted") || debug.contains("REDACTED")); +} + +#[rstest] +fn legacy_unknown_settings_do_not_change_resolution() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary", + &[], + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +#[case::configured( + "general_settings:\n coordination_redis:\n host: redis\n max_connections: 47", + &[], + "max_connections" +)] +#[case::from_the_environment( + "general_settings:\n coordination_redis:\n host: redis", + &[("REDIS_MAX_CONNECTIONS", "47")], + "REDIS_MAX_CONNECTIONS" +)] +fn ignored_settings_are_named_by_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let resolution = coordination_resolution(yaml, environment).unwrap(); + + let settings: Vec<&str> = resolution + .ignored + .iter() + .map(|ignored| ignored.setting.as_str()) + .collect(); + assert_eq!(settings, vec![source]); +} + +fn snapshot(values: &[(&str, &str)]) -> SecretSnapshot { + let values: HashMap<&str, &str> = values.iter().copied().collect(); + SecretSnapshot::read( + ConfiguredRedis::environment_only().secret_names(), + &|name: &str| values.get(name).map(|value| (*value).to_owned()), + ) +} + +#[rstest] +fn secret_manager_values_reach_the_resolver_through_the_snapshot() { + let settings = redis_cache_kwargs(json!({"host": "os.environ/COORD_HOST"})); + let secrets = SecretSnapshot::read(settings.secret_names(), &|name: &str| match name { + "COORD_HOST" => Some("from-the-manager".to_owned()), + "REDIS_PASSWORD" => Some("manager-secret".to_owned()), + _ => None, + }); + + let resolved = settings.resolve(&secrets, &|_: &str| None).unwrap(); + + assert!( + matches!(&resolved.target, Target::Standalone(host) if host.host == "from-the-manager") + ); + assert_eq!(resolved.auth, static_auth(None, Some("manager-secret"))); +} + +#[rstest] +#[case::from_the_environment(&[], &[("REDIS_USERNAME", "environment-user")], Some("environment-user"))] +#[case::never_from_the_secret_manager(&[("REDIS_USERNAME", "manager-user")], &[], None)] +fn the_azure_username_comes_from_the_process_environment_like_python( + #[case] secrets: &[(&str, &str)], + #[case] environment: &[(&str, &str)], + #[case] expected: Option<&str>, +) { + let settings = redis_cache_kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"})); + let environment: HashMap<&str, &str> = environment.iter().copied().collect(); + + let resolved = settings + .resolve(&snapshot(secrets), &|name: &str| { + environment.get(name).map(|value| (*value).to_owned()) + }) + .unwrap(); + + assert!( + matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected), + "{:?}", + resolved.auth + ); +} + +#[rstest] +#[case::coordination_drops_the_null( + "general_settings:\n coordination_redis:\n host: redis\n ssl: null", + false, + true +)] +#[case::cache_passes_an_unnamed_null_on( + "litellm_settings:\n cache_params:\n type: redis\n host: redis\n ssl: null", + true, + false +)] +#[case::cache_drops_a_null_it_names( + "litellm_settings:\n cache_params:\n type: redis\n host: null\n ssl: true", + true, + true +)] +fn a_null_overrides_the_environment_only_where_python_passes_it_on( + #[case] yaml: &str, + #[case] cache: bool, + #[case] tls: bool, +) { + let environment = [("REDIS_HOST", "environment"), ("REDIS_SSL", "true")]; + let resolved = match cache { + true => resolve_cache_in(yaml, &environment), + false => resolve(yaml, &environment), + } + .unwrap(); + + let Target::Standalone(host) = resolved.target else { + panic!("expected a standalone target, got {:?}", resolved.target); + }; + assert_eq!(host.tls.is_some(), tls); +} + +#[rstest] +#[case::empty_cluster_nodes_use_the_environment( + json!({"startup_nodes": []}), + &[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)], + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], +)] +#[case::empty_sentinel_nodes_use_the_environment( + json!({"sentinel_nodes": [], "service_name": "primary"}), + &[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)], + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES")], +)] +#[case::empty_service_name_uses_the_environment( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}), + &[("REDIS_SERVICE_NAME", "primary")], + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[("service_name", "REDIS_SERVICE_NAME")], +)] +#[case::environment_host_port_and_password( + json!({}), + &[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")], + json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}), + &[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")], +)] +#[case::padded_environment_numbers_parse_like_python_int( + json!({}), + &[("REDIS_HOST", "redis"), ("REDIS_PORT", "6_380"), ("REDIS_DB", "3\n")], + json!({"port": "6_380", "host": "redis", "db": "3\n"}), + &[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")], +)] +#[case::socket_timeout_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SOCKET_TIMEOUT", "0.75")], + json!({"socket_timeout": "0.75", "host": "redis"}), + &[("socket_timeout", "REDIS_SOCKET_TIMEOUT")], +)] +#[case::environment_url( + json!({}), + &[("REDIS_URL", "redis://environment:6379/1")], + json!({"url": "redis://environment:6379/1"}), + &[("url", "REDIS_URL")], +)] +#[case::environment_url_beats_the_environment_host( + json!({}), + &[("REDIS_HOST", "ignored"), ("REDIS_URL", "redis://from-url:6379")], + json!({"host": "ignored", "url": "redis://from-url:6379"}), + &[("host", "REDIS_HOST"), ("url", "REDIS_URL")], +)] +#[case::an_explicit_host_drops_the_environment_url( + json!({"host": "configured"}), + &[("REDIS_URL", "redis://environment:6379")], + json!({"host": "configured"}), + &[], +)] +#[case::environment_cluster_nodes( + json!({}), + &[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)], + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], +)] +#[case::sentinel_from_the_environment( + json!({}), + &[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")], + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")], +)] +#[case::an_explicit_credential_provider_beats_cloud_auth( + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[("REDIS_AZURE_AD_TOKEN", "true")], + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[], +)] +#[case::gcp_beats_azure_and_aws( + json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")], + json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")], +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[("REDIS_USERNAME", "environment-user")], + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[], +)] +#[case::an_empty_environment_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}), + &[("REDIS_USERNAME", "")], + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}), + &[("username", "REDIS_USERNAME")], +)] +#[case::aws_iam_over_tls_with_the_default_region( + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}), + &[("AWS_DEFAULT_REGION", "us-east-1")], + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}), + &[("aws_iam_region", "AWS_DEFAULT_REGION")], +)] +#[case::aws_region_beats_the_default_region( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}), + &[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")], + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}), + &[("aws_iam_region", "AWS_REGION")], +)] +#[case::a_null_cert_requirement_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[("REDIS_SSL_CERT_REQS", "required")], + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[], +)] +#[case::a_cert_requirement_spelled_like_python_none_fails( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CERT_REQS", "None")], + json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], +)] +#[case::tls_options_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")], + json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], +)] +#[case::environment_references_resolve( + json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}), + &[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")], + json!({"host": "referenced", "password": "referenced-secret"}), + &[], +)] +#[case::malformed_port_from_the_environment( + json!({}), + &[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")], + json!({"port": "abc", "host": "redis"}), + &[("host", "REDIS_HOST"), ("port", "REDIS_PORT")], +)] +#[case::retry_on_error_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_RETRY_ON_ERROR", "true")], + json!({"retry_on_error": true, "host": "redis"}), + &[("retry_on_error", "REDIS_RETRY_ON_ERROR")], +)] +#[case::ssl_ca_path_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")], + json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")], +)] +fn merges_like_python_redis( + #[case] kwargs: serde_json::Value, + #[case] environment: &[(&str, &str)], + #[case] merged: serde_json::Value, + #[case] sources: &[(&str, &str)], +) { + let configured = redis_cache_kwargs(kwargs); + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + + let kwargs = configured + .merged(&SecretSnapshot::read(configured.secret_names(), &lookup)) + .unwrap(); + + let values: serde_json::Map = kwargs + .iter() + .map(|(kwarg, entry)| (kwarg.clone(), entry.value.clone())) + .collect(); + let renamed: Vec<(&str, &str)> = kwargs + .iter() + .filter(|(kwarg, entry)| *kwarg != &entry.source) + .map(|(kwarg, entry)| (kwarg.as_str(), entry.source.as_str())) + .collect(); + assert_eq!(serde_json::Value::Object(values), merged); + assert_eq!(renamed, sources); +} diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index e36b7d70c28..fe6173dd100 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned { ) -> PyResult>; /// What the stream carries at hand-off, as the caller's stream receives it. + /// The unary response after the deployment transform, before success hooks run. + fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> { + Ok(()) + } + fn encode_stream_head( &mut self, py: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index dbf843a5295..052156cbc5d 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -399,8 +399,16 @@ where Ok(step) => step, Err(error) => return self.failure(error, FailureOrigin::Host).await, }; - match self.settle(step).await { - Ok(response) => self.succeeded(response).await, + let response = match self.settle(step).await { + Ok(response) => response, + Err(error) => return self.failure(error, FailureOrigin::Host).await, + }; + let delivered = { + let mut parts = self.parts(); + Python::attach(|py| parts.binding.deliver_response(py, response.bind(py))) + }; + match delivered { + Ok(()) => self.succeeded(response).await, Err(error) => self.failure(error, FailureOrigin::Host).await, } } diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8e37ab02669..a9ebae427d9 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -36,8 +36,8 @@ pub use marshal::{ pub use owned::PythonOwned; pub use runtime::{ ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value, - ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync, - run_sync_value, runtime_started, + ready_future, reserve_process_for_forking, run_async, run_async_value, run_isolated, run_sync, + run_sync_value, runtime_handle, runtime_started, }; pub use services::PythonHostCalls; diff --git a/litellm-rust/crates/host-python/src/runtime.rs b/litellm-rust/crates/host-python/src/runtime.rs index e8872c9340a..2a8fc598fc4 100644 --- a/litellm-rust/crates/host-python/src/runtime.rs +++ b/litellm-rust/crates/host-python/src/runtime.rs @@ -68,6 +68,23 @@ fn runtime() -> PyResult<&'static Runtime> { Ok(pyo3_async_runtimes::tokio::get_runtime()) } +/// Runs `future` on a current-thread runtime built and shut down inside the call, its blocking pool joined before +/// it returns. No thread outlives the call, so unlike every other door this one never claims the process in the fork +/// gate: a process that only used it can still fork. The future must not hold resources tied to another runtime +pub fn run_isolated(future: F) -> std::io::Result { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let output = runtime.block_on(future); + runtime.shutdown_timeout(Duration::from_secs(5)); + Ok(output) +} + +/// The native runtime, for work that must outlive the call that schedules it. +pub fn runtime_handle() -> PyResult { + Ok(runtime()?.handle().clone()) +} + #[expect(clippy::disallowed_methods, reason = "this is the gated door")] fn future_into_py(py: Python<'_>, future: F) -> PyResult> where @@ -813,4 +830,21 @@ asyncio.run(exercise()) .expect("async route contract should hold"); }); } + + #[rstest] + fn an_isolated_run_joins_blocking_work_it_left_running_before_returning() { + let finished = Arc::new(AtomicBool::new(false)); + let flag = Arc::clone(&finished); + let output = run_isolated(async move { + drop(tokio::task::spawn_blocking(move || { + thread::sleep(Duration::from_millis(50)); + flag.store(true, Ordering::SeqCst); + })); + 7 + }) + .unwrap(); + + assert_eq!(output, 7); + assert!(finished.load(Ordering::SeqCst)); + } } diff --git a/litellm-rust/crates/host-python/tests/isolated_run.rs b/litellm-rust/crates/host-python/tests/isolated_run.rs new file mode 100644 index 00000000000..7b8603229a7 --- /dev/null +++ b/litellm-rust/crates/host-python/tests/isolated_run.rs @@ -0,0 +1,9 @@ +use litellm_host_python::{reserve_process_for_forking, run_isolated, runtime_started}; + +#[test] +fn an_isolated_run_leaves_the_process_free_to_reserve_for_forking() { + run_isolated(async { tokio::task::spawn_blocking(|| ()).await.unwrap() }).unwrap(); + + assert!(!runtime_started()); + assert!(reserve_process_for_forking().is_ok()); +} diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 7078d4aba2a..9505aacf44b 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -16,6 +16,15 @@ pub enum TlsSource { #[error(transparent)] pub struct Error(Arc); +// A streamed body the connection ended early, after the response head arrived; `expected` is +// the declared `content-length`, `None` for a chunked body. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BodyCut { + pub received: u64, + pub expected: Option, + pub reset: bool, +} + #[derive(Debug, thiserror::Error)] pub(crate) enum Repr { #[error("could not read {}", path.display())] @@ -56,6 +65,12 @@ pub(crate) enum Repr { Connect(#[source] reqwest::Error), #[error("upstream network error")] Network(#[source] reqwest::Error), + #[error("upstream network error")] + BodyCut { + #[source] + source: reqwest::Error, + cut: BodyCut, + }, #[error("upstream request timed out")] Timeout(#[source] reqwest::Error), /// No bytes arrived within the read timeout, which the caller enforces per read; `body` says @@ -201,6 +216,37 @@ impl Error { matches!(&*self.0, Repr::ReadTimeout { body: true, .. }) } + pub fn body_read(error: reqwest::Error, received: u64, expected: Option) -> Self { + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(&error); + let kind = std::iter::from_fn(|| { + let current = source?; + source = current.source(); + Some(current) + }) + .find_map(|cause| cause.downcast_ref::().map(io::Error::kind)); + let reset = match kind { + Some(io::ErrorKind::UnexpectedEof) => false, + Some(io::ErrorKind::ConnectionReset) => true, + _ => return error.into(), + }; + Repr::BodyCut { + source: error.without_url(), + cut: BodyCut { + received, + expected, + reset, + }, + } + .into() + } + + pub fn body_cut(&self) -> Option { + match &*self.0 { + Repr::BodyCut { cut, .. } => Some(*cut), + _ => None, + } + } + pub fn computed_header(name: impl Into) -> Self { Repr::ComputedHeader(name.into()).into() } @@ -226,6 +272,7 @@ impl Error { pub fn dispatched(&self) -> bool { match &*self.0 { Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -285,6 +332,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -318,6 +366,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -355,6 +404,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -394,6 +444,7 @@ impl Error { | Repr::UnsignableHeader(_) | Repr::Signature(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -437,6 +488,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Status { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) @@ -473,6 +525,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::WebSocketUrl(_) @@ -513,6 +566,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::WebSocketUrl(_) | Repr::WebSocketScheme(_) | Repr::WebSocketSchemeRewrite(_) @@ -548,6 +602,7 @@ impl Error { | Repr::Signature(_) | Repr::Connect(_) | Repr::Network(_) + | Repr::BodyCut { .. } | Repr::Timeout(_) | Repr::ReadTimeout { .. } | Repr::Status { .. } @@ -742,4 +797,53 @@ mod tests { assert_eq!(error.upstream_response(), Some((429, body.as_str()))); assert_eq!(error.status(), Some(429)); } + + #[rstest] + #[case::chunked("transfer-encoding: chunked\r\n\r\n5\r\nhello\r\n", None)] + #[case::content_length("content-length: 100\r\n\r\nhello", Some(100))] + #[tokio::test] + async fn a_connection_that_ends_mid_body_is_a_body_cut( + #[case] tail: &'static str, + #[case] expected: Option, + ) { + use std::io::{Read as _, Write as _}; + + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("address"); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept"); + let mut request = [0_u8; 1024]; + let _ = stream.read(&mut request).expect("request"); + stream + .write_all(format!("HTTP/1.1 200 OK\r\n{tail}").as_bytes()) + .expect("response"); + }); + let mut response = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}/")) + .send() + .await + .expect("head"); + server.join().expect("server"); + let mut received = 0_u64; + let error = loop { + match response.chunk().await { + Ok(Some(chunk)) => received += chunk.len() as u64, + Ok(None) => panic!("the body must end with an error"), + Err(error) => break crate::Error::body_read(error, received, expected), + } + }; + assert_eq!( + error.body_cut(), + Some(super::BodyCut { + received: 5, + expected, + reset: false, + }) + ); + assert!(error.dispatched()); + assert_eq!(error.to_string(), "upstream network error"); + } } diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index 02129c62083..1aa11843b46 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -18,7 +18,7 @@ pub mod websocket; pub use client::Client; pub use config::{ClientIdentity, HttpClientConfig, Resolution, Verify}; -pub use error::{Error, TlsSource}; +pub use error::{BodyCut, Error, TlsSource}; pub use pool::{ClientVariant, HttpClientPool}; pub use proxy::EnvironmentProxies; pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive}; diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 3bdc63453ef..3b20bbda607 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -55,36 +55,13 @@ impl HttpClientPool { config: &HttpClientConfig, variant: ClientVariant, ) -> Result { - let effective = match variant { - ClientVariant::Media => HttpClientConfig { - client_certificate: None, - proxies: EnvironmentProxies::default(), - ..config.clone() - }, - ClientVariant::UnpinnedMedia => HttpClientConfig { - client_certificate: None, - ..config.clone() - }, - ClientVariant::WebSocket => HttpClientConfig { - http2: false, - verify: match &config.verify { - Verify::BuiltInRoots => Verify::NativeRoots, - verify => verify.clone(), - }, - ..config.clone() - }, - ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), - }; - let key = (effective, variant); + let key = (effective(config, variant), variant); if let Some(pooled) = self.lock().get(&key) && pooled.built_at.elapsed() < self.ttl { return Ok(Client::new(pooled.client.clone())); } - let client = self - .apply(variant, reqwest::ClientBuilder::try_from(&key.0)?) - .build() - .map_err(Repr::Client)?; + let client = self.build(&key.0, variant)?; self.lock().insert( key, PooledClient { @@ -95,6 +72,26 @@ impl HttpClientPool { Ok(Client::new(client)) } + /// A client nobody else holds, for a caller that drives it on a runtime of its own and drops both together + pub fn unpooled( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + self.build(&effective(config, variant), variant) + .map(Client::new) + } + + fn build( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + self.apply(variant, reqwest::ClientBuilder::try_from(config)?) + .build() + .map_err(|error| Error::from(Repr::Client(error))) + } + fn lock(&self) -> MutexGuard<'_, Clients> { self.clients.lock().unwrap_or_else(PoisonError::into_inner) } @@ -116,6 +113,29 @@ impl HttpClientPool { } } +fn effective(config: &HttpClientConfig, variant: ClientVariant) -> HttpClientConfig { + match variant { + ClientVariant::Media => HttpClientConfig { + client_certificate: None, + proxies: EnvironmentProxies::default(), + ..config.clone() + }, + ClientVariant::UnpinnedMedia => HttpClientConfig { + client_certificate: None, + ..config.clone() + }, + ClientVariant::WebSocket => HttpClientConfig { + http2: false, + verify: match &config.verify { + Verify::BuiltInRoots => Verify::NativeRoots, + verify => verify.clone(), + }, + ..config.clone() + }, + ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), + } +} + #[cfg(test)] mod tests { use std::{ diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index e5ece2837d8..d0f109c1e99 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -1,5 +1,5 @@ use litellm_core_utils::get_llm_provider_logic::LlmProviders; -use litellm_http::request::string_headers as shared_string_headers; +use litellm_llms::ErrorDetail; use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG; use litellm_llms::{ anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, @@ -19,8 +19,6 @@ use serde_json::{Map, Value}; use super::Error; -const HEADER_CONTEXT: &str = "messages"; - #[derive(Clone, Copy)] pub(crate) struct MessagesProvider { provider: LlmProviders, @@ -75,7 +73,52 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option>, ) -> Result, Error> { - shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from) + extra_headers + .unwrap_or_default() + .into_iter() + .map(|(name, value)| match value { + Value::String(value) => Ok((name, value)), + value => Err(Error::invalid_request(ErrorDetail::invalid( + "extra_headers", + NonStringHeader { + actual: python_type_name(&value), + name, + }, + ))), + }) + .collect() +} + +/// A header value the caller set to something other than a string, named by its Python type +/// since the caller wrote it in Python. +#[derive(Debug)] +struct NonStringHeader { + name: String, + actual: &'static str, +} + +impl std::fmt::Display for NonStringHeader { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "header {} must be a string, got {}", + self.name, self.actual + ) + } +} + +impl std::error::Error for NonStringHeader {} + +fn python_type_name(value: &Value) -> &'static str { + match value { + Value::Null => "NoneType", + Value::Bool(_) => "bool", + Value::Number(number) if number.is_f64() => "float", + Value::Number(_) => "int", + Value::String(_) => "str", + Value::Array(_) => "list", + Value::Object(_) => "dict", + } } #[cfg(test)] @@ -85,7 +128,6 @@ mod tests { use rstest::rstest; use super::{messages_provider, string_headers}; - use crate::Error; use litellm_core_utils::get_llm_provider_logic::LlmProviders; #[rstest] @@ -113,17 +155,28 @@ mod tests { assert!(messages_provider(provider, model).is_none()); } - #[test] - fn string_headers_rejects_non_string_values() { - let headers = json!({"x-count": 3}).as_object().unwrap().clone(); - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + #[rstest] + #[case::int(json!(12345), "int")] + #[case::float(json!(1.5), "float")] + #[case::bool(json!(true), "bool")] + #[case::null(json!(null), "NoneType")] + #[case::list(json!([1]), "list")] + #[case::dict(json!({"a": 1}), "dict")] + fn non_string_header_is_a_bad_request_naming_the_header_and_its_type( + #[case] value: serde_json::Value, + #[case] python_type: &str, + ) { + let headers = json!({"x-kept": "v", "x-count": value}) + .as_object() + .unwrap() + .clone(); + let error = string_headers(Some(headers)).expect_err("non-string header refused"); + assert!(error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); + let rejection = error.rejected_request().expect("a rejected request"); assert_eq!( - err, - Error::Headers(litellm_http::request::HeaderError { - context: "messages", - name: "x-count".to_string(), - actual: "number", - }) + litellm_error::Report::new(rejection).to_string(), + format!("invalid extra_headers: header x-count must be a string, got {python_type}") ); } } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index e0a7aa5adb8..c7a7d93bf98 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -26,10 +26,10 @@ use litellm_llms_types::formats::messages::{ }; use litellm_secrets::source::Secrets; use litellm_tracing::ByteChunk; -use serde_json::Value; +use serde_json::{Map, Value}; use super::{ - Error, MessagesCall, MessagesCallResponse, MessagesRoute, + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, common_utils::{MessagesProvider, merge_messages_header_sources, string_headers}, constants::MESSAGES_TIMEOUT_SECS, route::Messages, @@ -85,7 +85,10 @@ pub(super) async fn execute( ), }; let reported = match &output { - MessagesCallResponse::Complete(response) => provider.config().reported_cost(response), + MessagesCallResponse::Complete(response) => match response { + MessagesBody::Message(message) => provider.config().reported_cost(message), + MessagesBody::Unchecked(_) => None, + }, MessagesCallResponse::Stream { .. } => None, }; context @@ -130,9 +133,10 @@ async fn provider_call( )?; let transform_context = MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params); - let filtered = requested_optional_params(sanitized, provider, &transform_context)?; - let trimmed = - without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?; + let without_listed = + without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?; + let trimmed = requested_optional_params(without_listed, provider, &transform_context)?; + let requested_params = serde_json::to_value(&trimmed.params).map_err(serialize_failure)?; let scoped = get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str()); @@ -173,7 +177,7 @@ async fn provider_call( let request_context = RequestContext { model: body.model.clone(), custom_llm_provider: provider.as_str().to_string(), - optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?, + optional_params: requested_params, secret_fields: Vec::new(), api_key: api_key.map(SecretValue::new), }; @@ -346,8 +350,7 @@ async fn call_provider( // body it cannot parse, and the public mapping classifies it by that status. return Err(litellm_http::Error::upstream(status, text).into()); } - decode_response(config, &identity.model, &text) - .map(|message| MessagesCallResponse::Complete(Box::new(message))) + decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete) } fn serialize_failure(err: serde_json::Error) -> Error { @@ -383,13 +386,20 @@ fn timed_chunks( response: reqwest::Response, timeout: Duration, ) -> impl futures_util::Stream> + Send + 'static { - futures_util::stream::try_unfold(response, move |mut response| async move { - let chunk = tokio::time::timeout(timeout, response.chunk()) - .await - .map_err(litellm_http::Error::body_read_timeout)? - .map_err(litellm_http::Error::from)?; - Ok(chunk.map(|chunk| (chunk, response))) - }) + let expected = response.content_length(); + futures_util::stream::try_unfold( + (response, 0_u64), + move |(mut response, received)| async move { + let chunk = tokio::time::timeout(timeout, response.chunk()) + .await + .map_err(litellm_http::Error::body_read_timeout)? + .map_err(|error| litellm_http::Error::body_read(error, received, expected))?; + Ok(chunk.map(|chunk| { + let received = received + chunk.len() as u64; + (chunk, (response, received)) + })) + }, + ) } /// The provider's error, its body read under `timeout`: a stream has no overall deadline once @@ -411,15 +421,23 @@ fn decode_response( config: &dyn BaseMessagesConfig, model: &str, text: &str, -) -> Result { - let response = serde_json::from_str(text).map_err(|err| { - Error::invalid_response(litellm_llms::ErrorDetail::invalid( - "messages response JSON", - err, - )) - })?; +) -> Result { + let response = match serde_json::from_str::(text) { + Ok(response) => response, + Err(err) => { + let Ok(mut body) = serde_json::from_str::>(text) else { + return Err(Error::invalid_response(litellm_llms::ErrorDetail::invalid( + "messages response JSON", + err, + ))); + }; + body.entry("stop_sequence").or_insert(Value::Null); + return Ok(MessagesBody::Unchecked(body)); + } + }; config .transform_anthropic_messages_response(model, response) + .map(|message| MessagesBody::Message(Box::new(message))) .map_err(Error::from) .map_err(Error::after_send) } diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index 745442d0314..25ebeaacb49 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -17,8 +17,8 @@ use crate::common_utils::messages_provider; pub use common_utils::merge_messages_header_sources; pub use litellm_inference::Error; pub use types::{ - MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, - messages_body, + MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, + litellm_params, messages_body, }; #[derive(Clone)] diff --git a/litellm-rust/crates/inference-messages/src/route.rs b/litellm-rust/crates/inference-messages/src/route.rs index 7ff886689da..5f06d16d49f 100644 --- a/litellm-rust/crates/inference-messages/src/route.rs +++ b/litellm-rust/crates/inference-messages/src/route.rs @@ -5,11 +5,10 @@ use litellm_host::{ call::{HostedCompletion, HostedMachine, hosted_call}, protocol::Protocol, }; -use litellm_llms_types::formats::messages::MessagesResponse; -use super::{Error, MessagesCall}; +use super::{Error, MessagesBody, MessagesCall}; -pub type MessagesOutput = HostedCompletion>; +pub type MessagesOutput = HostedCompletion; /// The upstream response as the caller sees it at stream hand-off, before any chunk. pub struct MessagesStreamHead { @@ -19,7 +18,7 @@ pub struct MessagesStreamHead { pub struct Messages; impl Protocol for Messages { - type Response = Box; + type Response = MessagesBody; type Error = Error; type Request = MessagesCall; type HostCall = Infallible; diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index ae346259559..772ea6f7292 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -42,7 +42,16 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error { } pub type MessagesCallResponse = - CallOutput, super::route::MessagesStreamHead, Bytes, Error>; + CallOutput; + +// Python's `transform_anthropic_messages_response` returns any JSON object the provider sends +// with a 200, message or not, so one that is not a message passes through as sent. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MessagesBody { + Message(Box), + Unchecked(Map), +} #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct MessagesShaping { @@ -144,6 +153,7 @@ mod tests { supports_reasoning: true, supports_adaptive_thinking: true, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: true, supports_sampling_params: false, diff --git a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs index f91d5c79f7c..5207f4e9d58 100644 --- a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs +++ b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs @@ -1,7 +1,9 @@ use std::{sync::Arc, time::Duration}; use futures_util::TryStreamExt; -use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping}; +use litellm_inference_messages::{ + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesShaping, +}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse}; use rstest::{fixture, rstest}; @@ -74,7 +76,10 @@ async fn run(call: MessagesCall) -> Result { async fn run_message(call: MessagesCall) -> MessagesResponse { match run(call).await.unwrap() { - MessagesCallResponse::Complete(message) => *message, + MessagesCallResponse::Complete(message) => match message { + MessagesBody::Message(message) => *message, + MessagesBody::Unchecked(_) => panic!("expected a message"), + }, MessagesCallResponse::Stream { .. } => panic!("expected a completed message"), } } diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 663c87453a9..787501a8fed 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -50,42 +50,47 @@ async fn messages_cache_identity_includes_provider_native_parameters( #[case] original: Value, #[case] changed: Value, ) { - use litellm_inference_messages::route::Messages; + use litellm_inference_messages::{MessagesBody, route::Messages}; use litellm_llms_types::formats::messages::MessagesResponse; let calls = AtomicUsize::new(0); for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] { - let response = - execute_unary::( - CacheRequest::from_wire( - ProviderIdentity { - model: "test".into(), - provider: "anthropic".into(), - }, - Some(&WireRequest { - url: "https://example.test/v1/messages".into(), - headers: vec![], - body: json!({ - "model":"test", "messages":[{"role":"user","content":"hello"}], - "max_tokens":32, (field):value - }), - }), - ), - Some(cache.clone()), - Some(CacheOptions::new(CacheScope::Shared)), - &(), - None, - || async { - let call = calls.fetch_add(1, Ordering::SeqCst); - Ok(Box::new(serde_json::from_value::(json!({ - "id":call.to_string(), "type":"message", "role":"assistant", "model":"test", - "content":[{"type":"text","text":format!("answer {call}")}], - "stop_reason":"end_turn", "stop_sequence":null - })).unwrap())) + let response = execute_unary::( + CacheRequest::from_wire( + ProviderIdentity { + model: "test".into(), + provider: "anthropic".into(), }, - ) - .await - .unwrap(); + Some(&WireRequest { + url: "https://example.test/v1/messages".into(), + headers: vec![], + body: json!({ + "model":"test", "messages":[{"role":"user","content":"hello"}], + "max_tokens":32, (field):value + }), + }), + ), + Some(cache.clone()), + Some(CacheOptions::new(CacheScope::Shared)), + &(), + None, + || async { + let call = calls.fetch_add(1, Ordering::SeqCst); + Ok(MessagesBody::Message(Box::new( + serde_json::from_value::(json!({ + "id":call.to_string(), "type":"message", "role":"assistant", "model":"test", + "content":[{"type":"text","text":format!("answer {call}")}], + "stop_reason":"end_turn", "stop_sequence":null + })) + .unwrap(), + ))) + }, + ) + .await + .unwrap(); + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!(response.id, expected_call.to_string()); assert_eq!( response.content[0]["text"], @@ -253,7 +258,7 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ #[case] cost: Value, #[case] expected: Option, ) { - use litellm_inference_messages::{MessagesCall, MessagesCallResponse}; + use litellm_inference_messages::{MessagesBody, MessagesCall, MessagesCallResponse}; let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test", "content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null, "usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost}); @@ -277,6 +282,9 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ panic!("expected Messages responses"); }; assert_eq!(first, second); + let MessagesBody::Message(first) = first else { + panic!("expected a message"); + }; assert_eq!(first.extra.get("cost"), Some(&cost)); assert_eq!(support::received(&upstream).await.len(), 1); let facts = hooks.facts.lock().unwrap(); diff --git a/litellm-rust/crates/inference-messages/tests/live/support.rs b/litellm-rust/crates/inference-messages/tests/live/support.rs index 81e295e34ad..52684928f1f 100644 --- a/litellm-rust/crates/inference-messages/tests/live/support.rs +++ b/litellm-rust/crates/inference-messages/tests/live/support.rs @@ -4,7 +4,9 @@ use futures_util::{FutureExt, TryStreamExt}; use litellm_host::interceptors::{ ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest, }; -use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute}; +use litellm_inference_messages::{ + Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, +}; use litellm_inference_testing::live::{LiveResources, model, within_deadline}; use litellm_llms_types::formats::messages::MessagesResponse; use serde_json::{Value, json}; @@ -111,7 +113,10 @@ pub(super) fn call(provider: &str, payload: Value) -> MessagesCall { pub(super) fn complete(response: MessagesCallResponse) -> Box { match response { - MessagesCallResponse::Complete(message) => message, + MessagesCallResponse::Complete(message) => match message { + MessagesBody::Message(message) => message, + MessagesBody::Unchecked(_) => panic!("a complete request returned no message"), + }, MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"), } } diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index c3d06027c32..063e9d18738 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -6,9 +6,11 @@ use litellm_host::{ interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest}, lifecycle::CallEvent, }; -use litellm_inference_messages::{MessagesCallResponse, route::Messages}; +use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; -use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities; +use litellm_llms::base_llm::messages::context::{ + MessagesModelCapabilities as AnthropicModelCapabilities, SupportedEffortTiers, +}; use rstest::rstest; use super::*; @@ -338,7 +340,10 @@ async fn response_mode_follows_the_intercepted_request( } MessagesCallResponse::Complete(message) => { assert_eq!(expected_stream, Some(false)); - assert_eq!(*message, serde_json::from_value(message_body()).unwrap()); + assert_eq!( + message, + MessagesBody::Message(serde_json::from_value(message_body()).unwrap()) + ); } } Ok::<_, Error>(()) @@ -419,7 +424,7 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure( assert!(host.raw_responses().is_empty()); } -/// Python logs `optional_params` as what it is about to send, so a dropped param must +/// Python logs `optional_params` after its requested-param gate, so a dropped param must /// not resurface in callbacks. #[rstest] #[tokio::test] @@ -450,3 +455,82 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); assert_eq!(optional_params, json!({"max_tokens": 16})); } + +#[rstest] +#[tokio::test] +async fn the_request_context_leaves_out_a_listed_param_the_model_refuses(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + shaping: MessagesShaping { + settings: MessagesSettings { + additional_drop_params: vec!["temperature".into()], + ..MessagesSettings::default() + }, + capabilities: AnthropicModelCapabilities { + supports_sampling_params: false, + ..AnthropicModelCapabilities::default() + }, + }, + ..with_fields(call, json!({"temperature": 0.2})) + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); + assert_eq!(optional_params, json!({"max_tokens": 16})); +} + +#[rstest] +#[case::reasoning_effort_on_an_adaptive_model( + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_output_config: true, + effort_tiers: SupportedEffortTiers { low: true, ..SupportedEffortTiers::default() }, + ..AnthropicModelCapabilities::default() + }, + json!({"reasoning_effort": "low"}) +)] +#[case::adaptive_payload_on_a_legacy_model( + AnthropicModelCapabilities { + supports_reasoning: true, + ..AnthropicModelCapabilities::default() + }, + json!({"max_tokens": 3000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) +)] +#[tokio::test] +async fn the_request_context_carries_the_requested_reasoning_not_its_translation( + call: MessagesCall, + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, +) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + shaping: MessagesShaping { + capabilities, + ..MessagesShaping::default() + }, + ..with_fields(call, fields.clone()) + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); + assert_eq!( + optional_params, + Value::Object( + object(json!({"max_tokens": 16})) + .into_iter() + .chain(object(fields)) + .collect() + ) + ); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index a01e571ec75..fe203287766 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -5,7 +5,7 @@ use std::{ use litellm_http::{HttpSettings, Resolution}; use litellm_inference_messages::{ - Error, MessagesCall, MessagesSettings, MessagesShaping, + Error, MessagesBody, MessagesCall, MessagesSettings, MessagesShaping, route::{Messages, MessagesMachine, MessagesOutput}, }; use litellm_inference_testing::RecordingSecrets; @@ -119,7 +119,10 @@ async fn run(call: MessagesCall) -> Result { async fn run_message(call: MessagesCall) -> MessagesResponse { match run(call).await.expect("messages call succeeds") { - MessagesOutput::Complete(message) => *message, + MessagesOutput::Complete(message) => match message { + MessagesBody::Message(message) => *message, + MessagesBody::Unchecked(_) => panic!("a call returned no message"), + }, MessagesOutput::StreamEnded | MessagesOutput::Detached => { panic!("a non-streaming call returned a stream") } diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 216aa96fbe6..3784eb6cc68 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -530,6 +530,30 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i assert_eq!(sent["max_tokens"], 16); } +#[rstest] +#[tokio::test] +async fn additional_drop_params_clear_a_param_the_model_does_not_support(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities: sampling_removed(), + settings: MessagesSettings { + additional_drop_params: vec!["top_k".into()], + ..MessagesSettings::default() + }, + }, + ..with_fields(call, json!({"top_k": 5})) + }) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!(sent.get("top_k"), None); + assert_eq!(sent["max_tokens"], 16); +} + #[rstest] #[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))] #[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))] diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index 76813fe0a34..f978da5e995 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -2,7 +2,7 @@ use litellm_host::{ interceptors::{ExecutionFacts, ResultSource}, lifecycle::ExecutionEvent, }; -use litellm_inference_messages::{MessagesCallResponse, messages_body}; +use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body}; use litellm_inference_testing::{ RecordingSecrets, http_config, no_secrets, provider_http, resources, }; @@ -48,6 +48,9 @@ async fn calls_defer_execution_until_polled( let MessagesCallResponse::Complete(response) = future.await.unwrap() else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!( response.content, message_body()["content"].as_array().unwrap().as_slice() @@ -240,10 +243,42 @@ async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: Me assert_eq!(error.upstream_response(), Some((200, "not json"))); } +#[rstest] +#[case::error_object( + json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}), + json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}, + "stop_sequence": null}), +)] +#[case::message_missing_id( + json!({"type": "message", "content": [], "stop_sequence": "###"}), + json!({"type": "message", "content": [], "stop_sequence": "###"}), +)] +#[tokio::test] +async fn a_success_body_that_is_not_a_message_passes_through( + call: MessagesCall, + #[case] body: Value, + #[case] expected: Value, +) { + let upstream = upstream([json_response(body)]).await; + + let response = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .expect("any JSON object is returned"); + + let MessagesOutput::Complete(response) = response else { + panic!("expected a completed response"); + }; + assert_eq!(serde_json::to_value(response).unwrap(), expected); +} + #[rstest] #[tokio::test] async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) { - let upstream = upstream([json_response(json!({"unexpected": true}))]).await; + let upstream = upstream([json_response(json!(["unexpected"]))]).await; let error = run(MessagesCall { api_key: Some("sk".into()), @@ -320,6 +355,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes let MessagesCallResponse::Complete(message) = response else { panic!("a non-streaming request returns a message"); }; + let MessagesBody::Message(message) = message else { + panic!("expected a message"); + }; assert_eq!(message.id, "msg_1"); let sent = only_request(&upstream).await; assert_eq!(sent.header("x-api-key"), Some("sk-ant")); @@ -412,6 +450,9 @@ async fn route_uses_injected_dependencies_and_optional_cache( else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!( response.content, message_body()["content"].as_array().unwrap().as_slice() @@ -484,6 +525,9 @@ async fn cache_overrides_preserve_the_routes_isolated_scope(call: MessagesCall) else { panic!("expected a completed message"); }; + let MessagesBody::Message(response) = response else { + panic!("expected a message"); + }; assert_eq!(response.id, expected["id"].as_str().unwrap()); } assert_eq!(received(&upstream).await.len(), 2); diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index 841e0c53799..4135424d6ef 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -164,6 +164,14 @@ impl Error { } } + pub fn body_cut(&self) -> Option { + match self { + Self::Llm(error) => error.body_cut(), + Self::Http(error) => error.body_cut(), + _ => None, + } + } + /// `Some(connect)` when the call timed out, where `connect` says the connection was never /// established; `None` for every other error. pub fn timeout(&self) -> Option { diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 13653f438f7..0854820358b 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1913,6 +1913,7 @@ mod tests { supports_reasoning: false, supports_adaptive_thinking: false, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: false, supports_sampling_params: true, diff --git a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs index ae3488a21f9..0ddbc915154 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs @@ -2,8 +2,8 @@ use litellm_llms_types::{ formats::{ chat_completions::ReasoningEffort, messages::{ - AdaptiveThinking, EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig, - ThinkingConfig, ThinkingDisplay, + AdaptiveThinking, DisabledThinking, EffortLevel, MessagesOptionalParams, + MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay, }, }, recognized::Recognized, @@ -172,7 +172,11 @@ fn translate_reasoning_effort( ) else { return Ok(MessagesRequest { params: MessagesOptionalParams { - thinking: None, + thinking: (context.capabilities.thinking_on_by_default + && !context.capabilities.thinking_always_on) + .then(|| { + Recognized::Known(ThinkingConfig::Disabled(DisabledThinking::default())) + }), output_config: None, ..request.params }, @@ -748,6 +752,33 @@ mod tests { ); } + #[rstest] + #[case::default_thinking(true, false, Some(serde_json::json!({"type": "disabled"}))) ] + #[case::legacy(false, false, None)] + #[case::always_on(true, true, None)] + fn none_disables_only_optional_default_thinking( + #[case] thinking_on_by_default: bool, + #[case] thinking_always_on: bool, + #[case] thinking: Option, + ) { + let capabilities = MessagesModelCapabilities { + thinking_on_by_default, + thinking_always_on, + ..opus_4_7() + }; + let expected = match thinking { + Some(thinking) => serde_json::json!({"max_tokens": 1024, "thinking": thinking}), + None => serde_json::json!({"max_tokens": 1024}), + }; + assert_eq!( + translate( + capabilities, + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "none"}) + ), + Ok(request(expected)) + ); + } + #[rstest] #[case::adaptive_model(opus_4_7())] #[case::effort_capable_model(opus_4_5())] diff --git a/litellm-rust/crates/llms/src/base_llm/messages/context.rs b/litellm-rust/crates/llms/src/base_llm/messages/context.rs index b514e975aa4..3d8b0f41d8c 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/context.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/context.rs @@ -32,6 +32,8 @@ pub struct MessagesModelCapabilities { #[serde(default)] pub thinking_always_on: bool, #[serde(default)] + pub thinking_on_by_default: bool, + #[serde(default)] pub supports_legacy_thinking: bool, #[serde(default)] pub supports_output_config: bool, @@ -55,6 +57,7 @@ impl Default for MessagesModelCapabilities { supports_reasoning: false, supports_adaptive_thinking: false, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: false, supports_sampling_params: true, diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index 2844fbf3e01..2975ceee3cb 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -103,6 +103,13 @@ impl Error { matches!(self, Self::Http(error) if error.is_body_read_timeout()) } + pub fn body_cut(&self) -> Option { + match self { + Self::Http(error) => error.body_cut(), + _ => None, + } + } + pub fn is_timeout(&self) -> Option { match self { Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()), diff --git a/litellm-rust/crates/model-catalog-refresh/AGENTS.md b/litellm-rust/crates/model-catalog-refresh/AGENTS.md new file mode 100644 index 00000000000..96228716abf --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/AGENTS.md @@ -0,0 +1,9 @@ +# Model catalog refresh + +- `fetch` fetches one document over http(s) or `s3://` and never parses it +- Only 429, 500, 502, 503, 504 and transport errors are retried. S3 errors are final +- `Failed.retryable` means only the attempt limit ended the fetch +- A server's Retry-After overrides the backoff, capped at `max_delay` +- `Start::After(failed)` resumes the backoff schedule where another caller's first attempt left off +- Error reasons use the URL with userinfo, query and fragment stripped, since presigned URLs carry credentials +- Bodies over `MAX_BODY_BYTES` are rejected diff --git a/litellm-rust/crates/model-catalog-refresh/Cargo.toml b/litellm-rust/crates/model-catalog-refresh/Cargo.toml index 27219042c6c..76f4a9fb969 100644 --- a/litellm-rust/crates/model-catalog-refresh/Cargo.toml +++ b/litellm-rust/crates/model-catalog-refresh/Cargo.toml @@ -6,14 +6,19 @@ license.workspace = true repository.workspace = true [dependencies] +backon.workspace = true +bytes.workspace = true +litellm-error.workspace = true litellm-http.workspace = true -litellm-model-catalog.workspace = true +litellm-s3.workspace = true +litellm-tracing.workspace = true +reqwest.workspace = true thiserror.workspace = true +tokio = { workspace = true, features = ["time"] } url.workspace = true [dev-dependencies] -litellm-error.workspace = true +aws-credential-types = "1.3.0" litellm-http = { workspace = true, features = ["test-support"] } -reqwest.workspace = true rstest.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } diff --git a/litellm-rust/crates/model-catalog-refresh/src/attempt.rs b/litellm-rust/crates/model-catalog-refresh/src/attempt.rs new file mode 100644 index 00000000000..45d80e36a18 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/attempt.rs @@ -0,0 +1,153 @@ +use std::time::Duration; + +use bytes::Bytes; + +use litellm_http::Client; +use litellm_s3::S3Uri; +use litellm_s3::sdk::config::retry::RetryConfig; +use litellm_s3::sdk::config::timeout::TimeoutConfig; +use reqwest::header::{ETAG, HeaderMap, RETRY_AFTER}; +use url::Url; + +use crate::{Error, RETRYABLE_STATUSES, S3Access}; + +pub(crate) const MAX_BODY_BYTES: u64 = 100 * 1024 * 1024; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Fetched { + pub body: Vec, + pub etag: Option, +} + +#[derive(Debug)] +pub(crate) enum Cause { + Status(u16), + Transport(Error), +} + +#[derive(Debug)] +pub(crate) struct Rejected { + pub(crate) cause: Cause, + pub(crate) retryable: bool, + pub(crate) retry_after: Option, +} + +impl Rejected { + pub(crate) fn fatal(error: Error) -> Self { + Self { + cause: Cause::Transport(error), + retryable: false, + retry_after: None, + } + } +} + +/// Any status below 400 is the document +pub(crate) async fn http( + client: &Client, + url: Url, + timeout: Duration, +) -> Result { + let transport = |error: litellm_http::Error| Rejected { + cause: Cause::Transport(Error::from(error)), + retryable: true, + retry_after: None, + }; + let mut response = tokio::time::timeout(timeout, client.get(url).send()) + .await + .map_err(litellm_http::Error::read_timeout) + .and_then(|sent| sent.map_err(litellm_http::Error::from)) + .map_err(transport)?; + let status = response.status().as_u16(); + let etag = etag(response.headers()); + let retry_after = response + .headers() + .get(RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.trim().parse().ok()) + .map(Duration::from_secs); + if status >= 400 { + return Err(Rejected { + cause: Cause::Status(status), + retryable: RETRYABLE_STATUSES.contains(&status), + retry_after, + }); + } + if response + .content_length() + .is_some_and(|length| length > MAX_BODY_BYTES) + { + return Err(Rejected::fatal(Error::body_too_large(MAX_BODY_BYTES))); + } + let mut body = Vec::new(); + while let Some(chunk) = tokio::time::timeout(timeout, response.chunk()) + .await + .map_err(litellm_http::Error::body_read_timeout) + .and_then(|read| read.map_err(litellm_http::Error::from)) + .map_err(transport)? + { + append(&mut body, &chunk).map_err(Rejected::fatal)?; + } + Ok(Fetched { body, etag }) +} + +/// S3 failures are final: the SDK already classifies its own errors and a missing or denied object stays so +pub(crate) async fn s3( + client: &Client, + s3: &S3Access, + object: &S3Uri, + timeout: Duration, +) -> Result { + let fetched = async { + let config = match s3 { + S3Access::Environment { region } => { + litellm_s3::environment_config(client.clone(), region.clone()).await + } + S3Access::Config(config) => config.to_builder(), + } + .retry_config(RetryConfig::disabled()) + .timeout_config( + TimeoutConfig::builder() + .operation_attempt_timeout(timeout) + .build(), + ) + .build(); + let output = litellm_s3::sdk::Client::from_conf(config) + .get_object() + .bucket(&object.bucket) + .key(&object.key) + .send() + .await?; + if output + .content_length() + .is_some_and(|length| length.unsigned_abs() > MAX_BODY_BYTES) + { + return Err(Error::body_too_large(MAX_BODY_BYTES)); + } + let etag = output.e_tag().map(str::to_owned); + let mut stream = output.body; + let mut body = Vec::new(); + while let Some(chunk) = tokio::time::timeout(timeout, stream.try_next()).await?? { + append(&mut body, &chunk)?; + } + Ok::<_, Error>(Fetched { body, etag }) + }; + fetched.await.map_err(Rejected::fatal) +} + +fn append(body: &mut Vec, chunk: &Bytes) -> Result<(), Error> { + if (body.len() + chunk.len()) as u64 > MAX_BODY_BYTES { + return Err(Error::body_too_large(MAX_BODY_BYTES)); + } + body.extend_from_slice(chunk); + Ok(()) +} + +fn etag(headers: &HeaderMap) -> Option { + let values: Vec<&str> = headers + .get_all(ETAG) + .iter() + .filter_map(|value| value.to_str().ok()) + .collect(); + (!values.is_empty()).then(|| values.join(", ")) +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/error.rs b/litellm-rust/crates/model-catalog-refresh/src/error.rs index 656f32370cb..c79daeb4777 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/error.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/error.rs @@ -1,6 +1,9 @@ -use std::io; - +use litellm_s3::sdk::config::http::HttpResponse; +use litellm_s3::sdk::error::SdkError; +use litellm_s3::sdk::operation::get_object::GetObjectError; +use litellm_s3::sdk::primitives::ByteStreamError; use thiserror::Error; +use tokio::time::error::Elapsed; #[derive(Debug, Error)] #[error(transparent)] @@ -10,18 +13,48 @@ pub struct Error(Box); enum Repr { #[error(transparent)] Http(#[from] litellm_http::Error), - #[error("could not read the catalog")] - Read(#[source] io::Error), #[error("invalid catalog URL")] InvalidUrl(#[source] url::ParseError), #[error("unsupported catalog URL scheme {scheme}")] UnsupportedScheme { scheme: String }, + #[error("invalid S3 URL")] + S3Uri(#[source] litellm_s3::Error), + #[error("S3 GetObject failed")] + S3GetObject(#[source] Box>), + #[error("S3 GetObject body read failed")] + S3Body(#[source] ByteStreamError), + #[error("S3 GetObject timed out")] + S3Timeout(#[source] Elapsed), + #[error("response body is over the {limit} byte limit")] + BodyTooLarge { limit: u64 }, } impl Error { pub fn unsupported_scheme(scheme: String) -> Self { Self(Box::new(Repr::UnsupportedScheme { scheme })) } + + pub(crate) fn body_too_large(limit: u64) -> Self { + Self(Box::new(Repr::BodyTooLarge { limit })) + } +} + +impl From> for Error { + fn from(cause: SdkError) -> Self { + Self(Box::new(Repr::S3GetObject(Box::new(cause)))) + } +} + +impl From for Error { + fn from(cause: ByteStreamError) -> Self { + Self(Box::new(Repr::S3Body(cause))) + } +} + +impl From for Error { + fn from(cause: Elapsed) -> Self { + Self(Box::new(Repr::S3Timeout(cause))) + } } impl From for Error { @@ -30,9 +63,9 @@ impl From for Error { } } -impl From for Error { - fn from(cause: io::Error) -> Self { - Self(Box::new(Repr::Read(cause))) +impl From for Error { + fn from(cause: litellm_s3::Error) -> Self { + Self(Box::new(Repr::S3Uri(cause))) } } @@ -44,29 +77,18 @@ impl From for Error { #[cfg(test)] mod tests { + use std::io; + use litellm_error::{Report, find}; use rstest::rstest; use super::*; - use crate::{DoneReason, Reply, Step}; #[rstest] - fn exhausted_transport_keeps_its_typed_cause() { - let reply = Reply::Transport { - error: io::Error::from(io::ErrorKind::ConnectionReset).into(), - }; - let Reply::Transport { error } = reply else { - panic!("transport failure") - }; - let step = Step::Done { - reason: DoneReason::Exhausted(error), - }; - let Step::Done { - reason: DoneReason::Exhausted(error), - } = step - else { - panic!("exhausted attempts") - }; + fn an_s3_failure_keeps_its_typed_cause_once() { + let error = Error::from(SdkError::::timeout_error( + io::Error::from(io::ErrorKind::ConnectionReset), + )); let cause = find::(&error).expect("typed I/O cause"); assert_eq!(cause.kind(), io::ErrorKind::ConnectionReset); assert_eq!( diff --git a/litellm-rust/crates/model-catalog-refresh/src/fetch.rs b/litellm-rust/crates/model-catalog-refresh/src/fetch.rs new file mode 100644 index 00000000000..bd58827b546 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/fetch.rs @@ -0,0 +1,122 @@ +use std::borrow::Cow; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::Duration; + +use backon::Retryable; +use litellm_error::Report; +use litellm_http::Client; +use url::Url; + +use crate::attempt::{self, Cause, Fetched, Rejected}; +use crate::{RetryPolicy, Source}; + +/// How an `s3://` fetch reaches S3 +#[derive(Clone, Debug)] +pub enum S3Access { + /// What an AWS SDK finds in the environment and shared config files, loaded on each `s3://` attempt, with + /// `region` winning over the region it finds + Environment { region: Option }, + /// A config the caller built with its region, endpoint and credentials + Config(Box), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Failed { + pub reason: String, + /// Another attempt could succeed. Only the attempt limit ended the fetch + pub retryable: bool, + pub retry_after: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Start { + Fresh, + /// Continues a fetch whose first attempt failed retryably elsewhere, waiting before its next attempt + After(Failed), +} + +pub async fn fetch( + client: &Client, + s3: &S3Access, + url: &str, + policy: RetryPolicy, + start: Start, +) -> Result { + let shown = redacted(url); + let source = + Source::parse(url).map_err(|error| Failed::of(&shown, &Rejected::fatal(error), 1))?; + let mut backoff = policy.backoff(); + let made = AtomicU32::new(0); + if let Start::After(failed) = start { + let Some(wait) = backoff.next() else { + return Err(failed); + }; + tokio::time::sleep(policy.delay(wait, failed.retry_after)).await; + made.store(1, Ordering::Relaxed); + } + let attempt = || { + made.fetch_add(1, Ordering::Relaxed); + async { + match &source { + Source::Http(http) => attempt::http(client, http.clone(), policy.timeout).await, + Source::S3(object) => attempt::s3(client, s3, object, policy.timeout).await, + } + } + }; + attempt + .retry(backoff) + .sleep(tokio::time::sleep) + .when(|rejected| rejected.retryable) + .adjust(|rejected, wait| wait.map(|wait| policy.delay(wait, rejected.retry_after))) + .notify(|rejected, wait| { + litellm_tracing::warn!( + "LiteLLM: model cost map fetch attempt {}/{} failed ({}); retrying in {:.1}s", + made.load(Ordering::Relaxed), + policy.max_attempts, + reason(&shown, &rejected.cause), + wait.as_secs_f64(), + ); + }) + .await + .map_err(|rejected| Failed::of(&shown, &rejected, made.load(Ordering::Relaxed))) +} + +impl Failed { + fn of(url: &str, rejected: &Rejected, attempts: u32) -> Self { + let reason = reason(url, &rejected.cause); + Self { + reason: match rejected.retryable && attempts > 1 { + true => format!("{reason} (after {attempts} attempts)"), + false => reason, + }, + retryable: rejected.retryable, + retry_after: rejected.retry_after, + } + } +} + +fn reason(url: &str, cause: &Cause) -> String { + match cause { + Cause::Status(status) => format!("HTTP {status} from {url}"), + Cause::Transport(error) => format!("error fetching {url}: {}", Report::new(error)), + } +} + +/// Drops userinfo, query and fragment, where presigned URLs and tokens carry credentials +fn redacted(url: &str) -> Cow<'_, str> { + let Ok(mut parsed) = Url::parse(url) else { + return Cow::Borrowed(url); + }; + if parsed.username().is_empty() + && parsed.password().is_none() + && parsed.query().is_none() + && parsed.fragment().is_none() + { + return Cow::Borrowed(url); + } + let _ = parsed.set_username(""); + let _ = parsed.set_password(None); + parsed.set_query(None); + parsed.set_fragment(None); + Cow::Owned(parsed.into()) +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/lib.rs b/litellm-rust/crates/model-catalog-refresh/src/lib.rs index 006c4887fad..1742a669483 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/lib.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/lib.rs @@ -1,9 +1,11 @@ +mod attempt; mod error; +mod fetch; mod policy; mod source; -mod step; +pub use attempt::Fetched; pub use error::Error; +pub use fetch::*; pub use policy::*; pub use source::*; -pub use step::*; diff --git a/litellm-rust/crates/model-catalog-refresh/src/policy.rs b/litellm-rust/crates/model-catalog-refresh/src/policy.rs index bc32e40408e..1cdb737ecd8 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/policy.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/policy.rs @@ -1,20 +1,76 @@ use std::time::Duration; -use litellm_model_catalog::IntegrityLimits; +use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder}; -/// Main's defaults, 3 attempts, 30 s max wait, 5 s timeout, backoff base 2.0, jitter 0.0 to 1.0 s -#[derive(Clone, Copy, Debug, PartialEq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct RetryPolicy { pub max_attempts: u32, - pub max_wait: Duration, pub timeout: Duration, - pub backoff_base_seconds: f64, - pub jitter_min_seconds: f64, - pub jitter_max_seconds: f64, - /// Main's 50 model floor and 0.5 shrink ratio - pub limits: IntegrityLimits, + pub min_delay: Duration, + pub max_delay: Duration, } +pub const RETRYABLE_STATUSES: [u16; 5] = [429, 500, 502, 503, 504]; + /// Main's default `LITELLM_MODEL_COST_MAP_URL` pub const DEFAULT_REGISTRY_URL: &str = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"; + +impl RetryPolicy { + pub(crate) fn backoff(&self) -> ExponentialBackoff { + ExponentialBuilder::new() + .with_min_delay(self.min_delay) + .with_max_delay(self.max_delay) + .with_max_times(self.max_attempts.saturating_sub(1) as usize) + .with_jitter() + .build() + } + + pub(crate) fn delay(&self, backoff: Duration, retry_after: Option) -> Duration { + retry_after.map_or(backoff, |server| server.min(self.max_delay)) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + const POLICY: RetryPolicy = RetryPolicy { + max_attempts: 3, + timeout: Duration::from_secs(5), + min_delay: Duration::from_secs(1), + max_delay: Duration::from_secs(30), + }; + + #[rstest] + #[case::backoff_without_retry_after(None, 4)] + #[case::retry_after_wins(Some(7), 7)] + #[case::retry_after_zero(Some(0), 0)] + #[case::retry_after_capped(Some(120), 30)] + fn the_server_retry_after_overrides_the_backoff_up_to_max_delay( + #[case] retry_after: Option, + #[case] seconds: u64, + ) { + assert_eq!( + POLICY.delay(Duration::from_secs(4), retry_after.map(Duration::from_secs)), + Duration::from_secs(seconds) + ); + } + + #[rstest] + #[case::one_attempt(1, 0)] + #[case::three_attempts(3, 2)] + #[case::zero_attempts(0, 0)] + fn the_backoff_allows_one_retry_fewer_than_the_attempts( + #[case] max_attempts: u32, + #[case] retries: usize, + ) { + let policy = RetryPolicy { + max_attempts, + ..POLICY + }; + assert_eq!(policy.backoff().count(), retries); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/source.rs b/litellm-rust/crates/model-catalog-refresh/src/source.rs index 713e0386941..28849ba2841 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/source.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/source.rs @@ -1,17 +1,59 @@ -use std::path::PathBuf; - +use litellm_s3::S3Uri; use url::Url; +use crate::Error; + /// Where a fetched catalog comes from #[derive(Clone, Debug, PartialEq, Eq)] pub enum Source { Http(Url), - File(PathBuf), - S3(S3Object), + S3(S3Uri), } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct S3Object { - pub bucket: String, - pub key: String, +impl Source { + pub fn parse(url: &str) -> Result { + let parsed = Url::parse(url)?; + match parsed.scheme() { + "http" | "https" => Ok(Self::Http(parsed)), + "s3" => Ok(Self::S3(S3Uri::try_from(url)?)), + scheme => Err(Error::unsupported_scheme(scheme.to_owned())), + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::https("https://example.com/cost.json")] + #[case::http("http://127.0.0.1:8000/cost.json")] + fn http_urls_are_fetched_over_http(#[case] url: &str) { + assert_eq!( + Source::parse(url).unwrap(), + Source::Http(Url::parse(url).unwrap()) + ); + } + + #[rstest] + fn s3_urls_name_a_bucket_and_key() { + assert_eq!( + Source::parse("s3://bucket/prices/cost.json").unwrap(), + Source::S3(S3Uri { + bucket: "bucket".into(), + key: "prices/cost.json".into(), + }) + ); + } + + #[rstest] + #[case::s3_without_a_key("s3://bucket/")] + #[case::unsupported("ftp://example.com/cost.json")] + #[case::file("file:///tmp/cost.json")] + #[case::not_a_url("not a url")] + fn other_urls_are_refused(#[case] url: &str) { + assert!(Source::parse(url).is_err()); + } } diff --git a/litellm-rust/crates/model-catalog-refresh/src/step.rs b/litellm-rust/crates/model-catalog-refresh/src/step.rs deleted file mode 100644 index 822d1807916..00000000000 --- a/litellm-rust/crates/model-catalog-refresh/src/step.rs +++ /dev/null @@ -1,66 +0,0 @@ -use std::time::Duration; - -use litellm_model_catalog::Catalog; - -use crate::Error; -use crate::source::Source; - -/// One attempt. Where to fetch from is the `Source`, so the request carries no URL -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct FetchRequest { - pub if_none_match: Option, - pub timeout: Duration, -} - -#[derive(Debug)] -pub enum Reply { - /// An HTTP, S3 or file answer. The body is read only on 200 - Response { - status: u16, - etag: Option, - retry_after: Option, - body: Option>, - }, - /// No answer reached us - Transport { error: Error }, -} - -/// What the fetch state machine asks its driver to do next -#[derive(Debug)] -pub enum Step { - Fetch(FetchRequest), - Sleep(Duration), - /// A parsed document to install, with its rejected row errors - Adopt { - catalog: Box, - installed: Installed, - errors: Vec, - }, - /// 304, a final status, or retries exhausted - Done { - reason: DoneReason, - }, -} - -#[derive(Debug)] -pub enum DoneReason { - NotModified, - Status(u16), - Exhausted(Error), -} - -/// What is installed now, the input to the next revalidation -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct Installed { - pub origin: Origin, - pub etag: Option, - /// `git_blob_id` of the installed bytes - pub revision: String, - pub model_count: usize, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum Origin { - Bundled, - Remote { source: Source }, -} diff --git a/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs b/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs new file mode 100644 index 00000000000..2a4ec996bd4 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs @@ -0,0 +1,541 @@ +use std::io::{BufRead, BufReader, Write}; +use std::net::TcpListener; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use aws_credential_types::Credentials; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::{ProvideCredentials, future}; +use litellm_http::Client; +use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start, fetch}; +use rstest::rstest; + +const POLICY: RetryPolicy = RetryPolicy { + max_attempts: 3, + timeout: Duration::from_secs(5), + min_delay: Duration::from_millis(10), + max_delay: Duration::from_millis(50), +}; + +const OK: &str = + "HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}"; +const UNAVAILABLE: &str = "HTTP/1.1 503 Service Unavailable\r\nRetry-After: 0\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"; +const NOT_FOUND: &str = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"; + +const SLOW_BODY: &[u8] = b"{ }"; + +#[derive(Clone, Copy)] +enum Answer { + Raw(&'static str), + Stall, + StallBody, + SlowBody, + Oversized, +} + +/// A scripted upstream: the n-th request gets the n-th answer, repeating the last, and every request head is kept +struct Upstream { + base: String, + heads: Arc>>, +} + +impl Upstream { + fn serve(answers: Vec) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let heads = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&heads); + std::thread::spawn(move || { + for (index, stream) in listener.incoming().enumerate() { + let answer = answers[index.min(answers.len() - 1)]; + let recorded = Arc::clone(&recorded); + std::thread::spawn(move || { + let mut stream = stream.unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut head = String::new(); + loop { + let mut line = String::new(); + if reader.read_line(&mut line).unwrap() <= 2 { + break; + } + head.push_str(&line); + } + recorded.lock().unwrap().push(head); + match answer { + Answer::Raw(raw) => stream.write_all(raw.as_bytes()).unwrap(), + Answer::Stall => std::thread::sleep(Duration::from_secs(5)), + Answer::Oversized => { + let chunk = vec![b' '; 1024 * 1024]; + let written = stream + .write_all(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n") + .and_then(|()| { + (0..=100).try_for_each(|_| stream.write_all(&chunk)) + }); + drop(written); + } + Answer::StallBody => { + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{") + .unwrap(); + std::thread::sleep(Duration::from_secs(5)); + } + Answer::SlowBody => { + stream.set_nodelay(true).unwrap(); + stream + .write_all( + b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 8\r\nConnection: close\r\n\r\n", + ) + .unwrap(); + for byte in SLOW_BODY { + std::thread::sleep(Duration::from_millis(60)); + stream.write_all(&[*byte]).unwrap(); + } + } + } + }); + } + }); + Self { base, heads } + } + + fn url(&self) -> String { + format!("{}/cost.json", self.base) + } + + fn heads(&self) -> Vec { + self.heads.lock().unwrap().clone() + } +} + +fn header<'a>(head: &'a str, name: &str) -> Option<&'a str> { + head.lines().find_map(|line| { + let (key, value) = line.split_once(':')?; + key.eq_ignore_ascii_case(name).then(|| value.trim()) + }) +} + +#[derive(Clone, Debug)] +struct NoCredentials; + +impl ProvideCredentials for NoCredentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::ready(Err(CredentialsError::not_loaded("no credentials"))) + } +} + +fn keys() -> Credentials { + Credentials::new( + "AKIDMODELCOST", + "model-cost-secret", + None, + None, + "model-cost-test", + ) +} + +async fn fetched(url: &str, policy: RetryPolicy, start: Start) -> Result { + fetch( + &Client::for_test(), + &S3Access::Environment { region: None }, + url, + policy, + start, + ) + .await +} + +#[rstest] +#[case::every_etag_joined( + "HTTP/1.1 200 OK\r\nETag: \"a\"\r\nETag: \"b\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}", + Some("\"a\", \"b\"") +)] +#[case::no_etag( + "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}", + None +)] +#[tokio::test] +async fn a_200_returns_the_body_and_its_etags( + #[case] response: &'static str, + #[case] etag: Option<&str>, +) { + let upstream = Upstream::serve(vec![Answer::Raw(response)]); + assert_eq!( + fetched(&upstream.url(), POLICY, Start::Fresh).await, + Ok(Fetched { + body: b"{}".to_vec(), + etag: etag.map(str::to_owned), + }) + ); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test(flavor = "multi_thread")] +async fn a_fetch_runs_as_a_task_on_a_multi_thread_runtime() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let url = upstream.url(); + let fetched = tokio::spawn(async move { + fetch( + &Client::for_test(), + &S3Access::Environment { region: None }, + &url, + POLICY, + Start::Fresh, + ) + .await + }) + .await + .unwrap(); + assert!(fetched.is_ok()); +} + +#[rstest] +#[tokio::test] +async fn a_redirect_is_not_followed_and_its_body_is_the_document() { + let upstream = Upstream::serve(vec![ + Answer::Raw( + "HTTP/1.1 301 Moved Permanently\r\nLocation: /elsewhere.json\r\nContent-Length: 7\r\nConnection: close\r\n\r\n{\"a\":1}", + ), + Answer::Raw(OK), + ]); + assert_eq!( + fetched(&upstream.url(), POLICY, Start::Fresh).await, + Ok(Fetched { + body: b"{\"a\":1}".to_vec(), + etag: None, + }) + ); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_retry_after_shorter_than_the_backoff_is_honored() { + let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE), Answer::Raw(OK)]); + let policy = RetryPolicy { + min_delay: Duration::from_secs(10), + max_delay: Duration::from_secs(10), + ..POLICY + }; + let fetched = tokio::time::timeout( + Duration::from_secs(2), + fetched(&upstream.url(), policy, Start::Fresh), + ) + .await + .expect("Retry-After: 0 replaces the 10s backoff"); + assert!(fetched.is_ok(), "{fetched:?}"); + assert_eq!(upstream.heads().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn retryable_failures_stop_at_the_attempt_limit() { + let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]); + let url = upstream.url(); + assert_eq!( + fetched(&url, POLICY, Start::Fresh).await, + Err(Failed { + reason: format!("HTTP 503 from {url} (after 3 attempts)"), + retryable: true, + retry_after: Some(Duration::ZERO), + }) + ); + assert_eq!(upstream.heads().len(), 3); +} + +#[rstest] +#[tokio::test] +async fn a_404_fails_on_the_first_attempt() { + let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]); + let url = upstream.url(); + assert_eq!( + fetched(&url, POLICY, Start::Fresh).await, + Err(Failed { + reason: format!("HTTP 404 from {url}"), + retryable: false, + retry_after: None, + }) + ); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_stalled_response_times_out_and_is_retried() { + let upstream = Upstream::serve(vec![Answer::Stall]); + let url = upstream.url(); + let policy = RetryPolicy { + max_attempts: 2, + timeout: Duration::from_millis(100), + ..POLICY + }; + let failed = fetched(&url, policy, Start::Fresh).await.unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")) + && failed.reason.ends_with(" (after 2 attempts)"), + "{}", + failed.reason + ); + assert!(failed.retryable); + assert_eq!(upstream.heads().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn a_resumed_fetch_makes_only_the_attempts_left() { + let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]); + let url = upstream.url(); + let first = Failed { + reason: "first".into(), + retryable: true, + retry_after: Some(Duration::ZERO), + }; + let failed = fetched(&url, POLICY, Start::After(first)) + .await + .unwrap_err(); + assert_eq!( + failed.reason, + format!("HTTP 503 from {url} (after 3 attempts)") + ); + assert_eq!(upstream.heads().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn a_resumed_fetch_with_no_attempts_left_returns_the_failure_it_resumed() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let first = Failed { + reason: "first".into(), + retryable: true, + retry_after: None, + }; + let policy = RetryPolicy { + max_attempts: 1, + ..POLICY + }; + assert_eq!( + fetched(&upstream.url(), policy, Start::After(first.clone())).await, + Err(first) + ); + assert!(upstream.heads().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn a_failure_reason_drops_url_credentials() { + let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]); + let address = upstream.base.trim_start_matches("http://"); + let url = format!("http://user:hunter2@{address}/cost.json?X-Amz-Signature=sig-secret#frag"); + let failed = fetched(&url, POLICY, Start::Fresh).await.unwrap_err(); + assert_eq!(failed.reason, format!("HTTP 404 from {}", upstream.url())); +} + +#[rstest] +#[case::unsupported_scheme("ftp://example.com/cost.json")] +#[case::not_a_url("not a url")] +#[case::s3_without_a_key("s3://prices")] +#[tokio::test] +async fn an_unfetchable_url_fails_without_a_request(#[case] url: &str) { + let failed = fetched(url, POLICY, Start::Fresh).await.unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")), + "{}", + failed.reason + ); + assert!(!failed.retryable); +} + +fn s3(upstream: &Upstream, credentials: impl ProvideCredentials + 'static) -> S3Access { + S3Access::Config(Box::new( + litellm_s3::config( + "us-west-2".into(), + Some(&upstream.base), + Client::for_test(), + credentials, + ) + .build(), + )) +} + +async fn fetched_from_s3( + url: &str, + upstream: &Upstream, + credentials: impl ProvideCredentials + 'static, +) -> Result { + fetch( + &Client::for_test(), + &s3(upstream, credentials), + url, + POLICY, + Start::Fresh, + ) + .await +} + +#[rstest] +#[tokio::test] +async fn an_s3_fetch_sends_a_signed_get_object_to_the_endpoint() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + assert_eq!( + fetched_from_s3("s3://prices/maps/cost.json", &upstream, keys()).await, + Ok(Fetched { + body: b"{}".to_vec(), + etag: Some("\"v1\"".into()), + }) + ); + let heads = upstream.heads(); + let [head] = heads.as_slice() else { + panic!("expected one request, got {heads:?}"); + }; + assert!(head.starts_with("GET /prices/maps/cost.json"), "{head}"); + let authorization = header(head, "authorization").unwrap_or_default(); + assert!( + authorization.starts_with("AWS4-HMAC-SHA256 Credential=AKIDMODELCOST/") + && authorization.contains("/us-west-2/s3/aws4_request"), + "{authorization}" + ); +} + +#[rstest] +#[tokio::test] +async fn an_s3_body_that_stalls_after_the_headers_times_out() { + let upstream = Upstream::serve(vec![Answer::StallBody]); + let url = "s3://prices/cost.json"; + let policy = RetryPolicy { + timeout: Duration::from_millis(200), + ..POLICY + }; + let failed = tokio::time::timeout( + Duration::from_secs(2), + fetch( + &Client::for_test(), + &s3(&upstream, keys()), + url, + policy, + Start::Fresh, + ), + ) + .await + .expect("the attempt timeout bounds the body read") + .unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")), + "{}", + failed.reason + ); + assert!(!failed.retryable); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[case::http(|upstream: &Upstream| (S3Access::Environment { region: None }, upstream.url()))] +#[case::s3(|upstream: &Upstream| (s3(upstream, keys()), "s3://prices/cost.json".to_owned()))] +#[tokio::test] +async fn a_body_that_keeps_arriving_finishes_after_the_timeout( + #[case] source: fn(&Upstream) -> (S3Access, String), +) { + let upstream = Upstream::serve(vec![Answer::SlowBody]); + let (access, url) = source(&upstream); + let policy = RetryPolicy { + timeout: Duration::from_millis(300), + ..POLICY + }; + assert_eq!( + fetch(&Client::for_test(), &access, &url, policy, Start::Fresh).await, + Ok(Fetched { + body: SLOW_BODY.to_vec(), + etag: Some("\"v1\"".into()), + }) + ); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[case::missing_key(NOT_FOUND)] +#[case::denied("HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")] +#[tokio::test] +async fn an_s3_error_fails_without_a_retry(#[case] response: &'static str) { + let upstream = Upstream::serve(vec![Answer::Raw(response), Answer::Raw(OK)]); + let url = "s3://prices/cost.json"; + let failed = fetched_from_s3(url, &upstream, keys()).await.unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")), + "{}", + failed.reason + ); + assert!(!failed.retryable); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn an_s3_fetch_without_credentials_fails_without_a_request() { + let upstream = Upstream::serve(vec![Answer::Raw(OK)]); + let url = "s3://prices/cost.json"; + let failed = fetched_from_s3(url, &upstream, NoCredentials) + .await + .unwrap_err(); + assert!( + failed + .reason + .starts_with(&format!("error fetching {url}: ")), + "{}", + failed.reason + ); + assert!(!failed.retryable); + assert!(upstream.heads().is_empty()); +} + +const DECLARED_OVERSIZED: &str = + "HTTP/1.1 200 OK\r\nContent-Length: 104857601\r\nConnection: close\r\n\r\n"; + +#[rstest] +#[case::http_declared(Answer::Raw(DECLARED_OVERSIZED), false)] +#[case::http_streamed(Answer::Oversized, false)] +#[case::s3_declared(Answer::Raw(DECLARED_OVERSIZED), true)] +#[case::s3_streamed(Answer::Oversized, true)] +#[tokio::test] +async fn a_body_over_100_mib_fails_without_a_retry(#[case] answer: Answer, #[case] s3: bool) { + let upstream = Upstream::serve(vec![answer, Answer::Raw(OK)]); + let failed = match s3 { + true => fetched_from_s3("s3://prices/cost.json", &upstream, keys()).await, + false => fetched(&upstream.url(), POLICY, Start::Fresh).await, + } + .unwrap_err(); + assert!( + failed + .reason + .contains("response body is over the 104857600 byte limit"), + "{}", + failed.reason + ); + assert!(!failed.retryable); + assert_eq!(upstream.heads().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_body_of_exactly_100_mib_is_fetched() { + let upstream = Upstream::serve(vec![Answer::Raw(Box::leak( + format!( + "HTTP/1.1 200 OK\r\nContent-Length: 104857600\r\nConnection: close\r\n\r\n{}", + " ".repeat(100 * 1024 * 1024) + ) + .into_boxed_str(), + ))]); + let fetched = fetched(&upstream.url(), POLICY, Start::Fresh) + .await + .unwrap(); + assert_eq!(fetched.body.len(), 100 * 1024 * 1024); +} diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index abc5d198729..a9c473f2ca3 100644 --- a/litellm-rust/crates/model-catalog/AGENTS.md +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -46,7 +46,14 @@ Every price in a row lowers to one `Price { meter, channel, when }` and its `Uni `store/` is the catalog's in-memory store and the only place that indexes model names. Its rules are in [src/store/AGENTS.md](src/store/AGENTS.md) -`layers.rs`, `resolve/`, `model_info.rs` and `register.rs` hold types only. The functions over them land with their first caller +`cost_map.rs` holds a host's inputs, the installed document and the `register_model` overlay in `register.rs`. It serves the document in the key order and alias sharing Python's `litellm.model_cost` has, without the overlay, since where a registration lands is a lookup in the published table. It never holds the live table, which belongs to the host + +- `CostMap` is immutable behind an `Arc`. A new document is a new `CostMap` +- `Registrations` is generic over the host's value type, because Python registrations can hold `datetime` and other values JSON cannot. The host keeps its own objects (`Py` in python-bridge) and does the merging + - Entry order is latest-registration order. `landed_on` records which registered keys landed on each row key +- `write_pass` stops at the first `place` that returns `None` and reports `Pass::Swapped`. Retrying is the host's job + +`resolve/` and `model_info.rs` hold types only. The functions over them land with their first caller Catalog errors use one opaque `Error` with typed causes and JSON Pointer accessors (`path`, `key`, `field`). `Error::rejects` alone decides whether an error fails the document, rejects its row or neither, and `ErrorKind` is the stable name hosts report it by diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index 63b9ad14768..8cd8ee5760b 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -11,6 +11,7 @@ schema = ["dep:schemars", "litellm-cost/schema"] [dependencies] litellm-cost = { workspace = true, features = ["serde"] } litellm-error.workspace = true +litellm-tracing.workspace = true indexmap = { version = "2.14.0", features = ["serde"] } jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] } @@ -19,6 +20,7 @@ schemars = { workspace = true, optional = true } serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } serde_with.workspace = true +sha1.workspace = true strum.workspace = true thiserror.workspace = true time.workspace = true diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index 98af785b17e..6c11a39b0eb 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -12,7 +12,8 @@ use crate::store::{Rows, Store}; use crate::wire::{EntryWire, Fields}; const SAMPLE_SPEC: &str = "sample_spec"; -const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; +/// The top-level key of the rules block, which is never a model row +pub const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; #[derive(Clone, Copy, Debug, PartialEq)] pub struct IntegrityLimits { @@ -111,6 +112,11 @@ impl Rows for CatalogRows { } impl IntegrityLimits { + pub const UNCHECKED: Self = Self { + min_model_count: 0, + min_reference_ratio: 0.0, + }; + pub fn new(min_model_count: usize, min_reference_ratio: f64) -> Option { (0.0..=1.0).contains(&min_reference_ratio).then_some(Self { min_model_count, @@ -142,6 +148,7 @@ impl Catalog { let mut sections = Vec::new(); let mut entries = IndexMap::with_capacity(document.len()); let mut alias_lists = Vec::new(); + let mut pointer_rows = Vec::new(); for (index, (name, value)) in document.into_iter().enumerate() { let fields = match value { Value::Object(fields) if !is_section(&name) => fields, @@ -157,16 +164,34 @@ impl Catalog { continue; } }; - let (model, aliases, row_problems) = read_row(&fields); + let row = read_row(&fields); + let readable = row.problems.is_empty(); problems.extend( - row_problems + row.problems .into_iter() .map(|problem| problem.prepend(&name)), ); - if let Some(aliases) = aliases { + if let Some(target) = row.alias_of { + if readable { + pointer_rows.push((name.clone(), target)); + } + sections.push(Section { + index, + key: name, + value: Value::Object(fields), + }); + continue; + } + if let Some(aliases) = row.aliases { alias_lists.push((name.clone(), aliases)); } - entries.insert(name, ModelEntry { fields, model }); + entries.insert( + name, + ModelEntry { + fields, + model: row.model, + }, + ); } if entries.is_empty() { return Err(Error::new("", Repr::Empty)); @@ -190,6 +215,21 @@ impl Catalog { } } } + for (alias, target) in pointer_rows { + if !entries.contains_key(&target) { + problems.push(Error::from_key(&alias, Repr::AliasTargetMissing { target })); + } else if let Some(taken_by) = aliases.get(&alias) { + problems.push(Error::from_key( + &alias, + Repr::AliasTaken { + alias: alias.clone(), + taken_by: taken_by.clone(), + }, + )); + } else { + aliases.insert(alias, target); + } + } let fallback_generalizations = sections .iter() @@ -242,8 +282,8 @@ impl Catalog { } pub fn row_errors(model: &str, fields: &Fields) -> Vec { - let (_, _, errors) = read_row(fields); - errors + read_row(fields) + .problems .into_iter() .map(|error| error.prepend(model)) .collect() @@ -255,7 +295,13 @@ impl Catalog { (0..self.rows().entries.len() + self.sections.len()).map_while(move |index| match sections .next_if(|section| section.index == index) { - Some(section) => Some((section.key.as_str(), Written::Section(§ion.value))), + Some(section) => Some(( + section.key.as_str(), + match §ion.value { + Value::Object(fields) if !is_section(§ion.key) => Written::Entry(fields), + value => Written::Section(value), + }, + )), None => entries .next() .map(|(name, entry)| (name.as_str(), Written::Entry(&entry.fields))), @@ -378,6 +424,20 @@ impl Catalog { }) } + /// The row under an exact model or alias key, an alias reading its owner's row + pub fn get(&self, key: &str) -> Option<&Fields> { + self.rows().fields(key) + } + + /// A top-level value that is not a model row, such as `sample_spec` + pub fn section(&self, key: &str) -> Option<&Value> { + self.sections + .iter() + .filter(|section| is_section(§ion.key)) + .find(|section| section.key == key) + .map(|section| §ion.value) + } + pub fn model_count(&self) -> usize { self.rows().entries.len() } @@ -406,16 +466,41 @@ impl Catalog { } } -fn read_row(fields: &Fields) -> (Option, Option>, Vec) { +struct RowRead { + model: Option, + aliases: Option>, + alias_of: Option, + problems: Vec, +} + +fn read_row(fields: &Fields) -> RowRead { let wire = EntryWire::read(fields); + if let Some(target) = wire.attributes.alias_of.clone() { + let other_fields = fields + .keys() + .filter(|key| *key != "alias_of") + .map(|key| Error::from_key(key, Repr::AliasRowField)); + return RowRead { + model: None, + aliases: None, + alias_of: Some(target), + problems: wire.problems.into_iter().chain(other_fields).collect(), + }; + } let aliases = wire.attributes.aliases.clone(); let timezone_error = unknown_timezone(&wire); - match Model::try_from(wire) { - Ok(model) => (Some(model), aliases, timezone_error.into_iter().collect()), + let (model, problems) = match Model::try_from(wire) { + Ok(model) => (Some(model), timezone_error.into_iter().collect()), Err(mut errors) => { errors.extend(timezone_error); - (None, aliases, errors) + (None, errors) } + }; + RowRead { + model, + aliases, + alias_of: None, + problems, } } diff --git a/litellm-rust/crates/model-catalog/src/cost_map.rs b/litellm-rust/crates/model-catalog/src/cost_map.rs new file mode 100644 index 00000000000..bec417acae7 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/cost_map.rs @@ -0,0 +1,130 @@ +use std::sync::{Arc, LazyLock}; + +use indexmap::IndexMap; +use serde_json::Value; +use sha1::{Digest, Sha1}; + +use crate::bundled::BUNDLED_REGISTRY; +use crate::catalog::{Catalog, FALLBACK_GENERALIZATIONS, IntegrityLimits}; +use crate::error::Error; + +/// Which document is installed: the `git_blob_id` of its bytes and the ETag it was served with +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Provenance { + pub revision: String, + pub etag: Option, +} + +/// An installed cost map document. A host builds its live table from it +#[derive(Clone, Debug)] +pub struct CostMap { + base: Arc, + provenance: Provenance, +} + +#[derive(Debug)] +struct Base { + catalog: Catalog, + /// Each owner's names, the owner first, for owners with at least one alias + groups: IndexMap>, +} + +impl Base { + fn new(catalog: Catalog) -> Self { + let groups = catalog.aliases().iter().fold( + IndexMap::>::new(), + |mut groups, (alias, owner)| { + groups + .entry(owner.clone()) + .or_insert_with(|| vec![owner.clone()]) + .push(alias.clone()); + groups + }, + ); + Self { catalog, groups } + } +} + +static BUNDLED: LazyLock = LazyLock::new(|| CostMap { + base: Arc::new(Base::new( + Catalog::parse(BUNDLED_REGISTRY) + .expect("the bundled registry parses") + .value, + )), + provenance: Provenance { + revision: git_blob_id(BUNDLED_REGISTRY), + etag: None, + }, +}); + +impl CostMap { + pub fn bundled() -> Self { + BUNDLED.clone() + } + + /// Reads a fetched document. One that does not parse, or that has fewer models than `limits` allow against the + /// bundled one, is refused. Every alias problem it has is logged + pub fn read(body: &[u8], etag: Option, limits: IntegrityLimits) -> Result { + let loaded = Catalog::parse(body)?; + loaded + .value + .check_size(BUNDLED.base.catalog.model_count(), limits)?; + for problem in loaded.errors.iter().filter(|error| error.is_alias()) { + litellm_tracing::warn!( + "LiteLLM: model cost map alias problem: {}", + litellm_error::Report::new(problem) + ); + } + Ok(Self { + base: Arc::new(Base::new(loaded.value)), + provenance: Provenance { + revision: git_blob_id(body), + etag, + }, + }) + } + + pub fn catalog(&self) -> &Catalog { + &self.base.catalog + } + + pub fn provenance(&self) -> &Provenance { + &self.provenance + } + + /// Every name the row under `key` is served under, the owner first. Empty for a row without aliases + pub fn names(&self, key: &str) -> &[String] { + let owner = self + .catalog() + .aliases() + .get(key) + .map_or(key, String::as_str); + self.base + .groups + .get(owner) + .map(Vec::as_slice) + .unwrap_or_default() + } + + /// Each owner with its names, the owner first, in the order the document declared the aliases + pub fn alias_groups(&self) -> impl Iterator { + self.base + .groups + .iter() + .map(|(owner, names)| (owner.as_str(), names.as_slice())) + } + + /// The `fallback_generalizations` block as the document wrote it + pub fn fallback_generalizations(&self) -> Option<&Value> { + self.catalog().section(FALLBACK_GENERALIZATIONS) + } +} + +/// `git hash-object` of the bytes, which main reports as `source_revision` +pub fn git_blob_id(body: &[u8]) -> String { + let digest = Sha1::new() + .chain_update(format!("blob {}\0", body.len())) + .chain_update(body) + .finalize(); + format!("{digest:x}") +} diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index 89e27c05f46..a4b54b5a551 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -67,6 +67,12 @@ pub(crate) enum Repr { #[error("alias {alias:?} already belongs to {taken_by}")] #[strum_discriminants(strum(serialize = "alias_taken"))] AliasTaken { alias: String, taken_by: String }, + #[error("alias_of {target:?} is not a model name")] + #[strum_discriminants(strum(serialize = "alias_target_missing"))] + AliasTargetMissing { target: String }, + #[error("an alias_of row carries no other field")] + #[strum_discriminants(strum(serialize = "alias_row_field"))] + AliasRowField, #[error("{first} and {second} price the same meter under the same conditions")] #[strum_discriminants(strum(serialize = "duplicate_price"))] DuplicatePrice { first: String, second: String }, @@ -141,6 +147,8 @@ impl Error { | Repr::UnknownTimezone { .. } | Repr::AliasShadowsModel { .. } | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField | Repr::DuplicatePrice { .. } | Repr::SuffixInsideBlock { .. } | Repr::TierRange @@ -159,6 +167,74 @@ impl Error { } } +impl Error { + /// The document parsed but holds fewer models than the integrity limits trust + pub fn is_integrity(&self) -> bool { + match self.repr() { + Repr::Empty | Repr::BelowMinimum { .. } | Repr::Shrunk { .. } => true, + Repr::Json(_) + | Repr::NotAnObject + | Repr::Invalid(_) + | Repr::UnknownBlockKey + | Repr::UnknownTimezone { .. } + | Repr::AliasShadowsModel { .. } + | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField + | Repr::DuplicatePrice { .. } + | Repr::SuffixInsideBlock { .. } + | Repr::TierRange + | Repr::TierReversed + | Repr::TierGap { .. } + | Repr::MixedTierBasis + | Repr::NotPricedByResults + | Repr::BandsFromTwoSources + | Repr::NotAnOffPeakRate + | Repr::NoOffPeakPrices + | Repr::NoOffPeakHours + | Repr::EmptyList + | Repr::InvalidDate { .. } + | Repr::MissingProvider + | Repr::UnreadField + | Repr::NoFallbackRules => false, + } + } + + /// A problem with a row's aliases, which the served map skips + pub fn is_alias(&self) -> bool { + match self.repr() { + Repr::AliasShadowsModel { .. } + | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField => true, + Repr::Invalid(_) => matches!(self.field().as_deref(), Some("aliases" | "alias_of")), + Repr::Json(_) + | Repr::Empty + | Repr::BelowMinimum { .. } + | Repr::Shrunk { .. } + | Repr::NotAnObject + | Repr::UnknownBlockKey + | Repr::UnknownTimezone { .. } + | Repr::DuplicatePrice { .. } + | Repr::SuffixInsideBlock { .. } + | Repr::TierRange + | Repr::TierReversed + | Repr::TierGap { .. } + | Repr::MixedTierBasis + | Repr::NotPricedByResults + | Repr::BandsFromTwoSources + | Repr::NotAnOffPeakRate + | Repr::NoOffPeakPrices + | Repr::NoOffPeakHours + | Repr::EmptyList + | Repr::InvalidDate { .. } + | Repr::MissingProvider + | Repr::UnreadField + | Repr::NoFallbackRules => false, + } + } +} + #[derive(Debug)] pub struct Loaded { pub value: T, @@ -327,6 +403,8 @@ mod tests { #[case::unknown_timezone(Repr::UnknownTimezone { zone: "invalid".into() }, Rejects::Row)] #[case::alias_shadows_model(Repr::AliasShadowsModel { alias: "model".into() }, Rejects::Row)] #[case::alias_taken(Repr::AliasTaken { alias: "alias".into(), taken_by: "owner".into() }, Rejects::Row)] + #[case::alias_target_missing(Repr::AliasTargetMissing { target: "owner".into() }, Rejects::Row)] + #[case::alias_row_field(Repr::AliasRowField, Rejects::Row)] #[case::duplicate_price(Repr::DuplicatePrice { first: "a".into(), second: "b".into() }, Rejects::Row)] #[case::suffix_inside_block(Repr::SuffixInsideBlock { block: "tiered_pricing" }, Rejects::Row)] #[case::tier_range(Repr::TierRange, Rejects::Row)] diff --git a/litellm-rust/crates/model-catalog/src/layers.rs b/litellm-rust/crates/model-catalog/src/layers.rs deleted file mode 100644 index 881f0f6af9b..00000000000 --- a/litellm-rust/crates/model-catalog/src/layers.rs +++ /dev/null @@ -1,32 +0,0 @@ -use std::sync::Arc; - -use indexmap::IndexMap; -use serde_json::Value; - -use crate::catalog::Catalog; -use crate::register::Registrations; -use crate::wire::Fields; - -/// One field of a fetched row over the layer below, RFC 7386. An absent field inherits, JSON `null` is `Remove` -#[derive(Clone, Debug, PartialEq)] -pub enum Patch { - Set(Value), - Remove, -} - -/// The four layers served rows are composed from, lowest first -#[derive(Clone, Debug)] -#[expect( - dead_code, - reason = "02 and 07 add the functions that read these layers" -)] -pub struct Layers { - /// 0, the wheel's embedded map - bundled: Arc, - /// 1, the last installed fetch, rejected rows left out - fetched: Option>, - /// 2, `register(persist=False)`, cleared by every install and refilled by the router replay - transient: IndexMap, - /// 3, `register(persist=True)`, applied over every install - persisted: Registrations, -} diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index fe419dde054..9b869a078db 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -17,6 +17,9 @@ //! failing the document. //! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known //! families. They are data in the file, not entries. +//! - `CostMap`: an installed document with its provenance, which a host builds its live table from. +//! `Registrations` is the ordered `register_model` overlay over the host's own values, and +//! `write_pass` writes registrations into a host's table through `RegistrationTable`. //! - Validation: `Catalog::check_size` checks a catalog against `IntegrityLimits` (a remote copy //! that lost most of its models is refused), and `Catalog::row_errors` validates one row. //! @@ -37,9 +40,9 @@ macro_rules_attribute::attribute_alias! { mod bundled; mod catalog; +mod cost_map; mod error; mod fallback; -mod layers; mod lower; mod model; mod model_info; @@ -51,9 +54,9 @@ pub mod wire; pub use bundled::*; pub use catalog::*; +pub use cost_map::*; pub use error::*; pub use fallback::*; -pub use layers::*; pub use model::*; pub use model_info::*; pub use register::*; diff --git a/litellm-rust/crates/model-catalog/src/register.rs b/litellm-rust/crates/model-catalog/src/register.rs index ed34082ae57..7aaf2ed3978 100644 --- a/litellm-rust/crates/model-catalog/src/register.rs +++ b/litellm-rust/crates/model-catalog/src/register.rs @@ -1,26 +1,252 @@ -use std::collections::{BTreeSet, HashMap}; +use std::collections::BTreeMap; use indexmap::IndexMap; -use crate::wire::Fields; - -/// Where one `register` call writes and which row it inherits from -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RegistrationTarget { - /// Always the caller's key, never a case variant of it - pub write: String, - /// The first candidate that is an exact key whose row passes `provider_matches` - pub inherit_from: Option, +/// `register_model(persist_across_reloads=True)` calls, replayed above every installed document +#[derive(Clone, Debug, PartialEq)] +pub struct Registrations { + /// Each registered key's merged value, ordered by its latest registration + entries: IndexMap, + /// The registered keys whose write landed on each row key, in landing order + landed_on: IndexMap>, } -/// Python's inheritance candidates in main's order, `get_model_info`'s then `_resolve_builtin_model_cost_entry`'s -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RegistrationCandidates(pub Vec); - -/// Layer 3, and the fields `register` last wrote, which litellm owns again -#[derive(Clone, Debug, Default)] -#[expect(dead_code, reason = "07 adds the functions that read registrations")] -pub struct Registrations { - persisted: IndexMap, - registered_fields: HashMap>, +impl Default for Registrations { + fn default() -> Self { + Self { + entries: IndexMap::new(), + landed_on: IndexMap::new(), + } + } +} + +/// One registration written to a host's table: the key it was registered under and the row key it landed on +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Landed { + pub registered: String, + pub written: String, +} + +impl Registrations { + /// Keeps `value` as the key's registration, already merged with `get(key)` by the host, and makes it the latest + pub fn register(&mut self, key: &str, value: V) { + self.entries.shift_remove(key); + self.entries.insert(key.to_owned(), value); + } + + pub fn get(&self, key: &str) -> Option<&V> { + self.entries.get(key) + } + + pub fn record(&mut self, landed: &[Landed]) { + for Landed { + registered, + written, + } in landed + { + let names = self.landed_on.entry(written.clone()).or_default(); + if !names.contains(registered) { + names.push(registered.clone()); + } + } + } + + /// Every registration, oldest first + pub fn entries(&self) -> impl Iterator { + self.entries + .iter() + .map(|(key, value)| (key.as_str(), value)) + } + + /// The registrations that apply to the row under `key`, oldest first: `key`'s own and every one that landed on it + pub fn over<'a>(&'a self, key: &str) -> Vec<&'a V> { + let landed = self.landed_on.get(key).into_iter().flatten(); + let by_order: BTreeMap = landed + .map(String::as_str) + .chain([key]) + .filter_map(|name| self.entries.get_full(name)) + .map(|(order, _, value)| (order, value)) + .collect(); + by_order.into_values().collect() + } + + pub fn map(&self, mut convert: impl FnMut(&V) -> W) -> Registrations { + Registrations { + entries: self + .entries + .iter() + .map(|(key, value)| (key.clone(), convert(value))) + .collect(), + landed_on: self.landed_on.clone(), + } + } +} + +/// A host's live table that registrations are written into +pub trait RegistrationTable { + type Registration; + type Error; + + /// Writes `registration` where it lands, `None` when this table is no longer the one the host publishes + fn place( + &mut self, + key: &str, + registration: &Self::Registration, + ) -> Result, Self::Error>; +} + +/// Where a registration was written and whether a built-in row matched it +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Placed { + pub key: String, + pub matched_builtin: bool, +} + +#[derive(Debug, PartialEq)] +pub enum Pass<'a, R> { + /// Every entry was written. `unmatched` are those no built-in row matched + Finished { + landed: Vec, + unmatched: Vec<(&'a str, &'a R)>, + }, + /// The table stopped being the published one before every entry was written + Swapped { landed: Vec }, +} + +impl Pass<'_, R> { + pub fn landed(&self) -> &[Landed] { + match self { + Pass::Finished { landed, .. } | Pass::Swapped { landed } => landed, + } + } +} + +/// Writes `entries` into `table` in order +pub fn write_pass<'a, T: RegistrationTable>( + table: &mut T, + entries: &'a [(String, T::Registration)], +) -> Result, T::Error> { + let mut landed = Vec::with_capacity(entries.len()); + let mut unmatched = Vec::new(); + for (key, registration) in entries { + let Some(written) = table.place(key, registration)? else { + return Ok(Pass::Swapped { landed }); + }; + if !written.matched_builtin { + unmatched.push((key.as_str(), registration)); + } + landed.push(Landed { + registered: key.clone(), + written: written.key, + }); + } + Ok(Pass::Finished { landed, unmatched }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn landed(registered: &str, written: &str) -> Landed { + Landed { + registered: registered.to_owned(), + written: written.to_owned(), + } + } + + #[rstest] + fn reregistering_a_key_makes_it_the_latest() { + let mut registrations = Registrations::default(); + registrations.register("a", 1); + registrations.register("b", 2); + registrations.register("a", 3); + assert_eq!( + registrations.entries().collect::>(), + [("b", &2), ("a", &3)] + ); + } + + #[rstest] + fn over_applies_names_landed_on_the_row_in_registration_order() { + let mut registrations = Registrations::default(); + registrations.register("openai/o3-pro", 1); + registrations.register("o3-pro", 2); + registrations.register("other", 3); + registrations.record(&[ + landed("openai/o3-pro", "o3-pro"), + landed("o3-pro", "o3-pro"), + ]); + registrations.register("openai/o3-pro", 4); + assert_eq!(registrations.over("o3-pro"), [&2, &4]); + } + + #[rstest] + fn over_applies_a_keys_own_registration_that_landed_on_another_row() { + let mut registrations = Registrations::default(); + registrations.register("openai/o3-pro", 1); + registrations.record(&[landed("openai/o3-pro", "o3-pro")]); + assert_eq!( + ( + registrations.over("openai/o3-pro"), + registrations.over("o3-pro") + ), + (vec![&1], vec![&1]) + ); + } + + /// A table published until `swaps_after` writes, whose rows match a built-in only when the key starts with `b` + struct Table { + writes: Vec<(String, i32)>, + swaps_after: usize, + } + + impl RegistrationTable for Table { + type Registration = i32; + type Error = (); + + fn place(&mut self, key: &str, registration: &i32) -> Result, ()> { + if self.writes.len() >= self.swaps_after { + return Ok(None); + } + let written = key.to_uppercase(); + self.writes.push((written.clone(), *registration)); + Ok(Some(Placed { + matched_builtin: written.starts_with('B'), + key: written, + })) + } + } + + #[rstest] + #[case::published_throughout( + usize::MAX, + Pass::Finished { landed: vec![landed("a", "A"), landed("b", "B")], unmatched: vec![("a", &1)] }, + &[("A", 1), ("B", 2)], + )] + #[case::swapped_after_the_first_write( + 1, + Pass::Swapped { landed: vec![landed("a", "A")] }, + &[("A", 1)], + )] + fn a_pass_writes_entries_in_order_until_the_table_is_swapped( + #[case] swaps_after: usize, + #[case] expected: Pass<'static, i32>, + #[case] writes: &[(&str, i32)], + ) { + let entries = [("a".to_owned(), 1), ("b".to_owned(), 2)]; + let mut table = Table { + writes: Vec::new(), + swaps_after, + }; + assert_eq!(write_pass(&mut table, &entries), Ok(expected)); + assert_eq!( + table.writes, + writes + .iter() + .map(|(key, value)| ((*key).to_owned(), *value)) + .collect::>() + ); + } } diff --git a/litellm-rust/crates/model-catalog/src/schema.rs b/litellm-rust/crates/model-catalog/src/schema.rs index 6fba985a8bd..60c755fe94f 100644 --- a/litellm-rust/crates/model-catalog/src/schema.rs +++ b/litellm-rust/crates/model-catalog/src/schema.rs @@ -104,9 +104,18 @@ pub fn registry_json_schema() -> Schema { "title".into(), "LiteLLM model prices and context window registry".into(), ); + definitions.insert( + "aliasRow".into(), + json!({ + "type": "object", + "required": ["alias_of"], + "properties": {"alias_of": {"type": "string"}}, + "additionalProperties": false, + }), + ); root.insert( "additionalProperties".into(), - json!({"$ref": "#/$defs/modelEntry"}), + json!({"anyOf": [{"$ref": "#/$defs/modelEntry"}, {"$ref": "#/$defs/aliasRow"}]}), ); root.insert("$defs".into(), Value::Object(definitions)); Schema::from(root.clone()) @@ -153,6 +162,7 @@ fn decorate_model_entry(schema: &mut Value) { let rates = json!({FLAT_KEY_PATTERN: {"type": "number", "minimum": 0}}); let object = schema.as_object_mut().unwrap(); object.insert("required".into(), json!(["litellm_provider"])); + object.insert("not".into(), json!({"required": ["alias_of"]})); object.insert("additionalProperties".into(), Value::Bool(true)); object.insert( "patternProperties".into(), diff --git a/litellm-rust/crates/model-catalog/src/store/AGENTS.md b/litellm-rust/crates/model-catalog/src/store/AGENTS.md index f1dacd2e217..580798da895 100644 --- a/litellm-rust/crates/model-catalog/src/store/AGENTS.md +++ b/litellm-rust/crates/model-catalog/src/store/AGENTS.md @@ -24,7 +24,7 @@ | Source | Stamp | Blind spot | |---|---|---| | `Catalog`'s `CatalogRows` | `()`, every change goes through `Store::edit`, which drops the index | none | -| python-bridge's bound `litellm.model_cost` | the dict object and its `len` | a delete plus an add between two reads, or a `litellm_provider` rewrite on an existing key | +| python-bridge's bound `litellm.model_cost` | the dict object, its `len` and its last key | deletes plus adds between two reads that put the old last key back last, or a `litellm_provider` rewrite on an existing key | A source with a blind spot calls `Store::invalidate` when it knows keys changed. Until then a key added in the blind spot is not found, and a key moved to another provider stays listed under the old one. Removed keys are still never returned diff --git a/litellm-rust/crates/model-catalog/src/wire/attributes.rs b/litellm-rust/crates/model-catalog/src/wire/attributes.rs index 401634f8c62..52cd2e8525f 100644 --- a/litellm-rust/crates/model-catalog/src/wire/attributes.rs +++ b/litellm-rust/crates/model-catalog/src/wire/attributes.rs @@ -47,6 +47,8 @@ macro_rules! attributes { } attributes! { + /// Model this row is another name for; such a row carries no other key. + alias_of: String, /// Other names this entry is served under. aliases: Vec, audio_transcription_config: String, diff --git a/litellm-rust/crates/model-catalog/tests/catalog.rs b/litellm-rust/crates/model-catalog/tests/catalog.rs index 31f907689d3..679c67412c8 100644 --- a/litellm-rust/crates/model-catalog/tests/catalog.rs +++ b/litellm-rust/crates/model-catalog/tests/catalog.rs @@ -139,6 +139,18 @@ fn a_bad_entry_is_reported_by_model_and_key_and_kept_raw() { json!({"First": {"litellm_provider": "t", "aliases": ["shared"]}, "Second": {"litellm_provider": "t", "aliases": ["shared"]}}), r#"Second: alias "shared" already belongs to First"# )] +#[case::alias_of_target_missing( + json!({"Real": {"litellm_provider": "t"}, "ghost-latest": {"alias_of": "Ghost"}}), + r#"ghost-latest: alias_of "Ghost" is not a model name"# +)] +#[case::alias_of_row_with_other_fields( + json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real", "mode": "chat"}}), + "real-latest.mode: an alias_of row carries no other field" +)] +#[case::alias_of_claimed_by_a_list( + json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real"}, "Other": {"litellm_provider": "t", "aliases": ["real-latest"]}}), + r#"real-latest: alias "real-latest" already belongs to Other"# +)] #[case::aliases_not_a_list( json!({"a": {"litellm_provider": "t", "aliases": "bad"}}), "a.aliases: invalid value: invalid type: string \"bad\", expected a sequence" @@ -399,3 +411,70 @@ fn a_bad_value_keeps_its_json_error_as_the_source_and_names_its_key_and_field() assert_eq!(error.key().as_deref(), Some("open/router")); assert_eq!(error.field().as_deref(), Some("input_cost_per_token")); } + +#[rstest] +fn an_alias_of_row_is_served_as_its_target_and_written_as_itself() { + let catalog = catalog(json!({ + "Real": {"litellm_provider": "t", "input_cost_per_token": 1e-6}, + "real-latest": {"alias_of": "Real"}, + "After": {"litellm_provider": "t"} + })); + let found = catalog.lookup("REAL-LATEST").unwrap(); + assert_eq!( + (found.matched_key, found.canonical_key), + ("real-latest", "Real") + ); + assert_eq!( + found.entry.field("input_cost_per_token"), + Some(&json!(1e-6)) + ); + assert_eq!(catalog.model_count(), 2); + assert_eq!( + catalog.served().map(|(key, _)| key).collect::>(), + ["Real", "After", "real-latest"] + ); + assert_eq!( + written(&catalog), + r#"{"Real":{"litellm_provider":"t","input_cost_per_token":1e-6},"real-latest":{"alias_of":"Real"},"After":{"litellm_provider":"t"}}"# + ); +} + +#[rstest] +fn a_dangling_alias_of_row_is_rejected_and_dropped_from_the_written_document() { + let parsed = loaded(json!({ + "Real": {"litellm_provider": "t"}, + "ghost-latest": {"alias_of": "Ghost"} + })); + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(removed.len(), 1); + assert!(catalog.lookup("ghost-latest").is_none()); + assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#); +} + +#[rstest] +fn an_alias_of_row_with_other_fields_is_never_served() { + let parsed = loaded(json!({ + "Real": {"litellm_provider": "t"}, + "real-latest": {"alias_of": "Real", "mode": "chat"} + })); + assert!(parsed.value.lookup("real-latest").is_none()); + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(removed.len(), 1); + assert!(catalog.lookup("real-latest").is_none()); + assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#); +} + +#[rstest] +fn a_pointer_row_that_repeats_a_listed_alias_is_dropped_and_the_model_is_kept() { + let parsed = loaded(json!({ + "Real": {"litellm_provider": "t", "aliases": ["real-latest"]}, + "real-latest": {"alias_of": "Real"} + })); + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(removed.len(), 1); + assert_eq!(catalog.lookup("real-latest").unwrap().canonical_key, "Real"); + assert_eq!( + written(&catalog), + r#"{"Real":{"litellm_provider":"t","aliases":["real-latest"]}}"# + ); +} diff --git a/litellm-rust/crates/model-catalog/tests/cost_map.rs b/litellm-rust/crates/model-catalog/tests/cost_map.rs new file mode 100644 index 00000000000..04a04170df6 --- /dev/null +++ b/litellm-rust/crates/model-catalog/tests/cost_map.rs @@ -0,0 +1,53 @@ +use litellm_model_catalog::{ + BUNDLED_REGISTRY, CostMap, ErrorKind, IntegrityLimits, Provenance, git_blob_id, +}; +use rstest::rstest; +use serde_json::json; + +fn any_size() -> IntegrityLimits { + IntegrityLimits::new(0, 0.0).unwrap() +} + +#[rstest] +#[case::not_json(b"{".as_slice(), false, ErrorKind::Json)] +#[case::empty(b"{}".as_slice(), true, ErrorKind::Empty)] +fn a_refused_document_says_whether_it_failed_the_integrity_limits( + #[case] body: &[u8], + #[case] integrity: bool, + #[case] kind: ErrorKind, +) { + let error = CostMap::read(body, None, any_size()).unwrap_err(); + assert_eq!((error.is_integrity(), error.kind()), (integrity, kind)); +} + +#[rstest] +fn provenance_is_the_read_bytes_and_etag() { + let body = br#"{"m":{"litellm_provider":"p"}}"#; + let cost_map = CostMap::read(body, Some("\"v1\"".into()), any_size()).unwrap(); + assert_eq!( + cost_map.provenance(), + &Provenance { + revision: "eaa747f210a1527d5ea7138be0f768c47578ca68".to_owned(), + etag: Some("\"v1\"".into()), + } + ); + assert_eq!( + CostMap::bundled().provenance().revision, + git_blob_id(BUNDLED_REGISTRY) + ); +} + +#[rstest] +fn every_name_of_an_aliased_row_lists_its_owner_first( + #[values("owner", "short", "owner-latest")] key: &str, +) { + let body = json!({ + "owner": {"litellm_provider": "p", "aliases": ["short"]}, + "owner-latest": {"alias_of": "owner"}, + "plain": {"litellm_provider": "p"}, + }) + .to_string(); + let cost_map = CostMap::read(body.as_bytes(), None, any_size()).unwrap(); + assert_eq!(cost_map.names(key), ["owner", "short", "owner-latest"]); + assert!(cost_map.names("plain").is_empty()); +} diff --git a/litellm-rust/crates/model-catalog/tests/schema.rs b/litellm-rust/crates/model-catalog/tests/schema.rs index 464cf069cac..5fe5832cb58 100644 --- a/litellm-rust/crates/model-catalog/tests/schema.rs +++ b/litellm-rust/crates/model-catalog/tests/schema.rs @@ -35,6 +35,7 @@ fn generated_registry_schema_validates_checked_in_catalog() { #[case(json!({"example": {"litellm_provider": "test"}}))] #[case(json!({"example": {"litellm_provider": "test", "future_field": true}}))] #[case(json!({"sample_spec": {"litellm_provider": "placeholder"}}))] +#[case(json!({"example": {"litellm_provider": "test"}, "example-latest": {"alias_of": "example"}}))] fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) { assert!(registry_validator().is_valid(&document)); } @@ -54,6 +55,9 @@ fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) #[case::empty_windows(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}))] #[case::invalid_weekday(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}))] #[case::invalid_aliases(json!({"litellm_provider": "test", "aliases": "wrong"}))] +#[case::alias_of_on_a_model(json!({"litellm_provider": "test", "alias_of": "other"}))] +#[case::alias_row_with_other_fields(json!({"alias_of": "other", "mode": "chat"}))] +#[case::alias_of_not_a_string(json!({"alias_of": 4}))] #[case::non_object_model(json!(4))] fn generated_registry_schema_rejects_invalid_entries(#[case] entry: Value) { assert!(!registry_validator().is_valid(&json!({"example": entry}))); diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index ba892199dcf..48206d242e4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +rustls.workspace = true litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true @@ -30,6 +31,14 @@ bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true litellm-error.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-client.workspace = true +litellm-redis-connect.workspace = true +litellm-redis-store.workspace = true +litellm-redis-resolve.workspace = true +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true +tracing.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true @@ -51,8 +60,10 @@ litellm-router-types.workspace = true litellm-router-core.workspace = true litellm-host-python.workspace = true litellm-model-catalog.workspace = true +litellm-model-catalog-refresh.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } -pyo3.workspace = true +indexmap = "2.14.0" +pyo3 = { workspace = true, features = ["indexmap"] } prost.workspace = true pyo3-async-runtimes.workspace = true reqwest.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 4e3847951d8..fe9e06b3d23 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -6,6 +6,8 @@ pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled}; use pyo3::{exceptions::PyRuntimeError, prelude::*}; +const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache"; + #[derive(Debug, thiserror::Error)] pub(crate) enum HandleError { #[error("this inference route does not support configured caches")] diff --git a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md index 81a0997f02a..3b516251f43 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md @@ -4,6 +4,6 @@ This directory lets Rust inference use a selected Python cache. `service.rs` imp Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details -Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions +Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. The one exception is the asynchronous unary messages store: Python's `_async_set_cache` writes the returned response from `create_cache_write_task` after the router has stamped its `_hidden_params`, so `host.rs` defers that store until the deployment transform has returned the response the caller receives and hands that response to `litellm.rust_bridge.host.cache.store_response`, which schedules the same write task. Lookups, sync stores and stream stores stay in the caller's task. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs index 82b7d881df6..a8cbd0ce607 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/host.rs @@ -9,7 +9,7 @@ use pyo3::{ use serde_json::Value; use super::service::CacheCall; -use crate::cache::HandleError; +use crate::cache::{CACHE_HOST_MODULE, HandleError}; enum Pending { Lookup(Reply, Error>>), @@ -21,6 +21,7 @@ pub(crate) struct PythonCache { arguments: Option>, pending: Option, asynchronous: bool, + deferred: bool, } impl PythonCache { @@ -30,6 +31,7 @@ impl PythonCache { arguments: None, pending: None, asynchronous, + deferred: false, } } @@ -64,7 +66,14 @@ impl PythonCache { None, ) } - CacheCall::Store { value, reply } => { + CacheCall::Store { + unary: true, reply, .. + } if self.asynchronous => { + self.deferred = true; + reply.send(Ok(())); + return Ok(None); + } + CacheCall::Store { value, reply, .. } => { self.pending = Some(Pending::Store(reply)); ( if self.asynchronous { @@ -89,6 +98,19 @@ impl PythonCache { self.resume(py, result) } + pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> { + if !std::mem::take(&mut self.deferred) { + return Ok(()); + } + let (Some(cache), Some(arguments)) = (self.cache.as_ref(), self.arguments.as_ref()) else { + return Err(HandleError::Unconfigured.into()); + }; + py.import(CACHE_HOST_MODULE)? + .getattr("store_response")? + .call((cache, response), Some(&arguments.bind(py).copy()?))?; + Ok(()) + } + pub fn resume( &mut self, py: Python<'_>, @@ -124,6 +146,7 @@ impl PythonCache { pub fn close(&mut self) { self.pending = None; + self.deferred = false; self.cache = None; self.arguments = None; } diff --git a/litellm-rust/crates/python-bridge/src/cache/python/service.rs b/litellm-rust/crates/python-bridge/src/cache/python/service.rs index 38872b85d30..018c8647043 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/service.rs @@ -45,6 +45,7 @@ pub(crate) enum CacheCall { }, Store { value: Value, + unary: bool, reply: Reply>, }, } @@ -101,8 +102,13 @@ where ) -> Pin> + Send + 'a>> { Box::pin(async move { let value = to_python(value)?; + let unary = value.get(STREAM_EVENTS_KEY).is_none(); self.services - .call(|reply| CacheCall::Store { value, reply }) + .call(|reply| CacheCall::Store { + value, + unary, + reply, + }) .await .map_err(Error::unavailable)? }) diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index e96f079e5da..78c1606ea40 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -1,11 +1,10 @@ -use super::{HandleError, python}; +use super::{CACHE_HOST_MODULE, HandleError, python}; use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache}; use litellm_host::{ machine::{HostServices, MachineFault}, protocol::Protocol, }; use pyo3::{prelude::*, types::PyDict}; -const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache"; pub(crate) struct Cached

(std::marker::PhantomData

); diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 426cd861729..db5696d52ce 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -29,6 +29,8 @@ const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param"; /// The phase label (`Connection` or `Read`) of a timed-out call, for Python's timeout text, or /// `Body` for a streamed body read that timed out after the head arrived. const MESSAGES_TIMEOUT: &str = "messages_timeout"; +// `(received, expected, reset)` of a streamed body the connection ended early. +const MESSAGES_BODY_CUT: &str = "messages_body_cut"; /// No credential reached the provider, which Python raises as `AuthenticationError`. const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential"; /// Python refuses the request only once its logger knows the resolved model and provider (the @@ -134,6 +136,9 @@ pub(crate) fn route_failure( }; value.setattr(MESSAGES_TIMEOUT, phase)?; } + if let Some(cut) = error.body_cut() { + value.setattr(MESSAGES_BODY_CUT, (cut.received, cut.expected, cut.reset))?; + } if error.is_missing_credential() { value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?; } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index b4e20a8f750..c6a8b98747c 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -108,13 +108,35 @@ pub(crate) fn call_config( Some(snapshot) => configured(&snapshot)?, None => HttpSettingsLayer::default(), }, - ]) - .without_missing_files(&|path: &Path| path.exists()); - let resolution = Resolution::from(&settings); + ]); + Ok(resolved(py, settings)) +} + +/// The client config for a value of `litellm.rust_bridge.host.settings.http_settings()` the caller handed in, under +/// the process environment. `None` reads the environment alone +pub(crate) fn settings_config( + py: Python<'_>, + http_settings: Option<&Bound<'_, PyAny>>, +) -> PyResult { + let configured = match http_settings { + Some(value) => configured(&PythonSettings::Http.snapshot(value.clone()))?, + None => HttpSettingsLayer::default(), + }; + let settings = HttpSettings::from_layers([ + HttpSettingsLayer::from_environment(&ProcessEnvironment), + configured, + ]); + Ok(resolved(py, settings)) +} + +/// The client config `settings` resolve to, warning once per process about each setting the client cannot apply +fn resolved(py: Python<'_>, settings: HttpSettings) -> HttpClientConfig { + let resolution = + Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists())); for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) { crate::logger::capture(py).scope(|| litellm_tracing::warn!("{unsupported}")); } - Ok(resolution.config) + resolution.config } pub(crate) fn provider_client( diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f4a4f7c2b26..7edf13035c1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -9,9 +9,11 @@ mod execution; mod http; mod logger; mod marshal; +mod model_cost; mod model_registry; mod preflight; mod python_settings; +mod redis; mod router; mod routes; mod secrets; @@ -35,8 +37,39 @@ mod _native { #[pymodule_export] use crate::logger::NativeDiagnosticProcessor; #[pymodule_export] + use crate::model_cost::{ + NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched, NativeInstalled, + NativeModelCostMap, NativeModelCostMapState, NativeProvenance, NativeRegistered, + NativeRejected, NativeTableSwapped, + }; + #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states}; + #[pymodule_export] + use crate::redis::client::NativeRedisClient; + #[pymodule_export] + use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes}; + #[pymodule_export] + use crate::redis::store::rate_limits::{ + NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired, + PyGuardedReply, PyWindowCount, + }; + #[pymodule_export] + use crate::redis::store::request::{NativeRequestRedis, NativeSend}; + #[pymodule_export] + use crate::redis::store::router_usage::NativeRouterUsage; + #[pymodule_export] + use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins; + #[pymodule_export] + use crate::redis::store::spend_counters::{ + NativeSpendCounters, PyPostCallSettlement, PySpendRead, + }; + #[pymodule_export] + use crate::redis::store::{ + NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns, + }; + #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; @@ -66,6 +99,7 @@ mod _native { #[pymodule_init] fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { + let _ = rustls::crypto::ring::default_provider().install_default(); let py = module.py(); let dict = module.dict(); dict.set_item( @@ -111,11 +145,45 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", + "NativeRedisClient", + "NativeRedisBreaker", + "NativeAdmission", + "NativeRequestRedis", + "NativeSend", + "NativeAuthCache", + "NativeRateLimits", + "WindowCount", + "AtomicApplied", + "AtomicOverLimit", + "GuardedReply", + "GaugesAcquired", + "GaugeOverLimit", + "NativePending", + "NativeRouterCooldowns", + "NativeRouterUsage", + "NativeSensitiveRoutePins", + "NativeRedisStoreError", + "NativeSpendCounters", + "PostCallSettlement", + "SpendRead", + "LocalSizes", + "native_redis_breaker_states", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", "trace_encode_error", "bundled_model_cost_map", + "NativeModelCostMap", + "NativeModelCostMapState", + "NativeInstalled", + "NativeRejected", + "NativeRegistered", + "NativeTableSwapped", + "NativeProvenance", + "NativeCostMapSource", + "NativeFetchPolicy", + "NativeFetched", + "NativeFetchFailed", "TokenCounter", "Tokenizer", "RouterCoreHandle", @@ -140,4 +208,18 @@ mod tests { assert_eq!(public_names, expected); }); } + #[rstest::rstest] + fn module_import_installs_ring_as_the_process_crypto_provider() { + Python::initialize(); + Python::attach(|py| { + native_module(py); + }); + + let installed = rustls::crypto::CryptoProvider::get_default() + .expect("importing _native should install a process crypto provider"); + assert_eq!( + format!("{installed:?}"), + format!("{:?}", rustls::crypto::ring::default_provider()) + ); + } } diff --git a/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md b/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md new file mode 100644 index 00000000000..4ba23aac1ce --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md @@ -0,0 +1,16 @@ +# Model cost bridge + +- `NativeModelCostMap` holds inputs only, the installed `CostMap`, kept `Registrations` and the name index. Python's `litellm.model_cost` is the only copy of the live rows +- Every call that reads or writes rows takes the dict as an argument. Never store rows or build a second table here +- `LiveDict` holds a reference to that same dict and reads it under the GIL on every query, so a Python write is seen on the next read + - It needs the GIL, so Rust code running without the GIL cannot use it + - Its stamp is the dict's address, length and last key. [model-catalog/src/store/AGENTS.md](../../../model-catalog/src/store/AGENTS.md) says what that misses +- `host.rs` is the only code that calls Python, by name through `MODULE`. Calls happen with the GIL held and no `NativeModelCostMap` lock held + - `resolve` reaches `find_key` on this same object, so holding a lock across a host call deadlocks + - Test `a_host_lookup_that_reaches_find_key_on_the_same_map_finishes` guards this +- Rows and registrations cross as Python objects, never JSON, because router `model_info` holds `datetime` +- `table::build` serves the document as main's `_expand_model_aliases` did, so all names of a model share one row object made by `host.entry` +- `install` touches no state until `CostMap::read` succeeds. `NativeRejected` means nothing changed +- `register` keeps a registration only after its pass finished, so an exception in a host call keeps nothing +- `rows=None` writes each entry into whichever table is published at that moment and never reports a swap +- `fetch_blocking` runs on a throwaway runtime with an unpooled client so an importing process can still fork. `fetch` is for workers after the fork diff --git a/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs b/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs new file mode 100644 index 00000000000..77ec88cc1b9 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs @@ -0,0 +1,239 @@ +use std::time::Duration; + +use litellm_http::{ClientVariant, HttpClientConfig}; +use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +/// A cost map URL and how to connect to it. An `s3://` URL reaches S3 as an AWS SDK finds it in the environment, +/// with LiteLLM's `AWS_REGION_NAME` winning over the region it finds +#[pyclass( + frozen, + name = "NativeCostMapSource", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeCostMapSource { + url: String, + http: HttpClientConfig, +} + +#[pymethods] +impl NativeCostMapSource { + /// `http_settings` is `litellm.rust_bridge.host.settings.http_settings()`, or `None` for the environment alone + #[new] + #[pyo3(signature = (url, *, http_settings))] + fn new(py: Python<'_>, url: String, http_settings: Option>) -> PyResult { + Ok(Self { + url, + http: crate::http::settings_config(py, http_settings.as_ref())?, + }) + } + + /// Fetches on the shared runtime with a pooled client + fn fetch<'py>( + &self, + py: Python<'py>, + policy: PyRef<'_, NativeFetchPolicy>, + ) -> PyResult> { + let client = crate::http::pool() + .client(&self.http, ClientVariant::Provider) + .map_err(crate::http::client_error)?; + let url = self.url.clone(); + let policy = policy.0; + crate::execution::run_async_value(py, async move { + Ok(Outcome( + litellm_model_catalog_refresh::fetch(&client, &s3(), &url, policy, Start::Fresh) + .await, + )) + }) + } + + /// Fetches on a throwaway runtime with a client of its own. `resume_from` continues a fetch whose first attempt + /// failed retryably + #[pyo3(signature = (policy, resume_from = None))] + fn fetch_blocking( + &self, + py: Python<'_>, + policy: PyRef<'_, NativeFetchPolicy>, + resume_from: Option>, + ) -> PyResult> { + let start = resume_from.map_or(Start::Fresh, |failed| Start::After(failed.0.clone())); + let (config, url, policy) = (self.http.clone(), self.url.clone(), policy.0); + let logger = crate::logger::capture(py); + let fetched = litellm_host_python::release_gil(py, move || { + logger.scope(|| { + let client = crate::http::pool() + .unpooled(&config, ClientVariant::Provider) + .map_err(crate::http::client_error)?; + litellm_host_python::run_isolated(litellm_model_catalog_refresh::fetch( + &client, + &s3(), + &url, + policy, + start, + )) + .map_err(PyErr::from) + }) + })?; + Outcome(fetched).into_pyobject(py).map(Bound::unbind) + } +} + +fn s3() -> S3Access { + S3Access::Environment { + region: std::env::var("AWS_REGION_NAME").ok(), + } +} + +#[pyclass( + frozen, + name = "NativeFetchPolicy", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeFetchPolicy(RetryPolicy); + +#[pymethods] +impl NativeFetchPolicy { + #[new] + /// Retries wait an exponential backoff with jitter from `min_delay_seconds` up to `max_wait_seconds`, or the + /// server's `Retry-After` up to that cap + #[pyo3(signature = (*, max_attempts, max_wait_seconds, timeout_seconds, min_delay_seconds))] + fn new( + max_attempts: u32, + max_wait_seconds: f64, + timeout_seconds: f64, + min_delay_seconds: f64, + ) -> PyResult { + if max_attempts == 0 { + return Err(PyValueError::new_err("max_attempts must be at least 1")); + } + let seconds = |name: &str, value: f64| { + Duration::try_from_secs_f64(value).map_err(|_| { + PyValueError::new_err(format!("{name} must be a non-negative number of seconds")) + }) + }; + let min_delay = seconds("min_delay_seconds", min_delay_seconds)?; + let max_delay = seconds("max_wait_seconds", max_wait_seconds)?; + if min_delay > max_delay { + return Err(PyValueError::new_err( + "min_delay_seconds must not exceed max_wait_seconds", + )); + } + Ok(Self(RetryPolicy { + max_attempts, + timeout: seconds("timeout_seconds", timeout_seconds)?, + min_delay, + max_delay, + })) + } +} + +#[pyclass(frozen, name = "NativeFetched", module = "litellm.rust_bridge._native")] +pub(crate) struct NativeFetched { + #[pyo3(get)] + body: Py, + #[pyo3(get)] + etag: Option, +} + +#[pyclass( + frozen, + skip_from_py_object, + name = "NativeFetchFailed", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeFetchFailed(Failed); + +#[pymethods] +impl NativeFetchFailed { + #[getter] + fn reason(&self) -> &str { + &self.0.reason + } + + /// Another attempt could succeed. Only the attempt limit ended the fetch + #[getter] + fn retryable(&self) -> bool { + self.0.retryable + } + + #[getter] + fn retry_after(&self) -> Option { + self.0.retry_after.map(|wait| wait.as_secs_f64()) + } +} + +struct Outcome(Result); + +impl<'py> IntoPyObject<'py> for Outcome { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + match self.0 { + Ok(fetched) => Bound::new( + py, + NativeFetched { + body: PyBytes::new(py, &fetched.body).unbind(), + etag: fetched.etag, + }, + ) + .map(Bound::into_any), + Err(failed) => Bound::new(py, NativeFetchFailed(failed)).map(Bound::into_any), + } + } +} + +#[cfg(test)] +mod tests { + use std::io::{BufRead, BufReader, Write}; + use std::net::TcpListener; + + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::equal_bounds(2.0, 2.0, true)] + #[case::min_below_max(2.0, 30.0, true)] + #[case::min_above_max(31.0, 30.0, false)] + fn a_policy_accepts_a_min_delay_only_up_to_the_max_wait( + #[case] min_delay_seconds: f64, + #[case] max_wait_seconds: f64, + #[case] accepted: bool, + ) { + Python::initialize(); + let policy = NativeFetchPolicy::new(3, max_wait_seconds, 5.0, min_delay_seconds); + assert_eq!(policy.is_ok(), accepted); + } + + #[rstest] + fn a_blocking_fetch_hands_python_the_document_without_starting_the_shared_runtime() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/cost.json", listener.local_addr().unwrap()); + std::thread::spawn(move || { + let mut stream = listener.incoming().next().unwrap().unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut line = String::new(); + while reader.read_line(&mut line).unwrap() > 2 { + line.clear(); + } + stream + .write_all(b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}") + .unwrap(); + }); + Python::initialize(); + Python::attach(|py| { + let source = NativeCostMapSource::new(py, url, None).unwrap(); + let policy = + Bound::new(py, NativeFetchPolicy::new(3, 30.0, 5.0, 2.0).unwrap()).unwrap(); + let outcome = source.fetch_blocking(py, policy.borrow(), None).unwrap(); + let fetched = outcome.bind(py).cast::().unwrap().get(); + assert_eq!(fetched.body.bind(py).as_bytes(), b"{}"); + assert_eq!(fetched.etag.as_deref(), Some("\"v1\"")); + }); + assert!(!litellm_host_python::runtime_started()); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/host.rs b/litellm-rust/crates/python-bridge/src/model_cost/host.rs new file mode 100644 index 00000000000..25fbbefbb73 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/host.rs @@ -0,0 +1,81 @@ +//! Calls into `litellm.rust_bridge.host.model_cost` + +use pyo3::call::PyCallArgs; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +pub(crate) const MODULE: &str = "litellm.rust_bridge.host.model_cost"; + +/// Where a registration lands in the published table and the row it starts from +#[derive(FromPyObject)] +pub(crate) struct Landing<'py> { + #[pyo3(attribute)] + pub(crate) key: String, + #[pyo3(attribute)] + pub(crate) base: Bound<'py, PyAny>, + #[pyo3(attribute)] + pub(crate) matched_builtin: bool, +} + +fn call<'py>( + py: Python<'py>, + name: &str, + args: impl PyCallArgs<'py>, +) -> PyResult> { + py.import(MODULE)?.getattr(name)?.call1(args) +} + +/// The row object every name in `names` shares, the owner first +pub(crate) fn entry<'py>( + fields: Bound<'py, PyAny>, + names: &[String], +) -> PyResult> { + let py = fields.py(); + call(py, "entry", (fields, PyTuple::new(py, names)?)) +} + +/// Where `registration` lands in `rows`, `None` when `rows` is no longer the published table. `rows` of `None` is the +/// published table +pub(crate) fn resolve<'py>( + rows: Option<&Bound<'py, PyDict>>, + key: &str, + registration: &Bound<'py, PyAny>, + custom_llm_provider: Option<&str>, +) -> PyResult>> { + call( + registration.py(), + "resolve", + (rows, key, registration, custom_llm_provider), + )? + .extract() +} + +/// The row `registration` writes over `current` +pub(crate) fn merged<'py>( + current: Option<&Bound<'py, PyAny>>, + landing: &Landing<'py>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + call( + registration.py(), + "merged", + (current, &landing.base, registration), + ) +} + +/// The registration kept for a key: `registration` merged into the one kept before it +pub(crate) fn kept<'py>( + previous: Option<&Bound<'py, PyDict>>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + Ok(call(registration.py(), "kept", (previous, registration))?.cast_into::()?) +} + +pub(crate) fn published(py: Python<'_>) -> PyResult> { + Ok(call(py, "published", ())?.cast_into::()?) +} + +/// Writes `written` into `rows`, the published table, through the host, which tracks the table's version +pub(crate) fn write(rows: &Bound<'_, PyDict>, written: &Bound<'_, PyDict>) -> PyResult<()> { + call(rows.py(), "write", (rows, written)).map(drop) +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/live.rs b/litellm-rust/crates/python-bridge/src/model_cost/live.rs new file mode 100644 index 00000000000..0b7f071127d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/live.rs @@ -0,0 +1,176 @@ +use std::borrow::Cow; + +use litellm_model_catalog::Rows; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +/// The dict Python last handed in as `litellm.model_cost`, read under the GIL on every query +pub(crate) struct LiveDict(Py); + +/// One row exactly as the dict holds it +pub(crate) struct LiveRow(Py); + +impl Clone for LiveRow { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } +} + +impl LiveDict { + pub(crate) fn new(dict: &Bound<'_, PyDict>) -> Self { + Self(dict.clone().unbind()) + } + + pub(crate) fn is(&self, dict: &Bound<'_, PyDict>) -> bool { + self.0.bind(dict.py()).is(dict) + } +} + +impl Rows for LiveDict { + type Row = LiveRow; + /// The dict's address, length and last key + type Stamp = (usize, usize, Option); + + fn fold(name: &str) -> String { + name.to_lowercase() + } + + fn stamp(&self) -> Self::Stamp { + Python::attach(|py| { + let rows = self.0.bind(py); + let last = rows + .call_method0(intern!(py, "__reversed__")) + .and_then(|keys| keys.try_iter()) + .ok() + .and_then(|mut keys| keys.next()) + .and_then(Result::ok) + .and_then(|key| key.extract::().ok()); + (self.0.as_ptr() as usize, rows.len(), last) + }) + } + + fn contains(&self, key: &str) -> bool { + Python::attach(|py| self.0.bind(py).contains(key).unwrap_or(false)) + } + + fn row(&self, key: &str) -> Option> { + Python::attach(|py| { + let row = self.0.bind(py).get_item(key).ok().flatten()?; + Some(Cow::Owned(LiveRow(row.unbind()))) + }) + } + + /// Groups each key by its row's `litellm_provider` when that is a string + fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) { + Python::attach(|py| { + for (key, row) in self.0.bind(py).iter() { + let Ok(key) = key.extract::() else { + continue; + }; + let provider = row + .cast::() + .ok() + .and_then(|row| row.get_item("litellm_provider").ok().flatten()) + .and_then(|provider| provider.extract::().ok()); + visit(&key, provider.as_deref()); + } + }); + } +} + +#[cfg(test)] +mod tests { + use litellm_model_catalog::Store; + use pyo3::ffi::c_str; + use rstest::rstest; + + use super::*; + + fn dict<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + py.eval(source, None, None) + .unwrap() + .cast_into::() + .unwrap() + } + + #[rstest] + fn a_key_added_to_the_dict_is_found_on_the_next_query() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}")); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.find_key("GPT-4O"), Some("gpt-4o".to_owned())); + rows.set_item("Claude-X", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("claude-x"), Some("Claude-X".to_owned())); + }); + } + + #[rstest] + fn a_removed_key_is_never_returned() { + Python::initialize(); + Python::attach(|py| { + let rows = dict( + py, + c_str!("{'gpt-4o': {'litellm_provider': 'openai'}, 'o3': {}}"), + ); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.keys_in("openai"), ["gpt-4o"]); + rows.del_item("gpt-4o").unwrap(); + rows.set_item("o4", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("GPT-4O"), None); + assert!(store.keys_in("openai").is_empty()); + }); + } + + #[rstest] + fn a_key_swapped_in_for_a_removed_one_is_found() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {}, 'Lit-Gone': {}}")); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.find_key("lit-gone"), Some("Lit-Gone".to_owned())); + rows.del_item("Lit-Gone").unwrap(); + rows.set_item("Lit-New", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("lit-new"), Some("Lit-New".to_owned())); + assert_eq!(store.find_key("lit-gone"), None); + }); + } + + #[rstest] + fn two_keys_folding_alike_resolve_to_the_later_one() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'Model': {}, 'MODEL': {}}")); + assert_eq!( + Store::new(LiveDict::new(&rows)).find_key("model"), + Some("MODEL".to_owned()) + ); + }); + } + + #[rstest] + fn rows_are_the_objects_the_dict_holds() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}")); + let store = Store::new(LiveDict::new(&rows)); + let row = store.get("gpt-4o").unwrap().into_owned(); + assert!(row.0.bind(py).is(rows.get_item("gpt-4o").unwrap().unwrap())); + }); + } + + #[rstest] + fn providers_group_only_string_providers() { + Python::initialize(); + Python::attach(|py| { + let rows = dict( + py, + c_str!( + "{'a': {'litellm_provider': 'p'}, 'b': {'litellm_provider': None}, 'c': 1, 1: {'litellm_provider': 'p'}}" + ), + ); + assert_eq!(Store::new(LiveDict::new(&rows)).keys_in("p"), ["a"]); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/mod.rs b/litellm-rust/crates/python-bridge/src/model_cost/mod.rs new file mode 100644 index 00000000000..99451033122 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/mod.rs @@ -0,0 +1,497 @@ +mod fetch; +mod host; +mod live; +mod table; + +use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard}; + +use litellm_host_python::to_py; +use litellm_model_catalog::{ + CostMap, IntegrityLimits, Landed, Pass, Placed, Provenance, RegistrationTable, Registrations, + Store, write_pass, +}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +pub(crate) use fetch::{NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched}; +use live::LiveDict; + +/// The installed document and the `register_model` calls kept above it, each as the dict Python merged +struct Inputs { + cost_map: CostMap, + registrations: Registrations>, +} + +impl Inputs { + fn clone_ref(&self, py: Python<'_>) -> Self { + Self { + cost_map: self.cost_map.clone(), + registrations: self.registrations.map(|kept| kept.clone_ref(py)), + } + } +} + +/// The cost map's inputs and the name index over `litellm.model_cost` +#[pyclass( + frozen, + name = "NativeModelCostMap", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeModelCostMap { + inputs: RwLock, + limits: IntegrityLimits, + index: RwLock>>>, +} + +#[pyclass( + frozen, + name = "NativeInstalled", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeInstalled { + /// The new live table, the document without registrations + #[pyo3(get)] + rows: Py, + /// How many keys the document itself serves, aliases included and registrations left out + #[pyo3(get)] + model_count: usize, + /// The document's `fallback_generalizations` block as written, `None` when it has none + #[pyo3(get)] + fallback_generalizations: Py, +} + +#[pyclass( + frozen, + name = "NativeRejected", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeRejected { + #[pyo3(get)] + reason: String, + /// The document parsed but has fewer models than the integrity limits allow + #[pyo3(get)] + integrity: bool, +} + +/// A finished registration pass, with the registrations no built-in row matched, by registered key +#[pyclass( + frozen, + name = "NativeRegistered", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeRegistered { + #[pyo3(get)] + unmatched: Py, +} + +/// A registration pass that stopped because the table it was handed is no longer the published one +#[pyclass( + frozen, + name = "NativeTableSwapped", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeTableSwapped; + +/// The inputs at one moment, for a caller that has to put them back, such as test isolation +#[pyclass( + frozen, + name = "NativeModelCostMapState", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeModelCostMapState(Inputs); + +#[pyclass( + frozen, + name = "NativeProvenance", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeProvenance { + #[pyo3(get)] + revision: String, + #[pyo3(get)] + etag: Option, +} + +#[pymethods] +impl NativeModelCostMap { + /// Starts from the bundled document. A fetched document with fewer than `min_model_count` models, or fewer than + /// `max_shrink_ratio` times the bundled count, is rejected + #[staticmethod] + fn bundled(min_model_count: i64, max_shrink_ratio: f64) -> PyResult { + let min_model_count = usize::try_from(min_model_count).map_err(|_| { + PyValueError::new_err(format!( + "min_model_count must be 0 or more, got {min_model_count}" + )) + })?; + let limits = IntegrityLimits::new(min_model_count, max_shrink_ratio).ok_or_else(|| { + PyValueError::new_err(format!( + "max_shrink_ratio must be between 0 and 1, got {max_shrink_ratio}" + )) + })?; + Ok(Self { + inputs: RwLock::new(Inputs { + cost_map: CostMap::bundled(), + registrations: Registrations::default(), + }), + limits, + index: RwLock::new(None), + }) + } + + /// The live table built from a fetched document, or why it was rejected, in which case nothing changed + #[pyo3(signature = (body, etag = None))] + fn install(&self, py: Python<'_>, body: &[u8], etag: Option) -> PyResult> { + match crate::logger::capture(py).scope(|| CostMap::read(body, etag, self.limits)) { + Ok(cost_map) => self.install_document(py, cost_map), + Err(error) => rejected(py, &error), + } + } + + /// A caller's document as a table, read without size limits and installed nowhere + #[staticmethod] + fn finalize(py: Python<'_>, body: &[u8]) -> PyResult> { + match crate::logger::capture(py) + .scope(|| CostMap::read(body, None, IntegrityLimits::UNCHECKED)) + { + Ok(cost_map) => Ok(Py::new(py, built(py, &cost_map)?)?.into_any()), + Err(error) => rejected(py, &error), + } + } + + /// The live table built from the bundled document, which is never size checked + fn install_bundled(&self, py: Python<'_>) -> PyResult> { + self.install_document(py, CostMap::bundled()) + } + + /// Writes each entry into `rows`, or into the published table when `rows` is `None` + #[pyo3(signature = (rows, entries, persist, custom_llm_provider = None))] + fn register( + &self, + py: Python<'_>, + rows: Option<&Bound<'_, PyDict>>, + entries: &Bound<'_, PyDict>, + persist: bool, + custom_llm_provider: Option<&str>, + ) -> PyResult> { + let entries: Vec<(String, Bound<'_, PyAny>)> = entries + .iter() + .map(|(key, registration)| Ok((key.str()?.to_string(), registration))) + .collect::>()?; + let kept = match persist { + true => self.merged_with_kept(&entries)?, + false => Vec::new(), + }; + let pass = self.write_pass(rows, &entries, custom_llm_provider)?; + if persist { + self.keep(kept, pass.landed()); + } + into_py(py, pass) + } + + /// Writes every kept registration into `rows` again, oldest first + fn reapply(&self, rows: &Bound<'_, PyDict>) -> PyResult> { + let py = rows.py(); + let entries: Vec<(String, Bound<'_, PyAny>)> = self + .read_inputs() + .registrations + .entries() + .map(|(key, kept)| (key.to_owned(), kept.bind(py).clone().into_any())) + .collect(); + let pass = self.write_pass(Some(rows), &entries, None)?; + self.write_inputs().registrations.record(pass.landed()); + into_py(py, pass) + } + + /// The installed row under `key` with the kept registrations that apply to it, without anything the router wrote + fn registered_row<'py>( + &self, + py: Python<'py>, + key: &str, + ) -> PyResult>> { + let (cost_map, applied) = { + let inputs = self.read_inputs(); + let applied: Vec> = inputs + .registrations + .over(key) + .into_iter() + .map(|kept| kept.bind(py).clone()) + .collect(); + (inputs.cost_map.clone(), applied) + }; + table::registered_row(py, &cost_map, key, &applied) + } + + fn registered_keys(&self) -> Vec { + let inputs = self.read_inputs(); + table::registered_keys( + &inputs.cost_map, + inputs.registrations.entries().map(|(key, _)| key), + ) + } + + /// The installed document as a table, without any registration + fn loaded<'py>(&self, py: Python<'py>) -> PyResult> { + let cost_map = self.read_inputs().cost_map.clone(); + table::build(py, &cost_map) + } + + /// The key itself when `rows` has it, otherwise the key that differs from `name` only in case + fn find_key(&self, rows: &Bound<'_, PyDict>, name: &str) -> Option { + self.index_over(rows).find_key(name) + } + + /// Keys whose row's `litellm_provider` is exactly `provider`, in row order + fn keys_in(&self, rows: &Bound<'_, PyDict>, provider: &str) -> Vec { + self.index_over(rows).keys_in(provider) + } + + /// Drops the name index, for a caller that changed keys in a way the dict's length does not show + fn invalidate(&self) { + if let Some(index) = self + .index + .read() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + { + index.invalidate(); + } + } + + fn snapshot(&self, py: Python<'_>) -> NativeModelCostMapState { + NativeModelCostMapState(self.read_inputs().clone_ref(py)) + } + + /// Puts back the document, registrations and provenance a `snapshot` saw + fn restore(&self, py: Python<'_>, state: PyRef<'_, NativeModelCostMapState>) { + *self.write_inputs() = state.0.clone_ref(py); + self.invalidate(); + } + + #[getter] + fn provenance(&self) -> NativeProvenance { + let Provenance { revision, etag } = self.read_inputs().cost_map.provenance().clone(); + NativeProvenance { revision, etag } + } +} + +impl NativeModelCostMap { + fn read_inputs(&self) -> RwLockReadGuard<'_, Inputs> { + self.inputs.read().unwrap_or_else(PoisonError::into_inner) + } + + fn write_inputs(&self) -> RwLockWriteGuard<'_, Inputs> { + self.inputs.write().unwrap_or_else(PoisonError::into_inner) + } + + fn install_document(&self, py: Python<'_>, cost_map: CostMap) -> PyResult> { + let installed = built(py, &cost_map)?; + self.write_inputs().cost_map = cost_map; + *self.index.write().unwrap_or_else(PoisonError::into_inner) = + Some(Arc::new(Store::new(LiveDict::new(installed.rows.bind(py))))); + Ok(Py::new(py, installed)?.into_any()) + } + + /// Each entry merged into the registration kept for its key + fn merged_with_kept<'py>( + &self, + entries: &[(String, Bound<'py, PyAny>)], + ) -> PyResult)>> { + entries + .iter() + .map(|(key, registration)| { + let previous = self + .read_inputs() + .registrations + .get(key) + .map(|kept| kept.bind(registration.py()).clone()); + Ok((key.clone(), host::kept(previous.as_ref(), registration)?)) + }) + .collect() + } + + fn keep(&self, kept: Vec<(String, Bound<'_, PyDict>)>, landed: &[Landed]) { + let mut inputs = self.write_inputs(); + for (key, registration) in kept { + inputs.registrations.register(&key, registration.unbind()); + } + inputs.registrations.record(landed); + } + + fn write_pass<'a, 'py>( + &self, + rows: Option<&Bound<'py, PyDict>>, + entries: &'a [(String, Bound<'py, PyAny>)], + custom_llm_provider: Option<&str>, + ) -> PyResult>> { + let mut table = PublishedTable { + map: self, + rows, + custom_llm_provider, + }; + write_pass(&mut table, entries) + } + + /// The index over `rows`, rebuilt when `rows` is not the dict it was built over + fn index_over(&self, rows: &Bound<'_, PyDict>) -> Arc> { + let current = self + .index + .read() + .unwrap_or_else(PoisonError::into_inner) + .clone(); + if let Some(index) = current.filter(|index| index.rows().is(rows)) { + return index; + } + let index = Arc::new(Store::new(LiveDict::new(rows))); + *self.index.write().unwrap_or_else(PoisonError::into_inner) = Some(Arc::clone(&index)); + index + } +} + +/// The dict Python publishes as `litellm.model_cost`, written through the host +struct PublishedTable<'a, 'py> { + map: &'a NativeModelCostMap, + rows: Option<&'a Bound<'py, PyDict>>, + custom_llm_provider: Option<&'a str>, +} + +impl<'py> RegistrationTable for PublishedTable<'_, 'py> { + type Registration = Bound<'py, PyAny>; + type Error = PyErr; + + fn place(&mut self, key: &str, registration: &Self::Registration) -> PyResult> { + let Some(landing) = host::resolve(self.rows, key, registration, self.custom_llm_provider)? + else { + return Ok(None); + }; + let rows = match self.rows { + Some(rows) => rows.clone(), + None => host::published(registration.py())?, + }; + table::write(&rows, &landing, registration)?; + self.map.invalidate(); + Ok(Some(Placed { + key: landing.key, + matched_builtin: landing.matched_builtin, + })) + } +} + +fn built(py: Python<'_>, cost_map: &CostMap) -> PyResult { + let rows = table::build(py, cost_map)?; + let fallback_generalizations = match cost_map.fallback_generalizations() { + Some(block) => to_py(py, block)?, + None => py.None(), + }; + Ok(NativeInstalled { + model_count: rows.len(), + rows: rows.unbind(), + fallback_generalizations, + }) +} + +fn rejected(py: Python<'_>, error: &litellm_model_catalog::Error) -> PyResult> { + let rejected = NativeRejected { + reason: litellm_error::Report::new(error).to_string(), + integrity: error.is_integrity(), + }; + Ok(Py::new(py, rejected)?.into_any()) +} + +fn into_py<'py>(py: Python<'py>, pass: Pass<'_, Bound<'py, PyAny>>) -> PyResult> { + match pass { + Pass::Swapped { .. } => Ok(Py::new(py, NativeTableSwapped)?.into_any()), + Pass::Finished { unmatched, .. } => { + let table = PyDict::new(py); + for (key, registration) in unmatched { + table.set_item(key, registration)?; + } + let registered = NativeRegistered { + unmatched: table.unbind(), + }; + Ok(Py::new(py, registered)?.into_any()) + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::mpsc; + use std::time::Duration; + + use pyo3::ffi::c_str; + use rstest::rstest; + + use super::*; + + const HOST: &std::ffi::CStr = c_str!( + r#" +from types import SimpleNamespace + +class Row(dict): + names = () + +found = [] + +def resolve(rows, key, registration, custom_llm_provider): + found.append(cost_map.find_key(rows, key)) + return SimpleNamespace(key=key, base={}, matched_builtin=True) + +def merged(current, base, registration): + return Row({**(current or {}), **base, **registration}) + +def write(rows, written): + rows.update(written) +"# + ); + + #[rstest] + fn a_host_lookup_that_reaches_find_key_on_the_same_map_finishes() { + Python::initialize(); + let (sender, receiver) = mpsc::channel(); + std::thread::spawn(move || { + let outcome = Python::attach(|py| -> PyResult<(Vec>, f64)> { + let cost_map = Py::new(py, NativeModelCostMap::bundled(0, 0.0)?)?; + let module_name = std::ffi::CString::new(host::MODULE).unwrap(); + let host = + PyModule::from_code(py, HOST, c_str!("host_model_cost.py"), &module_name)?; + host.setattr("cost_map", &cost_map)?; + let modules = py.import("sys")?.getattr("modules")?; + let packages = ["litellm", "litellm.rust_bridge", "litellm.rust_bridge.host"]; + for package in packages { + modules.set_item(package, PyModule::new(py, package)?)?; + } + modules.set_item(host::MODULE, &host)?; + let rows = PyDict::new(py); + rows.set_item("Lit-Reentry", PyDict::new(py))?; + let entry = PyDict::new(py); + entry.set_item("input_cost_per_token", 1.5)?; + let entries = PyDict::new(py); + entries.set_item("lit-reentry", entry)?; + let registered = cost_map + .bind(py) + .call_method1("register", (&rows, &entries, false)); + for name in [host::MODULE].into_iter().chain(packages) { + modules.del_item(name)?; + } + registered?; + let price = rows + .get_item("lit-reentry")? + .expect("the registration landed") + .get_item("input_cost_per_token")? + .extract()?; + Ok((host.getattr("found")?.extract()?, price)) + }); + sender + .send(outcome.map_err(|error| error.to_string())) + .unwrap(); + }); + + let (found, price) = receiver + .recv_timeout(Duration::from_secs(30)) + .expect("register finished instead of deadlocking on its own lock") + .unwrap(); + assert_eq!(found, [Some("Lit-Reentry".to_owned())]); + assert!((price - 1.5).abs() < f64::EPSILON); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/table.rs b/litellm-rust/crates/python-bridge/src/model_cost/table.rs new file mode 100644 index 00000000000..67b5f581b49 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/table.rs @@ -0,0 +1,166 @@ +//! Builds `litellm.model_cost` from an installed document + +use std::borrow::Cow; +use std::collections::HashMap; + +use indexmap::IndexSet; +use litellm_host_python::to_py; +use litellm_model_catalog::wire::Fields; +use litellm_model_catalog::{CostMap, FALLBACK_GENERALIZATIONS, Written}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::Value; + +use super::host::{self, Landing}; + +const ALIASES: &str = "aliases"; + +enum Served<'a> { + Model, + AliasOf(&'a str), + /// A top-level value that is not a model row, such as `sample_spec` + Section(&'a Value), +} + +/// Every served key in dict order. An `alias_of` row repeats among the aliases and keeps its document position +fn served(cost_map: &CostMap) -> impl Iterator)> { + let catalog = cost_map.catalog(); + let documented = catalog + .written() + .filter(|(key, _)| *key != FALLBACK_GENERALIZATIONS) + .filter_map(|(key, written)| { + let served = match written { + Written::Section(value) => Served::Section(value), + Written::Entry(_) => match catalog.aliases().get(key) { + Some(owner) => Served::AliasOf(owner), + None => catalog.get(key).map(|_| Served::Model)?, + }, + }; + Some((key, served)) + }); + let listed = cost_map.alias_groups().flat_map(|(owner, names)| { + names + .iter() + .skip(1) + .map(move |alias| (alias.as_str(), Served::AliasOf(owner))) + }); + documented.chain(listed) +} + +/// The installed document as a new dict +pub(crate) fn build<'py>(py: Python<'py>, cost_map: &CostMap) -> PyResult> { + let table = PyDict::new(py); + let mut rows = HashMap::new(); + for (key, served) in served(cost_map) { + let value = match served { + Served::Model => row(py, cost_map, key, &mut rows)?, + Served::AliasOf(owner) => row(py, cost_map, owner, &mut rows)?, + Served::Section(value) => to_py(py, value)?.into_bound(py), + }; + table.set_item(key, value)?; + } + Ok(table) +} + +fn row<'a, 'py>( + py: Python<'py>, + cost_map: &'a CostMap, + owner: &'a str, + rows: &mut HashMap<&'a str, Bound<'py, PyAny>>, +) -> PyResult> { + if let Some(row) = rows.get(owner) { + return Ok(row.clone()); + } + let fields = cost_map + .catalog() + .get(owner) + .expect("a served model or alias owner is a catalog row"); + let fields = to_py(py, without_alias_list(fields).as_ref())?.into_bound(py); + let row = match cost_map.names(owner) { + [] => fields, + names => host::entry(fields, names)?, + }; + rows.insert(owner, row.clone()); + Ok(row) +} + +fn without_alias_list(fields: &Fields) -> Cow<'_, Fields> { + if fields.get(ALIASES).is_none_or(Value::is_null) { + return Cow::Borrowed(fields); + } + Cow::Owned( + fields + .iter() + .filter(|(name, _)| *name != ALIASES) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + ) +} + +/// The installed row under `key` with `applied` written over it in order, `None` when neither has one +pub(crate) fn registered_row<'py>( + py: Python<'py>, + cost_map: &CostMap, + key: &str, + applied: &[Bound<'py, PyDict>], +) -> PyResult>> { + let catalog = cost_map.catalog(); + let installed = catalog + .get(key) + .or_else(|| match catalog.section(key) { + Some(Value::Object(fields)) if key != FALLBACK_GENERALIZATIONS => Some(fields), + _ => None, + }) + .map(without_alias_list); + let row = match (installed, applied) { + (None, []) => return Ok(None), + (Some(fields), _) => to_py(py, fields.as_ref())? + .into_bound(py) + .cast_into::()?, + (None, _) => PyDict::new(py), + }; + for registration in applied { + row.update(registration.as_mapping())?; + } + Ok(Some(row)) +} + +/// The keys `registered_row` answers for, installed keys first +pub(crate) fn registered_keys<'a>( + cost_map: &'a CostMap, + registered: impl Iterator, +) -> Vec { + let served: IndexSet<&str> = served(cost_map).map(|(key, _)| key).collect(); + let only_registered: Vec<&str> = registered.filter(|key| !served.contains(key)).collect(); + served + .into_iter() + .chain(only_registered) + .map(str::to_owned) + .collect() +} + +/// Writes `registration` over the row under `landing.key`, and under every name that shares that row object +pub(crate) fn write<'py>( + rows: &Bound<'py, PyDict>, + landing: &Landing<'py>, + registration: &Bound<'py, PyAny>, +) -> PyResult<()> { + let current = rows.get_item(&landing.key)?; + let row = host::merged(current.as_ref(), landing, registration)?; + let names: Vec = row.getattr("names")?.extract()?; + let sharing: Vec<&String> = names + .iter() + .filter(|name| { + let served = rows.get_item(name.as_str()).ok().flatten(); + served + .zip(current.as_ref()) + .is_some_and(|(served, current)| served.is(current)) + }) + .collect(); + let written = PyDict::new(rows.py()); + written.set_item(&landing.key, &row)?; + sharing + .into_iter() + .try_for_each(|name| written.set_item(name, &row))?; + host::write(rows, &written) +} diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index eacec6859b5..93377abb7e9 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -13,6 +13,8 @@ pub(crate) enum PythonSettings { UrlPolicy, #[strum(serialize = "provider_defaults")] ProviderDefaults, + #[strum(serialize = "router_tiers")] + RouterTiers, #[strum(serialize = "secret_manager")] SecretManager, #[strum(serialize = "secret_manager_binding")] @@ -54,7 +56,7 @@ impl PythonSettings { } } - #[cfg(test)] + /// A snapshot of a value Python handed in, read with the same field rules as the accessor's pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> { Snapshot { group: self, value } } diff --git a/litellm-rust/crates/python-bridge/src/redis/breaker.rs b/litellm-rust/crates/python-bridge/src/redis/breaker.rs new file mode 100644 index 00000000000..733cf1bd969 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/breaker.rs @@ -0,0 +1,134 @@ +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::{Duration, Instant}, +}; + +use litellm_circuit_breaker::{Admission, Breaker, Config, Failure}; +use litellm_redis_client::RedisEvents; +use litellm_redis_types::Endpoint; +use pyo3::{exceptions::PyValueError, prelude::*}; + +use crate::redis::{ + client::{RedisSpec, resolved_endpoint}, + events::{ServiceLoggingEvents, state_label}, + hosts, +}; + +/// `litellm.redis.store_settings.RedisBreakerSettings`, filled from `litellm.constants`. +#[derive(FromPyObject)] +pub(crate) struct BreakerSettings { + failure_threshold: u32, + recovery_timeout: f64, + timeout_min_duration: f64, +} + +impl BreakerSettings { + pub(crate) fn config(&self) -> PyResult { + let seconds = |seconds: f64| { + Duration::try_from_secs_f64(seconds) + .map_err(|error| PyValueError::new_err(error.to_string())) + }; + Ok(Config { + failure_threshold: self.failure_threshold, + recovery_timeout: seconds(self.recovery_timeout)?, + timeout_min_duration: seconds(self.timeout_min_duration)?, + }) + } +} + +/// The one health signal for a Redis, shared by its `RedisCache` and its native client, so a Redis the +/// cache found wedged is fast-failed by every store too. +#[pyclass(frozen, name = "NativeRedisBreaker")] +pub(crate) struct NativeRedisBreaker(pub(crate) Breaker); + +impl NativeRedisBreaker { + fn build( + py: Python<'_>, + settings: &BreakerSettings, + endpoint: Option, + ) -> PyResult { + let config = settings.config()?; + let runtime = litellm_host_python::runtime_handle()?; + let events: Arc = + Arc::new(ServiceLoggingEvents::transitions_only(py, runtime)?); + Ok(Self(py.detach(|| { + hosts().breaker(endpoint, config, Some(events)) + }))) + } +} + +#[pymethods] +impl NativeRedisBreaker { + #[new] + fn new(py: Python<'_>, settings: BreakerSettings) -> PyResult { + Self::build(py, &settings, None) + } + + /// The breaker of the Redis server `redis_kwargs` names, which every cache and client on that server shares + /// whatever its db, credentials or namespace. The first caller's settings win. + #[staticmethod] + fn shared( + py: Python<'_>, + redis_kwargs: RedisSpec, + settings: BreakerSettings, + ) -> PyResult { + let endpoint = resolved_endpoint(py, &redis_kwargs)?; + Self::build(py, &settings, endpoint) + } + + fn admit(&self) -> Option { + self.0 + .admit(Instant::now()) + .ok() + .map(|admission| NativeAdmission(Mutex::new(Some(admission)))) + } + + fn failed(&self, timeout: bool) { + self.0.failed(failure(timeout), Instant::now()); + } + + #[getter] + fn state(&self) -> &'static str { + state_label(&self.0.state()) + } +} + +/// One admitted call. Settling it twice records nothing more, and dropping an unsettled probe reopens the breaker. +#[pyclass(frozen, name = "NativeAdmission")] +pub(crate) struct NativeAdmission(Mutex>); + +impl NativeAdmission { + fn take(&self) -> Option { + self.0.lock().unwrap_or_else(PoisonError::into_inner).take() + } +} + +#[pymethods] +impl NativeAdmission { + fn succeeded(&self) { + if let Some(admission) = self.take() { + admission.succeeded(); + } + } + + fn failed(&self, timeout: bool) { + if let Some(admission) = self.take() { + admission.failed(failure(timeout), Instant::now()); + } + } +} + +fn failure(timeout: bool) -> Failure { + match timeout { + true => Failure::Timeout, + false => Failure::Hard, + } +} + +#[pyfunction] +pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> { + py.detach(|| hosts().breakers()) + .iter() + .map(|breaker| state_label(&breaker.state())) + .collect() +} diff --git a/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs b/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs new file mode 100644 index 00000000000..21a37947c30 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs @@ -0,0 +1,47 @@ +use std::sync::Arc; + +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_redis_connect::{CredentialSource, CredentialStage, Error, Expiring, SourceFuture}; +use litellm_secret_value::SecretValue; +use pyo3::prelude::*; + +const MODULE: &str = "litellm.rust_bridge.host.redis_credentials"; + +/// A redis-py `CredentialProvider` the SDK caller passed in `cache_kwargs`, asked through the host +/// module the way redis-py's async connection asks it. +pub(crate) struct PythonCredentialSource { + provider: Arc>, + context: PythonContext, +} + +impl PythonCredentialSource { + pub(crate) fn new(provider: Py, context: PythonContext) -> Self { + Self { + provider: Arc::new(provider), + context, + } + } +} + +impl CredentialSource for PythonCredentialSource { + fn fetch(&self) -> SourceFuture<'_> { + let provider = Arc::clone(&self.provider); + let context = self.context.clone(); + Box::pin(async move { + let (username, password) = attach_blocking(context, move |py| { + py.import(MODULE)? + .getattr("credentials")? + .call1((provider.bind(py),))? + .extract::<(Option, String)>() + }) + .await + .and_then(|credentials| credentials) + .map_err(Error::credential(CredentialStage::Provider))?; + Ok(Expiring { + username, + password: SecretValue::new(password), + expires_at: None, + }) + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs new file mode 100644 index 00000000000..989f3dc9247 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs @@ -0,0 +1,337 @@ +mod credentials; + +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use litellm_circuit_breaker::Config as BreakerConfig; +use litellm_core_utils::settings::ProcessEnvironment; +use litellm_host_python::PythonContext; +use litellm_python_compat::number::{self, SecondsError}; +use litellm_redis_client::{RedisClient, RedisEvents}; +use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment}; +use litellm_redis_resolve::{Kwargs, PythonException}; +use litellm_redis_types::{Endpoint, ProviderHandle, Resolution}; +use pyo3::{ + exceptions::{PyOverflowError, PyTypeError, PyValueError}, + prelude::*, + types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple}, +}; + +use self::credentials::PythonCredentialSource; +use crate::redis::{breaker::BreakerSettings, events::ServiceLoggingEvents, hosts}; + +pyo3::import_exception!(redis.exceptions, RedisClusterException); + +/// The kwargs `_merge_redis_kwargs` returns for a `RedisCache`, which the resolver reads the way +/// redis-py's constructors do. A key whose value is a Python object, such as `retry`, reaches the +/// resolver by name only, so it can name the setting it cannot honor. A `credential_provider` object +/// is also kept, so the connector can ask it for credentials, and reaches the resolver as a handle. +pub(crate) struct RedisSpec { + kwargs: Kwargs, + credential_provider: Option>, +} + +impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec { + type Error = PyErr; + + fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult { + let py = kwargs.py(); + let source = PyDict::new(py); + let mapping = kwargs.cast::()?; + source.update(&mapping)?; + let plain = PyDict::new(py); + let mut host_objects = Vec::new(); + for (name, value) in source.iter() { + match value.is_none() || is_plain(&value) { + true => plain.set_item(&name, value)?, + false => host_objects.push(name.extract::()?), + } + } + let credential_provider = source + .get_item("credential_provider")? + .filter(|provider| !provider.is_none() && !is_plain(provider)) + .map(Bound::unbind); + let fields: BTreeMap = + litellm_host_python::from_py_argument(plain.as_any())?; + let objects = host_objects + .into_iter() + .map(|name| (name, serde_json::Value::Bool(true))); + Ok(Self { + kwargs: litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)), + credential_provider, + }) + } +} + +/// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built. +#[derive(FromPyObject)] +pub(crate) struct StoreSettings { + redis_kwargs: RedisSpec, + namespace: Option, + default_ttl: f64, + breaker: Option, +} + +fn is_plain(value: &Bound<'_, PyAny>) -> bool { + value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() +} + +/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool +/// and its breaker. A config redis-py refuses building its client raises the same class here. One it only fails on +/// once a command connects, where `RedisCache` logs every failed call, leaves the family on memory with an error. +fn shared_client( + py: Python<'_>, + spec: RedisSpec, + namespace: Option, + breaker: Option, +) -> PyResult> { + let environment: SharedEnvironment = Arc::new(ProcessEnvironment); + let logger = crate::logger::capture(py); + let resolved = match logger.scope(|| resolve_spec(py, &spec))? { + Ok(resolution) => { + logger.scope(|| { + for ignored in &resolution.ignored { + litellm_tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored: {}", + ignored.setting, + ignored.consequence + ); + } + }); + Resolution { + namespace, + ..resolution + } + } + Err(error) => return unresolved(py, &error), + }; + let runtime = litellm_host_python::runtime_handle()?; + let events: Arc = Arc::new(ServiceLoggingEvents::new(py, runtime.clone())?); + let host = HostServices { + events: Some(events), + http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?, + environment, + credential_provider: spec + .credential_provider + .as_ref() + .map(|provider| -> PyResult> { + Ok(Arc::new(PythonCredentialSource::new( + provider.clone_ref(py), + PythonContext::capture(py)?, + ))) + }) + .transpose()?, + }; + let connected = py.detach(|| { + let _runtime = runtime.enter(); + hosts().client(resolved, host, breaker) + }); + match connected { + Ok(client) => Ok(Some(client)), + Err(error) => Ok(unconnected(py, &error)), + } +} + +/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store +/// never sees the redis-py kwargs or the namespace. +#[pyclass(frozen, name = "NativeRedisClient")] +pub(crate) struct NativeRedisClient { + pub(crate) client: RedisClient, + /// `RedisCache.default_ttl`, the TTL a store's write takes when its caller gives none. + pub(crate) default_ttl: Duration, +} + +#[pymethods] +impl NativeRedisClient { + /// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory. + #[staticmethod] + fn connect(py: Python<'_>, settings: StoreSettings) -> PyResult> { + let StoreSettings { + redis_kwargs, + namespace, + default_ttl, + breaker, + } = settings; + let default_ttl = number::seconds(default_ttl).map_err(seconds_error)?; + let breaker = breaker.as_ref().map(BreakerSettings::config).transpose()?; + Ok( + shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self { + client, + default_ttl, + }), + ) + } + + #[getter] + fn default_ttl_seconds(&self) -> f64 { + self.default_ttl.as_secs_f64() + } +} + +/// The Redis server `spec` names, whatever its db, credentials or namespace, or none when the Rust client cannot reach it. +pub(crate) fn resolved_endpoint(py: Python<'_>, spec: &RedisSpec) -> PyResult> { + Ok(resolve_spec(py, spec)? + .ok() + .and_then(|resolution| resolution.target.endpoint())) +} + +// A live registry entry's client keeps its provider alive, so no live entry shares its address +fn resolve_spec( + py: Python<'_>, + spec: &RedisSpec, +) -> PyResult> { + let handle = spec + .credential_provider + .as_ref() + .map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64)); + resolve(&spec.kwargs, handle, &py.import("os")?.getattr("environ")?) +} + +/// What redis-py makes of the kwargs, with `REDIS_USERNAME` read from `environ` the way Python's Azure AD auth +/// reads `os.environ`. +fn resolve( + kwargs: &Kwargs, + provider: Option, + environ: &Bound<'_, PyAny>, +) -> PyResult> { + let redis_username: Option = environ + .call_method1("get", ("REDIS_USERNAME",))? + .extract()?; + Ok(litellm_redis_resolve::interpret( + kwargs, + provider, + redis_username.as_deref(), + )) +} + +fn unresolved( + py: Python<'_>, + error: &litellm_redis_resolve::Error, +) -> PyResult> { + let message = || litellm_error::Report::new(error).to_string(); + match error.raised_building_the_python_client() { + Some(PythonException::ValueError) => Err(PyValueError::new_err(message())), + Some(PythonException::TypeError) => Err(PyTypeError::new_err(message())), + Some(PythonException::RedisClusterException) => { + Err(RedisClusterException::new_err(message())) + } + None => Ok(memory_only(py, error)), + } +} + +// redis-py reads TLS files and fetches credentials only once a command connects, so none of these fails building +// its client +fn unconnected(py: Python<'_>, error: &litellm_redis_connect::Error) -> Option { + match error { + litellm_redis_connect::Error::Connect(_) + | litellm_redis_connect::Error::NoHostCredentialProvider + | litellm_redis_connect::Error::ProcessEnvironmentRequired { .. } + | litellm_redis_connect::Error::Credential { .. } + | litellm_redis_connect::Error::CredentialTimeout { .. } + | litellm_redis_connect::Error::ExpiredCredential + | litellm_redis_connect::Error::MissingCredential { .. } => memory_only(py, error), + } +} + +fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option { + let report = litellm_error::Report::new(error).to_string(); + crate::logger::capture(py).scope(|| { + litellm_tracing::error!( + "Redis is not usable from the Rust stores, their state stays in memory: {report}" + ) + }); + None +} + +fn seconds_error(error: SecondsError) -> PyErr { + match error { + SecondsError::Overflow => PyOverflowError::new_err(error.to_string()), + SecondsError::NotANumber => PyValueError::new_err(error.to_string()), + } +} + +#[cfg(test)] +mod tests { + use std::ffi::CStr; + + use litellm_redis_types::Auth; + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn kwargs(fields: serde_json::Value) -> Kwargs { + let fields: BTreeMap = serde_json::from_value(fields).unwrap(); + litellm_redis_resolve::kwargs(fields) + } + + #[rstest] + #[case::set(c"environ = {'REDIS_USERNAME': 'azure-user'}", Some("azure-user"))] + #[case::empty_is_unset(c"environ = {'REDIS_USERNAME': ''}", None)] + fn azure_auth_takes_its_username_from_the_environ_like_python( + #[case] source: &CStr, + #[case] expected: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + let environ = locals.get_item("environ").unwrap().unwrap(); + let kwargs = kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"})); + + let resolved = resolve(&kwargs, None, &environ).unwrap().unwrap(); + + assert!( + matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected), + "{:?}", + resolved.auth + ); + }); + } + + #[rstest] + #[case::no_target(json!({"port": 6379}), Some("ValueError"))] + #[case::unenforceable(json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), Some("TypeError"))] + #[case::fails_only_once_a_command_connects( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": "None"}), + None + )] + fn a_config_raises_where_redis_py_raises_building_its_client( + #[case] fields: serde_json::Value, + #[case] raised: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let error = litellm_redis_resolve::interpret(&kwargs(fields), None, None).unwrap_err(); + + let outcome = unresolved(py, &error) + .map(|client| client.is_some()) + .map_err(|error| error.get_type(py).name().unwrap().to_string()); + + assert_eq!( + outcome, + raised.map_or(Ok(false), |class| Err(class.to_owned())) + ); + }); + } + + #[rstest] + #[case::overflow(f64::INFINITY, "OverflowError")] + #[case::not_a_number(f64::NAN, "ValueError")] + fn a_default_ttl_python_cannot_make_an_int_raises_like_int( + #[case] seconds: f64, + #[case] raised: &str, + ) { + Python::initialize(); + Python::attach(|py| { + let error = number::seconds(seconds).map_err(seconds_error).unwrap_err(); + + assert_eq!(error.get_type(py).name().unwrap().to_string(), raised); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/events.rs b/litellm-rust/crates/python-bridge/src/redis/events.rs new file mode 100644 index 00000000000..a281d78d5a6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/events.rs @@ -0,0 +1,403 @@ +use std::{ + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_redis_client::{ + CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged, +}; +use pyo3::{prelude::*, types::PyModule}; +use tokio::runtime::Handle; + +/// The asyncio loop a Python caller ran on. `RedisCache` scheduled its service events there with +/// `asyncio.create_task`, so the native events go there too instead of a new loop per event. +#[derive(Clone)] +pub(crate) struct EventLoop(Arc>); + +impl EventLoop { + pub(crate) fn running(py: Python<'_>) -> Option { + pyo3_async_runtimes::get_running_loop(py) + .ok() + .map(|event_loop| Self(Arc::new(event_loop.unbind()))) + } +} + +/// The Python caller of one store operation. Its context carries `with_service_target` and the +/// ambient OTel span, and its `call_type` is the label `RedisCache` reports, so dashboards keep +/// their series. +pub(crate) struct RedisCall { + pub(crate) context: PythonContext, + pub(crate) event_loop: Option, + pub(crate) call_type: &'static str, +} + +impl RedisCall { + pub(crate) fn tag(self) -> CallTag { + Arc::new(self) + } + + pub(crate) async fn scope(self, operation: F) -> F::Output { + tagged(self.tag(), operation).await + } +} + +/// Reports every native Redis call and breaker transition to `ServiceLogging`. +pub(crate) struct ServiceLoggingEvents { + record_call: Arc>, + record_breaker: Arc>, + context: PythonContext, + /// The loop the client was built on, for calls whose caller had none running. + event_loop: Option, + runtime: Handle, +} + +impl ServiceLoggingEvents { + const MODULE: &str = "litellm.rust_bridge.host.redis_events"; + + pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult { + Self::from_module( + py, + &py.import(Self::MODULE)?, + runtime, + EventLoop::running(py), + ) + } + + /// For a breaker, whose transitions never run on the caller's loop, so building it never looks one up. + pub(crate) fn transitions_only(py: Python<'_>, runtime: Handle) -> PyResult { + Self::from_module(py, &py.import(Self::MODULE)?, runtime, None) + } + + fn from_module( + py: Python<'_>, + module: &Bound<'_, PyModule>, + runtime: Handle, + event_loop: Option, + ) -> PyResult { + Ok(Self { + record_call: Arc::new(module.getattr("record_redis_call")?.unbind()), + record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()), + context: PythonContext::capture(py)?, + event_loop, + runtime, + }) + } + + fn report(&self, context: PythonContext, record: Arc>, args: ReportArgs) { + self.runtime.spawn(async move { + let reported = attach_blocking(context, move |py| match args { + ReportArgs::Call(call) => record + .bind(py) + .call1(( + call.event_loop + .as_ref() + .map(|event_loop| event_loop.0.bind(py)), + call.call_type, + call.start_time, + call.end_time, + call.failure, + )) + .map(drop), + ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop), + }) + .await + .and_then(|reported| reported); + if let Err(error) = reported { + tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed"); + } + }); + } +} + +/// One call as `record_redis_call` takes it, its times in `time.time()` seconds. +struct CallReport { + event_loop: Option, + call_type: &'static str, + start_time: f64, + end_time: f64, + failure: Option<&'static str>, +} + +enum ReportArgs { + Call(CallReport), + Breaker(&'static str), +} + +/// The wall-clock start and end of a call that just finished after `duration`. +fn wall_clock_span(duration: Duration) -> (f64, f64) { + let end = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default(); + ( + end.saturating_sub(duration).as_secs_f64(), + end.as_secs_f64(), + ) +} + +fn operation_call_type(operation: Operation) -> &'static str { + match operation { + Operation::Command(name) => name, + Operation::Script => "eval", + Operation::Pipeline { .. } => "pipeline", + } +} + +fn failure_label(outcome: Outcome) -> Option<&'static str> { + match outcome { + Outcome::Succeeded => None, + Outcome::Failed(FailureKind::Timeout) => Some("timeout"), + Outcome::Failed(FailureKind::Transport) => Some("transport"), + Outcome::Failed(FailureKind::Server) => Some("server"), + Outcome::Failed(FailureKind::Decode) => Some("decode"), + Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"), + } +} + +pub(crate) fn state_label(state: &State) -> &'static str { + match state { + State::Closed { .. } => "closed", + State::Open { .. } => "open", + State::HalfOpen => "half_open", + } +} + +impl RedisEvents for ServiceLoggingEvents { + fn call(&self, event: &CallEvent) { + let scoped = current_tag().and_then(|tag| tag.downcast::().ok()); + let (context, event_loop, call_type) = match scoped { + Some(call) => ( + call.context.clone(), + call.event_loop.clone().or_else(|| self.event_loop.clone()), + call.call_type, + ), + None => ( + self.context.clone(), + self.event_loop.clone(), + operation_call_type(event.operation), + ), + }; + let (start_time, end_time) = wall_clock_span(event.duration); + self.report( + context, + self.record_call.clone(), + ReportArgs::Call(CallReport { + event_loop, + call_type, + start_time, + end_time, + failure: failure_label(event.outcome), + }), + ); + } + + fn breaker(&self, transition: &Transition) { + self.report( + self.context.clone(), + self.record_breaker.clone(), + ReportArgs::Breaker(state_label(&transition.to)), + ); + } +} + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use litellm_host_python::PythonContext; + use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; + use pyo3::{ + prelude::*, + types::{PyCFunction, PyModule, PyString}, + }; + use rstest::rstest; + use tokio::runtime::Handle; + + use super::{EventLoop, RedisCall, ServiceLoggingEvents}; + + const RECORDER: &str = " +import contextvars, time +target = contextvars.ContextVar('target', default=None) +seen = [] +def record_redis_call(event_loop, call_type, start_time, end_time, failure): + just_now = abs(time.time() - end_time) < 60 + seen.append((event_loop, call_type, round(end_time - start_time, 3), just_now, failure, target.get())) +def record_redis_breaker(state): + seen.append(('breaker', state)) +"; + + fn stand_in_loop(py: Python<'_>, name: &str) -> EventLoop { + EventLoop(Arc::new(PyString::new(py, name).into_any().unbind())) + } + + fn recorder(name: &str, client_loop: Option<&str>) -> (ServiceLoggingEvents, Py) { + Python::initialize(); + Python::attach(|py| { + let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap(); + let module = + PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name) + .unwrap(); + let events = ServiceLoggingEvents::from_module( + py, + &module, + Handle::current(), + client_loop.map(|client_loop| stand_in_loop(py, client_loop)), + ) + .unwrap(); + (events, module.unbind()) + }) + } + + async fn seen(module: &Py, count: usize) -> String { + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + loop { + let seen: (usize, String) = Python::attach(|py| { + let seen = module.bind(py).getattr("seen").unwrap(); + (seen.len().unwrap(), seen.repr().unwrap().to_string()) + }); + if seen.0 >= count || tokio::time::Instant::now() > deadline { + return seen.1; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + } + + fn scoped_call(namespace: &Py, caller_loop: Option<&str>) -> RedisCall { + Python::attach(|py| { + namespace + .bind(py) + .getattr("target") + .unwrap() + .call_method1("set", ("spend_counter",)) + .unwrap(); + RedisCall { + context: PythonContext::capture(py).unwrap(), + event_loop: caller_loop.map(|caller_loop| stand_in_loop(py, caller_loop)), + call_type: "async_increment", + } + }) + } + + #[rstest] + #[case::on_the_callers_loop(Some("caller loop"), "caller loop")] + #[case::on_the_clients_loop_when_the_caller_had_none(None, "client loop")] + #[tokio::test(flavor = "multi_thread")] + async fn a_scoped_call_reports_the_store_labels_and_its_times_in_the_callers_context( + #[case] caller_loop: Option<&str>, + #[case] expected_loop: &str, + ) { + let (events, namespace) = recorder( + expected_loop.replace(' ', "_").as_str(), + Some("client loop"), + ); + let call = scoped_call(&namespace, caller_loop); + let event = CallEvent { + operation: Operation::Command("INCRBYFLOAT"), + duration: Duration::from_millis(250), + outcome: Outcome::Failed(FailureKind::Timeout), + }; + + call.scope(async { events.call(&event) }).await; + + assert_eq!( + seen(&namespace, 1).await, + format!( + "[('{expected_loop}', 'async_increment', 0.25, True, 'timeout', 'spend_counter')]" + ) + ); + } + + #[rstest] + #[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")] + #[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")] + #[tokio::test(flavor = "multi_thread")] + async fn breaker_transitions_report_python_state_names( + #[case] to: litellm_circuit_breaker::State, + #[case] expected: &str, + ) { + let (events, namespace) = recorder(expected, None); + + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to, + generation: 1, + }); + + assert_eq!( + seen(&namespace, 1).await, + format!("[('breaker', '{expected}')]") + ); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn a_transition_on_a_thread_without_a_runtime_is_still_reported() { + let (events, namespace) = recorder("plain_thread", None); + + std::thread::spawn(move || { + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to: litellm_circuit_breaker::State::HalfOpen, + generation: 1, + }); + }) + .join() + .unwrap(); + + assert_eq!(seen(&namespace, 1).await, "[('breaker', 'half_open')]"); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn an_unscoped_call_falls_back_to_the_command_name_and_the_clients_loop() { + let (events, namespace) = recorder("unscoped", None); + + events.call(&CallEvent { + operation: Operation::Pipeline { ops: 3 }, + duration: Duration::from_millis(5), + outcome: Outcome::Succeeded, + }); + + assert_eq!( + seen(&namespace, 1).await, + "[(None, 'pipeline', 0.005, True, None, None)]" + ); + } + + #[rstest] + fn running_is_the_loop_the_caller_runs_on_and_none_outside_one() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::from_code( + py, + c"import asyncio\nasync def probe_inside(probe):\n return probe() is asyncio.get_running_loop()\n", + c"event_loop_probe", + c"event_loop_probe", + ) + .unwrap(); + let probe = PyCFunction::new_closure(py, None, None, |_, _| { + Python::attach(|py| -> PyResult>> { + Ok(EventLoop::running(py).map(|event_loop| event_loop.0.clone_ref(py))) + }) + }) + .unwrap(); + let coroutine = module + .getattr("probe_inside") + .unwrap() + .call1((probe,)) + .unwrap(); + + let inside: bool = py + .import("asyncio") + .unwrap() + .call_method1("run", (coroutine,)) + .unwrap() + .extract() + .unwrap(); + + assert!(inside); + assert!(EventLoop::running(py).is_none()); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/mod.rs b/litellm-rust/crates/python-bridge/src/redis/mod.rs new file mode 100644 index 00000000000..808f930dd20 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/mod.rs @@ -0,0 +1,15 @@ +use std::sync::LazyLock; + +use litellm_redis_connect::RedisHosts; + +pub(crate) mod breaker; +pub(crate) mod client; +mod events; +pub(crate) mod store; + +static HOSTS: LazyLock = LazyLock::new(RedisHosts::default); + +/// The Redis servers this process talks to, shared by every `RedisCache` and native store in it. +pub(crate) fn hosts() -> &'static RedisHosts { + &HOSTS +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs new file mode 100644 index 00000000000..8d7788b136e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs @@ -0,0 +1,265 @@ +use std::sync::Arc; + +use litellm_python_compat::Value; +use litellm_redis_store::auth_cache::{AuthCache, Settings}; +use pyo3::prelude::*; + +use super::{ + NativePending, PyNumber, kind_error, pending, redis_async, redis_sync, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, + value::PyValue, +}; +use crate::redis::client::NativeRedisClient; + +fn values(values: Vec>) -> Vec> { + values.into_iter().map(|value| value.map(PyValue)).collect() +} + +/// `UserApiKeyCache` with `enable_redis_auth_cache` on. The sync methods back `DualCache`'s sync +/// API and block with the GIL released, like redis-py's sync client. +#[pyclass(frozen, name = "NativeAuthCache")] +pub(crate) struct NativeAuthCache(Arc); + +#[pymethods] +impl NativeAuthCache { + #[new] + #[pyo3(signature = (client, *, memory_entries, memory_ttl_seconds = None))] + fn new( + client: Option>, + memory_entries: usize, + memory_ttl_seconds: Option, + ) -> PyResult { + Ok(Self(Arc::new(AuthCache::new( + client.map(|client| remote(&client)), + Settings { + memory_entries, + memory_ttl: ttl(memory_ttl_seconds)?, + }, + )))) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + /// `request` answers a key its pre-call reads already found absent without another `GET`. + #[pyo3(signature = (key, request = None))] + fn get<'py>( + &self, + py: Python<'py>, + key: String, + request: Option<&NativeRequestRedis>, + ) -> PyResult> { + let store = self.0.clone(); + let request = request.map(|request| request.0.clone()); + redis_async(py, "async_get_cache", kind_error, async move { + store + .get(request.as_ref(), &key) + .await + .map(|value| value.map(PyValue)) + }) + } + + fn get_sync(&self, py: Python<'_>, key: String) -> PyResult> { + let store = self.0.clone(); + redis_sync(py, "get_cache", kind_error, async move { + store.get(None, &key).await.map(|value| value.map(PyValue)) + }) + } + + /// `update_cache_ttl` and `update_in_memory_max_size`: both partitions keep what they hold. + #[pyo3(signature = (*, memory_entries, memory_ttl_seconds = None))] + fn reconfigure(&self, memory_entries: usize, memory_ttl_seconds: Option) -> PyResult<()> { + self.0.reconfigure(Settings { + memory_entries, + memory_ttl: ttl(memory_ttl_seconds)?, + }); + Ok(()) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set<'py>( + &self, + py: Python<'py>, + key: String, + value: PyValue, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.set(&key, value.0, ttl).await + }) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set_sync( + &self, + py: Python<'_>, + key: String, + value: PyValue, + ttl_seconds: Option, + ) -> PyResult<()> { + let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; + redis_sync(py, "set_cache", kind_error, async move { + store.set_str(&key, value.0, ttl).await + }) + } + + /// Memory now and the `SET` on the batch `send` picks. + #[pyo3(signature = (key, value, ttl_seconds, send))] + fn declare_set( + &self, + key: String, + value: PyValue, + ttl_seconds: Option, + send: &NativeSend, + ) -> PyResult { + let written = self + .0 + .declare_set(send.send(), &key, value.0, ttl(ttl_seconds)?) + .map_err(kind_error)?; + Ok(pending( + "async_set_cache", + kind_error, + written, + std::convert::identity, + )) + } + + fn declare_delete(&self, key: String, send: &NativeSend) -> PyResult { + let deleted = self + .0 + .declare_delete(send.send(), &key) + .map_err(kind_error)?; + Ok(pending( + "async_delete_cache", + kind_error, + deleted, + std::convert::identity, + )) + } + + #[pyo3(signature = (entries, ttl_seconds = None))] + fn set_many<'py>( + &self, + py: Python<'py>, + entries: Vec<(String, PyValue)>, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let entries = entries + .into_iter() + .map(|(key, value)| (key, value.0)) + .collect(); + let ttl = ttl(ttl_seconds)?; + redis_async(py, "async_set_cache_pipeline", kind_error, async move { + store.set_many(entries, ttl).await + }) + } + + /// On a request batch it returns before anything is sent, and `Now` is a pipeline of its own. + fn write_back<'py>( + &self, + py: Python<'py>, + entries: Vec<(String, PyValue, Option)>, + send: &NativeSend, + ) -> PyResult> { + let entries = entries + .into_iter() + .map(|(key, value, seconds)| Ok((key, value.0, ttl(seconds)?))) + .collect::>()?; + let written = self + .0 + .write_back(send.send(), entries) + .map_err(kind_error)?; + redis_async( + py, + "async_set_cache_pipeline_with_ttls", + kind_error, + written, + ) + } + + fn delete<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_delete_cache", kind_error, async move { + store.delete(&keys).await + }) + } + + fn delete_sync(&self, py: Python<'_>, keys: Vec) -> PyResult<()> { + let store = self.0.clone(); + redis_sync(py, "delete_cache", kind_error, async move { + store.delete(&keys).await + }) + } + + /// One `MGET` of keys the caller found missing in memory, on the batch `send` picks; `None` + /// without Redis. + fn get_many(&self, keys: Vec, send: &NativeSend) -> Option { + let read = self.0.get_many(send.send(), &keys)?; + Some(pending("async_batch_get_cache", kind_error, read, values)) + } + + fn increment_if_present<'py>( + &self, + py: Python<'py>, + key: String, + amount: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_increment_if_present", kind_error, async move { + store.increment_if_present(&key, amount.0).await + }) + } + + fn get_local(&self, key: &str) -> PyResult> { + Ok(self.0.get_local(key).map_err(kind_error)?.map(PyValue)) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set_local(&self, key: &str, value: PyValue, ttl_seconds: Option) -> PyResult<()> { + self.0 + .set_local(key, value.0, ttl(ttl_seconds)?) + .map_err(kind_error) + } + + fn delete_local(&self, key: &str) -> PyResult<()> { + self.0.delete_local(key).map_err(kind_error) + } + + fn flush_local(&self) -> PyResult<()> { + self.0.flush_local().map_err(kind_error) + } + + fn len_local(&self) -> PyResult { + let (objects, key_objects) = self.0.len_local().map_err(kind_error)?; + Ok(PyLocalSizes { + objects, + key_objects, + }) + } +} + +/// Entries the management-object and the user-key partitions hold in memory. +#[pyclass(frozen, eq, get_all, name = "LocalSizes")] +#[derive(PartialEq)] +pub(crate) struct PyLocalSizes { + objects: usize, + key_objects: usize, +} + +#[pymethods] +impl PyLocalSizes { + #[new] + fn new(objects: usize, key_objects: usize) -> Self { + Self { + objects, + key_objects, + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs new file mode 100644 index 00000000000..d9c3aabb814 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -0,0 +1,271 @@ +pub(crate) mod auth_cache; +mod policy; +pub(crate) mod rate_limits; +pub(crate) mod request; +pub(crate) mod router_cooldowns; +pub(crate) mod router_usage; +pub(crate) mod sensitive_route_pins; +pub(crate) mod spend_counters; +mod value; + +use value::PyValue; + +use std::{ + pin::Pin, + sync::{Mutex, PoisonError}, + time::Duration, +}; + +use litellm_host_python::PythonContext; +use litellm_python_compat::number::{SecondsError, seconds}; +use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote}; +use pyo3::{ + IntoPyObjectExt, + exceptions::{PyOverflowError, PyRuntimeError, PyTypeError, PyValueError}, + prelude::*, + types::{PyBool, PyFloat, PyInt}, +}; + +use crate::execution::{run_async_value, run_sync_value}; +use crate::redis::{ + client::NativeRedisClient, + events::{EventLoop, RedisCall}, +}; + +/// The labels `RedisCache` reports for this operation, in the caller's context so its +/// `service_target` reaches the service event, dispatched on the loop the caller runs on. +fn call(py: Python<'_>, call_type: &'static str) -> PyResult { + Ok(RedisCall { + context: PythonContext::capture(py)?, + event_loop: EventLoop::running(py), + call_type, + }) +} + +/// Runs one Redis operation as a Python awaitable, reported to `ServiceLogging` as `call_type`. +fn redis_async<'py, T, E: 'static>( + py: Python<'py>, + call_type: &'static str, + error: fn(E) -> PyErr, + operation: impl Future> + Send + 'static, +) -> PyResult> +where + T: for<'a> IntoPyObject<'a> + Send + 'static, +{ + let scope = call(py, call_type)?; + run_async_value( + py, + async move { scope.scope(operation).await.map_err(error) }, + ) +} + +/// Runs one Redis operation with the GIL released, for the router's sync paths. +fn redis_sync( + py: Python<'_>, + call_type: &'static str, + error: fn(E) -> PyErr, + operation: impl Future> + Send + 'static, +) -> PyResult +where + T: Send + 'static, +{ + let scope = call(py, call_type)?; + run_sync_value( + py, + async move { scope.scope(operation).await.map_err(error) }, + ) +} + +/// A TTL in seconds, raising what `int(seconds)` raises for infinity and NaN. +fn duration(value: f64) -> PyResult { + seconds(value).map_err(|error| match error { + SecondsError::Overflow => PyOverflowError::new_err(error.to_string()), + SecondsError::NotANumber => PyValueError::new_err(error.to_string()), + }) +} + +fn ttl(value: Option) -> PyResult> { + value.map(duration).transpose() +} + +fn remote(client: &NativeRedisClient) -> Remote { + Remote { + client: client.client.clone(), + default_ttl: client.default_ttl, + } +} + +type Resolve = Pin>> + Send>>; + +/// A Redis call declared on a request's batch. `resolve` flushes that batch if it has not gone +/// out and answers the call, once. +#[pyclass(frozen, name = "NativePending")] +pub(crate) struct NativePending { + resolve: Mutex>, + call_type: &'static str, +} + +impl NativePending { + pub(crate) fn new( + call_type: &'static str, + error: fn(E) -> PyErr, + reply: impl Future> + Send + 'static, + ) -> Self + where + T: for<'a> IntoPyObject<'a> + Send + 'static, + { + let resolve = async move { + let reply = reply.await.map_err(error)?; + Python::attach(|py| reply.into_py_any(py)) + }; + Self { + resolve: Mutex::new(Some(Box::pin(resolve))), + call_type, + } + } +} + +#[pymethods] +impl NativePending { + fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { + let resolve = self + .resolve + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .ok_or_else(|| PyRuntimeError::new_err("this Redis call was already resolved"))?; + redis_async(py, self.call_type, std::convert::identity, resolve) + } +} + +pyo3::create_exception!( + _native, + NativeRedisStoreError, + pyo3::exceptions::PyException, + "A native Redis store operation failed. Args are (kind, message), kind spelled as litellm-redis-store's ErrorKind." +); + +/// What every native store raises, its kind telling a caller which failures it may drop. +fn kind_error(error: Error) -> PyErr { + let kind: &'static str = error.kind().into(); + NativeRedisStoreError::new_err((kind, litellm_error::Report::new(&error).to_string())) +} + +/// A store op declared on the batch the caller's `NativeSend` picked, answered as `reply` shapes it. +fn pending( + call_type: &'static str, + error: fn(Error) -> PyErr, + declared: Pending, + reply: fn(T) -> R, +) -> NativePending +where + T: Send + 'static, + R: for<'a> IntoPyObject<'a> + Send + 'static, +{ + NativePending::new(call_type, error, async move { declared.await.map(reply) }) +} + +/// A Python `int` or `float` across the bridge, keeping which one it is. +#[derive(Clone, Copy, PartialEq)] +pub(crate) struct PyNumber(Number); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + if value.is_instance_of::() && !value.is_instance_of::() { + return Ok(Self(Number::Int(value.extract()?))); + } + if value.is_instance_of::() { + return Ok(Self(Number::Float(value.extract()?))); + } + Err(PyTypeError::new_err("expected an int or a float")) + } +} + +impl<'py> IntoPyObject<'py> for PyNumber { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Ok(match self.0 { + Number::Int(value) => value.into_pyobject(py)?.into_any(), + Number::Float(value) => value.into_pyobject(py)?.into_any(), + }) + } +} + +/// A store enum read from its wire spelling, so the bridge names no variant. +pub(crate) struct Spelled(pub(crate) T); + +impl<'a, 'py, T: std::str::FromStr> FromPyObject<'a, 'py> for Spelled { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let spelling: String = value.extract()?; + spelling + .parse() + .map(Self) + .map_err(|_| PyValueError::new_err(format!("unknown spelling {spelling:?}"))) + } +} + +/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue` +/// so the bridge never names a field. A shape the store refuses is `ValueError`. +pub(crate) struct PyData(pub(crate) T); + +impl<'a, 'py, T: FromValue> FromPyObject<'a, 'py> for PyData { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let PyValue(value) = value.extract()?; + T::from_value(&value) + .map(Self) + .map_err(|error| PyValueError::new_err(error.to_string())) + } +} + +impl<'py, T: IntoValue> IntoPyObject<'py> for PyData { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + PyValue(self.0.into_value()).into_pyobject(py) + } +} + +fn unwrap_all(values: Vec>) -> Vec { + values.into_iter().map(|value| value.0).collect() +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pyo3::prelude::*; + use rstest::rstest; + + use super::ttl; + + #[rstest] + #[case::infinity(f64::INFINITY, "OverflowError")] + #[case::nan(f64::NAN, "ValueError")] + fn a_ttl_int_refuses_raises_what_int_raises(#[case] seconds: f64, #[case] raised: &str) { + Python::initialize(); + Python::attach(|py| { + let error = ttl(Some(seconds)).unwrap_err(); + + assert_eq!(error.get_type(py).name().unwrap().to_string(), raised); + }); + } + + #[rstest] + #[case::fraction(Some(1.5), Some(Duration::from_millis(1500)))] + #[case::negative_clamps_to_zero(Some(-1.0), Some(Duration::ZERO))] + #[case::absent(None, None)] + fn a_finite_ttl_converts(#[case] seconds: Option, #[case] expected: Option) { + assert_eq!(ttl(seconds).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/policy.rs b/litellm-rust/crates/python-bridge/src/redis/store/policy.rs new file mode 100644 index 00000000000..ffdf7b97a2e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/policy.rs @@ -0,0 +1,151 @@ +use std::time::Duration; + +use litellm_redis_store::TierPolicy; +use pyo3::prelude::*; + +use crate::{ + coercion::{Field, FieldSpec, ProjectionError}, + python_settings::{PythonSettings, Snapshot}, +}; + +const COOLDOWN_REFRESH: FieldSpec = + FieldSpec::new("cooldown_redis_read_interval", seconds); +const USAGE_REFRESH: FieldSpec = FieldSpec::new("usage_redis_read_interval", seconds); +const TRACKED_READS: FieldSpec = FieldSpec::new("max_redis_batch_cache_size", count); +const MEMORY_TTL: FieldSpec> = + FieldSpec::new("default_in_memory_ttl", optional_seconds); + +/// The tier values this Python process runs with, from its environment and `litellm` globals. +#[derive(Debug)] +pub(crate) struct RouterTierPolicies { + pub(crate) cooldowns: TierPolicy, + pub(crate) usage: TierPolicy, +} + +pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult { + Ok(project(&PythonSettings::RouterTiers.read(py)?)?) +} + +fn project(snapshot: &Snapshot<'_>) -> Result { + let tracked_reads = snapshot.read(&TRACKED_READS)?; + let default_in_memory_ttl = snapshot.read(&MEMORY_TTL)?; + Ok(RouterTierPolicies { + cooldowns: TierPolicy { + refresh_after: snapshot.read(&COOLDOWN_REFRESH)?, + tracked_reads, + default_in_memory_ttl, + }, + usage: TierPolicy { + refresh_after: snapshot.read(&USAGE_REFRESH)?, + tracked_reads, + default_in_memory_ttl, + }, + }) +} + +fn seconds(field: &Field<'_>) -> Result { + match field.value().extract::() { + Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => { + Ok(Duration::from_secs_f64(seconds)) + } + _ => Err(field.invalid("a non-negative number of seconds")), + } +} + +fn optional_seconds(field: &Field<'_>) -> Result, ProjectionError> { + match field.value().is_none() { + true => Ok(None), + false => seconds(field).map(Some), + } +} + +fn count(field: &Field<'_>) -> Result { + field + .value() + .extract::() + .map_err(|_| field.invalid("a non-negative integer")) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pyo3::{prelude::*, types::PyDict}; + + use super::project; + use crate::python_settings::PythonSettings; + + fn projected(py: Python<'_>, fields: &str) -> Result { + let locals = PyDict::new(py); + let source = std::ffi::CString::new(format!( + "import types\nsettings = types.SimpleNamespace({fields})" + )) + .unwrap(); + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + let snapshot = + PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap()); + project(&snapshot).map_err(|error| PyErr::from(error).to_string()) + } + + #[test] + fn the_host_values_are_each_familys_refresh_and_tracked_reads() { + Python::initialize(); + Python::attach(|py| { + let policies = projected( + py, + "cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42, default_in_memory_ttl=None", + ) + .unwrap(); + assert_eq!( + policies.cooldowns.refresh_after, + Duration::from_millis(2500) + ); + assert_eq!(policies.usage.refresh_after, Duration::from_secs(30)); + assert_eq!(policies.cooldowns.tracked_reads, 42); + assert_eq!(policies.usage.tracked_reads, 42); + assert_eq!(policies.cooldowns.default_in_memory_ttl, None); + }); + } + + #[test] + fn a_set_default_in_memory_ttl_reaches_both_familys_backfill() { + Python::initialize(); + Python::attach(|py| { + let policies = projected( + py, + "cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=7.5", + ) + .unwrap(); + assert_eq!( + policies.cooldowns.default_in_memory_ttl, + Some(Duration::from_millis(7500)) + ); + assert_eq!( + policies.usage.default_in_memory_ttl, + Some(Duration::from_millis(7500)) + ); + }); + } + + #[test] + fn a_value_python_would_not_run_with_is_a_configuration_error() { + Python::initialize(); + Python::attach(|py| { + let negative = projected( + py, + "cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=None", + ) + .unwrap_err(); + let text = projected( + py, + "cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many', default_in_memory_ttl=None", + ) + .unwrap_err(); + assert!( + negative.contains("cooldown_redis_read_interval"), + "{negative}" + ); + assert!(text.contains("max_redis_batch_cache_size"), "{text}"); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs new file mode 100644 index 00000000000..ca95e1530f4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs @@ -0,0 +1,577 @@ +use std::{sync::Arc, time::Duration}; + +use indexmap::IndexMap; +use litellm_redis_store::{ + Pending, + rate_limits::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge, GuardedIncrement, + GuardedReply, Increment, Kind, Local, RateLimits, Scope, WindowCount, + }, +}; +use pyo3::{exceptions::PyValueError, prelude::*}; + +use super::{ + NativePending, PyNumber, Spelled, kind_error, redis_async, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, +}; +use crate::redis::client::NativeRedisClient; + +/// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`; +/// every Redis call is a future. +#[pyclass(frozen, name = "NativeRateLimits")] +pub(crate) struct NativeRateLimits(Arc); + +/// An `int` or `float` counter, a `str` window start, or a `{slot_id: acquired_at}` registry. +#[derive(FromPyObject)] +enum PyLocal { + Number(PyNumber), + Text(String), + Slots(IndexMap), +} + +impl From for Local { + fn from(local: PyLocal) -> Self { + match local { + PyLocal::Number(number) => Self::Number(number.0), + PyLocal::Text(text) => Self::Text(text), + PyLocal::Slots(slots) => Self::Slots(slots.into_iter().collect()), + } + } +} + +/// What the memory tier holds, back as the Python value it was written as. +struct LocalValue(Local); + +impl<'py> IntoPyObject<'py> for LocalValue { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Ok(match self.0 { + Local::Number(number) => PyNumber(number).into_pyobject(py)?, + Local::Text(text) => text.into_pyobject(py)?.into_any(), + Local::Slots(slots) => slots + .into_iter() + .collect::>() + .into_pyobject(py)? + .into_any(), + }) + } +} + +/// `litellm.redis.stores.rate_limits.RateLimitScope`. +#[derive(FromPyObject)] +struct PyScope { + descriptor_key: String, + descriptor_value: String, +} + +impl From for Scope { + fn from(scope: PyScope) -> Self { + Self { + descriptor: scope.descriptor_key, + value: scope.descriptor_value, + } + } +} + +/// `RateLimitCounter`. +#[derive(FromPyObject)] +struct PyCounter { + scope: PyScope, + kind: Spelled, +} + +impl From for Counter { + fn from(counter: PyCounter) -> Self { + Self { + scope: counter.scope.into(), + kind: counter.kind.0, + } + } +} + +/// `RateLimitWindow`, `RateLimitCounter` or `RateLimitGauge`, told apart by the key suffix each names. +#[derive(FromPyObject)] +struct PyEntry { + suffix: String, + scope: PyScope, +} + +impl TryFrom for Entry { + type Error = PyErr; + + fn try_from(entry: PyEntry) -> PyResult { + Self::new(&entry.suffix, entry.scope.into()) + .map_err(|_| PyValueError::new_err(format!("unknown key suffix {:?}", entry.suffix))) + } +} + +/// `CounterIncrement`. +#[derive(FromPyObject)] +struct PyIncrement { + counter: PyCounter, + amount: PyNumber, + ttl: Option, +} + +impl From for Increment { + fn from(increment: PyIncrement) -> Self { + Self { + counter: increment.counter.into(), + amount: increment.amount.0, + ttl: increment.ttl.map(Duration::from_secs), + } + } +} + +/// `AtomicCounter`. +#[derive(FromPyObject)] +struct PyAtomicCounter { + counter: PyCounter, + limit: i64, + increment: i64, + ttl: i64, + window_size: i64, +} + +impl From for AtomicCounter { + fn from(counter: PyAtomicCounter) -> Self { + Self { + counter: counter.counter.into(), + limit: counter.limit, + increment: counter.increment, + ttl_seconds: counter.ttl, + window_size_seconds: counter.window_size, + } + } +} + +fn all>(records: Vec

) -> Vec { + records.into_iter().map(T::from).collect() +} + +#[pyclass(frozen, eq, get_all, name = "WindowCount")] +#[derive(PartialEq)] +pub(crate) struct PyWindowCount { + window_start: String, + count: i64, +} + +#[pymethods] +impl PyWindowCount { + #[new] + fn new(window_start: String, count: i64) -> Self { + Self { + window_start, + count, + } + } +} + +/// Per counter, the count after the increment and the start of the window it landed in. +#[pyclass(frozen, eq, get_all, name = "AtomicApplied")] +#[derive(PartialEq)] +pub(crate) struct PyAtomicApplied { + counters: Vec<(i64, i64)>, +} + +#[pymethods] +impl PyAtomicApplied { + #[new] + fn new(counters: Vec<(i64, i64)>) -> Self { + Self { counters } + } +} + +#[pyclass(frozen, eq, get_all, name = "AtomicOverLimit")] +#[derive(PartialEq)] +pub(crate) struct PyAtomicOverLimit { + index: usize, + current: i64, + limit: i64, +} + +#[pymethods] +impl PyAtomicOverLimit { + #[new] + fn new(index: usize, current: i64, limit: i64) -> Self { + Self { + index, + current, + limit, + } + } +} + +#[derive(IntoPyObject)] +enum PyAtomicReply { + Applied(PyAtomicApplied), + OverLimit(PyAtomicOverLimit), +} + +#[pyclass(frozen, eq, get_all, name = "GuardedReply")] +#[derive(PartialEq)] +pub(crate) struct PyGuardedReply { + applied: bool, + count: i64, +} + +#[pymethods] +impl PyGuardedReply { + #[new] + fn new(applied: bool, count: i64) -> Self { + Self { applied, count } + } +} + +#[pyclass(frozen, eq, get_all, name = "GaugesAcquired")] +#[derive(PartialEq)] +pub(crate) struct PyGaugesAcquired { + in_flight: Vec, +} + +#[pymethods] +impl PyGaugesAcquired { + #[new] + fn new(in_flight: Vec) -> Self { + Self { in_flight } + } +} + +#[pyclass(frozen, eq, get_all, name = "GaugeOverLimit")] +#[derive(PartialEq)] +pub(crate) struct PyGaugeOverLimit { + index: usize, + in_flight: i64, + limit: i64, +} + +#[pymethods] +impl PyGaugeOverLimit { + #[new] + fn new(index: usize, in_flight: i64, limit: i64) -> Self { + Self { + index, + in_flight, + limit, + } + } +} + +#[derive(IntoPyObject)] +enum PyAcquireReply { + Acquired(PyGaugesAcquired), + OverLimit(PyGaugeOverLimit), +} + +fn window_counts(counts: Vec) -> Vec { + counts + .into_iter() + .map( + |WindowCount { + window_start, + count, + }| PyWindowCount { + window_start, + count, + }, + ) + .collect() +} + +fn atomic(reply: AtomicReply) -> PyAtomicReply { + match reply { + AtomicReply::Applied(applied) => PyAtomicReply::Applied(PyAtomicApplied { + counters: applied + .into_iter() + .map( + |Applied { + count, + window_start, + }| (count, window_start), + ) + .collect(), + }), + AtomicReply::OverLimit { + index, + current, + limit, + } => PyAtomicReply::OverLimit(PyAtomicOverLimit { + index, + current, + limit, + }), + } +} + +fn guarded(GuardedReply { applied, count }: GuardedReply) -> PyGuardedReply { + PyGuardedReply { applied, count } +} + +fn acquire(reply: AcquireReply) -> PyAcquireReply { + match reply { + AcquireReply::Acquired(in_flight) => { + PyAcquireReply::Acquired(PyGaugesAcquired { in_flight }) + } + AcquireReply::OverLimit { + index, + in_flight, + limit, + } => PyAcquireReply::OverLimit(PyGaugeOverLimit { + index, + in_flight, + limit, + }), + } +} + +/// One script call declared on the batch the caller's `NativeSend` picked; `None` without Redis. +fn script( + call_type: &'static str, + declared: Option>, + reply: fn(T) -> R, +) -> Option +where + T: Send + 'static, + R: for<'a> IntoPyObject<'a> + Send + 'static, +{ + Some(super::pending(call_type, kind_error, declared?, reply)) +} + +#[pymethods] +impl NativeRateLimits { + #[new] + #[pyo3(signature = (client, backfill_ttl_seconds = None))] + fn new( + client: Option>, + backfill_ttl_seconds: Option, + ) -> PyResult { + Ok(Self(Arc::new(RateLimits::new( + client.map(|client| remote(&client)), + ttl(backfill_ttl_seconds)?, + )))) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn slot_groups(&self, scopes: Vec, force: bool) -> Vec> { + self.0.slot_groups(&all(scopes), force) + } + + fn get_local(&self, entry: PyEntry) -> PyResult> { + Ok(self + .0 + .get_local(&entry.try_into()?) + .map_err(kind_error)? + .map(LocalValue)) + } + + #[pyo3(signature = (entry, value, ttl_seconds))] + fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option) -> PyResult<()> { + self.0 + .set_local(&entry.try_into()?, value.into(), ttl(ttl_seconds)?) + .map_err(kind_error) + } + + /// `request` answers a key its pre-call reads already found absent without another `GET`. + #[pyo3(signature = (entry, *, backfill_ttl_seconds = None, request = None))] + fn get<'py>( + &self, + py: Python<'py>, + entry: PyEntry, + backfill_ttl_seconds: Option, + request: Option<&NativeRequestRedis>, + ) -> PyResult> { + let store = self.0.clone(); + let entry = Entry::try_from(entry)?; + let backfill_ttl = ttl(backfill_ttl_seconds)?; + let request = request.map(|request| request.0.clone()); + redis_async(py, "async_get_cache", kind_error, async move { + store + .get(request.as_ref(), &entry, backfill_ttl) + .await + .map(|local| local.map(LocalValue)) + }) + } + + fn sliding_window( + &self, + counters: Vec, + now: i64, + window_size: i64, + send: &NativeSend, + ) -> Option { + script( + "batch_rate_limiter_script", + self.0 + .sliding_window(send.send(), &all(counters), now, window_size), + window_counts, + ) + } + + fn read_counters(&self, counters: Vec, send: &NativeSend) -> Option { + script( + "batch_counter_read_script", + self.0.read_counters(send.send(), &all(counters)), + std::convert::identity, + ) + } + + fn check_and_increment( + &self, + counters: Vec, + send: &NativeSend, + ) -> Option { + script( + "check_and_increment_by_n_script", + self.0.check_and_increment(send.send(), &all(counters)), + atomic, + ) + } + + fn guarded_increment( + &self, + counter: PyCounter, + expected_window_start: String, + amount: PyNumber, + ttl_seconds: Option, + send: &NativeSend, + ) -> Option { + let increment = GuardedIncrement { + counter: counter.into(), + expected_window_start, + amount: amount.0, + ttl: ttl_seconds.map(Duration::from_secs), + }; + script( + "window_guarded_token_increment_script", + self.0.guarded_increment(send.send(), &increment), + guarded, + ) + } + + fn parallel_acquire( + &self, + gauges: Vec<(PyScope, i64)>, + slot_ttl_seconds: u64, + slot_id: &str, + send: &NativeSend, + ) -> Option { + let gauges: Vec = gauges + .into_iter() + .map(|(scope, limit)| Gauge { + scope: scope.into(), + limit, + }) + .collect(); + script( + "parallel_acquire_script", + self.0.parallel_acquire( + send.send(), + &gauges, + Duration::from_secs(slot_ttl_seconds), + slot_id, + ), + acquire, + ) + } + + fn parallel_release( + &self, + gauges: Vec, + slot_id: &str, + send: &NativeSend, + ) -> Option { + script( + "parallel_release_script", + self.0.parallel_release(send.send(), &all(gauges), slot_id), + std::convert::identity, + ) + } + + fn parallel_count( + &self, + gauges: Vec, + slot_ttl_seconds: u64, + send: &NativeSend, + ) -> Option { + script( + "parallel_count_script", + self.0.parallel_count( + send.send(), + &all(gauges), + Duration::from_secs(slot_ttl_seconds), + ), + std::convert::identity, + ) + } + + fn token_increment( + &self, + increments: Vec, + send: &NativeSend, + ) -> Option { + script( + "token_increment_script", + self.0.token_increment(send.send(), &all(increments)), + std::convert::identity, + ) + } + + /// `False` without Redis. + fn token_increment_post_call( + &self, + request: &NativeRequestRedis, + increments: Vec, + ) -> bool { + self.0 + .token_increment_post_call(&request.0, all(increments)) + } + + /// `False` without Redis. + fn parallel_release_post_call( + &self, + request: &NativeRequestRedis, + gauges: Vec, + slot_id: &str, + ) -> bool { + self.0 + .parallel_release_post_call(&request.0, &all(gauges), slot_id) + } + + fn increment<'py>( + &self, + py: Python<'py>, + increment: PyIncrement, + ) -> PyResult> { + let store = self.0.clone(); + let increment = Increment::from(increment); + redis_async(py, "async_increment", kind_error, async move { + Ok(store.increment(&increment).await.map(PyNumber)) + }) + } + + fn increment_many<'py>( + &self, + py: Python<'py>, + increments: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let increments: Vec = all(increments); + redis_async(py, "async_increment_pipeline", kind_error, async move { + store.increment_many(&increments).await; + Ok(()) + }) + } + + fn increment_many_post_call(&self, request: &NativeRequestRedis, increments: Vec) { + self.0 + .increment_many_post_call(&request.0, &all(increments)); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/request.rs b/litellm-rust/crates/python-bridge/src/redis/store/request.rs new file mode 100644 index 00000000000..3436b26a16c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/request.rs @@ -0,0 +1,60 @@ +use litellm_redis_store::{RequestRedis, Send}; +use pyo3::prelude::*; + +use super::{call, duration}; +use crate::execution::run_async_value; + +/// One request's native Redis batches. Python builds one per request scope and passes it to each +/// request-aware store call. +#[pyclass(frozen, name = "NativeRequestRedis")] +pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis); + +#[pymethods] +impl NativeRequestRedis { + #[new] + fn new(py: Python<'_>, post_call_deadline_seconds: f64) -> PyResult { + let runtime = litellm_host_python::runtime_handle()?; + let deadline = duration(post_call_deadline_seconds)?; + let deadline_tag = call(py, "post_call_redis_batch")?.tag(); + Ok(Self(crate::logger::capture(py).scope(|| { + RequestRedis::with_deadline_tag(runtime, deadline, Some(deadline_tag)) + }))) + } + + fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "request_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_pre()).await; + Ok(()) + }) + } + + fn flush_post<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "post_call_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_post()).await; + Ok(()) + }) + } +} + +/// The Python face of `Send`: which batch a store op rides. Each variant carries exactly what it +/// needs, so there is no phase string and no phase without a request. +#[pyclass(frozen, name = "NativeSend")] +pub(crate) enum NativeSend { + Now(), + PreCall(Py), + PostCall(Py), +} + +impl NativeSend { + pub(crate) fn send(&self) -> Send<'_> { + match self { + Self::Now() => Send::Now, + Self::PreCall(request) => Send::PreCall(&request.get().0), + Self::PostCall(request) => Send::PostCall(&request.get().0), + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs new file mode 100644 index 00000000000..5d695ca7090 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs @@ -0,0 +1,182 @@ +use std::{collections::HashMap, sync::Arc}; + +use litellm_redis_store::{ + Send, + router_cooldowns::{Cooldown, RouterCooldowns}, +}; +use pyo3::{prelude::*, types::PyTuple}; + +use super::{ + NativePending, PyData, PyNumber, duration, kind_error, pending, policy::router_tier_policies, + redis_async, redis_sync, remote, request::NativeSend, +}; +use crate::redis::client::NativeRedisClient; + +/// One router's cooldown store. The router's sync completion path calls `put`, `active_sync`, +/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py; the +/// `async_` twins return awaitables for async callers. +#[pyclass(frozen, name = "NativeRouterCooldowns")] +pub(crate) struct NativeRouterCooldowns(Arc); + +/// `[(model_id, cooldown), ...]` the way `RedisCache` hands the router its batch read. +struct PyActive(Vec<(String, Cooldown)>); + +impl<'py> IntoPyObject<'py> for PyActive { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + active_to_py(py, self.0)? + .into_pyobject(py) + .map(Bound::into_any) + } +} + +fn active_to_py( + py: Python<'_>, + active: Vec<(String, Cooldown)>, +) -> PyResult>> { + active + .into_iter() + .map(|(model_id, cooldown)| { + PyTuple::new( + py, + [ + model_id.into_pyobject(py)?.into_any(), + PyData(cooldown).into_pyobject(py)?, + ], + ) + }) + .collect() +} + +#[pymethods] +impl NativeRouterCooldowns { + #[new] + fn new(py: Python<'_>, client: Option>) -> PyResult { + let remote = client.map(|client| remote(&client)); + let policy = router_tier_policies(py)?.cooldowns; + Ok(Self(Arc::new(RouterCooldowns::new(remote, policy)))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool { + self.0.attach(remote(&client)) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn put( + &self, + py: Python<'_>, + model_id: String, + cooldown: PyData, + ttl_seconds: f64, + ) -> PyResult<()> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_sync(py, "set_cache", kind_error, async move { + store.put_str(&model_id, &cooldown.0, ttl).await + }) + } + + fn increment_allowed_fails( + &self, + py: Python<'_>, + model_id: String, + suffix: Option, + ttl_seconds: f64, + ) -> PyResult { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_sync(py, "increment_cache", kind_error, async move { + store + .increment_allowed_fails_int(&model_id, suffix.as_deref(), ttl) + .await + }) + } + + fn async_put<'py>( + &self, + py: Python<'py>, + model_id: String, + cooldown: PyData, + ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.put(&model_id, &cooldown.0, ttl).await + }) + } + + fn async_increment_allowed_fails<'py>( + &self, + py: Python<'py>, + model_id: String, + suffix: Option, + ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_async(py, "async_increment", kind_error, async move { + store + .increment_allowed_fails(&model_id, suffix.as_deref(), ttl) + .await + }) + } + + /// The deployments still cooling down, read on the batch `send` picks. + fn active(&self, model_ids: Vec, send: &NativeSend) -> PyResult { + let active = self.0.active(send.send(), &model_ids).map_err(kind_error)?; + Ok(pending( + "async_batch_get_cache", + kind_error, + active, + PyActive, + )) + } + + /// `active` on its own with the GIL released, for the router's sync completion path. + fn active_sync<'py>( + &self, + py: Python<'py>, + model_ids: Vec, + ) -> PyResult>> { + let active = self.0.active(Send::Now, &model_ids).map_err(kind_error)?; + let active = redis_sync(py, "batch_get_cache", kind_error, active)?; + active_to_py(py, active) + } + + fn minimum(&self, py: Python<'_>, model_ids: Vec) -> PyResult> { + let store = self.0.clone(); + redis_sync(py, "batch_get_cache", kind_error, async move { + store + .minimum(&model_ids) + .await + .map(|minimum| minimum.map(PyNumber)) + }) + } + + fn delete_local(&self, model_id: &str) -> PyResult<()> { + self.0.delete_local(model_id).map_err(kind_error) + } + + fn flush_local_counters(&self) -> PyResult<()> { + self.0.flush_local_counters().map_err(kind_error) + } + + fn snapshot(&self) -> PyResult>> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, cooldown)| (key, PyData(cooldown))) + .collect()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs new file mode 100644 index 00000000000..48565028de3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -0,0 +1,182 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage}; +use pyo3::prelude::*; + +use super::{ + PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote, +}; +use crate::redis::client::NativeRedisClient; + +/// `litellm.redis.stores.router_usage.UsageCounter`. +#[derive(FromPyObject)] +struct PyCounter { + deployment_id: String, + model: Option, + minute: String, + kind: Spelled, +} + +impl From for Counter { + fn from(counter: PyCounter) -> Self { + Self { + deployment_id: counter.deployment_id, + model: counter.model, + minute: counter.minute, + kind: counter.kind.0, + } + } +} + +fn counters(counters: Vec) -> Vec { + counters.into_iter().map(Counter::from).collect() +} + +/// One router's global TPM/RPM counters. The `_sync` methods block with the GIL released, for the +/// Rust router's sync calls, which have no event loop to await on. +#[pyclass(frozen, name = "NativeRouterUsage")] +pub(crate) struct NativeRouterUsage(Arc); + +#[pymethods] +impl NativeRouterUsage { + #[new] + fn new(py: Python<'_>, client: Option>) -> PyResult { + let remote = client.map(|client| remote(&client)); + let policy = router_tier_policies(py)?.usage; + Ok(Self(Arc::new(RouterUsage::new(remote, policy)))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool { + self.0.attach(remote(&client)) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn increment<'py>( + &self, + py: Python<'py>, + counter: PyCounter, + value: PyNumber, + ttl_seconds: u64, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_async(py, "async_increment", kind_error, async move { + Ok(store + .increment(&counter, value.0, Duration::from_secs(ttl_seconds)) + .await + .map(PyNumber)) + }) + } + + /// `_increment_deployment`'s `async_increment_pipeline`, always on its own round trip. + fn increment_deployment<'py>( + &self, + py: Python<'py>, + deployment_id: String, + model: Option, + minute: String, + tokens: PyNumber, + requests: i64, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_increment_pipeline", kind_error, async move { + store + .increment_deployment( + &deployment_id, + model.as_deref(), + &minute, + tokens.0, + requests, + ) + .await + }) + } + + fn get_many<'py>( + &self, + py: Python<'py>, + counters: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let counters = self::counters(counters); + redis_async(py, "async_batch_get_cache", kind_error, async move { + let read = store.get_many(&counters).await?; + Ok(read + .into_iter() + .map(|count| count.map(PyNumber)) + .collect::>()) + }) + } + + fn increment_sync( + &self, + py: Python<'_>, + counter: PyCounter, + value: PyNumber, + ttl_seconds: u64, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_sync(py, "increment_cache", kind_error, async move { + Ok(store + .increment(&counter, value.0, Duration::from_secs(ttl_seconds)) + .await + .map(PyNumber)) + }) + } + + fn increment_deployment_sync( + &self, + py: Python<'_>, + deployment_id: String, + model: Option, + minute: String, + tokens: PyNumber, + requests: i64, + ) -> PyResult<()> { + let store = self.0.clone(); + redis_sync(py, "increment_pipeline", kind_error, async move { + store + .increment_deployment( + &deployment_id, + model.as_deref(), + &minute, + tokens.0, + requests, + ) + .await + }) + } + + fn get_many_sync( + &self, + py: Python<'_>, + counters: Vec, + ) -> PyResult>> { + let store = self.0.clone(); + let counters = self::counters(counters); + redis_sync(py, "batch_get_cache", kind_error, async move { + let read = store.get_many(&counters).await?; + Ok(read.into_iter().map(|count| count.map(PyNumber)).collect()) + }) + } + + fn flush_local_counters(&self) -> PyResult<()> { + self.0.flush_local_counters().map_err(kind_error) + } + + fn snapshot(&self) -> PyResult> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, count)| (key, PyNumber(count))) + .collect()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/sensitive_route_pins.rs b/litellm-rust/crates/python-bridge/src/redis/store/sensitive_route_pins.rs new file mode 100644 index 00000000000..8b1b82945fb --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/sensitive_route_pins.rs @@ -0,0 +1,49 @@ +use std::sync::Arc; + +use litellm_redis_store::sensitive_route_pins::SensitiveRoutePins; +use pyo3::prelude::*; + +use super::{duration, kind_error, redis_async, remote}; +use crate::redis::client::NativeRedisClient; + +#[pyclass(frozen, name = "NativeSensitiveRoutePins")] +pub(crate) struct NativeSensitiveRoutePins(Arc); + +#[pymethods] +impl NativeSensitiveRoutePins { + #[new] + fn new(client: Option>) -> Self { + Self(Arc::new(SensitiveRoutePins::new( + client.map(|client| remote(&client)), + ))) + } + + fn get<'py>( + &self, + py: Python<'py>, + tenant: String, + session_id: String, + default_ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let default_ttl = duration(default_ttl_seconds)?; + redis_async(py, "async_get_cache", kind_error, async move { + store.get(&tenant, &session_id, default_ttl).await + }) + } + + fn set<'py>( + &self, + py: Python<'py>, + tenant: String, + session_id: String, + model: String, + ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.set(&tenant, &session_id, &model, ttl).await + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs new file mode 100644 index 00000000000..9da6a0e00e5 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -0,0 +1,277 @@ +use std::{collections::HashMap, sync::Arc}; + +use litellm_redis_store::{ + Number, + spend_counters::{LostSeed, PostCallSettlement, Reading, SpendCounter, SpendCounters}, +}; +use pyo3::prelude::*; + +use super::{ + NativePending, PyData, PyNumber, duration, kind_error, pending, redis_async, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, unwrap_all, +}; +use crate::redis::client::NativeRedisClient; + +/// A counter's value and whether Redis answered it, a clean miss included. +#[pyclass(frozen, eq, get_all, name = "SpendRead")] +#[derive(PartialEq)] +pub(crate) struct PySpendRead { + value: Option, + from_redis: bool, +} + +#[pymethods] +impl PySpendRead { + #[new] + fn new(value: Option, from_redis: bool) -> Self { + Self { value, from_redis } + } +} + +/// Whether every increment one `increment_many_post_call` declared has its reply. +#[pyclass(frozen, name = "PostCallSettlement")] +pub(crate) struct PyPostCallSettlement(PostCallSettlement); + +#[pymethods] +impl PyPostCallSettlement { + #[getter] + fn settled(&self) -> bool { + self.0.is_settled() + } +} + +fn increments(increments: Vec<(PyData, PyNumber)>) -> Vec<(SpendCounter, Number)> { + increments + .into_iter() + .map(|(counter, amount)| (counter.0, amount.0)) + .collect() +} + +fn numbers(values: Option>>) -> Option>> { + values.map(|values| { + values + .into_iter() + .map(|value| value.map(PyNumber)) + .collect() + }) +} + +/// The proxy's spend counters, Redis first with this worker's memory copy behind it. +#[pyclass(frozen, name = "NativeSpendCounters")] +pub(crate) struct NativeSpendCounters(Arc); + +#[pymethods] +impl NativeSpendCounters { + #[new] + fn new(client: Option>) -> Self { + Self(Arc::new(SpendCounters::new( + client.map(|client| remote(&client)), + ))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool { + self.0.attach(remote(&client)) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + #[getter] + fn redis_ttl_seconds(&self) -> Option { + self.0.redis_ttl().map(|ttl| ttl.as_secs_f64()) + } + + /// The least TTL a Redis write gives a committed spend twin. + fn set_committed_ttl(&self, ttl_seconds: f64) -> PyResult<()> { + self.0.set_committed_ttl(duration(ttl_seconds)?); + Ok(()) + } + + /// `("redis" | "memory", value)`: where the answer came from, a Redis miss included. + fn read<'py>( + &self, + py: Python<'py>, + counter: PyData, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_get_cache", kind_error, async move { + let reading = store.read(&counter.0).await?; + Ok(PySpendRead { + value: reading.value().map(PyNumber), + from_redis: matches!(reading, Reading::Redis(_)), + }) + }) + } + + /// One `MGET` on the batch `send` picks; `None` without Redis. It resolves to `None` when the + /// read failed, so each caller reads on its own. + fn read_many( + &self, + counters: Vec>, + send: &NativeSend, + ) -> Option { + let read = self.0.read_many(send.send(), &unwrap_all(counters))?; + Some(pending("async_batch_get_cache", kind_error, read, numbers)) + } + + fn increment<'py>( + &self, + py: Python<'py>, + counter: PyData, + amount: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_increment", kind_error, async move { + store.increment(&counter.0, amount.0).await.map(PyNumber) + }) + } + + fn increment_many<'py>( + &self, + py: Python<'py>, + increments: Vec<(PyData, PyNumber)>, + ) -> PyResult> { + let store = self.0.clone(); + let increments = self::increments(increments); + redis_async(py, "async_increment_pipeline", kind_error, async move { + let totals = store.increment_many(&increments).await?; + Ok(totals.into_iter().map(PyNumber).collect::>()) + }) + } + + /// Declares the increments on the request's post-call batch; `None` when there is no Redis. + fn increment_many_post_call( + &self, + request: &NativeRequestRedis, + increments: Vec<(PyData, PyNumber)>, + ) -> Option { + self.0 + .increment_many_post_call(&request.0, &self::increments(increments)) + .map(PyPostCallSettlement) + } + + fn set_max<'py>( + &self, + py: Python<'py>, + counter: PyData, + spend: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_set_max", kind_error, async move { + store + .set_max(&counter.0, spend.0) + .await + .map(|stored| stored.map(PyNumber)) + }) + } + + #[pyo3(signature = (counter, spend, *, increment_if_lost))] + fn seed<'py>( + &self, + py: Python<'py>, + counter: PyData, + spend: PyNumber, + increment_if_lost: bool, + ) -> PyResult> { + let store = self.0.clone(); + let lost = if increment_if_lost { + LostSeed::Increment + } else { + LostSeed::Answer + }; + redis_async(py, "async_set_cache", kind_error, async move { + store.seed(&counter.0, spend.0, lost).await.map(PyNumber) + }) + } + + #[pyo3(signature = (counter, spend, ttl_seconds = None))] + fn set<'py>( + &self, + py: Python<'py>, + counter: PyData, + spend: PyNumber, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let ttl = ttl(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.set(&counter.0, spend.0, ttl).await + }) + } + + fn delete<'py>( + &self, + py: Python<'py>, + counters: Vec>, + ) -> PyResult> { + let store = self.0.clone(); + let counters = unwrap_all(counters); + redis_async(py, "async_delete_cache", kind_error, async move { + store.delete(&counters).await + }) + } + + fn refresh_ttl<'py>( + &self, + py: Python<'py>, + counter: PyData, + ) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_refresh_ttl", kind_error, async move { + store.refresh_ttl(&counter.0).await + }) + } + + fn get_local(&self, counter: PyData) -> PyResult> { + Ok(self + .0 + .get_local(&counter.0) + .map_err(kind_error)? + .map(PyNumber)) + } + + #[pyo3(signature = (counter, spend, ttl_seconds = None))] + fn set_local( + &self, + counter: PyData, + spend: PyNumber, + ttl_seconds: Option, + ) -> PyResult<()> { + self.0 + .set_local(&counter.0, spend.0, ttl(ttl_seconds)?) + .map_err(kind_error) + } + + fn delete_local(&self, counter: PyData) -> PyResult<()> { + self.0.delete_local(&counter.0).map_err(kind_error) + } + + fn increment_local( + &self, + counter: PyData, + amount: PyNumber, + ) -> PyResult { + self.0 + .increment_local(&counter.0, amount.0) + .map(PyNumber) + .map_err(kind_error) + } + + fn flush_local(&self) -> PyResult<()> { + self.0.flush_local().map_err(kind_error) + } + + fn snapshot(&self) -> PyResult> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, spend)| (key, PyNumber(spend))) + .collect()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/value.rs b/litellm-rust/crates/python-bridge/src/redis/store/value.rs new file mode 100644 index 00000000000..c86c9e9bb20 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/value.rs @@ -0,0 +1,125 @@ +use litellm_python_compat::{BigInt, MAX_DEPTH, Value}; +use pyo3::{ + exceptions::{PyTypeError, PyValueError}, + prelude::*, + types::{ + PyBool, PyBytes, PyComplex, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PySet, PyString, + PyTuple, + }, +}; + +/// A plain Python value (what `json.loads` or `ast.literal_eval` returns) across the bridge. +/// Anything else, a live object included, is refused with `TypeError`. +pub(crate) struct PyValue(pub(crate) Value); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyValue { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + to_value(&value, 0).map(Self) + } +} + +impl<'py> IntoPyObject<'py> for PyValue { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + to_python(py, self.0) + } +} + +fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult { + if depth > MAX_DEPTH { + return Err(PyValueError::new_err("value is nested too deeply")); + } + let nested = |items: Bound<'_, PyAny>| -> PyResult> { + items + .try_iter()? + .map(|item| to_value(&item?, depth + 1)) + .collect() + }; + if value.is_none() { + return Ok(Value::None); + } + if let Ok(flag) = value.cast::() { + return Ok(Value::Bool(flag.is_true())); + } + if value.is_instance_of::() { + return Ok(Value::Int(match value.extract::() { + Ok(small) => small.into(), + Err(_) => value + .str()? + .to_str()? + .parse::() + .map_err(|error| PyValueError::new_err(error.to_string()))?, + })); + } + if let Ok(number) = value.cast::() { + return Ok(Value::Float(number.value())); + } + if let Ok(text) = value.cast::() { + return Ok(Value::Str(text.to_str()?.to_owned())); + } + if let Ok(bytes) = value.cast::() { + return Ok(Value::Bytes(bytes.as_bytes().to_vec())); + } + if let Ok(complex) = value.cast::() { + return Ok(Value::Complex { + re: complex.real(), + im: complex.imag(), + }); + } + if let Ok(entries) = value.cast::() { + return entries + .iter() + .map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?))) + .collect::>() + .map(Value::Dict); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::List); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::Tuple); + } + if value.is_instance_of::() || value.is_instance_of::() { + return nested(value.clone()).map(Value::Set); + } + Err(PyTypeError::new_err(format!( + "{} is not a plain data value", + value.get_type().name()? + ))) +} + +fn to_python(py: Python<'_>, value: Value) -> PyResult> { + let all = |items: Vec| -> PyResult>> { + items.into_iter().map(|item| to_python(py, item)).collect() + }; + Ok(match value { + Value::None => py.None().into_bound(py), + Value::Bool(flag) => PyBool::new(py, flag).to_owned().into_any(), + Value::Int(number) => match i64::try_from(&number) { + Ok(small) => small.into_pyobject(py)?.into_any(), + Err(_) => py + .get_type::() + .call1((number.to_string(),))? + .into_any(), + }, + Value::Float(number) => number.into_pyobject(py)?.into_any(), + Value::Complex { re, im } => PyComplex::from_doubles(py, re, im).into_any(), + Value::Str(text) => text.into_pyobject(py)?.into_any(), + Value::Bytes(bytes) => PyBytes::new(py, &bytes).into_any(), + Value::Tuple(items) => PyTuple::new(py, all(items)?)?.into_any(), + Value::List(items) => PyList::new(py, all(items)?)?.into_any(), + Value::Set(items) => PySet::new(py, all(items)?)?.into_any(), + Value::Dict(entries) => { + let dict = PyDict::new(py); + for (key, item) in entries { + dict.set_item(to_python(py, key)?, to_python(py, item)?)?; + } + dict.into_any() + } + }) +} diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 6ccac4df659..27f9fe347e0 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -20,10 +20,10 @@ use std::{ use async_trait::async_trait; use litellm_router_core::{ - Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output, - Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, RouterEvent, RouterRequest, - RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness, - build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, + Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Deployment, + Error, Host, Output, Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, + RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, SelectStage, + ThreadRandomness, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, precall::PreCallFacts, }; use pyo3::{PyResult, Python, types::PyAnyMethods}; @@ -298,8 +298,9 @@ pub(crate) enum HostCallResponse { PreCall(PreCallFacts), /// The candidate models a `routing_plugins` select step reported. Models(Option>), - /// The ids a `callback_filters` select step left, or the error a filter raised. - Filtered(Result, AttemptError>), + /// The ids a `callback_filters` select step left with the one a filter prefers, or the error + /// a filter raised. + Filtered(Result<(Vec, Option), AttemptError>), /// What a pick's select or pre-route step, or a fallback check, raised on the host instead of /// answering: the walk takes it as the try's failure. Raised(AttemptError), @@ -352,9 +353,15 @@ impl HostCallResponse { got: self.kind(), }); } - if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) = - (self, call) - && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) + if let ( + Self::Select(ids) | Self::Filtered(Ok((ids, _))), + HostCall::Select { candidates, .. }, + ) = (self, call) + && let Some(id) = { + let known: std::collections::HashSet<&str> = + candidates.iter().map(String::as_str).collect(); + ids.iter().find(|id| !known.contains(id.as_str())) + } { return Err(DriveError::UnknownCandidate { id: id.clone() }); } @@ -685,14 +692,22 @@ fn shares_ids(candidates: &[Arc]) -> bool { /// The candidates a host answer names, in its order. An id names every candidate that carries /// it, so deployments sharing an id stay together through a filter, as Python's list keeps them. fn answered_by_id(candidates: &[Arc], ids: &[String]) -> Vec> { - let mut taken = vec![false; candidates.len()]; + let mut positions: std::collections::HashMap<&str, Vec> = + std::collections::HashMap::new(); + for (position, deployment) in candidates.iter().enumerate() { + positions + .entry(deployment.id.as_str()) + .or_default() + .push(position); + } let mut kept = Vec::new(); for id in ids { - for (position, deployment) in candidates.iter().enumerate() { - if !taken[position] && deployment.id == *id { - taken[position] = true; - kept.push(deployment.clone()); - } + if let Some(found) = positions.remove(id.as_str()) { + kept.extend( + found + .into_iter() + .map(|position| candidates[position].clone()), + ); } } kept @@ -890,7 +905,7 @@ impl Host for StepHost { &self, ctx: &RoutingContext<'_>, candidates: Vec>, - ) -> Result>, Error> { + ) -> Result { let (scope, group) = ctx.scope.wire(); let joined = |own: &[String], walked: &[String]| [own, walked].concat(); let call = HostCall::Select { @@ -909,8 +924,11 @@ impl Host for StepHost { let HostCallResponse::Filtered(result) = call_host(call).await else { unreachable!("{CHECKED}"); }; - let ids = result.map_err(Error::Attempt)?; - Ok(answered_by_id(&candidates, &ids)) + let (ids, preferred) = result.map_err(Error::Attempt)?; + Ok(CallbackFiltered { + kept: answered_by_id(&candidates, &ids), + preferred, + }) } async fn allow_fallback( @@ -1534,6 +1552,26 @@ mod tests { ); } + #[rstest] + fn an_answer_by_id_keeps_its_order_and_every_candidate_sharing_an_id() { + let candidates: Vec> = ["a", "b", "a", "c"] + .into_iter() + .map(|id| { + let spec: litellm_router_core::DeploymentSpec = + serde_json::from_value(deployment("g", id)).unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + }) + .collect(); + let ids = ["c", "a", "x", "a"].map(String::from); + + let kept = answered_by_id(&candidates, &ids); + + let positions: Vec = kept + .iter() + .map(|d| candidates.iter().position(|c| Arc::ptr_eq(c, d)).unwrap()) + .collect(); + assert_eq!(positions, [3, 0, 2]); + } #[rstest] fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() { let mut second = deployment("primary", "a1"); diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index 8157d0f5a32..6b58d63a714 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -1181,7 +1181,8 @@ fn host_call_response( stage: Stage::CallbackFilters { .. }, .. } => { - // The remaining ids, or a failure mapping as an attempt reports one. + // The remaining ids, alone or paired with the id a filter prefers, or a failure + // mapping as an attempt reports one. let outcome = outcome.ok_or_else(missing)?; if let Ok(failure) = outcome.cast::() { return match attempt_outcome(failure)? { @@ -1191,7 +1192,10 @@ fn host_call_response( )), }; } - Ok(HostCallResponse::Filtered(Ok(outcome.extract()?))) + if let Ok(answer) = outcome.cast::() { + return Ok(HostCallResponse::Filtered(Ok(answer.extract()?))); + } + Ok(HostCallResponse::Filtered(Ok((outcome.extract()?, None)))) } HostCall::Select { stage: Stage::Strategy { records, .. }, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index af614f17360..2a5f79ca393 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -239,9 +239,13 @@ impl PythonBinding for MessagesPythonHost { fn encode_response( &mut self, py: Python<'_>, - response: Box, + response: litellm_inference_messages::MessagesBody, ) -> PyResult> { - public_response(py, ROUTE_HOST_MODULE, response.as_ref()) + public_response(py, ROUTE_HOST_MODULE, &response) + } + + fn deliver_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> { + self.cache.store_response(py, response) } fn encode_stream_head( diff --git a/litellm-rust/crates/python-compat/benches/formats.rs b/litellm-rust/crates/python-compat/benches/formats.rs index f0b83cb1658..2f23b882e19 100644 --- a/litellm-rust/crates/python-compat/benches/formats.rs +++ b/litellm-rust/crates/python-compat/benches/formats.rs @@ -10,7 +10,7 @@ use std::{hint::black_box, time::Duration}; use criterion::{ - BatchSize, BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main, + BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main, measurement::WallTime, }; use litellm_python_compat::{Value, json, literal::literal_eval, pickle, repr::repr}; @@ -72,14 +72,9 @@ fn formats(c: &mut Criterion) { let mut group = c.benchmark_group("cached_completion"); text_formats(&mut group, &text, &value); binary_formats(&mut group, &value, &pickled); - // `from_json` consumes its input, so each iteration gets a freshly parsed one. group.throughput(Throughput::Bytes(dumped.len() as u64)); - group.bench_function("from_json", |bencher| { - bencher.iter_batched( - || serde_json::from_str::(&dumped).expect("dumps output parses"), - json::from_json, - BatchSize::SmallInput, - ) + group.bench_function("json_loads", |bencher| { + bencher.iter(|| json::loads(black_box(&dumped))) }); group.finish(); } diff --git a/litellm-rust/crates/python-compat/src/error.rs b/litellm-rust/crates/python-compat/src/error.rs index 6547907056e..db09925dde9 100644 --- a/litellm-rust/crates/python-compat/src/error.rs +++ b/litellm-rust/crates/python-compat/src/error.rs @@ -6,6 +6,8 @@ use crate::MAX_DEPTH; pub enum Error { #[error("malformed Python literal at byte {0}")] InvalidLiteral(usize), + #[error("malformed JSON at byte {0}")] + InvalidJson(usize), #[error("unhashable type: '{0}'")] Unhashable(&'static str), #[error("value nests deeper than {MAX_DEPTH} levels")] diff --git a/litellm-rust/crates/python-compat/src/json.rs b/litellm-rust/crates/python-compat/src/json.rs index d3436519ac4..37cc88fec78 100644 --- a/litellm-rust/crates/python-compat/src/json.rs +++ b/litellm-rust/crates/python-compat/src/json.rs @@ -1,14 +1,19 @@ -//! `json.dumps` with CPython's default options, and the JSON value `json.loads` returns. +//! `json.dumps` with CPython's default options, and `json.loads` with its default decoder. //! //! Defaults are `ensure_ascii=True`, `allow_nan=True`, separators `(", ", ": ")`, and no key -//! sorting. Python's `json.loads` is mapped by [`from_json`]; a `serde_json` number that does -//! not fit `i64` or `u64` arrives as a float, where Python would keep an `int`. +//! sorting on the way out, and `strict=True` with `NaN`, `Infinity` and `-Infinity` read as floats +//! on the way in. -use std::fmt::Write; +use std::{ + collections::{HashMap, hash_map::Entry}, + fmt::Write, + str::FromStr, +}; +use num_bigint::BigInt; use serde_json::{Map, Number}; -use crate::{Error, Value, repr::float_repr}; +use crate::{Error, MAX_DEPTH, Value, repr::float_repr}; /// `json.dumps(value)`. pub fn dumps(value: &Value) -> Result { @@ -44,6 +49,260 @@ pub fn from_json(value: serde_json::Value) -> Value { } } +/// `json.loads(text)`. Integers keep full precision, and a repeated key keeps its first position +/// with its last value, as `dict(pairs)` does. An escape that leaves a lone surrogate fails, since +/// a Rust `String` cannot hold one. +pub fn loads(text: &str) -> Result { + let mut decoder = Decoder { + text, + bytes: text.as_bytes(), + pos: 0, + depth: 0, + }; + decoder.skip_whitespace(); + let value = decoder.value()?; + decoder.skip_whitespace(); + if decoder.pos != decoder.bytes.len() { + return Err(decoder.error()); + } + Ok(value) +} + +/// CPython's `json.scanner` and `scanstring`, one pass over the bytes. +struct Decoder<'a> { + text: &'a str, + bytes: &'a [u8], + pos: usize, + depth: usize, +} + +impl Decoder<'_> { + fn error(&self) -> Error { + Error::InvalidJson(self.pos) + } + + fn peek(&self) -> Option { + self.bytes.get(self.pos).copied() + } + + fn eat(&mut self, token: &str) -> bool { + let found = self.bytes[self.pos..].starts_with(token.as_bytes()); + if found { + self.pos += token.len(); + } + found + } + + fn expect(&mut self, token: &str) -> Result<(), Error> { + if self.eat(token) { + Ok(()) + } else { + Err(self.error()) + } + } + + fn skip_whitespace(&mut self) { + while matches!(self.peek(), Some(b' ' | b'\t' | b'\n' | b'\r')) { + self.pos += 1; + } + } + + fn value(&mut self) -> Result { + match self.peek() { + Some(b'"') => self.string().map(Value::Str), + Some(b'[') => self.nested(Self::array), + Some(b'{') => self.nested(Self::object), + _ if self.eat("null") => Ok(Value::None), + _ if self.eat("true") => Ok(Value::Bool(true)), + _ if self.eat("false") => Ok(Value::Bool(false)), + _ if self.eat("NaN") => Ok(Value::Float(f64::NAN)), + _ if self.eat("Infinity") => Ok(Value::Float(f64::INFINITY)), + _ if self.eat("-Infinity") => Ok(Value::Float(f64::NEG_INFINITY)), + _ => self.number(), + } + } + + fn nested(&mut self, parse: fn(&mut Self) -> Result) -> Result { + self.depth += 1; + if self.depth > MAX_DEPTH { + return Err(Error::TooDeep); + } + let value = parse(self)?; + self.depth -= 1; + Ok(value) + } + + fn array(&mut self) -> Result { + self.expect("[")?; + self.skip_whitespace(); + let mut values = Vec::new(); + if self.eat("]") { + return Ok(Value::List(values)); + } + loop { + values.push(self.value()?); + self.skip_whitespace(); + if self.eat("]") { + return Ok(Value::List(values)); + } + self.expect(",")?; + self.skip_whitespace(); + } + } + + fn object(&mut self) -> Result { + self.expect("{")?; + self.skip_whitespace(); + let mut entries = Vec::new(); + let mut positions = HashMap::<_, usize>::new(); + if self.eat("}") { + return Ok(Value::Dict(entries)); + } + loop { + if self.peek() != Some(b'"') { + return Err(self.error()); + } + let key = self.string()?; + self.skip_whitespace(); + self.expect(":")?; + self.skip_whitespace(); + let value = self.value()?; + match positions.entry(key) { + Entry::Occupied(position) => entries[*position.get()].1 = value, + Entry::Vacant(position) => { + entries.push((Value::Str(position.key().clone()), value)); + position.insert(entries.len() - 1); + } + } + self.skip_whitespace(); + if self.eat("}") { + return Ok(Value::Dict(entries)); + } + self.expect(",")?; + self.skip_whitespace(); + } + } + + /// `scanstring` with `strict=True`: control characters must be escaped. + fn string(&mut self) -> Result { + self.expect("\"")?; + let mut out = String::new(); + loop { + let start = self.pos; + while self + .peek() + .is_some_and(|byte| byte != b'"' && byte != b'\\' && byte >= 0x20) + { + self.pos += 1; + } + out.push_str(&self.text[start..self.pos]); + match self.peek() { + Some(b'"') => { + self.pos += 1; + return Ok(out); + } + Some(b'\\') => { + self.pos += 1; + out.push(self.escape()?); + } + _ => return Err(self.error()), + } + } + } + + fn escape(&mut self) -> Result { + let escaped = match self.peek() { + Some(b'"') => '"', + Some(b'\\') => '\\', + Some(b'/') => '/', + Some(b'b') => '\u{8}', + Some(b'f') => '\u{c}', + Some(b'n') => '\n', + Some(b'r') => '\r', + Some(b't') => '\t', + Some(b'u') => { + self.pos += 1; + return self.unicode_escape(); + } + _ => return Err(self.error()), + }; + self.pos += 1; + Ok(escaped) + } + + /// `\uXXXX`, joining a high surrogate with the low one escaped right after it. + fn unicode_escape(&mut self) -> Result { + let unit = self.hex4()?; + let code = if (0xd800..0xdc00).contains(&unit) { + self.expect("\\u")?; + let low = self.hex4()?; + if !(0xdc00..0xe000).contains(&low) { + return Err(self.error()); + } + 0x10000 + ((unit - 0xd800) << 10) + (low - 0xdc00) + } else { + unit + }; + char::from_u32(code).ok_or_else(|| self.error()) + } + + fn hex4(&mut self) -> Result { + let digits = self + .text + .get(self.pos..self.pos + 4) + .filter(|digits| digits.bytes().all(|byte| byte.is_ascii_hexdigit())) + .ok_or_else(|| self.error())?; + self.pos += 4; + u32::from_str_radix(digits, 16).map_err(|_| self.error()) + } + + /// `NUMBER_RE`: `(-?(?:0|[1-9][0-9]*))(\.[0-9]+)?([eE][-+]?[0-9]+)?`, an `int` unless a + /// fraction or exponent matched. + fn number(&mut self) -> Result { + let start = self.pos; + self.eat("-"); + match self.peek() { + Some(b'0') => self.pos += 1, + Some(b'1'..=b'9') => self.digits(), + _ => return Err(Error::InvalidJson(start)), + } + let integer_end = self.pos; + if self.peek() == Some(b'.') && self.bytes.get(self.pos + 1).is_some_and(u8::is_ascii_digit) + { + self.pos += 1; + self.digits(); + } + if matches!(self.peek(), Some(b'e' | b'E')) { + let exponent = self.pos; + self.pos += 1; + if matches!(self.peek(), Some(b'+' | b'-')) { + self.pos += 1; + } + if self.peek().is_some_and(|byte| byte.is_ascii_digit()) { + self.digits(); + } else { + self.pos = exponent; + } + } + let text = &self.text[start..self.pos]; + if self.pos == integer_end { + BigInt::from_str(text) + .map(Value::Int) + .map_err(|_| Error::InvalidJson(start)) + } else { + f64::from_str(text) + .map(Value::Float) + .map_err(|_| Error::InvalidJson(start)) + } + } + + fn digits(&mut self) { + while self.peek().is_some_and(|byte| byte.is_ascii_digit()) { + self.pos += 1; + } + } +} + /// `json.loads(json.dumps(value))` as a `serde_json` value: tuples become arrays and dict /// keys are coerced to strings as `json.dumps` does. Non-finite floats, which Python writes /// as `NaN` and `Infinity`, have no `serde_json` form and fail with [`Error::NonFiniteFloat`]. @@ -166,3 +425,86 @@ fn write_string(out: &mut String, text: &str) { } out.push('"'); } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::{dumps, loads}; + use crate::{Error, Value, repr::repr}; + + // Expected values are the repr of CPython 3.14 json.loads(text) + #[rstest] + #[case::nan("NaN", "nan")] + #[case::infinity("Infinity", "inf")] + #[case::negative_infinity("-Infinity", "-inf")] + #[case::non_finite_in_a_list("[NaN, Infinity, -Infinity]", "[nan, inf, -inf]")] + #[case::beyond_u64("18446744073709551616", "18446744073709551616")] + #[case::exponent_overflow("1e400", "inf")] + #[case::upper_exponent("1E+2", "100.0")] + #[case::negative_zero_int("-0", "0")] + #[case::negative_zero_float("-0.0", "-0.0")] + #[case::repeated_key(r#"{"a": 1, "b": 2, "a": 3}"#, "{'a': 3, 'b': 2}")] + #[case::surrogate_pair(r#""😀""#, "'😀'")] + #[case::upper_hex_surrogate_pair(r#""😀""#, "'😀'")] + #[case::escapes( + r#"{"k": [true, false, null, 1.5, "é\/"]}"#, + "{'k': [True, False, None, 1.5, 'é/']}" + )] + #[case::delete_is_not_a_control_character("\"\x7f\"", "'\\x7f'")] + #[case::surrounding_whitespace(" \n[1]\t", "[1]")] + #[case::empty_containers("[{}, []]", "[{}, []]")] + fn loads_reads_like_python(#[case] text: &str, #[case] expected: &str) { + assert_eq!(repr(&loads(text).unwrap()), expected); + } + + // CPython 3.14 json.loads raises JSONDecodeError on each + #[rstest] + #[case::lowercase_nan("nan")] + #[case::signed_nan("-NaN")] + #[case::plus_sign("+1")] + #[case::bare_minus("-")] + #[case::leading_zero("01")] + #[case::bare_point("1.")] + #[case::bare_exponent("1e")] + #[case::trailing_data("Infinity1")] + #[case::trailing_comma("[1,]")] + #[case::missing_comma("[1 2]")] + #[case::missing_colon(r#"{"a" 1}"#)] + #[case::unquoted_key("{1: 2}")] + #[case::raw_newline_in_string("\"a\nb\"")] + #[case::unknown_escape(r#""\x41""#)] + #[case::empty("")] + #[case::byte_order_mark("\u{feff}1")] + fn loads_rejects_what_python_rejects(#[case] text: &str) { + assert!(matches!(loads(text), Err(Error::InvalidJson(_)))); + } + + #[rstest] + fn a_lone_surrogate_is_a_known_limit() { + assert!(matches!(loads(r#""\ud800""#), Err(Error::InvalidJson(_)))); + } + + #[rstest] + fn loads_reads_back_the_non_finite_floats_dumps_writes() { + let written = dumps(&Value::List(vec![ + Value::Float(f64::NAN), + Value::Float(f64::INFINITY), + Value::Float(f64::NEG_INFINITY), + ])) + .unwrap(); + + assert_eq!(repr(&loads(&written).unwrap()), "[nan, inf, -inf]"); + } + + #[rstest] + fn loads_rejects_nesting_past_max_depth() { + let text = format!( + "{}{}", + "[".repeat(crate::MAX_DEPTH + 1), + "]".repeat(crate::MAX_DEPTH + 1) + ); + + assert!(matches!(loads(&text), Err(Error::TooDeep))); + } +} diff --git a/litellm-rust/crates/python-compat/src/lib.rs b/litellm-rust/crates/python-compat/src/lib.rs index b23aff409f2..2cc0eeac397 100644 --- a/litellm-rust/crates/python-compat/src/lib.rs +++ b/litellm-rust/crates/python-compat/src/lib.rs @@ -7,9 +7,10 @@ //! |---|---| //! | [`literal`] | `ast.literal_eval(text)` | //! | [`repr`] | `repr(value)` and `str(value)` | -//! | [`json`] | `json.dumps(value)`, and `json.loads(json.dumps(value))` as a JSON value | +//! | [`json`] | `json.dumps(value)` and `json.loads(text)` | //! | [`pickle`] | `pickle.loads(data)` and `pickle.dumps(value)` for plain data | //! | [`truthy`] | `bool(value)` | +//! | [`number`] | `int(text)`, `float(text)`, and `int(seconds)` as a `Duration` | //! //! [`Value`] is the closed data model these formats share. Live Python objects //! (descriptors, `__bool__`, `__str__`, callbacks) are out of scope: those belong to the @@ -20,10 +21,12 @@ //! (string hashes are randomized per process). Sets keep their literal order. //! - `str` values are Rust `String`s, so lone surrogates cannot be represented. //! - [`pickle::loads`] decodes `tuple`, `set`, and `frozenset` as lists. +//! - [`number`] reads ASCII digits only, while Python's `int` and `float` accept any Unicode decimal digit. mod error; pub mod json; pub mod literal; +pub mod number; pub mod pickle; pub mod repr; pub mod truthy; diff --git a/litellm-rust/crates/python-compat/src/number.rs b/litellm-rust/crates/python-compat/src/number.rs new file mode 100644 index 00000000000..6e2f067ae25 --- /dev/null +++ b/litellm-rust/crates/python-compat/src/number.rs @@ -0,0 +1,121 @@ +use std::{str::FromStr, time::Duration}; + +use num_bigint::BigInt; + +/// `int(text)` in base 10, or `None` where Python raises `ValueError`. +pub fn int(text: &str) -> Option { + parse(text) +} + +/// `float(text)`, or `None` where Python raises `ValueError`. +pub fn float(text: &str) -> Option { + parse(text) +} + +/// Why [`seconds`] refused a value, as `int(value)` would. Messages are CPython's. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum SecondsError { + /// `OverflowError`. Finite values past `Duration::MAX` land here too, where Python builds a big int. + #[error("cannot convert float infinity to integer")] + Overflow, + /// `ValueError`. + #[error("cannot convert float NaN to integer")] + NotANumber, +} + +/// A float of seconds as a `Duration`, refusing what `int(value)` refuses. Finite negative values clamp to zero. +pub fn seconds(value: f64) -> Result { + match value { + value if value.is_nan() => Err(SecondsError::NotANumber), + value if value.is_infinite() => Err(SecondsError::Overflow), + value => Duration::try_from_secs_f64(value.max(0.0)).map_err(|_| SecondsError::Overflow), + } +} + +fn parse(text: &str) -> Option { + let text = text.trim(); + let digit_at = |index: Option| { + index + .and_then(|index| text.as_bytes().get(index)) + .is_some_and(u8::is_ascii_digit) + }; + let underscores_between_digits = text + .match_indices('_') + .all(|(index, _)| digit_at(index.checked_sub(1)) && digit_at(Some(index + 1))); + if !underscores_between_digits { + return None; + } + text.replace('_', "").parse().ok() +} + +#[cfg(test)] +mod tests { + use num_bigint::BigInt; + use rstest::rstest; + + use std::time::Duration; + + use super::{SecondsError, float, int, seconds}; + + // Expected values are CPython 3.14 int(text) + #[rstest] + #[case::trailing_newline("2\n", Some(2))] + #[case::surrounding_spaces_and_underscore(" 6_380 ", Some(6380))] + #[case::plus_sign("+5", Some(5))] + #[case::negative_zero("-0", Some(0))] + #[case::leading_zeros("007", Some(7))] + #[case::unicode_spaces("\u{a0}7\u{2003}", Some(7))] + #[case::beyond_i64("99999999999999999999", Some(99_999_999_999_999_999_999_i128))] + #[case::double_underscore("1__0", None)] + #[case::leading_underscore("_1", None)] + #[case::trailing_underscore("1_", None)] + #[case::decimal_point("5.0", None)] + #[case::information_separator("\x1c7", None)] + #[case::empty("", None)] + #[case::sign_only("+", None)] + #[case::non_ascii_digit_is_a_known_limit("\u{667}", None)] + fn int_parses_like_python(#[case] text: &str, #[case] expected: Option) { + assert_eq!(int(text), expected.map(BigInt::from)); + } + + // Expected values are CPython 3.14 float(text) + #[rstest] + #[case::underscore_and_newline(" 1_0.5\n", Some(10.5))] + #[case::exponent("1e3", Some(1000.0))] + #[case::leading_point(".5", Some(0.5))] + #[case::trailing_point("5.", Some(5.0))] + #[case::plus_sign("+1.5", Some(1.5))] + #[case::infinity_word("Infinity", Some(f64::INFINITY))] + #[case::infinity_with_space("infinity ", Some(f64::INFINITY))] + #[case::underscore_before_exponent("1_e3", None)] + #[case::underscore_after_exponent("1e_3", None)] + #[case::underscore_after_point("1._5", None)] + #[case::underscore_after_word("inf_", None)] + #[case::hexadecimal("0x10", None)] + #[case::bare_exponent("1.5e", None)] + fn float_parses_like_python(#[case] text: &str, #[case] expected: Option) { + assert_eq!(float(text), expected); + } + + #[rstest] + fn float_reads_a_signed_nan_like_python() { + assert!(float("-NaN").is_some_and(f64::is_nan)); + } + + // CPython 3.14: int(inf) and int(-inf) raise OverflowError, int(nan) raises ValueError + #[rstest] + #[case::fraction(1.5, Ok(Duration::from_millis(1500)))] + #[case::zero(0.0, Ok(Duration::ZERO))] + #[case::negative_clamps_to_zero(-3.0, Ok(Duration::ZERO))] + #[case::infinity(f64::INFINITY, Err(SecondsError::Overflow))] + #[case::negative_infinity(f64::NEG_INFINITY, Err(SecondsError::Overflow))] + #[case::past_duration_max(1e20, Err(SecondsError::Overflow))] + #[case::nan(f64::NAN, Err(SecondsError::NotANumber))] + #[case::negative_nan(-f64::NAN, Err(SecondsError::NotANumber))] + fn seconds_refuses_what_int_refuses( + #[case] value: f64, + #[case] expected: Result, + ) { + assert_eq!(seconds(value), expected); + } +} diff --git a/litellm-rust/crates/python-compat/tests/fixtures.rs b/litellm-rust/crates/python-compat/tests/fixtures.rs index 9c202fdc5f1..3f9c6fd7010 100644 --- a/litellm-rust/crates/python-compat/tests/fixtures.rs +++ b/litellm-rust/crates/python-compat/tests/fixtures.rs @@ -76,6 +76,11 @@ const KNOWN: &[(&str, &str, &str)] = &[ "literal_eval", "deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it", ), + ( + "nested_150", + "json.loads", + "deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it", + ), ( "nested_150", "pickle.loads", @@ -341,3 +346,19 @@ fn to_json_matches_python_json_round_trip(fixtures: &Fixtures) { } mismatches.finish(); } + +/// CPython gives back the same text from `json.dumps(json.loads(text))` for every text its +/// `json.dumps` wrote here, non-finite floats included, so `loads` must too. +#[rstest] +fn loads_reads_back_what_python_json_dumps_wrote(fixtures: &Fixtures) { + let mut mismatches = Mismatches::default(); + for row in &fixtures.rows { + let Some(text) = &row.json else { continue }; + let actual = match json::loads(text).and_then(|value| json::dumps(&value)) { + Ok(text) => text, + Err(error) => format!("error: {error}"), + }; + mismatches.check(row, "json.loads", text, &actual); + } + mismatches.finish(); +} diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index c73a4ed8c4d..6135919c539 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -1,12 +1,12 @@ - The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace - It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default - Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route -- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`) +- Every connection sets TCP_NODELAY, like redis-py's `socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)`. Set every redis-rs timeout and retry explicitly, citing the redis-py 5.3.1 or `litellm/` source next to each constant. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's. `Timeouts` takes `socket_timeout` and `socket_connect_timeout`, and the connect timeout falls back to the socket timeout like redis-py. Standalone reconnects on the next command with no retry of the failed one, matching redis-py's `Retry(NoBackoff(), 0)`. redis-rs keeps a socket whose command timed out, so `Renewing` in `connection.rs` swaps in a fresh `ConnectionManager`, or a fresh cluster connection like redis-py's cluster `aclose`, after a timeout. Without it a socket that went silent times out every command until TCP gives up on it. Cluster keeps redis-py's 3 `cluster_error_retry_attempts` and waits a flat 250ms between them, which is redis-py's CLUSTERDOWN sleep. redis-py's jittered per-node reconnect backoff has no redis-rs equivalent. The retry set also differs, redis-rs retries LOADING and READONLY internally, while redis-py 5.3.1 raises them. Do not claim exact retry parity - Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values ## Errors - `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it -- Classify in `From for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error +- Classify with `Error::classify(error, deployment)`, never a bare `From`, because redis-py (5.3.1) decides by topology. `Transport` is what Python's `_redis_health_error_types` counts: LOADING everywhere, CLUSTERDOWN and MASTERDOWN only on cluster (only `ClusterParser` maps them), READONLY only under Sentinel (where `SentinelManagedConnection` turns it into `ConnectionError` and rediscovers), plus NOAUTH, WRONGPASS and the `ERR` messages redis-py maps to `ConnectionError`, matched exactly like `parse_error`. TRYAGAIN is `Server`, since the cluster client already retried it. Each code gets its own rstest case built from a parsed RESP error - A `Server` error never drops the connection - `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome - Never swallow an error or turn it into a miss. The caller picks the fallback @@ -17,27 +17,61 @@ - Every public call (each command, `eval`, one `run` of a pipeline) takes one admission through `guarded` / `admit` in `client.rs`. Never call the connection around it - `Transport` counts as a failure (`Timeout` when `is_timeout()`, else `Hard`). `Server` and `Decode` count as nothing, since Redis answered and the request was bad. A probe that gets one still reopens, through the dropped-probe rule in `circuit-breaker`. A pipeline counts one failure if any op came back `Transport`, else one success - An open breaker returns `Error::BreakerOpen` without touching the socket -- Python doesn't count MASTERDOWN, TRYAGAIN or READONLY as health failures. Here they are `Transport` and do count ## Namespace - Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace` - The rule never changes, because Python and Rust pods share keys in the same Redis ## Commands -- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- A command is a value, one constructor in `commands.rs` returning `Command`: its alone and pipelined forms and its reply decode. `RedisClient::send`, `Pipeline::with` and `RedisBatch::add` are the only entry points, and each applies its namespace to the keys once. A new command is one constructor, never a method on an executor. `ping` and `subscribe` stay methods. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- A store attaches its value codec with `Command::map`, never a keyspace codec inside this crate - On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test -- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`) +- Every command is the one `RedisCache` sends, argument for argument, so `MONITOR` on a mixed fleet shows one shape. A difference from main's Python is a defect even when the result is the same +- TTLs are `Duration` and are sent in whole seconds (`SET ... EX`, `EXPIRE`, `TTL`), like redis-py's `int(ttl)`. `set_nx` sends `EX` before `NX`, the order redis-py builds +- Numbers go out as `Amount`, written the way redis-py encodes them: an `int` with `str()` and a `float` with `repr()` from `python-compat` +- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` alone is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. Pipelined it is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when an expiry is given, `KeepExisting` lowered to `Refresh`. `Plan` keeps only the cases whose alone and pipelined forms differ ## Scripts and pipelines - A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production -- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async` +- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline a script is sent as `EVALSHA`, like `RedisBatch`. A `NOSCRIPT` reply means that op ran nothing, so it alone runs again after the pipeline through `invoke_async` (`EVALSHA`, `SCRIPT LOAD`, `EVALSHA`), like `RedisBatch`'s `_settle_alone`. It then lands after the ops around it, as it does in Python +- `Plan::Script` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice +- `RedisBatch::pipelined` is one `Pipeline` whose ops answer as typed `Declared`s. A store that sends its own pipeline uses it rather than decoding `run`'s raw `Value`s - A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip -- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order -- Pipeline ops take one key each. A multi-key op would span slots inside one group -- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once +- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key and runs the slots concurrently, returning results in input order. Within a slot, runs of idempotent ops (GET, SET, DEL, EXPIRE) share a request and every script op gets its own, sent in order. redis-rs retries a cluster request whole, so a script op batched with others would be applied twice when only a neighbour failed +- Pipeline ops take one key each. A multi-key op would span slots inside one group. An op can send more than one command for that key (`INCRBYFLOAT` then `EXPIRE`), and settles from the first reply, failing if any of its replies failed +- Command constructors live in `commands.rs` and every executor sends them, so a command is written once -## Tests +## Pub/sub +- Classic `PUBLISH` / `SUBSCRIBE`, never sharded, since redis-py publishes and subscribes through one plain client and cluster broadcasts classic messages to every node. Payloads are raw bytes, and the codec belongs to the store that owns the channel +- Channel names never get the namespace. Python builds `{namespace}:{channel}` itself, so the store does too +- `publish` is a command like any other, guarded by the breaker and reported to events +- `subscribe` returns a `Stream` that opens its own RESP2 `PubSub` connection when first polled, like redis-py's `PubSub`, to the endpoint, the first reachable cluster node, or the master Sentinel names at that moment. It never touches the breaker or events, like Python's subscriber, which reads off its own connection +- After a lost connection or a failed subscribe it waits and resubscribes, 5s doubling to 60s and reset by a successful subscribe, like `AuthCacheInvalidationSubscriber._run`. Messages published while it is down are lost. The connection has no idle ping, so a socket that dies silently is only noticed by TCP +- redis-rs's `PubSub` takes no credentials provider, so a provider is asked for its current credentials once when the subscription connects, which is when redis-py's pubsub connection asks it. The connect is bounded by the connect timeout + +## Batches +- `RedisBatch` is litellm's `RedisBatch` (`litellm/caching/redis_batch.py`). It keeps each declared op as data and, on a standalone or Sentinel Redis, sends them later as one `Pipeline`, so it inherits every pipeline rule above. On a cluster it never pipelines and runs each op on its own, concurrently, the way `_settle_alone` calls the `RedisCache` method behind it. An increment there keeps a TTL the key already has, as `async_increment` does, where the pipelined one refreshes it like `_Increment` +- `mget` is one `MGET` of the keys in order with repeats read once, like `_MGet`, and answers one value per key it was given, in that order. Use it instead of one `get` per key whenever Python reads several keys in one call +- It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local +- `RedisBatch::alone` runs every op on its own like the cluster path, so a fresh alone batch is "send it now" and `Declared` is the one reply handle on every route +- Each op returns a `Declared` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once +- An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis +- `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client - Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes +- Tests that only need the public API live in `tests/.rs`, with client builders in `tests/support`. Inline `#[cfg(test)]` modules are only for private items (`error.rs`, `namespace.rs`) +- Observe connections from the server side (`RedisDeployment::clients`, `cli_on_every_node`), never by reaching into `RedisClient`'s connection - Every command test runs on each topology through `#[values(Topology::Standalone, Topology::Cluster)]`. A test is single-topology only when the behavior exists on one (a `db` number has no cluster meaning) -- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables +- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables. A socket that never answers again comes from `TcpProxy::blackhole_open_connections`, since `thaw` revives the same socket - Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline + +## TLS, credentials and Sentinel +- `TlsOptions` follows redis-py 5.3.1. It verifies the certificate by default, skips the hostname check by default (`ssl_check_hostname=False`), and trusts the system roots plus `ca_certs` +- rustls only reads unencrypted keys, so a PKCS#8 `ENCRYPTED PRIVATE KEY` is decrypted with `key_password` before it reaches redis-rs. A legacy `Proc-Type: 4,ENCRYPTED` PEM key is refused with `KeyError::LegacyEncryption` rather than decrypted by hand-rolled OpenSSL key derivation +- redis-rs loads the system trust store for every TLS connection and fails if that fails, even with `ca_certs` set. Where the macOS keychain is blocked (agent sandboxes), run the TLS tests with `SSL_CERT_FILE=/etc/ssl/cert.pem` +- `Credentials::Static` are sent at connect, and URL userinfo wins over them, like redis-py's `from_url` +- A url's `?protocol=` is ignored and every connection speaks RESP2, which every reply decode assumes +- `Credentials::Provider` wraps a redis-rs `StreamingCredentialsProvider`. redis-rs takes the first item at connect and re-`AUTH`s live connections on every later item, on each cluster node. Providers push, redis-rs never polls. With a provider, redis-rs overwrites URL userinfo at connect, which matches Python stripping `user:pass@` from the URL +- Sentinel discovers the master on each new connection and drops it on a transport failure or a READONLY reply, to one command or inside a pipeline, like redis-py's `SentinelConnectionPool` +- redis-rs sends `ROLE` to the master a Sentinel names before using it, which redis-py never does. That probe gets the data credentials (a provider's current pair) and TLS, so the Sentinel client is built on each discovery. redis-rs gives the probe no way to skip the hostname check, and drops the client certificate when the certificate is not verified, so a TLS master whose certificate does not name the announced address, or that requires a client certificate while verification is off, fails discovery where Python connects +- Sentinel monitors get `sentinel_password`, the data username when that password is set, and the TLS mode and certificates. Python's `_sentinel_auth_kwargs` forwards every other connection kwarg too, and it sends the username even without a sentinel password +- A `Debug` impl never shows a secret. Follow `.agents/skills/rust-redaction` diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index 3e9447a6bfa..d3d3226abc6 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -6,13 +6,28 @@ license.workspace = true repository.workspace = true [dependencies] +tracing.workspace = true litellm-circuit-breaker.workspace = true litellm-secret-value.workspace = true +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true futures-util.workspace = true -redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } +redis = { workspace = true, features = [ + "cluster-async", + "connection-manager", + "sentinel", + "tls-rustls-insecure", + "token-based-authentication", + "tokio-rustls-comp", +] } +pem = "4.0.0" +pkcs8 = { version = "0.10.2", features = ["3des", "encryption", "std"] } +rustls-native-certs.workspace = true +socket2 = "0.6.5" thiserror.workspace = true +tokio = { workspace = true, features = ["rt", "sync", "time"] } [dev-dependencies] litellm-redis-testing.workspace = true +rustls.workspace = true rstest.workspace = true -tokio.workspace = true diff --git a/litellm-rust/crates/redis-client/src/batch.rs b/litellm-rust/crates/redis-client/src/batch.rs new file mode 100644 index 00000000000..84fc1643424 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/batch.rs @@ -0,0 +1,147 @@ +use std::{ + io, + sync::{Arc, Mutex, MutexGuard}, +}; + +use futures_util::future::join_all; +use redis::{RedisError, Value}; +use tokio::sync::oneshot::{self, error::TryRecvError}; + +use crate::{Command, Error, Pipeline, RedisClient, commands::Op}; + +/// Ops declared by many callers and sent together, like litellm's `RedisBatch`: one pipeline on a +/// standalone or Sentinel Redis, and each op on its own on a cluster or in an alone batch. +#[derive(Clone)] +pub struct RedisBatch(Arc); + +#[derive(Clone, Copy)] +enum Mode { + /// litellm's `RedisBatch`: one pipeline, or each op on its own on a cluster. + Batch, + /// Each op on its own, the way the `RedisCache` method behind it sends it. + Alone, + /// One redis-py pipeline on every topology, split by slot on a cluster. + Pipeline, +} + +struct Shared { + client: RedisClient, + mode: Mode, + pending: Mutex, + flushing: tokio::sync::Mutex<()>, +} + +#[derive(Default)] +struct Pending { + ops: Vec, + replies: Vec>>, +} + +type Decode = Box Result + Send>; + +#[must_use] +pub struct Declared { + batch: RedisBatch, + reply: oneshot::Receiver>, + decode: Decode, +} + +impl RedisBatch { + pub(crate) fn new(client: RedisClient) -> Self { + Self::build(client, Mode::Batch) + } + + /// A batch whose every op goes out on its own, the way the `RedisCache` method behind it sends + /// it, so a fresh one is "send it now": the first `resolve` sends it. + pub fn alone(client: RedisClient) -> Self { + Self::build(client, Mode::Alone) + } + + /// A batch sent as one redis-py pipeline, like `async_increment_pipeline`, on every topology. + pub fn pipelined(client: RedisClient) -> Self { + Self::build(client, Mode::Pipeline) + } + + fn build(client: RedisClient, mode: Mode) -> Self { + Self(Arc::new(Shared { + client, + mode, + pending: Mutex::new(Pending::default()), + flushing: tokio::sync::Mutex::new(()), + })) + } + + pub fn add(&self, command: Command) -> Declared { + let (op, decode) = command.prepare(self.0.client.namespace()); + let (sender, reply) = oneshot::channel(); + let mut pending = self.pending(); + pending.ops.push(op); + pending.replies.push(sender); + Declared { + batch: self.clone(), + reply, + decode, + } + } + + pub async fn flush(&self) { + let _flushing = self.0.flushing.lock().await; + let Pending { ops, replies } = std::mem::take(&mut *self.pending()); + if replies.is_empty() { + return; + } + let client = &self.0.client; + let alone = match self.0.mode { + Mode::Alone => true, + Mode::Batch => client.is_cluster(), + Mode::Pipeline => false, + }; + let results = match alone { + true => join_all(ops.into_iter().map(|op| client.run_alone(op))).await, + false => { + client + .run(&Pipeline::of(client.namespace().clone(), ops)) + .await + } + }; + for (reply, result) in replies.into_iter().zip(results) { + let _ = reply.send(result); + } + } + + fn pending(&self) -> MutexGuard<'_, Pending> { + self.0 + .pending + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +impl Declared { + pub async fn resolve(mut self) -> Result { + match self.reply.try_recv() { + Ok(reply) => return reply.and_then(self.decode), + Err(TryRecvError::Closed) => return Err(abandoned()), + Err(TryRecvError::Empty) => {} + } + self.batch.flush().await; + self.settled().await + } + + pub async fn settled(self) -> Result { + let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned())); + reply.and_then(self.decode) + } + + /// Every reply in order, flushing their batch once if nobody has yet. + pub async fn all(declared: impl IntoIterator) -> Vec> { + join_all(declared.into_iter().map(Self::resolve)).await + } +} + +fn abandoned() -> Error { + Error::Transport(RedisError::from(io::Error::new( + io::ErrorKind::Interrupted, + "the batch flush carrying this op was cancelled", + ))) +} diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 48dbc2b2dfb..d9a8b3af7e8 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -1,21 +1,43 @@ use std::{ future::Future, + sync::{Arc, Mutex, PoisonError, Weak}, time::{Duration, Instant}, }; +use futures_util::stream::BoxStream; use litellm_circuit_breaker::{Admission, Breaker, Failure}; -use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value}; +use redis::{Cmd, ErrorKind, RedisError, ServerErrorKind, Value}; use crate::{ - ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands, + CallEvent, Command, ConnectError, Error, FailureKind, Namespace, Operation, Outcome, Pipeline, + RedisBatch, RedisConfig, RedisEvents, + commands::{self, Op, Plan, Trips}, connection::Connection, + pubsub::Subscriber, }; +/// One Redis client. Clones share the connection, the breaker and the subscriber. #[derive(Clone)] -pub struct RedisClient { - pub(crate) connection: Connection, +pub struct RedisClient(Arc); + +/// A handle that does not keep the client alive, for a registry that must not outlive its users. +#[derive(Clone)] +pub struct WeakRedisClient(Weak); + +impl WeakRedisClient { + pub fn upgrade(&self) -> Option { + self.0.upgrade().map(RedisClient) + } +} + +struct Shared { + connection: Connection, + subscriber: Subscriber, namespace: Namespace, breaker: Option, + events: Option>, + health_check: Option, + last_used: Mutex, } #[derive(Debug, PartialEq, Eq)] @@ -27,122 +49,152 @@ pub enum Ttl { impl RedisClient { pub fn new(config: RedisConfig) -> Result { - let connection = match &config.endpoint { - Endpoint::Url(url) => { - Connection::standalone(url.expose().into_connection_info()?, config.timeout)? - } - Endpoint::Host(host) => { - Connection::standalone(host.connection_info()?, config.timeout)? - } - Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout)?, - }; - Ok(Self { + let connection = Connection::new( + &config.target, + &config.credentials, + config.timeouts, + config.keepalive, + )?; + Ok(Self(Arc::new(Shared { + subscriber: Subscriber::new( + &config.target, + &config.credentials, + config.timeouts, + config.keepalive, + &connection, + )?, connection, namespace: config.namespace, breaker: config.breaker, - }) + events: config.events, + health_check: config.health_check, + last_used: Mutex::new(Instant::now()), + }))) + } + + pub fn downgrade(&self) -> WeakRedisClient { + WeakRedisClient(Arc::downgrade(&self.0)) } pub async fn ping(&self) -> Result<(), Error> { - self.query(&commands::ping()).await - } - - pub async fn get(&self, key: &str) -> Result, Error> { - self.query(&commands::get(self.key(key))).await - } - - pub async fn mget(&self, keys: &[&str]) -> Result>, Error> { - if keys.is_empty() { - return Ok(Vec::new()); - } - self.query(&commands::mget(self.keys(keys))).await - } - - pub async fn set( - &self, - key: &str, - value: impl ToSingleRedisArg, - ttl: Duration, - ) -> Result<(), Error> { - self.query(&commands::set(self.key(key), value, ttl)).await - } - - pub async fn set_nx( - &self, - key: &str, - value: impl ToSingleRedisArg, - ttl: Duration, - ) -> Result { - let reply: Option = self - .query(&commands::set_nx(self.key(key), value, ttl)) - .await?; - Ok(reply.is_some()) - } - - pub async fn del(&self, keys: &[&str]) -> Result { - if keys.is_empty() { - return Ok(0); - } - self.query(&commands::del(self.keys(keys))).await - } - - pub async fn expire(&self, key: &str, ttl: Duration) -> Result { - self.query(&commands::expire(self.key(key), ttl)).await - } - - pub async fn ttl(&self, key: &str) -> Result { - let reply: i64 = self.query(&commands::ttl(self.key(key))).await?; - Ok(match reply { - -2 => Ttl::Missing, - -1 => Ttl::Persistent, - millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())), - }) - } - - pub async fn incr_by(&self, key: &str, delta: i64) -> Result { - self.query(&commands::incr_by(self.key(key), delta)).await - } - - pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result { - self.query(&commands::incr_by_float(self.key(key), delta)) - .await - } - - pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result { - self.query(&commands::rpush(self.key(key), values)).await - } - - pub async fn eval( - &self, - script: &Script, - keys: &[&str], - args: impl ToRedisArgs, - ) -> Result { - let keys = self.keys(keys); - let args = args.to_redis_args(); - self.guarded(async { - Ok(script - .invocation(&keys, &args) - .invoke_async(&mut self.connection.clone()) - .await?) + self.guarded(Operation::Command("PING"), async { + commands::ping() + .query_async(&mut self.0.connection.clone()) + .await }) .await } + /// Sends one command on its own, the way the `RedisCache` method behind it sends it. + pub async fn send(&self, command: Command) -> Result { + let (op, decode) = command.prepare(&self.0.namespace); + self.run_alone(op).await.and_then(decode) + } + + pub(crate) async fn run_alone(&self, op: Op) -> Result { + let Op { + operation, plan, .. + } = op; + self.guarded(operation, async { + let mut connection = self.0.connection.clone(); + match plan { + Plan::Reply(value) => Ok(value), + Plan::Commands(commands) => match <[Cmd; 1]>::try_from(commands) { + Ok([command]) => command.query_async(&mut connection).await, + Err(commands) => { + let mut pipeline = redis::pipe(); + for command in commands { + pipeline.add_command(command); + } + let replies: Vec = pipeline.query_async(&mut connection).await?; + Ok(replies.into_iter().next().unwrap_or(Value::Nil)) + } + }, + Plan::Script(call) => { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut connection) + .await + } + Plan::KeepingTtl { + increment, + key, + ttl, + trips, + } => { + let (value, current): (Value, i64) = match trips { + Trips::Shared => { + redis::pipe() + .add_command(increment) + .add_command(commands::ttl_cmd(&key)) + .query_async(&mut connection) + .await? + } + Trips::Separate => ( + increment.query_async(&mut connection).await?, + commands::ttl_cmd(&key).query_async(&mut connection).await?, + ), + }; + if current == -1 { + redis::cmd("EXPIRE") + .arg(&key) + .arg(ttl.as_secs()) + .exec_async(&mut connection) + .await?; + } + Ok(value) + } + } + }) + .await + } + + /// Messages published on `channel`, on a connection of its own. After the connection is lost + /// it reconnects and resubscribes, and messages published meanwhile are not delivered. + pub fn subscribe(&self, channel: &str) -> BoxStream<'static, Vec> { + self.0.subscriber.clone().messages(channel.to_owned()) + } + pub fn pipeline(&self) -> Pipeline { - Pipeline::new(self.namespace.clone()) + Pipeline::new(self.0.namespace.clone()) + } + + pub(crate) fn namespace(&self) -> &Namespace { + &self.0.namespace + } + + pub fn is_cluster(&self) -> bool { + matches!(self.0.connection, Connection::Cluster(_)) + } + + pub fn same_client(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } + + pub fn batch(&self) -> RedisBatch { + RedisBatch::new(self.clone()) } pub async fn run(&self, pipeline: &Pipeline) -> Vec> { + let operation = Operation::Pipeline { + ops: pipeline.len(), + }; + let started = Instant::now(); let admission = match self.admit() { Ok(admission) => admission, Err(_) => { + self.observe( + operation, + started, + Outcome::Failed(FailureKind::BreakerOpen), + ); return (0..pipeline.len()) .map(|_| Err(Error::BreakerOpen)) .collect(); } }; - let results = pipeline.run(&self.connection).await; + self.check_health().await; + let results = pipeline.run(&self.0.connection).await; let transport = results.iter().find_map(|result| { result .as_ref() @@ -150,31 +202,76 @@ impl RedisClient { .filter(|error| matches!(error, Error::Transport(_))) }); settle(admission, transport); + let failure = transport.or_else(|| results.iter().find_map(|result| result.as_ref().err())); + let outcome = failure.map_or(Outcome::Succeeded, |error| { + Outcome::Failed(FailureKind::of(error)) + }); + self.observe(operation, started, outcome); results } - fn key(&self, key: &str) -> String { - self.namespace.apply(key) - } - - fn keys(&self, keys: &[&str]) -> Vec { - keys.iter().map(|key| self.key(key)).collect() - } - - async fn query(&self, command: &Cmd) -> Result { - self.guarded(async { Ok(command.query_async(&mut self.connection.clone()).await?) }) - .await - } - - async fn guarded(&self, call: impl Future>) -> Result { - let admission = self.admit()?; - let result = call.await; - settle(admission, result.as_ref().err()); + async fn guarded( + &self, + operation: Operation, + call: impl Future>, + ) -> Result { + let started = Instant::now(); + let result = match self.admit() { + Ok(admission) => { + self.check_health().await; + let result = call + .await + .map_err(|error| Error::classify(error, self.0.connection.deployment())); + settle(admission, result.as_ref().err()); + result + } + Err(error) => Err(error), + }; + self.observe(operation, started, Outcome::of(&result)); result } + fn observe(&self, operation: Operation, started: Instant, outcome: Outcome) { + *self + .0 + .last_used + .lock() + .unwrap_or_else(PoisonError::into_inner) = Instant::now(); + if let Some(events) = &self.0.events { + events.call(&CallEvent { + operation, + duration: started.elapsed(), + outcome, + }); + } + } + + /// redis-py pings a connection that sat idle longer than `health_check_interval` before + /// sending the next command on it. A ping that times out or finds the socket dropped leaves + /// that connection behind, as a command would, so the command then goes out on a fresh one, + /// where redis-py 5.3.1 disconnects and fails the command with the ping's error. The ping is + /// not a call of the caller's and never reaches the breaker or the events. + async fn check_health(&self) { + let Some(interval) = self.0.health_check else { + return; + }; + let idle = self + .0 + .last_used + .lock() + .unwrap_or_else(PoisonError::into_inner) + .elapsed(); + if idle < interval { + return; + } + if let Err(error) = self.0.connection.ping_every_node().await { + tracing::debug!(%error, "idle Redis connection failed its health check, replacing it"); + } + } + fn admit(&self) -> Result, Error> { - self.breaker + self.0 + .breaker .as_ref() .map(|breaker| { breaker @@ -195,458 +292,15 @@ fn settle(admission: Option, error: Option<&Error>) { admission.failed(Failure::Timeout, Instant::now()) } Some(Error::Transport(_)) => admission.failed(Failure::Hard, Instant::now()), + Some(Error::Server(error)) if proves_redis_answers(error) => admission.succeeded(), Some(Error::Server(_) | Error::Decode(_) | Error::BreakerOpen) => drop(admission), } } -#[cfg(test)] -mod tests { - use std::time::{Duration, Instant}; - - use litellm_redis_testing::RedisServer; - use litellm_secret_value::SecretValue; - use rstest::rstest; - - use litellm_circuit_breaker::{Breaker, Config, State}; - - use super::{RedisClient, Ttl}; - use crate::{ - DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig, - test_support::{Server, Topology, client_id}, - }; - - #[rstest] - #[tokio::test] - async fn namespaces_isolate_keys_on_the_server( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let a = server.client("a", DEFAULT_TIMEOUT).await; - let b = server.client("b", DEFAULT_TIMEOUT).await; - let bare = server.client("", DEFAULT_TIMEOUT).await; - - a.set("x", "1", Duration::from_secs(60)).await.unwrap(); - - assert_eq!(a.get::("x").await.unwrap().as_deref(), Some("1")); - assert_eq!(a.get::("a:x").await.unwrap().as_deref(), Some("1")); - assert_eq!( - bare.get::("a:x").await.unwrap().as_deref(), - Some("1") - ); - assert_eq!(b.get::("x").await.unwrap(), None); - } - - #[rstest] - #[tokio::test] - async fn a_rejected_command_keeps_the_connection( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.rpush("list", "item").await.unwrap(); - let before = client_id(&client, "list").await; - - let rejected = client.incr_by("list", 1).await; - - assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); - assert_eq!(client.rpush("list", "next").await.unwrap(), 2); - assert_eq!(client_id(&client, "list").await, before); - } - - #[rstest] - #[tokio::test] - async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client - .set("float", "15.0", Duration::from_secs(60)) - .await - .unwrap(); - - let decoded = client.get::("float").await; - - assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); - assert_eq!(client.get::("missing").await.unwrap(), None); - } - - #[rstest] - #[tokio::test] - async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - client - .set("first", "1", Duration::from_secs(60)) - .await - .unwrap(); - client - .set("second", "2", Duration::from_secs(60)) - .await - .unwrap(); - - server.freeze().await; - let timed_out = client.get::("first").await; - server.thaw().await; - - assert!( - matches!(timed_out, Err(Error::Transport(_))), - "{timed_out:?}" - ); - assert_eq!( - client.get::("second").await.unwrap().as_deref(), - Some("2") - ); - } - - #[rstest] - #[tokio::test] - async fn the_same_client_recovers_after_redis_restarts( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - - server.restart().await; - - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.incr_by("after_restart", 1).await { - Ok(value) => break Some(value), - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(200)).await - } - Err(_) => break None, - } - }; - assert_eq!(recovered, Some(1)); - } - - #[rstest] - #[tokio::test] - async fn a_client_created_while_redis_is_down_works_once_redis_is_up( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - server.stop().await; - - let client = server.client("", Duration::from_millis(500)).await; - let while_down = client.ping().await; - server.resume().await; - - assert!( - matches!(while_down, Err(Error::Transport(_))), - "{while_down:?}" - ); - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.ping().await { - Ok(()) => break true, - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(200)).await - } - Err(_) => break false, - } - }; - assert!(recovered); - } - - #[rstest] - #[tokio::test] - async fn empty_batches_succeed_without_asking_redis( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - assert_eq!( - client.mget::(&[]).await.unwrap(), - Vec::>::new() - ); - assert_eq!(client.del(&[]).await.unwrap(), 0); - } - - #[rstest] - #[tokio::test] - async fn mget_and_del_span_slots_in_input_order( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.set("a", "1", Duration::from_secs(60)).await.unwrap(); - client.set("c", "3", Duration::from_secs(60)).await.unwrap(); - - let values = client.mget::(&["c", "b", "a"]).await.unwrap(); - let deleted = client.del(&["a", "b", "c"]).await.unwrap(); - - assert_eq!( - values, - vec![Some("3".to_owned()), None, Some("1".to_owned())] - ); - assert_eq!(deleted, 2); - assert_eq!( - client.mget::(&["a", "c"]).await.unwrap(), - vec![None, None] - ); - } - - #[rstest] - #[tokio::test] - async fn set_nx_only_writes_a_missing_key( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - let first = client - .set_nx("lock", "owner-1", Duration::from_secs(60)) - .await - .unwrap(); - let second = client - .set_nx("lock", "owner-2", Duration::from_secs(60)) - .await - .unwrap(); - - assert!(first); - assert!(!second); - assert_eq!( - client.get::("lock").await.unwrap().as_deref(), - Some("owner-1") - ); - } - - #[rstest] - #[tokio::test] - async fn ttl_reports_missing_persistent_and_remaining_time( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.incr_by("persistent", 1).await.unwrap(); - client - .set("expiring", "1", Duration::from_secs(60)) - .await - .unwrap(); - client.incr_by_float("expired_later", 0.5).await.unwrap(); - assert!( - client - .expire("expired_later", Duration::from_secs(30)) - .await - .unwrap() - ); - - assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing); - assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent); - assert!(matches!( - client.ttl("expiring").await.unwrap(), - Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) - )); - assert!(matches!( - client.ttl("expired_later").await.unwrap(), - Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) - )); - assert!( - !client - .expire("missing", Duration::from_secs(30)) - .await - .unwrap() - ); - } - - #[rstest] - #[tokio::test] - async fn incr_by_float_accumulates_fractions( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - client.incr_by_float("spend", 0.25).await.unwrap(); - - assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75); - } - - #[rstest] - #[tokio::test] - async fn concurrent_commands_on_one_cloned_client_all_apply( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - let tasks: Vec<_> = (0..100) - .map(|_| { - let client = client.clone(); - tokio::spawn(async move { - for _ in 0..10 { - client.incr_by("shared", 1).await.unwrap(); - } - }) - }) - .collect(); - for task in tasks { - task.await.unwrap(); - } - - assert_eq!(client.get::("shared").await.unwrap(), Some(1000)); - } - - fn breaker(failure_threshold: u32) -> Breaker { - Breaker::new( - Config { - failure_threshold, - recovery_timeout: Duration::from_millis(200), - timeout_min_duration: Duration::ZERO, - }, - Box::new(|_| {}), - ) - } - - #[rstest] - #[tokio::test] - async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let breaker = breaker(2); - let client = server - .guarded_client("", Duration::from_millis(300), Some(breaker.clone())) - .await; - - server.freeze().await; - let first = client.get::("k").await; - let second = client.get::("k").await; - let started = Instant::now(); - let third = client.get::("k").await; - let rejected_in = started.elapsed(); - server.thaw().await; - - assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); - assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); - assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); - assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); - - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.get::("k").await { - Ok(value) => break Some(value), - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(50)).await - } - Err(_) => break None, - } - }; - assert_eq!(recovered, Some(None)); - assert!( - matches!(breaker.state(), State::Closed { .. }), - "{:?}", - breaker.state() - ); - } - - #[rstest] - #[tokio::test] - async fn rejected_commands_never_open_the_breaker( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let breaker = breaker(2); - let client = server - .guarded_client("", DEFAULT_TIMEOUT, Some(breaker.clone())) - .await; - client.rpush("list", "item").await.unwrap(); - - let rejected = [ - client.incr_by("list", 1).await, - client.incr_by("list", 1).await, - client.incr_by("list", 1).await, - ]; - - assert!( - matches!(rejected[2], Err(Error::Server(_))), - "{:?}", - rejected[2] - ); - assert!( - matches!(breaker.state(), State::Closed { hard: 0, .. }), - "{:?}", - breaker.state() - ); - } - - #[rstest] - #[tokio::test] - async fn an_open_breaker_fails_every_pipeline_op_without_sending( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server - .guarded_client("", Duration::from_millis(300), Some(breaker(1))) - .await; - let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); - - server.freeze().await; - let failed = client.run(&pipeline).await; - let rejected = client.run(&pipeline).await; - server.thaw().await; - - assert!( - matches!(failed[0], Err(Error::Transport(_))), - "{:?}", - failed[0] - ); - assert!( - matches!(rejected[0], Err(Error::BreakerOpen)), - "{:?}", - rejected[0] - ); - assert!( - matches!(rejected[1], Err(Error::BreakerOpen)), - "{:?}", - rejected[1] - ); - } - - #[rstest] - #[tokio::test] - async fn a_host_endpoint_connects_to_its_database() { - let server = RedisServer::start().await.unwrap(); - let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); - let connect = |endpoint| { - RedisClient::new(RedisConfig { - endpoint, - namespace: Namespace::default(), - timeout: DEFAULT_TIMEOUT, - breaker: None, - }) - }; - let in_db_3 = connect(Endpoint::Host(HostEndpoint { - host: "127.0.0.1".to_owned(), - port, - username: None, - password: None, - tls: false, - db: 3, - })) - .unwrap(); - let in_db_0 = connect(Endpoint::Url(SecretValue::new(server.url().to_owned()))).unwrap(); - let via_url = connect(Endpoint::Url(SecretValue::new(format!( - "{}/3", - server.url() - )))) - .unwrap(); - - in_db_3 - .set("k", "3", Duration::from_secs(60)) - .await - .unwrap(); - - assert_eq!(in_db_0.get::("k").await.unwrap(), None); - assert_eq!( - via_url.get::("k").await.unwrap().as_deref(), - Some("3") - ); - } +// A replica refusing a write is the old primary mid-failover, which must not reset a run of failures +fn proves_redis_answers(error: &RedisError) -> bool { + !matches!( + error.kind(), + ErrorKind::InvalidClientConfig | ErrorKind::Server(ServerErrorKind::ReadOnly) + ) } diff --git a/litellm-rust/crates/redis-client/src/commands.rs b/litellm-rust/crates/redis-client/src/commands.rs index fb8a4a28c94..5f8c1eff3f7 100644 --- a/litellm-rust/crates/redis-client/src/commands.rs +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -1,62 +1,407 @@ use std::time::Duration; -use redis::{Cmd, ToRedisArgs, ToSingleRedisArg}; +use litellm_python_compat::repr::float_repr; +use redis::{ + Cmd, ErrorKind, FromRedisValue, RedisError, RedisWrite, ToRedisArgs, ToSingleRedisArg, Value, +}; + +use crate::{Error, Namespace, Operation, Script, Ttl, script::ScriptCall}; + +/// How an increment treats the key's TTL, the two ways `RedisCache` increments with one. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Expiry { + /// `async_increment(refresh_ttl=True)` and every pipelined increment: `EXPIRE` after each write. + Refresh(Duration), + /// `async_increment(refresh_ttl=False)`: read the `TTL` and `EXPIRE` only a key that has none. + /// Pipelined it is lowered to [`Expiry::Refresh`], like `RedisBatch.increment`. + KeepExisting(Duration), +} + +/// A number written the way redis-py encodes it: an `int` with `str()`, a `float` with `repr()`. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Amount { + Int(i64), + Float(f64), +} + +impl ToRedisArgs for Amount { + fn write_redis_args(&self, out: &mut W) { + match self { + Self::Int(value) => out.write_arg(value.to_string().as_bytes()), + Self::Float(value) => out.write_arg(float_repr(*value).as_bytes()), + } + } +} + +impl ToSingleRedisArg for Amount {} + +type Build = Box) -> Plan + Send>; +type Decode = Box Result + Send>; + +/// One Redis command as a value: what it sends alone and pipelined, and how its reply decodes. Every +/// executor (`RedisClient::send`, `Pipeline::with`, `RedisBatch::add`) applies its namespace to the +/// keys once. +#[must_use] +pub struct Command { + operation: Operation, + keys: Vec, + idempotent: bool, + build: Build, + decode: Decode, +} + +/// What a command sends once its keys carry the namespace. Only the cases whose alone and pipelined +/// forms differ get a variant of their own. +pub(crate) enum Plan { + /// Answered without a round trip, like an empty `MGET` or `DEL`. + Reply(Value), + /// Sent together; the reply is the first command's, failing if any failed. + Commands(Vec), + /// Alone through `invoke_async`, pipelined as `EVALSHA` resent once on `NOSCRIPT`. + Script(ScriptCall), + /// Alone the increment reads the `TTL` and sends `EXPIRE` only to a key without one, pipelined it + /// is the increment then `EXPIRE`. + KeepingTtl { + increment: Cmd, + key: String, + ttl: Duration, + trips: Trips, + }, +} + +#[derive(Clone, Copy)] +pub(crate) enum Trips { + /// `async_increment`: the increment and its `TTL` in one round trip. + Shared, + /// `increment_cache`: the increment and its `TTL` each on its own round trip. + Separate, +} + +/// A command with the namespace applied, ready for an executor to send. +pub(crate) struct Op { + pub(crate) operation: Operation, + pub(crate) key: String, + pub(crate) idempotent: bool, + pub(crate) plan: Plan, +} + +impl Plan { + /// What a pipeline sends for this op. + pub(crate) fn pipelined(&self) -> Vec { + match self { + Self::Reply(_) => Vec::new(), + Self::Commands(commands) => commands.clone(), + Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)], + Self::KeepingTtl { + increment, + key, + ttl, + .. + } => vec![increment.clone(), expire_cmd(key, *ttl)], + } + } +} + +impl Command { + fn new( + operation: Operation, + keys: Vec, + idempotent: bool, + build: impl FnOnce(Vec) -> Plan + Send + 'static, + decode: impl FnOnce(Value) -> Result + Send + 'static, + ) -> Self { + Self { + operation, + keys, + idempotent, + build: Box::new(build), + decode: Box::new(decode), + } + } + + /// Attaches a codec over the decoded reply, such as a store's value format. + pub fn map(self, decode: fn(T) -> Result) -> Command + where + T: 'static, + { + let inner = self.decode; + Command { + operation: self.operation, + keys: self.keys, + idempotent: self.idempotent, + build: self.build, + decode: Box::new(move |value| inner(value).and_then(decode)), + } + } + + pub(crate) fn prepare(self, namespace: &Namespace) -> (Op, Decode) { + let keys: Vec = self.keys.iter().map(|key| namespace.apply(key)).collect(); + let key = keys.first().cloned().unwrap_or_default(); + let op = Op { + operation: self.operation, + key, + idempotent: self.idempotent, + plan: (self.build)(keys), + }; + (op, self.decode) + } +} + +fn decode(value: Value) -> Result { + T::from_redis_value(value).map_err(|error| Error::Decode(error.into())) +} + +fn owned(keys: &[&str]) -> Vec { + keys.iter().map(|key| (*key).to_owned()).collect() +} + +fn single(keys: Vec) -> String { + keys.into_iter().next().unwrap_or_default() +} + +fn expire_cmd(key: &str, ttl: Duration) -> Cmd { + redis::cmd("EXPIRE").arg(key).arg(ttl.as_secs()).to_owned() +} + +pub(crate) fn ttl_cmd(key: &str) -> Cmd { + redis::cmd("TTL").arg(key).to_owned() +} + +/// A reply whose shape a codec over it refuses, as `what` describes. +pub fn unexpected_reply(what: &'static str) -> Error { + Error::Decode(RedisError::from((ErrorKind::UnexpectedReturnType, what))) +} + +pub fn get(key: &str) -> Command> { + Command::new( + Operation::Command("GET"), + vec![key.to_owned()], + true, + |keys| Plan::Commands(vec![redis::cmd("GET").arg(single(keys)).to_owned()]), + decode, + ) +} + +/// One `MGET` of `keys` in order, repeated keys read once, like `RedisBatch.mget`, answering one +/// value per key of `keys` in their order. No keys is no round trip, since Redis rejects an empty +/// `MGET`. Pipelined, every key must share a hash slot on a cluster, and a batch never pipelines +/// there. +pub fn mget(keys: &[&str]) -> Command>> { + let unique = keys.iter().fold(Vec::::new(), |mut unique, key| { + if !unique.iter().any(|seen| seen == key) { + unique.push((*key).to_owned()); + } + unique + }); + let positions: Vec = keys + .iter() + .filter_map(|key| unique.iter().position(|seen| seen == key)) + .collect(); + Command::new( + Operation::Command("MGET"), + unique, + true, + |keys| match keys.is_empty() { + true => Plan::Reply(Value::Array(Vec::new())), + false => Plan::Commands(vec![redis::cmd("MGET").arg(keys).to_owned()]), + }, + move |reply| { + let values: Vec = decode(reply)?; + let in_key_order = positions + .iter() + .map(|&position| values.get(position).cloned()) + .collect::>>() + .ok_or_else(|| unexpected_reply("MGET answered fewer values than keys"))?; + decode(Value::Array(in_key_order)) + }, + ) +} + +// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())` +pub fn set(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<()> { + let value = value.to_redis_args().concat(); + Command::new( + Operation::Command("SET"), + vec![key.to_owned()], + true, + move |keys| { + Plan::Commands(vec![ + redis::cmd("SET") + .arg(single(keys)) + .arg(value) + .arg("EX") + .arg(ttl.as_secs()) + .to_owned(), + ]) + }, + decode, + ) +} + +/// `SET ... EX ... NX`, `EX` before `NX` as redis-py builds it. `true` when this call wrote the key. +pub fn set_nx(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command { + let value = value.to_redis_args().concat(); + Command::new( + Operation::Command("SET"), + vec![key.to_owned()], + false, + move |keys| { + Plan::Commands(vec![ + redis::cmd("SET") + .arg(single(keys)) + .arg(value) + .arg("EX") + .arg(ttl.as_secs()) + .arg("NX") + .to_owned(), + ]) + }, + |value| decode::>(value).map(|reply| reply.is_some()), + ) +} + +/// One multi-key `DEL`, no round trip for no keys. Pipelined, every key must share a hash slot. +pub fn del(keys: &[&str]) -> Command { + Command::new( + Operation::Command("DEL"), + owned(keys), + true, + |keys| match keys.is_empty() { + true => Plan::Reply(Value::Int(0)), + false => Plan::Commands(vec![redis::cmd("DEL").arg(keys).to_owned()]), + }, + decode, + ) +} + +pub fn expire(key: &str, ttl: Duration) -> Command { + Command::new( + Operation::Command("EXPIRE"), + vec![key.to_owned()], + true, + move |keys| Plan::Commands(vec![expire_cmd(&single(keys), ttl)]), + decode, + ) +} + +pub fn ttl(key: &str) -> Command { + Command::new( + Operation::Command("TTL"), + vec![key.to_owned()], + true, + |keys| Plan::Commands(vec![ttl_cmd(&single(keys))]), + |value| { + Ok(match decode::(value)? { + -2 => Ttl::Missing, + -1 => Ttl::Persistent, + seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())), + }) + }, + ) +} + +/// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no TTL, +/// each its own round trip. +pub fn incr_by(key: &str, delta: i64, ttl: Duration) -> Command { + Command::new( + Operation::Command("INCRBY"), + vec![key.to_owned()], + false, + move |keys| { + let key = single(keys); + Plan::KeepingTtl { + increment: redis::cmd("INCRBY").arg(&key).arg(delta).to_owned(), + key, + ttl, + trips: Trips::Separate, + } + }, + decode, + ) +} + +/// `INCRBYFLOAT`, with the TTL handled as `expiry` says, or left alone without one. Alone it is +/// `RedisCache.async_increment`, pipelined `RedisBatch.increment`. +pub fn incr_by_float(key: &str, amount: Amount, expiry: Option) -> Command { + Command::new( + Operation::Command("INCRBYFLOAT"), + vec![key.to_owned()], + false, + move |keys| { + let key = single(keys); + let increment = redis::cmd("INCRBYFLOAT").arg(&key).arg(amount).to_owned(); + match expiry { + None => Plan::Commands(vec![increment]), + Some(Expiry::Refresh(ttl)) => { + Plan::Commands(vec![increment, expire_cmd(&key, ttl)]) + } + Some(Expiry::KeepExisting(ttl)) => Plan::KeepingTtl { + increment, + key, + ttl, + trips: Trips::Shared, + }, + } + }, + decode, + ) +} + +/// A registered script, like redis-py's `Script`. Every key it touches must be in `keys`. +pub fn eval( + script: &Script, + keys: &[&str], + args: impl ToRedisArgs, +) -> Command { + let script = script.clone(); + let args = args.to_redis_args(); + Command::new( + Operation::Script, + owned(keys), + false, + move |keys| Plan::Script(ScriptCall { script, keys, args }), + decode, + ) +} + +/// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the scripts +/// `RedisCache` runs that way rather than through `register_script`. +pub fn eval_source( + source: &'static str, + keys: &[&str], + args: impl ToRedisArgs, +) -> Command { + let args = args.to_redis_args(); + Command::new( + Operation::Script, + owned(keys), + false, + move |keys| { + Plan::Commands(vec![ + redis::cmd("EVAL") + .arg(source) + .arg(keys.len()) + .arg(keys) + .arg(args) + .to_owned(), + ]) + }, + decode, + ) +} + +/// The channel is sent as given, without the namespace, like redis-py's `publish`. +pub fn publish(channel: &str, payload: &[u8]) -> Command<()> { + let command = redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned(); + Command::new( + Operation::Command("PUBLISH"), + Vec::new(), + false, + move |_| Plan::Commands(vec![command]), + |_| Ok(()), + ) +} pub(crate) fn ping() -> Cmd { redis::cmd("PING") } - -pub(crate) fn get(key: String) -> Cmd { - redis::cmd("GET").arg(key).to_owned() -} - -pub(crate) fn mget(keys: Vec) -> Cmd { - redis::cmd("MGET").arg(keys).to_owned() -} - -pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { - redis::cmd("SET") - .arg(key) - .arg(value) - .arg("PX") - .arg(millis(ttl)) - .to_owned() -} - -pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { - redis::cmd("SET") - .arg(key) - .arg(value) - .arg("NX") - .arg("PX") - .arg(millis(ttl)) - .to_owned() -} - -pub(crate) fn del(keys: Vec) -> Cmd { - redis::cmd("DEL").arg(keys).to_owned() -} - -pub(crate) fn expire(key: String, ttl: Duration) -> Cmd { - redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned() -} - -pub(crate) fn ttl(key: String) -> Cmd { - redis::cmd("PTTL").arg(key).to_owned() -} - -pub(crate) fn incr_by(key: String, delta: i64) -> Cmd { - redis::cmd("INCRBY").arg(key).arg(delta).to_owned() -} - -pub(crate) fn incr_by_float(key: String, delta: f64) -> Cmd { - redis::cmd("INCRBYFLOAT").arg(key).arg(delta).to_owned() -} - -pub(crate) fn rpush(key: String, values: impl ToRedisArgs) -> Cmd { - redis::cmd("RPUSH").arg(key).arg(values).to_owned() -} - -fn millis(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index c40b9eb6b36..f43ab4feb85 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -1,118 +1,458 @@ -use std::time::Duration; +use std::{fmt, path::PathBuf, pin::Pin, sync::Arc, time::Duration}; -use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; +use futures_util::{Stream, StreamExt}; + +use redis::{ + BasicAuth, ClientTlsConfig, ConnectionAddr, ConnectionInfo, ErrorKind, IntoConnectionInfo, + ProtocolVersion, RedisConnectionInfo, RedisError, RedisResult, StreamingCredentialsProvider, + TlsCertificates, TlsMode, io::tcp::TcpSettings, +}; use litellm_circuit_breaker::Breaker; use litellm_secret_value::SecretValue; -use crate::Namespace; +use pkcs8::EncryptedPrivateKeyInfo; +use crate::{ConnectError, KeyError, Namespace, RedisEvents}; +use litellm_redis_types::{Cluster, RedisUrl, Target, TlsOptions}; + +// litellm/constants.py REDIS_SOCKET_TIMEOUT, which _coerce_redis_kwargs_types sets when socket_timeout is unset pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); +// below the idle timeouts of common cloud load balancers (AWS NLB 350s, Azure Load Balancer 240s), which the +// OS default of two hours never probes within +const KEEPALIVE_IDLE: Duration = Duration::from_secs(60); pub struct RedisConfig { - pub endpoint: Endpoint, + pub target: Target, + pub credentials: Credentials, pub namespace: Namespace, - pub timeout: Duration, + pub timeouts: Timeouts, pub breaker: Option, + pub events: Option>, + /// Ping every node before the next command once the client sat idle this long, like + /// redis-py's `health_check_interval`, so a connection the server dropped silently is + /// replaced before a real command pays for it. + pub health_check: Option, + /// Probe a connection with TCP keepalive once it sat idle, like redis-py's `socket_keepalive`. + pub keepalive: bool, } -#[derive(Clone, Debug, PartialEq)] -pub enum Endpoint { - Url(SecretValue), - Host(HostEndpoint), - Cluster(ClusterEndpoint), +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Timeouts { + pub response: Duration, + pub connect: Duration, } -#[derive(Clone, Debug, PartialEq)] -pub struct HostEndpoint { - pub host: String, - pub port: u16, - pub username: Option, - pub password: Option, - pub tls: bool, - pub db: i64, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct ClusterEndpoint { - pub nodes: Vec, - pub username: Option, - pub password: Option, - pub tls: bool, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct ClusterNode { - pub host: String, - pub port: u16, -} - -impl ClusterEndpoint { - pub(crate) fn addrs(&self) -> impl Iterator + '_ { - self.nodes - .iter() - .map(|node| addr(&node.host, node.port, self.tls)) +impl Timeouts { + pub fn new(socket_timeout: Duration, connect_timeout: Option) -> Self { + Self { + response: socket_timeout, + connect: connect_timeout.unwrap_or(socket_timeout), + } } } -impl HostEndpoint { - pub(crate) fn connection_info(&self) -> RedisResult { - let settings = RedisConnectionInfo::default().set_db(self.db); - let settings = match &self.username { - Some(username) => settings.set_username(username.expose()), - None => settings, - }; - let settings = match &self.password { - Some(password) => settings.set_password(password.expose()), - None => settings, - }; - Ok(addr(&self.host, self.port, self.tls) - .into_connection_info()? - .set_redis_settings(settings)) +impl Default for Timeouts { + fn default() -> Self { + Self::new(DEFAULT_TIMEOUT, None) } } -fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr { +#[derive(Clone, Debug)] +pub enum Credentials { + Static { + username: Option, + password: Option, + }, + Provider(CredentialsProvider), +} + +#[derive(Clone)] +pub struct CredentialsProvider(pub Arc); + +impl fmt::Debug for CredentialsProvider { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CredentialsProvider") + .finish_non_exhaustive() + } +} + +impl StreamingCredentialsProvider for CredentialsProvider { + fn subscribe(&self) -> Pin> + Send + 'static>> { + self.0.subscribe() + } +} + +impl Credentials { + pub const NONE: Self = Self::Static { + username: None, + password: None, + }; +} + +pub(crate) fn tls_mode(tls: &TlsOptions) -> TlsMode { + match tls.verify_certificate { + true => TlsMode::Secure, + false => TlsMode::Insecure, + } +} + +pub(crate) fn tls_certificates(tls: &TlsOptions) -> Result { + let client_tls = match (&tls.certfile, &tls.keyfile) { + (Some(certfile), keyfile) => Some(ClientTlsConfig { + client_cert: read(certfile)?, + client_key: client_key( + keyfile.as_ref().unwrap_or(certfile), + tls.key_password.as_ref(), + )?, + }), + (None, _) => None, + }; + let from_file: Option> = tls.ca_certs.as_ref().map(read).transpose()?; + let from_directory: Vec> = tls + .ca_path + .as_ref() + .map(read_directory) + .transpose()? + .unwrap_or_default(); + let configured: Vec> = from_file + .into_iter() + .chain(tls.ca_data.iter().map(|pem| pem.clone().into_bytes())) + .chain(from_directory) + .collect(); + Ok(TlsCertificates { + client_tls, + root_cert: (!configured.is_empty()).then(|| with_system_roots(configured.join(&b'\n'))), + }) +} + +/// Every regular file in `path`, in name order, as redis-py hands `ssl_ca_path` to +/// `load_verify_locations(capath=...)` +fn read_directory(path: &PathBuf) -> Result>, ConnectError> { + let tls_file = |source| ConnectError::TlsFile { + path: path.display().to_string(), + source, + }; + let mut files: Vec = std::fs::read_dir(path) + .map_err(tls_file)? + .map(|entry| entry.map(|entry| entry.path())) + .collect::>() + .map_err(tls_file)?; + files.sort(); + files + .iter() + .filter(|file| file.is_file()) + .map(read) + .collect() +} + +pub(crate) fn apply_tls( + tls: &TlsOptions, + info: ConnectionInfo, +) -> Result { + let tls_info = redis::Client::build_with_tls(info, tls_certificates(tls)?)? + .get_connection_info() + .clone(); + let mut addr = match tls_info.addr().clone() { + ConnectionAddr::TcpTls { + host, + port, + tls_params, + .. + } => ConnectionAddr::TcpTls { + host, + port, + insecure: !tls.verify_certificate, + tls_params, + }, + addr => addr, + }; + addr.set_danger_accept_invalid_hostnames(!tls.check_hostname); + Ok(tls_info.set_addr(addr)) +} + +fn with_system_roots(configured: Vec) -> Vec { + let system = rustls_native_certs::load_native_certs(); + let system_pem = system + .certs + .iter() + .map(|cert| pem::encode(&pem::Pem::new("CERTIFICATE", cert.to_vec()))); + configured + .into_iter() + .chain(b"\n".iter().copied()) + .chain(system_pem.flat_map(String::into_bytes)) + .collect() +} + +/// The key as unencrypted PEM, decrypting a PKCS#8 `ENCRYPTED PRIVATE KEY` with `ssl_password` the way +/// `load_cert_chain(password=...)` does, since rustls only reads unencrypted keys. Like OpenSSL, a password next to +/// an unencrypted key is unused +fn client_key(path: &PathBuf, password: Option<&SecretValue>) -> Result, ConnectError> { + let contents = read(path)?; + let key_error = |reason| ConnectError::TlsKey { + path: path.display().to_string(), + reason, + }; + let Ok(sections) = pem::parse_many(&contents) else { + return Ok(contents); + }; + let encrypted = sections + .iter() + .find(|section| section.tag() == "ENCRYPTED PRIVATE KEY"); + match (encrypted, password) { + (Some(_), None) => Err(key_error(KeyError::MissingPassword)), + (Some(section), Some(password)) => { + let decrypted = EncryptedPrivateKeyInfo::try_from(section.contents()) + .and_then(|info| info.decrypt(password.expose())) + .map_err(|source| key_error(KeyError::Decrypt(source)))?; + Ok(pem::encode(&pem::Pem::new("PRIVATE KEY", decrypted.as_bytes())).into_bytes()) + } + (None, _) if sections.iter().any(is_legacy_encrypted) => { + Err(key_error(KeyError::LegacyEncryption)) + } + (None, _) => Ok(contents), + } +} + +fn is_legacy_encrypted(section: &pem::Pem) -> bool { + section + .headers() + .get("Proc-Type") + .is_some_and(|proc_type| proc_type.contains("ENCRYPTED")) +} + +fn read(path: &PathBuf) -> Result, ConnectError> { + std::fs::read(path).map_err(|source| ConnectError::TlsFile { + path: path.display().to_string(), + source, + }) +} + +pub(crate) fn addr(host: &str, port: u16, tls: Option<&TlsOptions>) -> ConnectionAddr { match tls { - true => ConnectionAddr::TcpTls { + Some(tls) => ConnectionAddr::TcpTls { host: host.to_owned(), port, - insecure: false, + insecure: !tls.verify_certificate, tls_params: None, }, - false => ConnectionAddr::Tcp(host.to_owned(), port), + None => ConnectionAddr::Tcp(host.to_owned(), port), } } +pub(crate) fn cluster_addrs(cluster: &Cluster) -> impl Iterator + '_ { + cluster + .nodes + .iter() + .map(|node| addr(&node.host, node.port, cluster.tls.as_ref())) +} + +pub(crate) fn tcp_settings(keepalive: bool) -> TcpSettings { + let settings = TcpSettings::default().set_nodelay(true); + match keepalive { + true => settings.set_keepalive(socket2::TcpKeepalive::new().with_time(KEEPALIVE_IDLE)), + false => settings, + } +} + +pub(crate) fn connection_info( + target: &Target, + credentials: &Credentials, + keepalive: bool, +) -> Result { + let (info, tls) = match target { + Target::Url(url) => (url_info(url, credentials)?, url.tls.as_ref()), + Target::Standalone(host) => { + let info = addr(&host.host, host.port, host.tls.as_ref()).into_connection_info()?; + let settings = info.redis_settings().clone().set_db(host.db); + (info.set_redis_settings(settings), host.tls.as_ref()) + } + Target::Unix(socket) => { + let info = ConnectionAddr::Unix(socket.path.clone()).into_connection_info()?; + let settings = info.redis_settings().clone().set_db(socket.db); + (info.set_redis_settings(settings), None) + } + Target::Cluster(_) | Target::Sentinel(_) => { + unreachable!("cluster and sentinel endpoints build their own connection info") + } + }; + let info = with_static_credentials(info, credentials).set_tcp_settings(tcp_settings(keepalive)); + match (tls, info.addr()) { + (Some(tls), ConnectionAddr::TcpTls { .. }) => apply_tls(tls, info), + _ => Ok(info), + } +} + +// parse_url reads ?db= over the path, and the url's userinfo over ?username= and ?password=. A credential +// provider supersedes both, as it does in redis-py. ?protocol=3 would switch redis-rs to RESP3, which no +// reply decode here expects, so the protocol is ignored as the resolver warns +fn url_info(url: &RedisUrl, credentials: &Credentials) -> Result { + let info = url.url.expose().into_connection_info()?; + let settings = info + .redis_settings() + .clone() + .set_protocol(ProtocolVersion::RESP2); + let settings = match url.db { + Some(db) => settings.set_db(db), + None => settings, + }; + let info = info.set_redis_settings(settings); + Ok(match credentials { + Credentials::Static { .. } => with_static_credentials( + info, + &Credentials::Static { + username: url.username.clone(), + password: url.password.clone(), + }, + ), + Credentials::Provider(_) => info, + }) +} + +fn with_static_credentials(info: ConnectionInfo, credentials: &Credentials) -> ConnectionInfo { + let settings = static_credentials(info.redis_settings().clone(), credentials); + info.set_redis_settings(settings) +} + +fn static_credentials( + settings: RedisConnectionInfo, + credentials: &Credentials, +) -> RedisConnectionInfo { + let Credentials::Static { username, password } = credentials else { + return settings; + }; + let settings = match (username, settings.username()) { + (Some(username), None) => settings.set_username(username.expose()), + _ => settings, + }; + match (password, settings.password()) { + (Some(password), None) => settings.set_password(password.expose()), + _ => settings, + } +} + +/// `settings` with the credentials a connection opened now sends. A provider is asked once for its +/// current pair, which is when redis-py asks its `credential_provider`. +pub(crate) async fn authenticated( + settings: RedisConnectionInfo, + credentials: &Credentials, +) -> RedisResult { + let Credentials::Provider(provider) = credentials else { + return Ok(static_credentials(settings, credentials)); + }; + let current = provider.subscribe().next().await.ok_or_else(|| { + RedisError::from(( + ErrorKind::AuthenticationFailed, + "the Redis credential provider ended before giving credentials", + )) + })??; + Ok(settings + .set_username(current.username()) + .set_password(current.password())) +} + #[cfg(test)] mod tests { - use litellm_secret_value::SecretValue; + use litellm_redis_types::{RedisUrl, UnixSocket}; use rstest::rstest; - use super::{ClusterEndpoint, Endpoint, HostEndpoint}; + use super::*; #[rstest] - #[case::url(Endpoint::Url(SecretValue::new("redis://credential-value@cache:6379/0")))] - #[case::host( - Endpoint::Host(HostEndpoint { - host: "cache".to_owned(), - port: 6379, - username: Some(SecretValue::new("credential-value")), - password: Some(SecretValue::new("credential-value")), - tls: false, - db: 0, - }) - )] - #[case::cluster(Endpoint::Cluster(ClusterEndpoint { - nodes: Vec::new(), - username: Some(SecretValue::new("credential-value")), - password: Some(SecretValue::new("credential-value")), - tls: false, + #[case::host(Target::Standalone(crate::Standalone { host: "redis".into(), port: 6379, tls: None, db: 0 }))] + #[case::url(Target::Url(RedisUrl { + url: SecretValue::new("redis://cache:6379".to_owned()), + db: None, + username: None, + password: None, + tls: None, }))] - fn debug_output_never_shows_the_secret(#[case] endpoint: Endpoint) { - let debug = format!("{endpoint:?}"); + fn keepalive_probes_an_idle_socket_only_when_enabled(#[case] target: Target) { + let probed = |keepalive| { + connection_info(&target, &Credentials::NONE, keepalive) + .unwrap() + .tcp_settings() + .keepalive() + .is_some() + }; - assert!(!debug.contains("credential-value"), "{debug}"); - assert!(debug.contains("[REDACTED]"), "{debug}"); + assert!(probed(true)); + assert!(!probed(false)); + } + + #[test] + fn a_unix_socket_replaces_host_port_and_tls_like_redis_py() { + let target = Target::Unix(UnixSocket { + path: "/tmp/redis.sock".into(), + db: 3, + }); + + let info = connection_info(&target, &Credentials::NONE, true).unwrap(); + + assert_eq!(*info.addr(), ConnectionAddr::Unix("/tmp/redis.sock".into())); + assert_eq!(info.redis_settings().db(), 3); + } + + #[rstest] + #[case::the_query_db_wins_over_the_path("redis://cache:6379/1", Some(3), None, 3, None)] + #[case::the_path_db_without_a_query_db("redis://cache:6379/1", None, None, 1, None)] + #[case::the_query_password_fills_a_missing_userinfo( + "redis://cache:6379", + None, + Some("query"), + 0, + Some("query") + )] + #[case::the_userinfo_wins_over_the_query_password( + "redis://:userinfo@cache:6379", + None, + Some("query"), + 0, + Some("userinfo") + )] + fn a_url_reads_its_query_like_parse_url( + #[case] url: &str, + #[case] db: Option, + #[case] password: Option<&str>, + #[case] expected_db: i64, + #[case] expected_password: Option<&str>, + ) { + let target = Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db, + username: None, + password: password.map(SecretValue::new), + }); + let static_password = Credentials::Static { + username: None, + password: Some(SecretValue::new("static")), + }; + + let info = connection_info(&target, &static_password, true).unwrap(); + + assert_eq!(info.redis_settings().db(), expected_db); + assert_eq!( + info.redis_settings().password(), + expected_password.or(Some("static")) + ); + } + + #[rstest] + #[case::by_number("redis://cache:6379/?protocol=3")] + #[case::by_name("redis://cache:6379/?protocol=resp3")] + fn a_url_never_switches_the_protocol_from_resp2(#[case] url: &str) { + let target = Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db: None, + username: None, + password: None, + }); + + let info = connection_info(&target, &Credentials::NONE, true).unwrap(); + + assert_eq!(info.redis_settings().protocol(), ProtocolVersion::RESP2); } } diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index 41729bfd107..c1f82490b3b 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -1,53 +1,161 @@ -use std::time::Duration; - -use redis::{ - Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value, - aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig}, - cluster::{ClusterClientBuilder, ClusterConfig}, - cluster_async::ClusterConnection, +use std::{ + sync::{ + Arc, Mutex, PoisonError, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, }; -use crate::ClusterEndpoint; +use redis::{ + AsyncConnectionConfig, Cmd, ConnectionAddr, ConnectionInfo, ErrorKind, Pipeline, + RedisConnectionInfo, RedisError, RedisFuture, RedisResult, ServerErrorKind, TlsCertificates, + TlsMode, Value, + aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig, MultiplexedConnection}, + cluster::{ClusterClientBuilder, ClusterConfig}, + cluster_async::ClusterConnection, + cluster_routing::{MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo}, + sentinel::{SentinelClient, SentinelClientBuilder, SentinelServerType}, +}; -const RETRIES: u32 = 3; +use crate::{ + Cluster, ConnectError, Credentials, Sentinel, Target, Timeouts, + config::{self, addr}, + error::Deployment, +}; + +// redis-py 5.3.1 redis/asyncio/connection.py: without retry options a connection gets Retry(NoBackoff(), 0) +const RECONNECT_RETRIES: usize = 0; +// redis-py 5.3.1 redis/asyncio/cluster.py: cluster_error_retry_attempts=3 +const CLUSTER_RETRIES: u32 = 3; +// redis-py 5.3.1 redis/asyncio/cluster.py sleeps a flat 0.25s before retrying after CLUSTERDOWN. +// redis-rs draws the wait from [min, max), so a 1ms window pins it to 250ms +const CLUSTER_RETRY_WAIT_MS: u64 = 250; #[derive(Clone)] pub(crate) enum Connection { - Standalone(ConnectionManager), - Cluster(ClusterConnection), + Standalone(Renewing), + Cluster(Renewing), + Sentinel(SentinelConnection), } impl Connection { - pub(crate) fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult { - let config = ConnectionManagerConfig::new() - .set_connection_timeout(Some(timeout)) - .set_response_timeout(Some(timeout)) - .set_number_of_retries(RETRIES as usize); - let manager = ConnectionManager::new_lazy_with_config(redis::Client::open(info)?, config)?; - Ok(Self::Standalone(manager)) + pub(crate) fn new( + target: &Target, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + match target { + Target::Url(_) | Target::Standalone(_) | Target::Unix(_) => Self::standalone( + config::connection_info(target, credentials, keepalive)?, + credentials, + timeouts, + ), + Target::Cluster(cluster) => Self::cluster(cluster, credentials, timeouts, keepalive), + Target::Sentinel(sentinel) => Ok(Self::Sentinel(SentinelConnection::new( + sentinel, + credentials, + timeouts, + keepalive, + )?)), + } } - pub(crate) fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult { - let builder = ClusterClientBuilder::new(endpoint.addrs()) - .connection_timeout(timeout) - .response_timeout(timeout) - .retries(RETRIES); - let builder = match &endpoint.username { - Some(username) => builder.username(username.expose()), + pub(crate) fn deployment(&self) -> Deployment { + match self { + Self::Standalone(_) => Deployment::Standalone, + Self::Cluster(_) => Deployment::Cluster, + Self::Sentinel(_) => Deployment::Sentinel, + } + } + + fn standalone( + info: ConnectionInfo, + credentials: &Credentials, + timeouts: Timeouts, + ) -> Result { + let config = ConnectionManagerConfig::new() + .set_connection_timeout(Some(timeouts.connect)) + .set_response_timeout(Some(timeouts.response)) + .set_number_of_retries(RECONNECT_RETRIES); + let config = match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + }; + let client = redis::Client::open(info)?; + Ok(Self::Standalone(Renewing::new(move || { + ConnectionManager::new_lazy_with_config(client.clone(), config.clone()) + })?)) + } + + fn cluster( + endpoint: &Cluster, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + let builder = ClusterClientBuilder::new(config::cluster_addrs(endpoint)) + .tcp_settings(config::tcp_settings(keepalive)) + .connection_timeout(timeouts.connect) + .response_timeout(timeouts.response) + .retries(CLUSTER_RETRIES) + .retry_wait_formula(0, 2) + .min_retry_wait(CLUSTER_RETRY_WAIT_MS) + .max_retry_wait(CLUSTER_RETRY_WAIT_MS + 1); + let builder = match &endpoint.tls { + Some(tls) => builder + .tls(config::tls_mode(tls)) + .certs(config::tls_certificates(tls)?) + .danger_accept_invalid_hostnames(!tls.check_hostname), None => builder, }; - let builder = match &endpoint.password { - Some(password) => builder.password(password.expose()), - None => builder, + let builder = match credentials { + Credentials::Static { username, password } => { + let builder = match username { + Some(username) => builder.username(username.expose()), + None => builder, + }; + match password { + Some(password) => builder.password(password.expose()), + None => builder, + } + } + Credentials::Provider(provider) => builder.set_credentials_provider(provider.clone()), }; - let config = ClusterConfig::new() - .set_connection_timeout(timeout) - .set_response_timeout(timeout); - Ok(Self::Cluster( - builder - .build()? - .get_pending_async_connection_with_config(config), - )) + // The builder's response timeout bounds the whole request, but with_config overwrites each node's + // response timeout with this one even when unset, so it has to be repeated here + let config = ClusterConfig::new().set_response_timeout(timeouts.response); + let client = builder.build()?; + Ok(Self::Cluster(Renewing::new(move || { + Ok(client.get_pending_async_connection_with_config(config.clone())) + })?)) + } +} + +impl Connection { + /// `PING` on the one multiplexed connection, or on every cluster node, so a dead socket is + /// replaced before a real command uses it. + pub(crate) async fn ping_every_node(&self) -> RedisResult<()> { + let ping = redis::cmd("PING"); + match self { + Self::Cluster(renewing) => { + let (generation, mut connection) = renewing.current(); + let reply = connection + .route_command( + ping, + RoutingInfo::MultiNode(( + MultipleNodeRoutingInfo::AllNodes, + Some(ResponsePolicy::AllSucceeded), + )), + ) + .await; + renewing.settle(generation, &reply); + reply.map(drop) + } + Self::Standalone(_) | Self::Sentinel(_) => { + ping.query_async::(&mut self.clone()).await.map(drop) + } + } } } @@ -56,6 +164,7 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.req_packed_command(cmd), Self::Cluster(connection) => connection.req_packed_command(cmd), + Self::Sentinel(connection) => connection.req_packed_command(cmd), } } @@ -68,6 +177,7 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.req_packed_commands(pipeline, offset, count), Self::Cluster(connection) => connection.req_packed_commands(pipeline, offset, count), + Self::Sentinel(connection) => connection.req_packed_commands(pipeline, offset, count), } } @@ -75,6 +185,441 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.get_db(), Self::Cluster(connection) => connection.get_db(), + Self::Sentinel(connection) => connection.get_db(), } } } + +/// A connection swapped for a fresh one once a command on it times out, like redis-py disconnecting +/// a connection whose read timed out. redis-rs keeps the socket after a timeout and the cluster +/// client retries on it, so a socket that went silent would time out every command until TCP gives +/// up on it. +#[derive(Clone)] +pub(crate) struct Renewing(Arc>); + +type Open = Box RedisResult + Send + Sync>; + +struct RenewingState { + open: Open, + current: Mutex<(u64, C)>, +} + +impl Renewing { + fn new(open: impl Fn() -> RedisResult + Send + Sync + 'static) -> RedisResult { + let first = open()?; + Ok(Self(Arc::new(RenewingState { + open: Box::new(open), + current: Mutex::new((0, first)), + }))) + } + + fn current(&self) -> (u64, C) { + self.0 + .current + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + /// Replaces the connection of `generation` when its reply timed out, unless another timeout + /// already replaced it. + fn settle(&self, generation: u64, reply: &RedisResult) { + if !reply.as_ref().is_err_and(RedisError::is_timeout) { + return; + } + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(PoisonError::into_inner); + if current.0 != generation { + return; + } + match (self.0.open)() { + Ok(fresh) => *current = (generation + 1, fresh), + Err(error) => { + tracing::warn!(%error, "replacing a timed-out Redis connection failed") + } + } + } +} + +impl ConnectionLike for Renewing { + fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { + Box::pin(async move { + let (generation, mut connection) = self.current(); + let reply = connection.req_packed_command(cmd).await; + self.settle(generation, &reply); + reply + }) + } + + fn req_packed_commands<'a>( + &'a mut self, + pipeline: &'a Pipeline, + offset: usize, + count: usize, + ) -> RedisFuture<'a, Vec> { + Box::pin(async move { + let (generation, mut connection) = self.current(); + let replies = connection + .req_packed_commands(pipeline, offset, count) + .await; + self.settle(generation, &replies); + replies + }) + } + + fn get_db(&self) -> i64 { + self.current().1.get_db() + } +} + +#[derive(Clone)] +pub(crate) struct SentinelConnection(Arc); + +struct SentinelState { + discovery: tokio::sync::Mutex<()>, + current: Mutex>, + generations: AtomicU64, + endpoint: Sentinel, + credentials: Credentials, + tls: Option<(TlsMode, TlsCertificates)>, + config: AsyncConnectionConfig, + connect_timeout: Duration, + keepalive: bool, +} + +impl SentinelConnection { + fn new( + endpoint: &Sentinel, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + let tls = match &endpoint.tls { + Some(tls) => Some((config::tls_mode(tls), config::tls_certificates(tls)?)), + None => None, + }; + let state = SentinelState { + discovery: tokio::sync::Mutex::new(()), + current: Mutex::new(None), + generations: AtomicU64::new(0), + endpoint: endpoint.clone(), + credentials: credentials.clone(), + tls, + config: async_config(credentials, timeouts), + connect_timeout: timeouts.connect, + keepalive, + }; + state.sentinel_client(&RedisConnectionInfo::default())?; + Ok(Self(Arc::new(state))) + } + + fn cached(&self) -> Option<(u64, MultiplexedConnection)> { + self.0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + } + + async fn connection(&self) -> RedisResult<(u64, MultiplexedConnection)> { + if let Some(current) = self.cached() { + return Ok(current); + } + let _discovering = self.0.discovery.lock().await; + if let Some(current) = self.cached() { + return Ok(current); + } + let connection = self + .discover() + .await? + .get_multiplexed_async_connection_with_config(&self.0.config) + .await?; + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let generation = self.0.generations.fetch_add(1, Ordering::Relaxed); + *current = Some((generation, connection.clone())); + Ok((generation, connection)) + } + + pub(crate) async fn master(&self) -> RedisResult { + self.discover().await + } + + async fn discover(&self) -> RedisResult { + let master = tokio::time::timeout(self.0.connect_timeout, self.ask_sentinels()) + .await + .map_err(|_| timed_out("asking the Sentinels for the master timed out"))??; + self.master_client(master.get_connection_info().addr()) + } + + /// redis-rs sends `ROLE` to the master a Sentinel names before trusting it, so that probe gets + /// the data credentials and TLS a connection opened now would use. The Sentinel client is built + /// per discovery since a provider's credentials change after `new`. + async fn ask_sentinels(&self) -> RedisResult { + let data_auth = + config::authenticated(RedisConnectionInfo::default(), &self.0.credentials).await?; + self.0.sentinel_client(&data_auth)?.async_get_client().await + } + + fn master_client(&self, master: &ConnectionAddr) -> Result { + let (host, port) = match master { + ConnectionAddr::Tcp(host, port) | ConnectionAddr::TcpTls { host, port, .. } => { + (host.as_str(), *port) + } + _ => { + return Err(RedisError::from(( + ErrorKind::InvalidClientConfig, + "Sentinel named a master that is not a TCP address", + ))); + } + }; + let endpoint = Target::Standalone(crate::Standalone { + host: host.to_owned(), + port, + tls: self.0.endpoint.tls.clone(), + db: self.0.endpoint.db, + }); + let info = config::connection_info(&endpoint, &self.0.credentials, self.0.keepalive) + .map_err(|error| match error { + ConnectError::Redis(error) => error, + ConnectError::TlsFile { path, .. } => RedisError::from(( + ErrorKind::InvalidClientConfig, + "reading a Redis TLS file", + path, + )), + ConnectError::TlsKey { path, reason } => RedisError::from(( + ErrorKind::InvalidClientConfig, + "reading the Redis TLS key", + format!("{path}: {reason}"), + )), + })?; + redis::Client::open(info) + } + + fn forget(&self, generation: u64) { + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if current + .as_ref() + .is_some_and(|(cached, _)| *cached == generation) + { + *current = None; + } + } +} + +impl SentinelState { + fn sentinel_client(&self, data_auth: &RedisConnectionInfo) -> RedisResult { + let endpoint = &self.endpoint; + let sentinels = endpoint + .monitors + .iter() + .map(|node| addr(&node.host, node.port, endpoint.tls.as_ref())); + let builder = SentinelClientBuilder::new( + sentinels, + &endpoint.service_name, + SentinelServerType::Master, + )?; + let (username, password) = match &endpoint.monitor_auth { + Some(auth) => (auth.username.as_ref(), auth.password.as_ref()), + None => (None, None), + }; + let builder = match username { + Some(username) => builder.set_client_to_sentinel_username(username.expose()), + None => builder, + }; + let builder = match password { + Some(password) => builder.set_client_to_sentinel_password(password.expose()), + None => builder, + }; + let builder = match data_auth.username() { + Some(username) => builder.set_client_to_redis_username(username), + None => builder, + }; + let builder = match data_auth.password() { + Some(password) => builder.set_client_to_redis_password(password), + None => builder, + }; + let builder = match &self.tls { + Some((mode, certificates)) => builder + .set_client_to_sentinel_tls_mode(*mode) + .set_client_to_sentinel_certificates(certificates.clone()) + .set_client_to_redis_tls_mode(*mode) + .set_client_to_redis_certificates(certificates.clone()), + None => builder, + }; + builder.build() + } +} + +pub(crate) fn async_config(credentials: &Credentials, timeouts: Timeouts) -> AsyncConnectionConfig { + let config = AsyncConnectionConfig::new() + .set_connection_timeout(Some(timeouts.connect)) + .set_response_timeout(Some(timeouts.response)); + match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + } +} + +fn timed_out(message: &'static str) -> RedisError { + RedisError::from(std::io::Error::new(std::io::ErrorKind::TimedOut, message)) +} + +fn needs_rediscovery(error: &RedisError) -> bool { + matches!( + crate::Error::classify(error.clone(), Deployment::Sentinel), + crate::Error::Transport(_) + ) +} + +fn replied_read_only(values: &[Value]) -> bool { + values.iter().any(|value| { + matches!(value, Value::ServerError(error) if error.kind() == Some(ServerErrorKind::ReadOnly)) + }) +} + +impl ConnectionLike for SentinelConnection { + fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { + Box::pin(async move { + let (generation, mut connection) = self.connection().await?; + let reply = connection.req_packed_command(cmd).await; + let rediscover = match &reply { + Ok(value) => replied_read_only(std::slice::from_ref(value)), + Err(error) => needs_rediscovery(error), + }; + if rediscover { + self.forget(generation); + } + reply + }) + } + + fn req_packed_commands<'a>( + &'a mut self, + pipeline: &'a Pipeline, + offset: usize, + count: usize, + ) -> RedisFuture<'a, Vec> { + Box::pin(async move { + let (generation, mut connection) = self.connection().await?; + let replies = connection + .req_packed_commands(pipeline, offset, count) + .await; + let rediscover = match &replies { + Ok(values) => replied_read_only(values), + Err(error) => needs_rediscovery(error), + }; + if rediscover { + self.forget(generation); + } + replies + }) + } + + fn get_db(&self) -> i64 { + self.0.endpoint.db + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Node; + use litellm_redis_testing::{RedisSentinel, SERVICE_NAME}; + + fn endpoint(server: &RedisSentinel) -> Sentinel { + let (host, port) = server.sentinel(); + Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + } + } + + #[rstest::rstest] + #[tokio::test] + async fn a_late_failure_does_not_discard_the_reconnected_sentinel_master() { + let server = RedisSentinel::start().await.unwrap(); + let mut client = SentinelConnection::new( + &endpoint(&server), + &Credentials::NONE, + Timeouts::default(), + true, + ) + .unwrap(); + let (old_generation, _) = client.connection().await.unwrap(); + let original: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + client.forget(old_generation); + let replacement: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + assert_ne!(original, replacement); + + client.forget(old_generation); + let after_late_failure: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + + assert_eq!(replacement, after_late_failure); + } + + #[rstest::rstest] + #[tokio::test] + async fn a_read_only_reply_to_one_command_rediscovers_the_master() { + let server = RedisSentinel::start().await.unwrap(); + let mut client = SentinelConnection::new( + &endpoint(&server), + &Credentials::NONE, + Timeouts::default(), + true, + ) + .unwrap(); + let replica = redis::Client::open(format!("redis://127.0.0.1:{}", server.replica_port())) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let stale = client.0.generations.fetch_add(1, Ordering::Relaxed); + *client.0.current.lock().unwrap() = Some((stale, replica)); + + let on_replica = redis::cmd("SET") + .arg("k") + .arg("1") + .exec_async(&mut client) + .await; + let on_master = redis::cmd("SET") + .arg("k") + .arg("1") + .exec_async(&mut client) + .await; + + assert_eq!( + on_replica.unwrap_err().kind(), + ErrorKind::Server(ServerErrorKind::ReadOnly) + ); + on_master.unwrap(); + } +} diff --git a/litellm-rust/crates/redis-client/src/error.rs b/litellm-rust/crates/redis-client/src/error.rs index ed3b2c74004..54e55615e60 100644 --- a/litellm-rust/crates/redis-client/src/error.rs +++ b/litellm-rust/crates/redis-client/src/error.rs @@ -13,32 +13,93 @@ pub enum Error { } #[derive(Debug, thiserror::Error)] -#[error("connecting to Redis")] -pub struct ConnectError(#[from] RedisError); +pub enum ConnectError { + #[error("connecting to Redis")] + Redis(#[from] RedisError), + #[error("reading the Redis TLS file {path}")] + TlsFile { + path: String, + #[source] + source: std::io::Error, + }, + #[error("reading the Redis TLS key {path}")] + TlsKey { + path: String, + #[source] + reason: KeyError, + }, +} -impl From for Error { - fn from(error: RedisError) -> Self { - match error.kind() { - ErrorKind::Server( - ServerErrorKind::BusyLoading - | ServerErrorKind::ClusterDown - | ServerErrorKind::MasterDown - | ServerErrorKind::TryAgain - | ServerErrorKind::ReadOnly, - ) => Self::Transport(error), - ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error), - ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error), +#[derive(Debug, thiserror::Error)] +pub enum KeyError { + #[error("the key is encrypted, so ssl_password is required")] + MissingPassword, + #[error("ssl_password does not decrypt the key, or the key uses an unsupported cipher")] + Decrypt(#[source] pkcs8::Error), + #[error( + "the key uses legacy OpenSSL PEM encryption, which Rust cannot read. Convert it to PKCS#8 with `openssl pkcs8 -topk8`" + )] + LegacyEncryption, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Deployment { + Standalone, + Cluster, + Sentinel, +} + +// redis-py 5.3.1 redis/_parsers/base.py EXCEPTION_CLASSES["ERR"] entries that subclass ConnectionError +const CONNECTION_ERR_MESSAGES: &[&str] = &[ + "max number of clients reached", + "invalid password", + "Client sent AUTH, but no password is set", + "AUTH called without any password configured for the default user. Are you sure your configuration is correct?", +]; + +impl Error { + pub(crate) fn classify(error: RedisError, deployment: Deployment) -> Self { + match (error.kind(), deployment) { + (ErrorKind::Parse | ErrorKind::UnexpectedReturnType, _) => Self::Decode(error), + (ErrorKind::Server(kind), _) if is_unavailable(kind, deployment) => { + Self::Transport(error) + } + (ErrorKind::Server(ServerErrorKind::ResponseError), _) + if error + .detail() + .is_some_and(|detail| CONNECTION_ERR_MESSAGES.contains(&detail)) => + { + Self::Transport(error) + } + (ErrorKind::Extension, _) if matches!(error.code(), Some("NOAUTH" | "WRONGPASS")) => { + Self::Transport(error) + } + (ErrorKind::Server(_) | ErrorKind::Extension | ErrorKind::InvalidClientConfig, _) => { + Self::Server(error) + } _ => Self::Transport(error), } } } +fn is_unavailable(kind: ServerErrorKind, deployment: Deployment) -> bool { + matches!( + (kind, deployment), + (ServerErrorKind::BusyLoading, _) + | ( + ServerErrorKind::ClusterDown | ServerErrorKind::MasterDown, + Deployment::Cluster + ) + | (ServerErrorKind::ReadOnly, Deployment::Sentinel) + ) +} + #[cfg(test)] mod tests { use redis::{RedisError, Value}; use rstest::rstest; - use super::Error; + use super::{Deployment, Error}; fn reply(resp: &str) -> RedisError { match redis::parse_redis_value(resp.as_bytes()).unwrap() { @@ -47,23 +108,101 @@ mod tests { } } + fn is_transport(resp: &str, deployment: Deployment) -> bool { + match Error::classify(reply(resp), deployment) { + Error::Transport(_) => true, + Error::Server(_) => false, + other => panic!("{resp:?} classified as {other:?}"), + } + } + #[rstest] #[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")] #[case::no_script("-NOSCRIPT No matching script\r\n")] #[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")] #[case::generic("-ERR value is not an integer or out of range\r\n")] #[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")] - fn rejections_are_server_errors(#[case] resp: &str) { - assert!(matches!(Error::from(reply(resp)), Error::Server(_))); + #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] + #[case::out_of_memory("-OOM command not allowed when used memory > 'maxmemory'.\r\n")] + #[case::auth_arity("-ERR wrong number of arguments for 'auth' command\r\n")] + #[case::max_clients_with_suffix("-ERR max number of clients reached, try later\r\n")] + fn rejections_are_server_errors_on_every_deployment( + #[case] resp: &str, + #[values(Deployment::Standalone, Deployment::Cluster, Deployment::Sentinel)] + deployment: Deployment, + ) { + assert!(!is_transport(resp, deployment)); } #[rstest] #[case::loading("-LOADING Redis is loading the dataset in memory\r\n")] - #[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")] - #[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")] - #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] - #[case::read_only("-READONLY You can't write against a read only replica.\r\n")] - fn unavailable_nodes_are_transport_errors(#[case] resp: &str) { - assert!(matches!(Error::from(reply(resp)), Error::Transport(_))); + #[case::no_auth("-NOAUTH Authentication required.\r\n")] + #[case::wrong_pass("-WRONGPASS invalid username-password pair or user is disabled.\r\n")] + #[case::max_clients("-ERR max number of clients reached\r\n")] + #[case::invalid_password("-ERR invalid password\r\n")] + #[case::no_password_set("-ERR Client sent AUTH, but no password is set\r\n")] + #[case::no_default_password( + "-ERR AUTH called without any password configured for the default user. Are you sure your configuration is correct?\r\n" + )] + fn unavailable_nodes_are_transport_errors_on_every_deployment( + #[case] resp: &str, + #[values(Deployment::Standalone, Deployment::Cluster, Deployment::Sentinel)] + deployment: Deployment, + ) { + assert!(is_transport(resp, deployment)); + } + + #[rstest] + #[case::master_down_on_standalone( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Standalone, + false + )] + #[case::master_down_on_cluster( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Cluster, + true + )] + #[case::cluster_down_on_standalone( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Standalone, + false + )] + #[case::cluster_down_on_cluster( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Cluster, + true + )] + #[case::read_only_on_standalone( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Standalone, + false + )] + #[case::read_only_on_cluster( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Cluster, + false + )] + #[case::master_down_on_sentinel( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Sentinel, + false + )] + #[case::cluster_down_on_sentinel( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Sentinel, + false + )] + #[case::read_only_on_sentinel( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Sentinel, + true + )] + fn topology_decides_like_redis_py( + #[case] resp: &str, + #[case] deployment: Deployment, + #[case] transport: bool, + ) { + assert_eq!(is_transport(resp, deployment), transport); } } diff --git a/litellm-rust/crates/redis-client/src/events.rs b/litellm-rust/crates/redis-client/src/events.rs new file mode 100644 index 00000000000..9ea9ff4defa --- /dev/null +++ b/litellm-rust/crates/redis-client/src/events.rs @@ -0,0 +1,78 @@ +use std::{any::Any, future::Future, sync::Arc, time::Duration}; + +use litellm_circuit_breaker::{State, Transition}; + +use crate::Error; + +/// What the host attached to the calls made inside `tagged`, opaque to the client. +pub type CallTag = Arc; + +tokio::task_local! { + static TAG: CallTag; +} + +pub async fn tagged(tag: CallTag, operation: F) -> F::Output { + TAG.scope(tag, operation).await +} + +/// The tag of the `tagged` scope this task runs in. `RedisEvents::call` runs on the calling task, +/// so a sink reads its call's tag here. +pub fn current_tag() -> Option { + TAG.try_with(Arc::clone).ok() +} + +pub trait RedisEvents: Send + Sync { + fn call(&self, event: &CallEvent); + + fn breaker(&self, transition: &Transition); +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Operation { + Command(&'static str), + Script, + Pipeline { ops: usize }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Outcome { + Succeeded, + Failed(FailureKind), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FailureKind { + Timeout, + Transport, + Server, + Decode, + BreakerOpen, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallEvent { + pub operation: Operation, + pub duration: Duration, + pub outcome: Outcome, +} + +impl FailureKind { + pub fn of(error: &Error) -> Self { + match error { + Error::Transport(error) if error.is_timeout() => Self::Timeout, + Error::Transport(_) => Self::Transport, + Error::Server(_) => Self::Server, + Error::Decode(_) => Self::Decode, + Error::BreakerOpen => Self::BreakerOpen, + } + } +} + +impl Outcome { + pub(crate) fn of(result: &Result) -> Self { + match result { + Ok(_) => Self::Succeeded, + Err(error) => Self::Failed(FailureKind::of(error)), + } + } +} diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 47ff8677e06..2754666c329 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -1,20 +1,30 @@ +mod batch; mod client; -mod commands; +pub mod commands; mod config; mod connection; mod error; +mod events; mod namespace; mod pipeline; +mod pubsub; mod script; -#[cfg(test)] -mod test_support; -pub use client::{RedisClient, Ttl}; -pub use config::{ - ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig, +pub use batch::{Declared, RedisBatch}; +pub use client::{RedisClient, Ttl, WeakRedisClient}; +pub use commands::{Amount, Command, Expiry}; +pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts}; +pub use error::{ConnectError, Error, KeyError}; +pub use events::{ + CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged, +}; +pub use litellm_redis_types::{ + Cluster, Node, RedisUrl, Sentinel, Standalone, StaticAuth, Target, TlsOptions, UnixSocket, }; -pub use error::{ConnectError, Error}; pub use namespace::Namespace; pub use pipeline::Pipeline; -pub use redis::Value; +pub use redis::{ + BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value, + cluster_routing::Slot, +}; pub use script::Script; diff --git a/litellm-rust/crates/redis-client/src/pipeline.rs b/litellm-rust/crates/redis-client/src/pipeline.rs index e54c65e9605..0b6bb6a6673 100644 --- a/litellm-rust/crates/redis-client/src/pipeline.rs +++ b/litellm-rust/crates/redis-client/src/pipeline.rs @@ -1,30 +1,20 @@ -use std::{collections::HashMap, time::Duration}; +use std::collections::HashMap; use futures_util::future::join_all; -use redis::{ - Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value, - cluster_routing::Slot, -}; +use redis::{ErrorKind, RedisError, RedisResult, ServerErrorKind, Value, cluster_routing::Slot}; -use crate::{Error, Namespace, Script, commands, connection::Connection}; +use crate::{ + Command, Error, Namespace, + commands::{Op, Plan}, + connection::Connection, + script::ScriptCall, +}; pub struct Pipeline { namespace: Namespace, ops: Vec, } -struct Op { - command: Cmd, - key: String, - script: Option, -} - -struct ScriptCall { - script: Script, - keys: Vec, - args: Vec>, -} - type Settled = (usize, Result); impl Pipeline { @@ -35,60 +25,13 @@ impl Pipeline { } } - pub fn get(self, key: &str) -> Self { - let key = self.namespace.apply(key); - self.command(commands::get(key.clone()), key) + pub(crate) fn of(namespace: Namespace, ops: Vec) -> Self { + Self { namespace, ops } } - pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self { - let key = self.namespace.apply(key); - self.command(commands::set(key.clone(), value, ttl), key) - } - - pub fn del(self, key: &str) -> Self { - let key = self.namespace.apply(key); - self.command(commands::del(vec![key.clone()]), key) - } - - pub fn expire(self, key: &str, ttl: Duration) -> Self { - let key = self.namespace.apply(key); - self.command(commands::expire(key.clone(), ttl), key) - } - - pub fn incr_by(self, key: &str, delta: i64) -> Self { - let key = self.namespace.apply(key); - self.command(commands::incr_by(key.clone(), delta), key) - } - - pub fn incr_by_float(self, key: &str, delta: f64) -> Self { - let key = self.namespace.apply(key); - self.command(commands::incr_by_float(key.clone(), delta), key) - } - - pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self { - let keys: Vec = keys.iter().map(|key| self.namespace.apply(key)).collect(); - let args = args.to_redis_args(); - let op = Op { - command: script.evalsha(&keys, &args), - key: keys.first().cloned().unwrap_or_default(), - script: Some(ScriptCall { - script: script.clone(), - keys, - args, - }), - }; - self.push(op) - } - - fn command(self, command: Cmd, key: String) -> Self { - self.push(Op { - command, - key, - script: None, - }) - } - - fn push(mut self, op: Op) -> Self { + /// Appends one command in its pipelined form; `RedisClient::run` answers it in order, undecoded. + pub fn with(mut self, command: Command) -> Self { + let (op, _) = command.prepare(&self.namespace); self.ops.push(op); self } @@ -102,7 +45,7 @@ impl Pipeline { let settled = join_all( groups .iter() - .map(|group| self.run_group(connection.clone(), group)), + .map(|batches| self.run_in_order(connection, batches)), ) .await; let mut settled: Vec = settled.into_iter().flatten().collect(); @@ -110,183 +53,131 @@ impl Pipeline { settled.into_iter().map(|(_, result)| result).collect() } - fn groups(&self, connection: &Connection) -> Vec> { + async fn run_in_order(&self, connection: &Connection, batches: &[Vec]) -> Vec { + let mut settled = Vec::with_capacity(batches.iter().map(Vec::len).sum()); + for batch in batches { + settled.extend(self.run_group(connection.clone(), batch).await); + } + settled + } + + fn groups(&self, connection: &Connection) -> Vec>> { match connection { _ if self.ops.is_empty() => Vec::new(), - Connection::Standalone(_) => vec![(0..self.ops.len()).collect()], + Connection::Standalone(_) | Connection::Sentinel(_) => { + vec![vec![(0..self.ops.len()).collect()]] + } Connection::Cluster(_) => { let mut by_slot: HashMap> = HashMap::new(); - for (index, op) in self.ops.iter().enumerate() { + for index in 0..self.ops.len() { by_slot - .entry(Slot::for_key(&op.key)) + .entry(Slot::for_key(&self.ops[index].key)) .or_default() .push(index); } - by_slot.into_values().collect() + by_slot + .into_values() + .map(|indices| self.retry_safe_batches(indices)) + .collect() } } } + fn retry_safe_batches(&self, indices: Vec) -> Vec> { + indices + .chunk_by(|&left, &right| self.ops[left].idempotent && self.ops[right].idempotent) + .map(<[usize]>::to_vec) + .collect() + } + + /// A script the server had not cached ran nothing, so it alone runs again, the way litellm's + /// `RedisBatch` reruns it through redis-py's `Script`: `EVALSHA`, then `SCRIPT LOAD` and + /// `EVALSHA` on `NOSCRIPT`. + async fn settle_script( + &self, + connection: Connection, + index: usize, + reply: RedisResult, + ) -> (usize, RedisResult) { + match (&self.ops[index].plan, reply) { + (Plan::Script(call), Err(error)) if is_missing_script(&error) => { + (index, rerun_missing_script(connection, call).await) + } + (_, reply) => (index, reply), + } + } + async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec { + let deployment = connection.deployment(); let mut pipeline = redis::pipe(); pipeline.ignore_errors(); - for &index in indices { - pipeline.add_command(self.ops[index].command.clone()); - } - let replies: RedisResult>> = - pipeline.query_async(&mut connection).await; + let widths: Vec = indices + .iter() + .map(|&index| { + let commands = self.ops[index].plan.pipelined(); + let width = commands.len(); + for command in commands { + pipeline.add_command(command); + } + width + }) + .collect(); + let replies: RedisResult>> = match pipeline.is_empty() { + true => Ok(Vec::new()), + false => pipeline.query_async(&mut connection).await, + }; let replies = match replies { Ok(replies) => replies, Err(error) => { return indices .iter() - .map(|&index| (index, Err(Error::from(error.clone())))) + .map(|&index| (index, Err(Error::classify(error.clone(), deployment)))) .collect(); } }; - let mut settled = Vec::with_capacity(indices.len()); - for (&index, reply) in indices.iter().zip(replies) { - settled.push((index, self.settle(&mut connection, index, reply).await)); - } - settled - } - - async fn settle( - &self, - connection: &mut Connection, - index: usize, - reply: RedisResult, - ) -> Result { - match (reply, &self.ops[index].script) { - (Err(error), Some(call)) - if error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) => - { - Ok(call - .script - .invocation(&call.keys, &call.args) - .invoke_async(connection) - .await?) - } - (reply, _) => Ok(reply?), - } + let mut replies = replies.into_iter(); + let per_op: Vec<(usize, RedisResult)> = indices + .iter() + .zip(widths) + .map(|(&index, width)| match &self.ops[index].plan { + Plan::Reply(value) => (index, Ok(value.clone())), + _ => (index, first_reply(replies.by_ref().take(width))), + }) + .collect(); + join_all( + per_op + .into_iter() + .map(|(index, reply)| self.settle_script(connection.clone(), index, reply)), + ) + .await + .into_iter() + .map(|(index, reply)| { + ( + index, + reply.map_err(|error| Error::classify(error, deployment)), + ) + }) + .collect() } } -#[cfg(test)] -mod tests { - use std::time::{Duration, Instant}; - - use redis::Value; - use rstest::rstest; - - use crate::{ - DEFAULT_TIMEOUT, Error, Script, - test_support::{Server, Topology, flush_scripts}, - }; - - #[rstest] - #[tokio::test] - async fn each_op_settles_on_its_own_in_input_order( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client - .set("text", "not a number", Duration::from_secs(60)) - .await - .unwrap(); - let pipeline = client - .pipeline() - .incr_by("a", 1) - .incr_by("text", 1) - .incr_by("c", 5) - .set("b", "2", Duration::from_secs(60)) - .get("b") - .incr_by("a", 1) - .get("missing"); - - let results = client.run(&pipeline).await; - - assert_eq!(results.len(), 7); - assert_eq!(results[0].as_ref().unwrap(), &Value::Int(1)); - assert!( - matches!(results[1], Err(Error::Server(_))), - "{:?}", - results[1] - ); - assert_eq!(results[2].as_ref().unwrap(), &Value::Int(5)); - assert_eq!(results[3].as_ref().unwrap(), &Value::Okay); - assert_eq!( - results[4].as_ref().unwrap(), - &Value::BulkString(b"2".to_vec()) - ); - assert_eq!(results[5].as_ref().unwrap(), &Value::Int(2)); - assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); - } - - #[rstest] - #[tokio::test] - async fn a_flushed_script_in_a_pipeline_reruns_alone( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("ns", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.eval::(&incr, &["counter"], 1).await.unwrap(); - flush_scripts(&client).await; - let pipeline = client - .pipeline() - .eval(&incr, &["counter"], 2) - .incr_by("other", 7); - - let results = client.run(&pipeline).await; - - assert_eq!(results[0].as_ref().unwrap(), &Value::Int(3)); - assert_eq!(results[1].as_ref().unwrap(), &Value::Int(7)); - assert_eq!(client.get::("counter").await.unwrap(), Some(3)); - } - - #[rstest] - #[tokio::test] - async fn an_unreachable_server_fails_every_op_as_transport( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); - - server.freeze().await; - let results = client.run(&pipeline).await; - server.thaw().await; - - assert_eq!(results.len(), 2); - assert!( - matches!(results[0], Err(Error::Transport(_))), - "{:?}", - results[0] - ); - assert!( - matches!(results[1], Err(Error::Transport(_))), - "{:?}", - results[1] - ); - } - - #[rstest] - #[tokio::test] - async fn an_empty_pipeline_returns_without_asking_redis( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_secs(2)).await; - - server.freeze().await; - let started = Instant::now(); - let results = client.run(&client.pipeline()).await; - let elapsed = started.elapsed(); - server.thaw().await; - - assert!(results.is_empty()); - assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); - } +async fn rerun_missing_script(mut connection: Connection, call: &ScriptCall) -> RedisResult { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut connection) + .await +} + +fn is_missing_script(error: &RedisError) -> bool { + error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) +} + +// RedisBatch settles an op from its first reply and fails it on any reply that failed +fn first_reply(replies: impl Iterator>) -> RedisResult { + replies + .reduce(|first, next| match (first, next) { + (Err(error), _) | (Ok(_), Err(error)) => Err(error), + (first, Ok(_)) => first, + }) + .unwrap_or(Ok(Value::Nil)) } diff --git a/litellm-rust/crates/redis-client/src/pubsub.rs b/litellm-rust/crates/redis-client/src/pubsub.rs new file mode 100644 index 00000000000..b6b18d3e481 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/pubsub.rs @@ -0,0 +1,225 @@ +use std::{sync::Arc, time::Duration}; + +use futures_util::{StreamExt, stream::BoxStream}; +use redis::{ConnectionInfo, RedisResult, aio::PubSubStream}; + +use crate::{ + ConnectError, Credentials, Standalone, Target, Timeouts, config, + connection::{Connection, SentinelConnection}, +}; + +// litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py: _BACKOFF_INITIAL_SECONDS, _BACKOFF_MAX_SECONDS +const BACKOFF_INITIAL: Duration = Duration::from_secs(5); +const BACKOFF_MAX: Duration = Duration::from_secs(60); + +/// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends +/// `SUBSCRIBE`, authenticated once when it connects. +#[derive(Clone)] +pub(crate) struct Subscriber { + source: Source, + credentials: Credentials, + connect_timeout: Duration, +} + +#[derive(Clone)] +enum Source { + Nodes(Arc<[ConnectionInfo]>), + Sentinel(SentinelConnection), +} + +impl Subscriber { + pub(crate) fn new( + endpoint: &Target, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + connection: &Connection, + ) -> Result { + let source = match (endpoint, connection) { + (_, Connection::Sentinel(sentinel)) => Source::Sentinel(sentinel.clone()), + (Target::Cluster(cluster), _) => Source::Nodes( + cluster + .nodes + .iter() + .map(|node| { + let node = Target::Standalone(Standalone { + host: node.host.clone(), + port: node.port, + tls: cluster.tls.clone(), + db: 0, + }); + config::connection_info(&node, credentials, keepalive) + }) + .collect::>()?, + ), + _ => Source::Nodes(Arc::new([config::connection_info( + endpoint, + credentials, + keepalive, + )?])), + }; + Ok(Self { + source, + credentials: credentials.clone(), + connect_timeout: timeouts.connect, + }) + } + + pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec> { + futures_util::stream::unfold( + (self, channel, None::, BACKOFF_INITIAL), + |(subscriber, channel, live, backoff)| async move { + let mut live = live; + let mut backoff = backoff; + loop { + let Some(current) = live.as_mut() else { + match subscriber.subscribe(&channel).await { + Ok(subscribed) => { + live = Some(subscribed); + backoff = BACKOFF_INITIAL; + } + Err(error) => { + tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(BACKOFF_MAX); + } + } + continue; + }; + match current.next().await { + Some(message) => { + let payload = message.get_payload_bytes().to_vec(); + return Some((payload, (subscriber, channel, live, backoff))); + } + None => live = None, + } + tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(BACKOFF_MAX); + } + }, + ) + .boxed() + } + + async fn subscribe(&self, channel: &str) -> RedisResult { + match &self.source { + Source::Sentinel(sentinel) => { + let master = sentinel.master().await?; + self.subscribe_on(master.get_connection_info().clone(), channel) + .await + } + Source::Nodes(nodes) => { + let mut last = None; + for node in nodes.iter() { + match self.subscribe_on(node.clone(), channel).await { + Ok(live) => return Ok(live), + Err(error) => last = Some(error), + } + } + Err(last.expect("an endpoint always has at least one node")) + } + } + } + + async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult { + let settings = + config::authenticated(info.redis_settings().clone(), &self.credentials).await?; + let client = redis::Client::open(info.set_redis_settings(settings))?; + let mut pubsub = tokio::time::timeout(self.connect_timeout, client.get_async_pubsub()) + .await + .map_err(|_| { + redis::RedisError::from(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "connecting the Redis subscription timed out", + )) + })??; + subscribe_connected(&mut pubsub, channel, self.connect_timeout).await?; + Ok(pubsub.into_on_message()) + } +} + +async fn subscribe_connected( + pubsub: &mut redis::aio::PubSub, + channel: &str, + connect_timeout: Duration, +) -> RedisResult<()> { + tokio::time::timeout(connect_timeout, pubsub.subscribe(channel)) + .await + .map_err(|_| { + redis::RedisError::from(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "subscribing to the Redis channel timed out", + )) + })? +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use futures_util::StreamExt; + use litellm_redis_testing::{RedisServer, TcpProxy}; + use rstest::rstest; + + use super::subscribe_connected; + + #[rstest] + #[tokio::test] + async fn a_stalled_pubsub_subscribe_times_out_and_a_fresh_pubsub_recovers() { + let server = RedisServer::start().await.unwrap(); + let proxy = TcpProxy::start(server.port()).await.unwrap(); + let client = redis::Client::open(format!("redis://127.0.0.1:{}", proxy.port())).unwrap(); + let mut pubsub = tokio::time::timeout(Duration::from_secs(2), client.get_async_pubsub()) + .await + .expect("connecting the initial pubsub should finish") + .unwrap(); + tokio::time::timeout( + Duration::from_secs(2), + subscribe_connected(&mut pubsub, "existing", Duration::from_secs(1)), + ) + .await + .expect("the initial subscription acknowledgement should arrive") + .unwrap(); + + proxy.blackhole_open_connections(); + let result = tokio::time::timeout( + Duration::from_secs(1), + subscribe_connected(&mut pubsub, "stalled", Duration::from_millis(100)), + ) + .await + .expect("the stalled subscription should return before the safety deadline"); + assert!( + matches!(&result, Err(error) if error.is_timeout()), + "{result:?}" + ); + drop(pubsub); + + let mut recovered = tokio::time::timeout(Duration::from_secs(2), client.get_async_pubsub()) + .await + .expect("a fresh pubsub connection should finish") + .unwrap(); + tokio::time::timeout( + Duration::from_secs(2), + subscribe_connected(&mut recovered, "recovered", Duration::from_secs(1)), + ) + .await + .expect("the fresh subscription acknowledgement should arrive") + .unwrap(); + + let publisher = redis::Client::open(server.url()).unwrap(); + let mut connection = publisher.get_multiplexed_async_connection().await.unwrap(); + let subscriber_count: i64 = redis::cmd("PUBLISH") + .arg("recovered") + .arg("payload") + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(subscriber_count, 1); + let message = tokio::time::timeout(Duration::from_secs(2), recovered.on_message().next()) + .await + .expect("the fresh subscription should receive the published payload") + .unwrap(); + assert_eq!(message.get_payload::().unwrap(), "payload"); + } +} diff --git a/litellm-rust/crates/redis-client/src/script.rs b/litellm-rust/crates/redis-client/src/script.rs index c5d73133ecc..6dbdbafe9c9 100644 --- a/litellm-rust/crates/redis-client/src/script.rs +++ b/litellm-rust/crates/redis-client/src/script.rs @@ -1,13 +1,18 @@ +use std::sync::Arc; + use redis::{Cmd, ScriptInvocation}; +/// A Lua script, cheap to clone into every call that runs it. #[derive(Clone)] -pub struct Script(redis::Script); +pub struct Script(Arc); impl Script { pub fn new(source: &str) -> Self { - Self(redis::Script::new(source)) + Self(Arc::new(redis::Script::new(source))) } + /// The `EVALSHA` a pipeline sends, like redis-py. A `NOSCRIPT` reply is settled by + /// [`Script::invocation`] alone, which loads the script and sends it again. pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec]) -> Cmd { redis::cmd("EVALSHA") .arg(self.0.get_hash()) @@ -24,65 +29,10 @@ impl Script { } } -#[cfg(test)] -mod tests { - use std::time::Duration; - - use rstest::rstest; - - use super::Script; - use crate::{ - DEFAULT_TIMEOUT, Error, - test_support::{Server, Topology, flush_scripts}, - }; - - #[rstest] - #[tokio::test] - async fn script_keys_are_namespaced( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("a", DEFAULT_TIMEOUT).await; - let bare = server.client("", DEFAULT_TIMEOUT).await; - let read = Script::new("return redis.call('GET', KEYS[1])"); - bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap(); - - let value: Option = client - .eval(&read, &["x"], Vec::::new()) - .await - .unwrap(); - - assert_eq!(value.as_deref(), Some("1")); - } - - #[rstest] - #[tokio::test] - async fn a_failing_script_is_a_server_error_and_keeps_the_connection( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.rpush("list", "item").await.unwrap(); - - let failed = client.eval::(&incr, &["list"], 1).await; - - assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); - assert_eq!(client.eval::(&incr, &["counter"], 3).await.unwrap(), 3); - } - - #[rstest] - #[tokio::test] - async fn a_flushed_script_is_loaded_again( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.eval::(&incr, &["counter"], 1).await.unwrap(); - - flush_scripts(&client).await; - - assert_eq!(client.eval::(&incr, &["counter"], 1).await.unwrap(), 2); - } +/// A script call in a pipeline, kept apart from plain commands because it is the one entry a +/// pipeline may send again. +pub(crate) struct ScriptCall { + pub(crate) script: Script, + pub(crate) keys: Vec, + pub(crate) args: Vec>, } diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs deleted file mode 100644 index e503739a942..00000000000 --- a/litellm-rust/crates/redis-client/src/test_support.rs +++ /dev/null @@ -1,126 +0,0 @@ -use std::time::Duration; - -use litellm_circuit_breaker::Breaker; -use litellm_redis_testing::{RedisCluster, RedisServer}; -use litellm_secret_value::SecretValue; -use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; - -use crate::{ - ClusterEndpoint, ClusterNode, Endpoint, Namespace, RedisClient, RedisConfig, - connection::Connection, -}; - -#[derive(Clone, Copy, Debug)] -pub(crate) enum Topology { - Standalone, - Cluster, -} - -pub(crate) enum Server { - Standalone(RedisServer), - Cluster(RedisCluster), -} - -impl Server { - pub(crate) async fn start(topology: Topology) -> Self { - match topology { - Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()), - Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()), - } - } - - pub(crate) fn endpoint(&self) -> Endpoint { - match self { - Self::Standalone(server) => Endpoint::Url(SecretValue::new(server.url().to_owned())), - Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { - nodes: cluster - .nodes() - .map(|(host, port)| ClusterNode { - host: host.to_owned(), - port, - }) - .collect(), - username: None, - password: None, - tls: false, - }), - } - } - - pub(crate) async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient { - self.guarded_client(namespace, timeout, None).await - } - - pub(crate) async fn guarded_client( - &self, - namespace: &str, - timeout: Duration, - breaker: Option, - ) -> RedisClient { - RedisClient::new(RedisConfig { - endpoint: self.endpoint(), - namespace: Namespace::new(namespace), - timeout, - breaker, - }) - .unwrap() - } - - pub(crate) async fn restart(&self) { - match self { - Self::Standalone(server) => server.restart().await.unwrap(), - Self::Cluster(cluster) => cluster.restart().await.unwrap(), - } - } - - pub(crate) async fn stop(&self) { - match self { - Self::Standalone(server) => server.stop().await.unwrap(), - Self::Cluster(cluster) => cluster.stop().await.unwrap(), - } - } - - pub(crate) async fn resume(&self) { - match self { - Self::Standalone(server) => server.resume().await.unwrap(), - Self::Cluster(cluster) => cluster.resume().await.unwrap(), - } - } - - pub(crate) async fn freeze(&self) { - match self { - Self::Standalone(server) => server.freeze().await.unwrap(), - Self::Cluster(cluster) => cluster.freeze().await.unwrap(), - } - } - - pub(crate) async fn thaw(&self) { - match self { - Self::Standalone(server) => server.thaw().await.unwrap(), - Self::Cluster(cluster) => cluster.thaw().await.unwrap(), - } - } -} - -pub(crate) async fn client_id(client: &RedisClient, key: &str) -> i64 { - let command = redis::cmd("CLIENT").arg("ID").to_owned(); - match client.connection.clone() { - Connection::Standalone(mut connection) => { - command.query_async(&mut connection).await.unwrap() - } - Connection::Cluster(mut connection) => { - let node = Route::with_key(key, SlotAddr::Master); - let routing = RoutingInfo::SingleNode(SingleNodeRoutingInfo::SpecificNode(node)); - redis::from_redis_value(connection.route_command(command, routing).await.unwrap()) - .unwrap() - } - } -} - -pub(crate) async fn flush_scripts(client: &RedisClient) { - redis::cmd("SCRIPT") - .arg("FLUSH") - .query_async::<()>(&mut client.connection.clone()) - .await - .unwrap(); -} diff --git a/litellm-rust/crates/redis-client/tests/batch.rs b/litellm-rust/crates/redis-client/tests/batch.rs new file mode 100644 index 00000000000..45ca419951e --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/batch.rs @@ -0,0 +1,330 @@ +mod support; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_redis_client::commands; +use litellm_redis_client::{ + Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Declared, Error, Expiry, Namespace, Operation, + RedisBatch, RedisClient, RedisConfig, RedisEvents, Script, Timeouts, Ttl, +}; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, endpoint}; + +const MINUTE: Option = Some(Expiry::KeepExisting(Duration::from_secs(60))); + +// RedisBatch pipelines a standalone flush and runs each op on its own on a cluster +fn expected_trips(topology: Topology, alone: &[Operation]) -> Vec { + match topology { + Topology::Standalone => vec![Operation::Pipeline { ops: alone.len() }], + Topology::Cluster => alone.to_vec(), + } +} + +#[derive(Default)] +struct RoundTrips(Mutex>); + +impl RoundTrips { + fn seen(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +impl RedisEvents for RoundTrips { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(event.operation); + } + + fn breaker(&self, _transition: &Transition) {} +} + +fn observed(server: &RedisDeployment) -> (RedisClient, Arc) { + let trips = Arc::new(RoundTrips::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(server), + credentials: Credentials::NONE, + namespace: Namespace::new("ns"), + timeouts: Timeouts::default(), + breaker: None, + events: Some(trips.clone()), + health_check: None, + keepalive: true, + }) + .unwrap(); + (client, trips) +} + +#[rstest] +#[tokio::test] +async fn declared_ops_wait_for_the_flush_and_share_one_round_trip( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(3), MINUTE)); + let total = batch.add(commands::incr_by_float("total", Amount::Float(1.5), MINUTE)); + let written = batch.add(commands::set("other", "v", Duration::from_secs(60))); + let read = batch.add(commands::get::("value")); + + let outside = connect(&server, "ns", DEFAULT_TIMEOUT); + assert_eq!( + outside.send(commands::get::("counter")).await.unwrap(), + None + ); + assert!(trips.seen().is_empty()); + + batch.flush().await; + + let expected = expected_trips( + topology, + &[ + Operation::Command("INCRBYFLOAT"), + Operation::Command("INCRBYFLOAT"), + Operation::Command("SET"), + Operation::Command("GET"), + ], + ); + let mut seen = trips.seen(); + seen.sort_by_key(|operation| format!("{operation:?}")); + let mut sorted = expected.clone(); + sorted.sort_by_key(|operation| format!("{operation:?}")); + assert_eq!(seen, sorted); + assert_eq!(counter.resolve().await.unwrap(), 3.0); + assert_eq!(total.resolve().await.unwrap(), 1.5); + written.resolve().await.unwrap(); + assert_eq!(read.resolve().await.unwrap(), None); + assert_eq!(trips.seen().len(), expected.len()); +} + +#[rstest] +#[tokio::test] +async fn resolving_before_a_flush_sends_everything_declared_so_far( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let first = batch.add(commands::incr_by_float("a", Amount::Int(1), MINUTE)); + let second = batch.add(commands::incr_by_float("b", Amount::Int(2), MINUTE)); + let increment = Operation::Command("INCRBYFLOAT"); + + assert_eq!(second.resolve().await.unwrap(), 2.0); + let late = batch.add(commands::incr_by_float("a", Amount::Int(10), MINUTE)); + assert_eq!(first.resolve().await.unwrap(), 1.0); + assert_eq!( + trips.seen(), + expected_trips(topology, &[increment, increment]) + ); + + assert_eq!(late.resolve().await.unwrap(), 11.0); + let both: Vec = expected_trips(topology, &[increment, increment]) + .into_iter() + .chain(expected_trips(topology, &[increment])) + .collect(); + assert_eq!(trips.seen(), both); + batch.flush().await; + assert_eq!(trips.seen().len(), both.len()); +} + +#[rstest] +#[tokio::test] +async fn a_pipelined_batch_is_one_pipeline_on_every_topology( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = RedisBatch::pipelined(client); + let declared = + ["a", "b", "c"].map(|key| batch.add(commands::incr_by_float(key, Amount::Int(2), MINUTE))); + + let totals = Declared::all(declared).await; + + assert_eq!( + totals.into_iter().map(Result::unwrap).collect::>(), + [2.0; 3] + ); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 3 }]); +} + +#[rstest] +#[tokio::test] +async fn each_declared_op_gets_its_own_typed_result( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set( + "text", + "not a number", + Duration::from_secs(60), + )) + .await + .unwrap(); + let batch = client.batch(); + let rejected = batch.add(commands::incr_by_float("text", Amount::Int(1), MINUTE)); + let undecodable = batch.add(commands::get::("text")); + let missing = batch.add(commands::get::("missing")); + let script = batch.add(commands::eval::>( + &Script::new("return {KEYS[1], ARGV[1]}"), + &["k"], + "arg", + )); + let removed = batch.add(commands::del(&["text"])); + + batch.flush().await; + + assert!(matches!(rejected.resolve().await, Err(Error::Server(_)))); + assert!(matches!(undecodable.resolve().await, Err(Error::Decode(_)))); + assert_eq!(missing.resolve().await.unwrap(), None); + assert_eq!(script.resolve().await.unwrap(), ["k", "arg"]); + assert_eq!(removed.resolve().await.unwrap(), 1); +} + +#[rstest] +#[tokio::test] +async fn an_op_whose_flush_was_cancelled_resolves_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(5)); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE)); + + server.freeze().await.unwrap(); + let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; + server.thaw().await.unwrap(); + + assert!(cancelled.is_err()); + assert!(matches!(counter.resolve().await, Err(Error::Transport(_)))); +} + +#[rstest] +#[tokio::test] +async fn a_settled_op_waits_for_someone_else_to_flush_and_sends_nothing_itself( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let counter = tokio::spawn( + batch + .add(commands::incr_by_float("counter", Amount::Int(4), MINUTE)) + .settled(), + ); + + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(!counter.is_finished()); + assert!(trips.seen().is_empty()); + + batch.flush().await; + + assert_eq!(counter.await.unwrap().unwrap(), 4.0); + assert_eq!( + trips.seen(), + expected_trips(topology, &[Operation::Command("INCRBYFLOAT")]) + ); +} + +#[rstest] +#[tokio::test] +async fn a_settled_op_whose_flush_was_cancelled_resolves_as_transport() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(5)); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE)); + + server.freeze().await.unwrap(); + let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; + server.thaw().await.unwrap(); + + assert!(cancelled.is_err()); + assert!(matches!(counter.settled().await, Err(Error::Transport(_)))); +} + +#[rstest] +#[tokio::test] +async fn clones_are_the_same_client_and_a_second_connection_is_not() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let twin = connect(&server, "", DEFAULT_TIMEOUT); + + assert!(client.same_client(&client.clone())); + assert!(!client.same_client(&twin)); +} + +#[rstest] +#[tokio::test] +async fn mget_reads_each_key_once_and_answers_every_key_in_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("a", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("c", "3", Duration::from_secs(60))) + .await + .unwrap(); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + let batch = client.batch(); + + let values = batch.add(commands::mget::(&["c", "missing", "a", "c"])); + batch.flush().await; + + assert_eq!( + values.resolve().await.unwrap(), + [ + Some("3".to_owned()), + None, + Some("1".to_owned()), + Some("3".to_owned()) + ] + ); + let stats = support::command_stats(&server).await; + assert!(support::calls(&stats, "mget") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "get"), 0, "{stats}"); + assert_eq!(support::keys_read(&server).await, 3); +} + +// RedisBatch's _Increment refreshes the TTL in a pipeline, while its cluster path calls +// async_increment, which keeps a TTL the key already has +#[rstest] +#[case::pipelined(Topology::Standalone, false)] +#[case::alone_on_a_cluster(Topology::Cluster, true)] +#[tokio::test] +async fn a_batched_increment_treats_the_ttl_like_redis_batch( + #[case] topology: Topology, + #[case] keeps_the_ttl: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("counter", "1", Duration::from_secs(30))) + .await + .unwrap(); + let batch = client.batch(); + + let counter = batch.add(commands::incr_by_float( + "counter", + Amount::Int(1), + Some(Expiry::KeepExisting(Duration::from_secs(600))), + )); + batch.flush().await; + + assert_eq!(counter.resolve().await.unwrap(), 2.0); + let Ttl::Expires(left) = client.send(commands::ttl("counter")).await.unwrap() else { + panic!("expected an expiring counter") + }; + assert_eq!(left <= Duration::from_secs(30), keeps_the_ttl, "{left:?}"); +} diff --git a/litellm-rust/crates/redis-client/tests/client.rs b/litellm-rust/crates/redis-client/tests/client.rs new file mode 100644 index 00000000000..e9c6cd7e712 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/client.rs @@ -0,0 +1,1114 @@ +mod support; + +use std::{ + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use litellm_redis_testing::{RedisCluster, RedisDeployment, RedisServer, Topology, eventually}; +use rstest::rstest; + +use litellm_circuit_breaker::{Breaker, Config, State}; + +use litellm_redis_client::commands; +use litellm_redis_client::{ + Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Error, Expiry, FailureKind, Namespace, + Operation, Outcome, RedisClient, RedisConfig, RedisEvents, RedisUrl, Script, Standalone, + Target, Timeouts, Ttl, Value, +}; +use litellm_secret_value::SecretValue; +use support::{KEEP, connect, connect_guarded, connections, endpoint}; + +async fn persist(client: &RedisClient, key: &str) { + let set = Script::new("return redis.call('SET', KEYS[1], '0')"); + client + .send(commands::eval::<()>(&set, &[key], Vec::::new())) + .await + .unwrap(); +} + +#[derive(Clone, Copy, Debug)] +enum Counter { + IncrBy, + IncrByFloat, +} + +const MINUTE: Duration = Duration::from_secs(60); + +async fn write(client: &RedisClient, counter: Counter, key: &str, expiry: Expiry) { + match (counter, expiry) { + (Counter::IncrBy, Expiry::KeepExisting(ttl)) => { + client.send(commands::incr_by(key, 1, ttl)).await.unwrap(); + } + (Counter::IncrBy, Expiry::Refresh(_)) => unreachable!("increment_cache never refreshes"), + (Counter::IncrByFloat, expiry) => { + client + .send(commands::incr_by_float( + key, + Amount::Float(0.5), + Some(expiry), + )) + .await + .unwrap(); + } + } +} + +async fn persist_as(client: &RedisClient, key: &str) { + let source = "return redis.call('SET', KEYS[1], '1')"; + client + .send(commands::eval::( + &Script::new(source), + &[key], + Vec::::new(), + )) + .await + .unwrap(); +} + +fn remaining(ttl: Ttl) -> Duration { + match ttl { + Ttl::Expires(left) => left, + other => panic!("expected an expiring key, got {other:?}"), + } +} + +#[rstest] +#[case::increment_cache(Counter::IncrBy, Expiry::KeepExisting(MINUTE))] +#[case::async_increment(Counter::IncrByFloat, Expiry::KeepExisting(MINUTE))] +#[case::async_increment_refreshing(Counter::IncrByFloat, Expiry::Refresh(MINUTE))] +#[tokio::test] +async fn a_counter_write_expires_a_new_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[case] counter: Counter, + #[case] expiry: Expiry, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + write(&client, counter, "fresh", expiry).await; + + let left = remaining(client.send(commands::ttl("fresh")).await.unwrap()); + assert!(left > Duration::from_secs(50) && left <= MINUTE, "{left:?}"); +} + +#[rstest] +#[tokio::test] +async fn keep_existing_leaves_a_running_ttl_and_refresh_resets_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let short = Expiry::Refresh(Duration::from_secs(30)); + let long = Duration::from_secs(600); + write(&client, Counter::IncrByFloat, "kept", short).await; + write(&client, Counter::IncrByFloat, "refreshed", short).await; + write( + &client, + Counter::IncrBy, + "kept_int", + Expiry::KeepExisting(Duration::from_secs(30)), + ) + .await; + + write( + &client, + Counter::IncrByFloat, + "kept", + Expiry::KeepExisting(long), + ) + .await; + write( + &client, + Counter::IncrByFloat, + "refreshed", + Expiry::Refresh(long), + ) + .await; + write( + &client, + Counter::IncrBy, + "kept_int", + Expiry::KeepExisting(long), + ) + .await; + + assert!( + remaining(client.send(commands::ttl("kept")).await.unwrap()) <= Duration::from_secs(30) + ); + assert!( + remaining(client.send(commands::ttl("kept_int")).await.unwrap()) <= Duration::from_secs(30) + ); + assert!( + remaining(client.send(commands::ttl("refreshed")).await.unwrap()) + > Duration::from_secs(500) + ); +} + +#[rstest] +#[tokio::test] +async fn keep_existing_expires_a_key_that_had_no_ttl( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(Counter::IncrBy, Counter::IncrByFloat)] counter: Counter, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + persist_as(&client, "persistent").await; + + write(&client, counter, "persistent", KEEP).await; + + let left = remaining(client.send(commands::ttl("persistent")).await.unwrap()); + assert!(left > Duration::from_secs(50), "{left:?}"); +} + +#[rstest] +#[tokio::test] +async fn counters_send_the_commands_redis_cache_sends( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + client + .send(commands::incr_by("int", 1, MINUTE)) + .await + .unwrap(); + client + .send(commands::incr_by_float("kept", Amount::Int(1), Some(KEEP))) + .await + .unwrap(); + client + .send(commands::incr_by_float( + "refreshed", + Amount::Float(0.5), + Some(Expiry::Refresh(MINUTE)), + )) + .await + .unwrap(); + client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + "piped", + Amount::Float(0.25), + Some(Expiry::Refresh(MINUTE)), + )) + .with(commands::incr_by_float("bare", Amount::Float(0.25), None)), + ) + .await; + + let stats = support::command_stats(&server).await; + let calls = |command: &str| support::calls(&stats, command); + assert_eq!(calls("incrby"), 1); + assert_eq!(calls("incrbyfloat"), 4); + assert_eq!(calls("ttl"), 2); + assert_eq!(calls("expire"), 4); + assert_eq!( + calls("eval") + calls("evalsha") + calls("pttl") + calls("pexpire"), + 0 + ); +} + +#[derive(Default)] +struct Recorded(Mutex>); + +#[rstest] +#[case::beyond_double_precision(9_007_199_254_740_992_i64, 1_i64)] +#[case::negative_beyond_double_precision(-9_007_199_254_740_992_i64, -1_i64)] +#[case::maximum_integer(i64::MAX - 1, 1_i64)] +#[case::minimum_integer(i64::MIN + 1, -1_i64)] +#[tokio::test] +async fn integer_increments_preserve_every_bit( + #[case] initial: i64, + #[case] delta: i64, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("integer", initial, Duration::from_secs(60))) + .await + .unwrap(); + let expected = initial + delta; + let result = client + .send(commands::incr_by("integer", delta, MINUTE)) + .await + .unwrap(); + assert_eq!(result, expected); + assert_eq!( + client.send(commands::get::("integer")).await.unwrap(), + Some(expected) + ); +} + +#[rstest] +#[tokio::test] +async fn outage_events_report_the_failure_kind( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(false, true)] timeout: bool, + #[values(false, true)] pipeline: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(300), None), + breaker: None, + events: Some(recorded.clone()), + credentials: Credentials::NONE, + health_check: None, + keepalive: true, + }) + .unwrap(); + client.ping().await.unwrap(); + if timeout { + server.freeze().await.unwrap(); + } else { + server.stop().await.unwrap(); + } + let result = if pipeline { + client + .run(&client.pipeline().with(commands::get::("missing"))) + .await + .remove(0) + .map(drop) + } else { + client + .send(commands::get::("missing")) + .await + .map(drop) + }; + if timeout { + server.thaw().await.unwrap(); + } else { + server.resume().await.unwrap(); + } + assert!(matches!(result, Err(Error::Transport(_)))); + let expected = if timeout { + FailureKind::Timeout + } else { + FailureKind::Transport + }; + assert_eq!( + recorded.0.lock().unwrap().last().unwrap().outcome, + Outcome::Failed(expected) + ); +} + +impl RedisEvents for Recorded { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(*event); + } + + fn breaker(&self, _transition: &litellm_circuit_breaker::Transition) {} +} + +#[rstest] +#[tokio::test] +async fn every_call_reports_its_operation_and_outcome( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: Some(recorded.clone()), + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let _ = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + let _ = client.send(commands::get::("missing")).await; + let _ = client + .run( + &client + .pipeline() + .with(commands::get::("a")) + .with(commands::get::("b")), + ) + .await; + let failures = client + .run(&client.pipeline().with(commands::incr_by_float( + "list", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + ))) + .await; + assert!(matches!(failures[0], Err(Error::Server(_)))); + + let seen: Vec<(Operation, Outcome)> = recorded + .0 + .lock() + .unwrap() + .iter() + .map(|event| (event.operation, event.outcome)) + .collect(); + assert_eq!( + seen, + vec![ + (Operation::Command("SET"), Outcome::Succeeded), + ( + Operation::Command("INCRBY"), + Outcome::Failed(FailureKind::Server) + ), + (Operation::Command("GET"), Outcome::Succeeded), + (Operation::Pipeline { ops: 2 }, Outcome::Succeeded), + ( + Operation::Pipeline { ops: 1 }, + Outcome::Failed(FailureKind::Server) + ), + ] + ); +} + +#[rstest] +#[tokio::test] +async fn namespaces_isolate_keys_on_the_server( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let a = connect(&server, "a", DEFAULT_TIMEOUT); + let b = connect(&server, "b", DEFAULT_TIMEOUT); + let bare = connect(&server, "", DEFAULT_TIMEOUT); + + a.send(commands::set("x", "1", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + a.send(commands::get::("x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!( + a.send(commands::get::("a:x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!( + bare.send(commands::get::("a:x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!(b.send(commands::get::("x")).await.unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_command_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + let before = connections(&server).await; + + let rejected = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + + assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); + assert_eq!( + client + .send(commands::get::("list")) + .await + .unwrap() + .as_deref(), + Some("item") + ); + assert_eq!(connections(&server).await, before); +} + +#[rstest] +#[tokio::test] +async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("float", "15.0", Duration::from_secs(60))) + .await + .unwrap(); + + let decoded = client.send(commands::get::("float")).await; + + assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); + assert_eq!( + client.send(commands::get::("missing")).await.unwrap(), + None + ); +} + +#[rstest] +#[tokio::test] +async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + client + .send(commands::set("first", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("second", "2", Duration::from_secs(60))) + .await + .unwrap(); + + server.freeze().await.unwrap(); + let timed_out = client.send(commands::get::("first")).await; + server.thaw().await.unwrap(); + + assert!( + matches!(timed_out, Err(Error::Transport(_))), + "{timed_out:?}" + ); + assert_eq!( + client + .send(commands::get::("second")) + .await + .unwrap() + .as_deref(), + Some("2") + ); +} + +#[rstest] +#[tokio::test] +async fn the_same_client_recovers_after_redis_restarts( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + + server.restart().await.unwrap(); + + let recovered = eventually(async || { + client + .send(commands::incr_by( + "after_restart", + 1, + Duration::from_secs(60), + )) + .await + .ok() + }) + .await; + assert_eq!(recovered, 1); +} + +#[rstest] +#[tokio::test] +async fn a_client_created_while_redis_is_down_works_once_redis_is_up( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server.stop().await.unwrap(); + + let client = connect(&server, "", Duration::from_millis(500)); + let started = Instant::now(); + let while_down = client.ping().await; + let failed_in = started.elapsed(); + server.resume().await.unwrap(); + + assert!(failed_in < Duration::from_secs(2), "{failed_in:?}"); + + assert!( + matches!(while_down, Err(Error::Transport(_))), + "{while_down:?}" + ); + let recovered = eventually(async || client.ping().await.ok()).await; + assert_eq!(recovered, ()); +} + +#[rstest] +#[tokio::test] +async fn empty_batches_succeed_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + assert_eq!( + client.send(commands::mget::(&[])).await.unwrap(), + Vec::>::new() + ); + assert_eq!(client.send(commands::del(&[])).await.unwrap(), 0); +} + +#[rstest] +#[tokio::test] +async fn mget_and_del_span_slots_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("a", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("c", "3", Duration::from_secs(60))) + .await + .unwrap(); + + let values = client + .send(commands::mget::(&["c", "b", "a"])) + .await + .unwrap(); + let deleted = client.send(commands::del(&["a", "b", "c"])).await.unwrap(); + + assert_eq!( + values, + vec![Some("3".to_owned()), None, Some("1".to_owned())] + ); + assert_eq!(deleted, 2); + assert_eq!( + client + .send(commands::mget::(&["a", "c"])) + .await + .unwrap(), + vec![None, None] + ); +} + +#[rstest] +#[tokio::test] +async fn set_nx_only_writes_a_missing_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + let first = client + .send(commands::set_nx("lock", "owner-1", Duration::from_secs(60))) + .await + .unwrap(); + let second = client + .send(commands::set_nx("lock", "owner-2", Duration::from_secs(60))) + .await + .unwrap(); + + assert!(first); + assert!(!second); + assert_eq!( + client + .send(commands::get::("lock")) + .await + .unwrap() + .as_deref(), + Some("owner-1") + ); +} + +#[rstest] +#[tokio::test] +async fn ttl_reports_missing_persistent_and_remaining_time( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + persist(&client, "persistent").await; + client + .send(commands::set("expiring", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::incr_by_float( + "expired_later", + Amount::Float(0.5), + Some(KEEP), + )) + .await + .unwrap(); + assert!( + client + .send(commands::expire("expired_later", Duration::from_secs(30))) + .await + .unwrap() + ); + + assert_eq!( + client.send(commands::ttl("missing")).await.unwrap(), + Ttl::Missing + ); + assert_eq!( + client.send(commands::ttl("persistent")).await.unwrap(), + Ttl::Persistent + ); + assert!(matches!( + client.send(commands::ttl("expiring")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); + assert!(matches!( + client.send(commands::ttl("expired_later")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) + )); + assert!( + !client + .send(commands::expire("missing", Duration::from_secs(30))) + .await + .unwrap() + ); +} + +#[rstest] +#[tokio::test] +async fn incr_by_float_accumulates_fractions( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + client + .send(commands::incr_by_float( + "spend", + Amount::Float(0.25), + Some(KEEP), + )) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::incr_by_float( + "spend", + Amount::Float(0.5), + Some(KEEP) + )) + .await + .unwrap(), + 0.75 + ); +} + +#[rstest] +#[tokio::test] +async fn concurrent_commands_on_one_cloned_client_all_apply( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + let tasks: Vec<_> = (0..100) + .map(|_| { + let client = client.clone(); + tokio::spawn(async move { + for _ in 0..10 { + client + .send(commands::incr_by("shared", 1, Duration::from_secs(60))) + .await + .unwrap(); + } + }) + }) + .collect(); + for task in tasks { + task.await.unwrap(); + } + + assert_eq!( + client.send(commands::get::("shared")).await.unwrap(), + Some(1000) + ); +} + +fn breaker(failure_threshold: u32) -> Breaker { + Breaker::new( + Config { + failure_threshold, + recovery_timeout: Duration::from_millis(200), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ) +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + + server.freeze().await.unwrap(); + let first = client.send(commands::get::("k")).await; + let second = client.send(commands::get::("k")).await; + let started = Instant::now(); + let third = client.send(commands::get::("k")).await; + let rejected_in = started.elapsed(); + server.thaw().await.unwrap(); + + assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); + assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); + assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + + let recovered = eventually(async || client.send(commands::get::("k")).await.ok()).await; + assert_eq!(recovered, None); + assert!( + matches!(breaker.state(), State::Closed { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn rejected_commands_never_open_the_breaker( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded(&server, "", DEFAULT_TIMEOUT, Some(breaker.clone())); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let rejected = [ + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + ]; + + assert!( + matches!(rejected[2], Err(Error::Server(_))), + "{:?}", + rejected[2] + ); + assert!( + matches!(breaker.state(), State::Closed { hard: 0, .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn a_probe_redis_rejects_still_proves_redis_answers_and_closes_the_breaker( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + assert!( + matches!(breaker.state(), State::Open { .. }), + "{:?}", + breaker.state() + ); + tokio::time::sleep(Duration::from_millis(250)).await; + + let probe = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + + assert!(matches!(probe, Err(Error::Server(_))), "{probe:?}"); + assert!( + matches!(breaker.state(), State::Closed { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn a_replica_refusing_a_write_does_not_reset_a_run_of_failures() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + client.send(commands::get::("k")).await.unwrap(); + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + server + .cli_on_every_node(&["REPLICAOF", "127.0.0.1", "1"]) + .await + .unwrap(); + + let refused = client + .send(commands::set("k", "v", Duration::from_secs(60))) + .await; + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + + assert!(matches!(refused, Err(Error::Server(_))), "{refused:?}"); + assert!( + matches!(breaker.state(), State::Open { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_fails_every_pipeline_op_without_sending( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker(1))); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )); + + server.freeze().await.unwrap(); + let failed = client.run(&pipeline).await; + let rejected = client.run(&pipeline).await; + server.thaw().await.unwrap(); + + assert!( + matches!(failed[0], Err(Error::Transport(_))), + "{:?}", + failed[0] + ); + assert!( + matches!(rejected[0], Err(Error::BreakerOpen)), + "{:?}", + rejected[0] + ); + assert!( + matches!(rejected[1], Err(Error::BreakerOpen)), + "{:?}", + rejected[1] + ); +} + +#[rstest] +#[tokio::test] +async fn a_host_endpoint_connects_to_its_database() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let connect = |target| { + RedisClient::new(RedisConfig { + target, + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + }; + let url = |url: String| { + Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db: None, + username: None, + password: None, + }) + }; + let in_db_3 = connect(Target::Standalone(Standalone { + host: "127.0.0.1".to_owned(), + port, + tls: None, + db: 3, + })) + .unwrap(); + let in_db_0 = connect(url(server.url().to_owned())).unwrap(); + let via_url = connect(url(format!("{}/3", server.url()))).unwrap(); + + in_db_3 + .send(commands::set("k", "3", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + in_db_0.send(commands::get::("k")).await.unwrap(), + None + ); + assert_eq!( + via_url + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("3") + ); +} + +async fn ping_calls(cluster: &RedisCluster) -> Vec { + cluster + .cli_outputs(&["INFO", "commandstats"]) + .await + .unwrap() + .iter() + .map(|info| { + info.lines() + .find_map(|line| line.strip_prefix("cmdstat_ping:calls=")) + .and_then(|rest| rest.split(',').next()) + .and_then(|calls| calls.trim().parse().ok()) + .unwrap_or(0) + }) + .collect() +} + +fn health_checked( + server: &RedisDeployment, + interval: Duration, + timeout: Duration, + breaker: Option, + events: Option>, +) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(server), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(timeout, None), + breaker, + events: events.map(|events| events as Arc), + health_check: Some(interval), + keepalive: true, + }) + .unwrap() +} + +// redis-py 5.3.1 redis/asyncio/connection.py check_health: a connection idle longer than +// health_check_interval gets a PING before the next command, a busy one does not +#[rstest] +#[tokio::test] +async fn an_idle_client_pings_every_node_before_the_next_command_like_redis_py() { + let server = RedisDeployment::start(Topology::Cluster).await.unwrap(); + let RedisDeployment::Cluster(cluster) = &server else { + unreachable!("started a cluster") + }; + let client = health_checked(&server, Duration::from_secs(3), DEFAULT_TIMEOUT, None, None); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = ping_calls(cluster).await; + + client + .send(commands::set("k", "2", Duration::from_secs(60))) + .await + .unwrap(); + let busy = ping_calls(cluster).await; + tokio::time::sleep(Duration::from_millis(3500)).await; + client + .send(commands::set("k", "3", Duration::from_secs(60))) + .await + .unwrap(); + let idle = ping_calls(cluster).await; + + assert_eq!(busy, before); + assert_eq!( + idle, + before.iter().map(|calls| calls + 1).collect::>(), + "before {before:?}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_failed_health_check_is_neither_a_call_nor_a_breaker_failure() { + let server = RedisDeployment::start(Topology::Cluster).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let breaker = breaker(2); + let client = health_checked( + &server, + Duration::from_millis(200), + Duration::from_millis(300), + Some(breaker.clone()), + Some(recorded.clone()), + ); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + + server.freeze().await.unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + let failed = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + + assert!(matches!(failed, Err(Error::Transport(_))), "{failed:?}"); + assert!( + matches!(breaker.state(), State::Closed { timeouts: 1, .. }), + "{:?}", + breaker.state() + ); + assert_eq!(recorded.0.lock().unwrap().len(), 2); +} + +#[rstest] +#[case::int(Amount::Int(1), "1")] +#[case::float_keeps_its_point(Amount::Float(1.0), "1.0")] +#[case::small_float(Amount::Float(1e-5), "1e-05")] +#[case::large_float(Amount::Float(1e16), "1e+16")] +#[case::fraction(Amount::Float(0.1 + 0.2), "0.30000000000000004")] +fn amounts_are_written_like_redis_py_encodes_them(#[case] amount: Amount, #[case] expected: &str) { + assert_eq!( + redis::ToRedisArgs::to_redis_args(&amount), + vec![expected.as_bytes().to_vec()] + ); +} diff --git a/litellm-rust/crates/redis-client/tests/config.rs b/litellm-rust/crates/redis-client/tests/config.rs new file mode 100644 index 00000000000..101c6529e94 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/config.rs @@ -0,0 +1,290 @@ +use rstest::rstest; + +use litellm_redis_client::{Credentials, RedisUrl, Sentinel, StaticAuth, Target, TlsOptions}; +use litellm_secret_value::SecretValue; + +#[rstest] +#[case::url(format!("{:?}", Target::Url(RedisUrl { url: SecretValue::new("redis://user:credential-value@cache:6379/0"), tls: None, db: None, username: None, password: None })))] +#[case::static_password(format!("{:?}", Credentials::Static { username: None, password: Some(SecretValue::new("credential-value")) }))] +#[case::sentinel_password(format!("{:?}", Target::Sentinel(Sentinel { + monitors: Vec::new(), + service_name: "primary".to_owned(), + monitor_auth: Some(StaticAuth { username: None, password: Some(SecretValue::new("credential-value")) }), + tls: None, + db: 0, +})))] +fn debug_output_never_shows_the_secret(#[case] debug: String) { + assert!(!debug.contains("credential-value"), "{debug}"); + assert!(debug.contains("[REDACTED]"), "{debug}"); +} + +#[derive(Clone, Copy, Debug)] +enum Ca { + Trusted, + TrustedData, + TrustedDirectory, + SystemRoots, +} + +fn directory_holding(ca: &std::path::Path) -> std::path::PathBuf { + let directory = std::env::temp_dir().join(format!("litellm-ca-dir-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + std::fs::copy(ca, directory.join("ca.pem")).unwrap(); + directory +} + +#[derive(Clone, Copy, Debug)] +enum ClientKey { + Absent, + Plain, + Encrypted { password: Option<&'static str> }, +} + +async fn tls_server(require_client_certificate: bool) -> litellm_redis_testing::RedisTlsServer { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + litellm_redis_testing::RedisTlsServer::start(require_client_certificate) + .await + .unwrap() +} + +fn tls_client( + server: &litellm_redis_testing::RedisTlsServer, + host: &str, + tls: TlsOptions, +) -> Result { + litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: host.to_owned(), + port: server.port(), + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) +} + +fn tls_options( + files: &litellm_redis_testing::TlsFiles, + ca: Ca, + verify_certificate: bool, + check_hostname: bool, + client_key: ClientKey, +) -> TlsOptions { + let (keyfile, key_password) = match client_key { + ClientKey::Absent => (None, None), + ClientKey::Plain => (Some(files.client_key.clone()), None), + ClientKey::Encrypted { password } => ( + Some(files.encrypted_client_key.clone()), + password.map(SecretValue::new), + ), + }; + TlsOptions { + ca_certs: matches!(ca, Ca::Trusted).then(|| files.ca.clone()), + ca_data: matches!(ca, Ca::TrustedData).then(|| std::fs::read_to_string(&files.ca).unwrap()), + ca_path: matches!(ca, Ca::TrustedDirectory).then(|| directory_holding(&files.ca)), + certfile: keyfile.is_some().then(|| files.client_cert.clone()), + keyfile, + key_password, + verify_certificate, + check_hostname, + } +} + +async fn ping_over_tls( + require_client_certificate: bool, + host: &str, + ca: Ca, + verify_certificate: bool, + check_hostname: bool, + client_key: ClientKey, +) -> Result<(), litellm_redis_client::Error> { + let server = tls_server(require_client_certificate).await; + let tls = tls_options( + server.files(), + ca, + verify_certificate, + check_hostname, + client_key, + ); + tls_client(&server, host, tls).unwrap().ping().await +} + +#[rstest] +#[case::trusted_ca_without_hostname_check_like_redis_py( + "127.0.0.1", + Ca::Trusted, + true, + false, + true +)] +#[case::trusted_ca_as_pem_text_like_ssl_ca_data("127.0.0.1", Ca::TrustedData, true, false, true)] +#[case::trusted_ca_in_a_directory_like_ssl_ca_path( + "127.0.0.1", + Ca::TrustedDirectory, + true, + false, + true +)] +#[case::hostname_check_with_a_matching_name("localhost", Ca::Trusted, true, true, true)] +#[case::hostname_check_against_an_unlisted_address("127.0.0.1", Ca::Trusted, true, true, false)] +#[case::an_untrusted_ca_is_refused("127.0.0.1", Ca::SystemRoots, true, false, false)] +#[case::cert_reqs_none_skips_verification("127.0.0.1", Ca::SystemRoots, false, false, true)] +#[tokio::test] +async fn tls_options_decide_which_servers_are_trusted( + #[case] host: &str, + #[case] ca: Ca, + #[case] verify_certificate: bool, + #[case] check_hostname: bool, + #[case] connects: bool, +) { + let result = ping_over_tls( + false, + host, + ca, + verify_certificate, + check_hostname, + ClientKey::Absent, + ) + .await; + + assert_eq!(result.is_ok(), connects, "{result:?}"); + if let Err(error) = result { + assert!( + matches!(error, litellm_redis_client::Error::Transport(_)), + "{error:?}" + ); + } +} + +#[rstest] +#[case::with_the_client_certificate(ClientKey::Plain, true)] +#[case::with_an_encrypted_key_and_its_ssl_password( + ClientKey::Encrypted { password: Some(litellm_redis_testing::CLIENT_KEY_PASSWORD) }, + true +)] +#[case::without_a_client_certificate(ClientKey::Absent, false)] +#[tokio::test] +async fn mutual_tls_sends_the_configured_client_certificate( + #[case] client_key: ClientKey, + #[case] connects: bool, +) { + let result = ping_over_tls(true, "127.0.0.1", Ca::Trusted, true, false, client_key).await; + + assert_eq!(result.is_ok(), connects, "{result:?}"); +} + +#[rstest] +#[case::without_ssl_password(None, "ssl_password is required")] +#[case::with_the_wrong_ssl_password(Some("wrong-password"), "does not decrypt the key")] +#[tokio::test] +async fn an_encrypted_key_that_cannot_be_decrypted_fails_before_connecting( + #[case] password: Option<&'static str>, + #[case] reason: &str, +) { + let server = tls_server(true).await; + let tls = tls_options( + server.files(), + Ca::Trusted, + true, + false, + ClientKey::Encrypted { password }, + ); + + let error = tls_client(&server, "127.0.0.1", tls).err().unwrap(); + + let litellm_redis_client::ConnectError::TlsKey { + path, + reason: cause, + } = &error + else { + panic!("expected a TLS key error, got {error:?}"); + }; + assert_eq!( + *path, + server.files().encrypted_client_key.display().to_string() + ); + assert!(cause.to_string().contains(reason), "{cause}"); +} + +#[rstest] +fn a_missing_tls_file_names_the_path() { + let tls = TlsOptions { + ca_certs: Some("/nonexistent/ca.pem".into()), + ..TlsOptions::verified() + }; + + let error = litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: "127.0.0.1".to_owned(), + port: 6379, + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .err() + .unwrap(); + + assert!(error.to_string().contains("/nonexistent/ca.pem"), "{error}"); +} + +#[rstest] +fn a_legacy_encrypted_pem_key_is_refused_with_a_conversion_hint() { + let directory = std::env::temp_dir().join(format!("litellm-legacy-key-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + let keyfile = directory.join("client.key"); + std::fs::write( + &keyfile, + "-----BEGIN RSA PRIVATE KEY-----\nProc-Type: 4,ENCRYPTED\nDEK-Info: AES-256-CBC,00112233445566778899AABBCCDDEEFF\n\nAAAA\n-----END RSA PRIVATE KEY-----\n", + ) + .unwrap(); + let tls = TlsOptions { + certfile: Some(keyfile.clone()), + keyfile: Some(keyfile), + key_password: Some(SecretValue::new("legacy-password")), + ..TlsOptions::verified() + }; + + let error = litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: "127.0.0.1".to_owned(), + port: 6379, + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .err() + .unwrap(); + + assert!( + matches!( + &error, + litellm_redis_client::ConnectError::TlsKey { + reason: litellm_redis_client::KeyError::LegacyEncryption, + .. + } + ), + "{error:?}" + ); +} diff --git a/litellm-rust/crates/redis-client/tests/connection.rs b/litellm-rust/crates/redis-client/tests/connection.rs new file mode 100644 index 00000000000..4efa6038cf0 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/connection.rs @@ -0,0 +1,329 @@ +mod support; + +use std::{collections::BTreeSet, sync::Arc, time::Duration}; + +use litellm_redis_testing::{ + ClientConnection, RedisDeployment, RedisSentinel, RedisServer, SERVICE_NAME, Topology, + eventually, +}; +use rstest::rstest; + +use litellm_redis_client::commands; +use litellm_redis_client::{ + Credentials, CredentialsProvider, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, + Standalone, Target, Timeouts, +}; +use litellm_secret_value::SecretValue; +use support::{Rotating, connect_as, connections, proxied}; + +#[rstest] +#[tokio::test] +async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect_uses_the_latest( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = Arc::new(Rotating::new("alice", "first")); + let client = connect_as( + &server, + Credentials::Provider(CredentialsProvider(provider.clone())), + ); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = connections(&server).await; + + provider.rotate("bob", "second"); + let rotated = eventually(async || { + let now = connections(&server).await; + (users(&now) == BTreeSet::from(["bob"])).then_some(now) + }) + .await; + server + .cli_on_every_node(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + server + .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let after_reconnect = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + + let reconnected = connections(&server).await; + + assert_eq!(users(&before), BTreeSet::from(["alice"])); + assert_eq!(ids(&rotated), ids(&before)); + assert_eq!(after_reconnect, "1"); + assert!(!reconnected.is_empty()); + assert_eq!(users(&reconnected), BTreeSet::from(["bob"])); +} + +fn users(connections: &[ClientConnection]) -> BTreeSet<&str> { + connections + .iter() + .map(|connection| connection.user.as_str()) + .collect() +} + +fn ids(connections: &[ClientConnection]) -> Vec<(u16, u64)> { + connections + .iter() + .map(|connection| (connection.node, connection.id)) + .collect() +} + +#[rstest] +#[tokio::test] +async fn an_out_of_range_database_is_a_server_rejection() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let client = RedisClient::new(RedisConfig { + target: Target::Standalone(Standalone { + host: "127.0.0.1".into(), + port, + tls: None, + db: 999, + }), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(2), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + let result = client.ping().await; + assert!(matches!(result, Err(Error::Server(_))), "{result:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_credential_is_a_transport_failure() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + server + .cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let with_password = |password: &str| { + connect_as( + &server, + Credentials::Static { + username: None, + password: Some(SecretValue::new(password)), + }, + ) + }; + with_password("right").ping().await.unwrap(); + + let rejected = with_password("wrong").ping().await; + + assert!(matches!(rejected, Err(Error::Transport(_))), "{rejected:?}"); +} + +fn sentinel_client(sentinel: &RedisSentinel, credentials: Credentials) -> RedisClient { + let (host, port) = sentinel.sentinel(); + RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(1), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +#[rstest] +#[tokio::test] +async fn sentinel_writes_follow_the_master_through_a_failover() { + let sentinel = RedisSentinel::start().await.unwrap(); + let client = sentinel_client(&sentinel, Credentials::NONE); + client + .send(commands::set("before", "1", Duration::from_secs(600))) + .await + .unwrap(); + eventually(async || { + let replica = + redis::Client::open(format!("redis://127.0.0.1:{}", sentinel.replica_port())).ok()?; + let mut replica = replica.get_multiplexed_async_connection().await.ok()?; + let value: Option = redis::cmd("GET") + .arg("before") + .query_async(&mut replica) + .await + .ok()?; + value + }) + .await; + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(600))) + .await + .ok() + }) + .await; + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!( + client + .send(commands::get::("before")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + let promoted = redis::Client::open(format!("redis://127.0.0.1:{promoted}")).unwrap(); + let mut promoted = promoted.get_multiplexed_async_connection().await.unwrap(); + let after: Option = redis::cmd("GET") + .arg("after") + .query_async(&mut promoted) + .await + .unwrap(); + assert_eq!(after.as_deref(), Some("2")); +} + +const DATA_PASSWORD: &str = "data-secret"; + +#[rstest] +#[tokio::test] +async fn sentinel_reaches_a_master_that_requires_a_password() { + let sentinel = RedisSentinel::start_with_password(DATA_PASSWORD) + .await + .unwrap(); + let client = sentinel_client( + &sentinel, + Credentials::Static { + username: None, + password: Some(SecretValue::new(DATA_PASSWORD)), + }, + ); + + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("1") + ); +} + +#[rstest] +#[tokio::test] +async fn sentinel_rediscovers_with_the_credentials_the_provider_holds_now() { + let sentinel = RedisSentinel::start_with_password(DATA_PASSWORD) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = Arc::new(Rotating::new("alice", "first")); + let client = sentinel_client( + &sentinel, + Credentials::Provider(CredentialsProvider(provider.clone())), + ); + client + .send(commands::set("k", "1", Duration::from_secs(600))) + .await + .unwrap(); + + provider.rotate("bob", "second"); + sentinel + .cli_on_data_nodes(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let after_rediscovery = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + + assert_eq!(after_rediscovery, "1"); +} + +#[rstest] +#[tokio::test] +async fn a_timed_out_command_replaces_a_connection_that_stopped_answering( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let proxies = proxied(&server).await; + let client = RedisClient::new(RedisConfig { + target: proxies.target.clone(), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(500), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("k", "1", Duration::from_secs(600))) + .await + .unwrap(); + + proxies.blackhole_open_connections(); + let stalled = client.send(commands::get::("k")).await; + let next = client.send(commands::get::("k")).await; + + assert!( + matches!(&stalled, Err(Error::Transport(error)) if error.is_timeout()), + "{stalled:?}" + ); + assert_eq!(next.unwrap().as_deref(), Some("1")); +} diff --git a/litellm-rust/crates/redis-client/tests/pipeline.rs b/litellm-rust/crates/redis-client/tests/pipeline.rs new file mode 100644 index 00000000000..5aed70b42dd --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pipeline.rs @@ -0,0 +1,362 @@ +mod support; + +use std::time::{Duration, Instant}; + +use litellm_redis_client::Value; +use litellm_redis_client::commands; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; + +use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Expiry, Script}; +use support::connect; + +async fn begin_migration(cluster: &litellm_redis_testing::RedisCluster, key: &str) { + let (_, first) = cluster.nodes().next().unwrap(); + let slot = cluster + .cli_on_node(first, &["CLUSTER", "KEYSLOT", key]) + .await + .unwrap(); + let slot = slot.trim(); + let number: u16 = slot.parse().unwrap(); + let nodes = cluster + .cli_on_node(first, &["CLUSTER", "NODES"]) + .await + .unwrap(); + let owner = nodes.lines().find(|line| line.split_whitespace().skip(8).any(|range| { + let (start, end) = range.split_once('-').unwrap_or((range, range)); + matches!((start.parse::(), end.parse::()), (Ok(start), Ok(end)) if (start..=end).contains(&number)) + })).unwrap(); + let target = nodes.lines().find(|line| *line != owner).unwrap(); + let address = |line: &str| { + let value = line.split_whitespace().nth(1).unwrap(); + value + .split_once('@') + .unwrap() + .0 + .rsplit_once(':') + .unwrap() + .1 + .parse::() + .unwrap() + }; + let source_port = address(owner); + let target_port = address(target); + let source_id = owner.split_whitespace().next().unwrap(); + let target_id = target.split_whitespace().next().unwrap(); + assert_eq!( + cluster + .cli_on_node(source_port, &["GET", key]) + .await + .unwrap() + .trim(), + "0" + ); + assert_eq!( + cluster + .cli_on_node( + target_port, + &["CLUSTER", "SETSLOT", slot, "IMPORTING", source_id] + ) + .await + .unwrap() + .trim(), + "OK" + ); + assert_eq!( + cluster + .cli_on_node( + source_port, + &["CLUSTER", "SETSLOT", slot, "MIGRATING", target_id] + ) + .await + .unwrap() + .trim(), + "OK" + ); + let state = cluster + .cli_on_node(source_port, &["CLUSTER", "NODES"]) + .await + .unwrap(); + assert!( + state.contains(&format!("[{slot}->-{target_id}]")), + "{state}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() { + let cluster = litellm_redis_testing::RedisCluster::start().await.unwrap(); + let server = RedisDeployment::Cluster(cluster); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let counter = "{migration}:counter"; + client + .send(commands::set(counter, 0, Duration::from_secs(60))) + .await + .unwrap(); + let RedisDeployment::Cluster(cluster) = &server else { + unreachable!() + }; + begin_migration(cluster, counter).await; + let both_keys = Script::new("return redis.call('MGET', KEYS[1], KEYS[2])"); + + let results = client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + counter, + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::eval::( + &both_keys, + &[counter, "{migration}:missing"], + Vec::::new(), + )), + ) + .await; + + assert_eq!( + client.send(commands::get::(counter)).await.unwrap(), + Some(1) + ); + assert_eq!( + results[0].as_ref().unwrap(), + &Value::BulkString(b"1".to_vec()) + ); + assert!( + matches!(&results[1], Err(Error::Server(error)) if error.code() == Some("TRYAGAIN")), + "{:?}", + results[1] + ); +} + +#[rstest] +#[tokio::test] +async fn each_op_settles_on_its_own_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set( + "text", + "not a number", + Duration::from_secs(60), + )) + .await + .unwrap(); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "text", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(5), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::set("b", "2", Duration::from_secs(60))) + .with(commands::get::("b")) + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::get::("missing")); + + let results = client.run(&pipeline).await; + + assert_eq!(results.len(), 7); + assert_eq!( + results[0].as_ref().unwrap(), + &Value::BulkString(b"1".to_vec()) + ); + assert!( + matches!(results[1], Err(Error::Server(_))), + "{:?}", + results[1] + ); + assert_eq!( + results[2].as_ref().unwrap(), + &Value::BulkString(b"5".to_vec()) + ); + assert_eq!(results[3].as_ref().unwrap(), &Value::Okay); + assert_eq!( + results[4].as_ref().unwrap(), + &Value::BulkString(b"2".to_vec()) + ); + assert_eq!( + results[5].as_ref().unwrap(), + &Value::BulkString(b"2".to_vec()) + ); + assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); +} + +#[rstest] +#[case::cached(true, Value::BulkString(b"2".to_vec()))] +#[case::lost_by_the_server(false, Value::Nil)] +#[tokio::test] +async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py( + #[case] cached: bool, + #[case] read_in_the_same_pipeline: Value, +) { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::eval::(&incr, &["warmup"], 1)) + .await + .unwrap(); + if !cached { + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + } + let pipeline = client + .pipeline() + .with(commands::eval::(&incr, &["counter"], 2)) + .with(commands::get::("counter")); + + let results = client.run(&pipeline).await; + + assert_eq!(results[0].as_ref().unwrap(), &Value::Int(2)); + assert_eq!(results[1].as_ref().unwrap(), &read_in_the_same_pipeline); + assert_eq!( + client.send(commands::get::("counter")).await.unwrap(), + Some(2) + ); +} + +#[rstest] +#[tokio::test] +async fn an_unreachable_server_fails_every_op_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )); + + server.freeze().await.unwrap(); + let results = client.run(&pipeline).await; + server.thaw().await.unwrap(); + + assert_eq!(results.len(), 2); + assert!( + matches!(results[0], Err(Error::Transport(_))), + "{:?}", + results[0] + ); + assert!( + matches!(results[1], Err(Error::Transport(_))), + "{:?}", + results[1] + ); +} + +#[rstest] +#[tokio::test] +async fn an_empty_pipeline_returns_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(2)); + + server.freeze().await.unwrap(); + let started = Instant::now(); + let results = client.run(&client.pipeline()).await; + let elapsed = started.elapsed(); + server.thaw().await.unwrap(); + + assert!(results.is_empty()); + assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_lost( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&deployment, "scripts", DEFAULT_TIMEOUT); + let count = Script::new("return redis.call('INCR', KEYS[1])"); + deployment + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + deployment + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + let first = client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + "{a}spend", + Amount::Float(0.5), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::eval::( + &count, + &["{a}"], + Vec::::new(), + )) + .with(commands::eval::( + &count, + &["b"], + Vec::::new(), + )), + ) + .await; + let second = client + .run(&client.pipeline().with(commands::eval::( + &count, + &["{a}"], + Vec::::new(), + ))) + .await; + + let replies: Vec = first + .into_iter() + .chain(second) + .map(Result::unwrap) + .collect(); + assert_eq!( + replies, + [ + Value::BulkString(b"0.5".to_vec()), + Value::Int(1), + Value::Int(1), + Value::Int(2) + ] + ); + assert_eq!( + client.send(commands::get::("{a}spend")).await.unwrap(), + Some(0.5) + ); + let stats = support::command_stats(&deployment).await; + assert!(support::calls(&stats, "evalsha") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "eval"), 0, "{stats}"); +} diff --git a/litellm-rust/crates/redis-client/tests/pubsub.rs b/litellm-rust/crates/redis-client/tests/pubsub.rs new file mode 100644 index 00000000000..2ad7cb38a5a --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pubsub.rs @@ -0,0 +1,189 @@ +mod support; + +use std::time::{Duration, Instant}; + +use futures_util::StreamExt; +use litellm_circuit_breaker::{Breaker, Config}; +use litellm_redis_client::commands; +use litellm_redis_client::{ + Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, + Target, Timeouts, +}; +use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually}; +use rstest::rstest; +use support::{connect, connect_guarded}; +use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + +const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation"; + +fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver> { + let (sender, receiver) = unbounded_channel(); + let mut messages = client.subscribe(channel); + tokio::spawn(async move { + while let Some(message) = messages.next().await { + if sender.send(message).is_err() { + return; + } + } + }); + receiver +} + +async fn delivered( + publisher: &RedisClient, + channel: &str, + received: &mut UnboundedReceiver>, + payload: &[u8], +) { + eventually(async || { + publisher + .send(commands::publish(channel, payload)) + .await + .ok()?; + let deadline = tokio::time::Instant::now() + Duration::from_millis(200); + while let Ok(Some(message)) = tokio::time::timeout_at(deadline, received.recv()).await { + if message == payload { + return Some(()); + } + } + None + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn a_subscriber_receives_raw_bytes_on_the_channel_name_without_the_namespace( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let publisher = connect(&server, "ns", DEFAULT_TIMEOUT); + let plain = connect(&server, "", DEFAULT_TIMEOUT); + let mut received = listen(&plain, CHANNEL); + + delivered( + &publisher, + CHANNEL, + &mut received, + b"{\"cache_key\":\"k\"}\xff", + ) + .await; +} + +#[rstest] +#[tokio::test] +async fn a_subscriber_resubscribes_after_redis_restarts( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let mut received = listen(&client, CHANNEL); + delivered(&client, CHANNEL, &mut received, b"before").await; + + server.restart().await.unwrap(); + + delivered(&client, CHANNEL, &mut received, b"after").await; +} + +#[rstest] +#[tokio::test] +async fn a_sentinel_subscriber_receives_from_the_master() { + let server = RedisSentinel::start().await.unwrap(); + let (host, port) = server.sentinel(); + let client = RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + let mut received = listen(&client, CHANNEL); + + delivered(&client, CHANNEL, &mut received, b"payload").await; +} + +async fn publish_calls(server: &RedisDeployment) -> Vec { + let replies = match server { + RedisDeployment::Standalone(server) => { + vec![server.cli(&["INFO", "commandstats"]).await.unwrap()] + } + RedisDeployment::Cluster(cluster) => cluster + .cli_outputs(&["INFO", "commandstats"]) + .await + .unwrap(), + }; + replies + .iter() + .flat_map(|reply| reply.lines()) + .filter(|line| line.starts_with("cmdstat_publish:")) + .map(str::to_owned) + .collect() +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_rejects_a_publish_without_sending_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = Breaker::new( + Config { + failure_threshold: 2, + recovery_timeout: Duration::from_secs(60), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ); + let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker)); + + server.freeze().await.unwrap(); + let opened = [client.ping().await, client.ping().await]; + let started = Instant::now(); + let rejected = client.send(commands::publish(CHANNEL, b"payload")).await; + let rejected_in = started.elapsed(); + server.thaw().await.unwrap(); + + assert!( + opened + .iter() + .all(|ping| matches!(ping, Err(Error::Transport(_)))), + "{opened:?}" + ); + assert!(matches!(rejected, Err(Error::BreakerOpen)), "{rejected:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + assert_eq!(publish_calls(&server).await, Vec::::new()); +} + +// redis-py's PubSub subscribes on a RESP2 connection, so it never sends HELLO +#[rstest] +#[tokio::test] +async fn a_subscription_speaks_resp2_like_redis_py( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let mut received = listen(&client, CHANNEL); + + delivered(&client, CHANNEL, &mut received, b"evict").await; + + let stats = support::command_stats(&server).await; + assert!(support::calls(&stats, "subscribe") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "hello"), 0, "{stats}"); +} diff --git a/litellm-rust/crates/redis-client/tests/script.rs b/litellm-rust/crates/redis-client/tests/script.rs new file mode 100644 index 00000000000..ea6d1c5ff78 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/script.rs @@ -0,0 +1,133 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; + +use litellm_redis_client::Script; +use litellm_redis_client::commands; +use litellm_redis_client::{DEFAULT_TIMEOUT, Error, Ttl}; +use support::connect; + +#[rstest] +#[tokio::test] +async fn script_keys_are_namespaced( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "a", DEFAULT_TIMEOUT); + let bare = connect(&server, "", DEFAULT_TIMEOUT); + let read = Script::new("return redis.call('GET', KEYS[1])"); + bare.send(commands::set("a:x", "1", Duration::from_secs(60))) + .await + .unwrap(); + + let value: Option = client + .send(commands::eval(&read, &["x"], Vec::::new())) + .await + .unwrap(); + + assert_eq!(value.as_deref(), Some("1")); +} + +#[rstest] +#[tokio::test] +async fn a_failing_script_is_a_server_error_and_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let failed = client + .send(commands::eval::(&incr, &["list"], 1)) + .await; + + assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); + assert_eq!( + client + .send(commands::eval::(&incr, &["counter"], 3)) + .await + .unwrap(), + 3 + ); +} + +#[rstest] +#[tokio::test] +async fn a_flushed_script_is_loaded_again( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::eval::(&incr, &["counter"], 1)) + .await + .unwrap(); + + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::eval::(&incr, &["counter"], 1)) + .await + .unwrap(), + 2 + ); +} + +const INCREMENT_WITH_FLOOR: &str = "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) \ +if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end \ +if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end \ +return count"; + +#[rstest] +#[tokio::test] +async fn pythons_increment_with_floor_script_floors_at_zero_and_expires_in_one_call( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(false, true)] pipelined: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + let floor = Script::new(INCREMENT_WITH_FLOOR); + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + + let count = match pipelined { + false => client + .send(commands::eval::(&floor, &["slots"], (-5, 60))) + .await + .unwrap(), + true => { + let results = client + .run(&client.pipeline().with(commands::eval::( + &floor, + &["slots"], + (-5, 60), + ))) + .await; + redis::from_redis_value(results[0].as_ref().unwrap().clone()).unwrap() + } + }; + + assert_eq!(count, 0); + assert_eq!( + client.send(commands::get::("slots")).await.unwrap(), + Some(0) + ); + assert!(matches!( + client.send(commands::ttl("slots")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); +} diff --git a/litellm-rust/crates/redis-client/tests/support/mod.rs b/litellm-rust/crates/redis-client/tests/support/mod.rs new file mode 100644 index 00000000000..9cf209dce4a --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/support/mod.rs @@ -0,0 +1,206 @@ +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::{pin::Pin, time::Duration}; + +use futures_util::Stream; +use litellm_circuit_breaker::Breaker; +use litellm_redis_client::{ + BasicAuth, Cluster, Credentials, Expiry, Namespace, Node, RedisClient, RedisConfig, RedisError, + RedisUrl, Standalone, StreamingCredentialsProvider, Target, Timeouts, +}; +use litellm_redis_testing::{ClientConnection, RedisDeployment, TcpProxy}; +use litellm_secret_value::SecretValue; +use tokio::sync::watch; + +pub const KEEP: Expiry = Expiry::KeepExisting(Duration::from_secs(60)); + +pub fn endpoint(deployment: &RedisDeployment) -> Target { + match deployment { + RedisDeployment::Standalone(server) => Target::Url(RedisUrl { + url: SecretValue::new(server.url()), + tls: None, + db: None, + username: None, + password: None, + }), + RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster { + nodes: cluster + .nodes() + .map(|(host, port)| Node { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + } +} + +pub fn connect(deployment: &RedisDeployment, namespace: &str, timeout: Duration) -> RedisClient { + connect_guarded(deployment, namespace, timeout, None) +} + +pub fn connect_guarded( + deployment: &RedisDeployment, + namespace: &str, + timeout: Duration, + breaker: Option, +) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(deployment), + credentials: Credentials::NONE, + namespace: Namespace::new(namespace), + timeouts: Timeouts::new(timeout, None), + breaker, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +pub fn connect_as(deployment: &RedisDeployment, credentials: Credentials) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(deployment), + credentials, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(2), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +pub async fn connections(deployment: &RedisDeployment) -> Vec { + deployment.clients().await.unwrap() +} + +/// A `TcpProxy` in front of every node, and a target that reaches the nodes only through them. A +/// cluster announces the proxy ports, so its slot map and redirects lead there too. +pub struct Proxied { + proxies: Vec, + pub target: Target, +} + +impl Proxied { + pub fn blackhole_open_connections(&self) { + self.proxies + .iter() + .for_each(TcpProxy::blackhole_open_connections); + } +} + +pub async fn proxied(deployment: &RedisDeployment) -> Proxied { + match deployment { + RedisDeployment::Standalone(server) => { + let proxy = TcpProxy::start(server.port()).await.unwrap(); + let target = Target::Standalone(Standalone { + host: "127.0.0.1".to_owned(), + port: proxy.port(), + tls: None, + db: 0, + }); + Proxied { + proxies: vec![proxy], + target, + } + } + RedisDeployment::Cluster(cluster) => { + let mut proxies = Vec::new(); + for (_, port) in cluster.nodes() { + proxies.push(TcpProxy::start(port).await.unwrap()); + } + let ports: [u16; 3] = std::array::from_fn(|node| proxies[node].port()); + cluster.announce_ports(ports).await.unwrap(); + let target = Target::Cluster(Cluster { + nodes: ports + .iter() + .map(|&port| Node { + host: "127.0.0.1".to_owned(), + port, + }) + .collect(), + tls: None, + }); + Proxied { proxies, target } + } + } +} + +/// `INFO

` from every node, concatenated. +async fn info(deployment: &RedisDeployment, section: &str) -> String { + match deployment { + RedisDeployment::Standalone(server) => server.cli(&["INFO", section]).await.unwrap(), + RedisDeployment::Cluster(cluster) => { + let mut stats = String::new(); + for (_, port) in cluster.nodes() { + stats.push_str(&cluster.cli_on_node(port, &["INFO", section]).await.unwrap()); + } + stats + } + } +} + +/// `INFO commandstats` from every node, so a test can count what was actually sent. +pub async fn command_stats(deployment: &RedisDeployment) -> String { + info(deployment, "commandstats").await +} + +/// Keys looked up on every node since the last `CONFIG RESETSTAT`, found or not. +pub async fn keys_read(deployment: &RedisDeployment) -> u64 { + let stats = info(deployment, "stats").await; + ["keyspace_hits:", "keyspace_misses:"] + .iter() + .flat_map(|field| { + stats + .lines() + .filter_map(move |line| line.trim().strip_prefix(field)?.parse::().ok()) + }) + .sum() +} + +/// How many times `command` reached a server, summed across the nodes in `stats`. +pub fn calls(stats: &str, command: &str) -> u64 { + let prefix = format!("cmdstat_{command}:calls="); + stats + .lines() + .filter_map(|line| line.trim().strip_prefix(&prefix)) + .filter_map(|rest| rest.split(',').next()?.parse::().ok()) + .sum() +} + +pub struct Rotating(watch::Sender); + +impl Rotating { + pub fn new(username: &str, password: &str) -> Self { + Self(watch::Sender::new(BasicAuth::new( + username.to_owned(), + password.to_owned(), + ))) + } + + pub fn rotate(&self, username: &str, password: &str) { + self.0 + .send_replace(BasicAuth::new(username.to_owned(), password.to_owned())); + } +} + +impl StreamingCredentialsProvider for Rotating { + fn subscribe( + &self, + ) -> Pin> + Send + 'static>> { + let receiver = self.0.subscribe(); + Box::pin(futures_util::stream::unfold( + (receiver, true), + |(mut receiver, first)| async move { + if !first { + receiver.changed().await.ok()?; + } + let credentials = receiver.borrow_and_update().clone(); + Some((Ok(credentials), (receiver, false))) + }, + )) + } +} diff --git a/litellm-rust/crates/redis-connect/AGENTS.md b/litellm-rust/crates/redis-connect/AGENTS.md new file mode 100644 index 00000000000..8e9d1302eaf --- /dev/null +++ b/litellm-rust/crates/redis-connect/AGENTS.md @@ -0,0 +1,11 @@ +- Turns a `litellm_redis_types::Resolution` into a `RedisClient`, the same way for every host. It depends on `litellm-redis-types`, never on `litellm-config`, so a host that resolves Redis some other way connects the same way. `redis-client` stays cloud-agnostic, so credential providers and cloud token sources are built here +- The host passes only what it alone can build in `HostServices`: the event sink, the HTTP client and the environment lookup. The breaker config comes from `Resolution::breaker`, which `config` reads from `REDIS_CIRCUIT_BREAKER_*` with Python's fallback, so a bad value keeps the default instead of failing connect. Each cloud source asks its auth crate for one explicit identity, the way the Python SDK builds it: `AwsAuthService::default_chain_credentials` (botocore's default chain), `CredentialSource::application_default` (`google.auth.default`) and `AzureIdentity` (the class `_build_azure_credential` picks). LiteLLM's provider settings such as `AWS_ROLE_NAME`, `AZURE_AD_TOKEN` or `VERTEXAI_CREDENTIALS` belong to Bedrock, Azure OpenAI and Vertex and never reach those constructors. `Auth::Provider` (a handle for a Python `credential_provider` object) fails at connect with an error naming the key, since no host can supply one yet +- `RedisHosts` is the one place a host shares Redis state. Breakers are keyed by `Endpoint`, so every caller on a server sees one health signal whatever its db or credentials. A client takes only breaker settings and admits through its server's breaker, so no caller can pair a client with a breaker of its own, and clients are keyed by `Resolution` alone. Both are held weakly, so a breaker or client is freed with its last user, and the lock is held while a client connects so each is built once. A caller holding the Python GIL releases it first. Hosts own one instance (python-bridge per process, the gateway per `connect`) and never keep a registry of their own +- Every expiring credential goes through the one refresh loop in `refresh.rs`. A cloud source only fetches. Never add a second loop or let a source sleep. The expiry deadline lives only in the published `Latest::Ready`, so expiring it leaves `Ready` and cannot fire again. Never keep a copy of the deadline in the loop +- AWS SDK identity settings must be in the process environment. `aws-config` cannot take the host lookup, so config-only overrides for profiles, role/web-identity settings, credentials, files and metadata sources fail before connecting rather than silently using another identity. Equal process/config values are allowed. Bedrock-only settings remain unrelated to Redis +- redis-rs subscribes once for the connection handshake and again for re-authentication, so a new connection sends a redundant initial AUTH. Unchanged refreshes do not trigger another AUTH. Do not skip a provider's first or alternating subscriptions, because concurrent connections and retries each need the current token. Removing the initial duplicate requires an upstream connection-level fix +- The loop starts on the first `subscribe()`, which redis-rs calls inside its runtime, so `connect` is safe off a runtime +- A token-only credential (Azure without `REDIS_USERNAME`, GCP) is sent as `AUTH default `, because redis-rs always sends a username with a credentials provider. On Redis 6+ that is the same as Python's `AUTH ` +- Tests for the refresh loop use a fake source and real Redis ACL users. Cloud sources are proven against the live service, unit tests only pin request shapes we own +- Rotation tests must retain the same server-side connection IDs across re-AUTH, then force a reconnect after removing the old user and verify the new identity still works +- Live cloud tests live in `tests/live/`. Set `LITELLM_REDIS_AWS_LIVE_CONFIG`, `LITELLM_REDIS_AZURE_LIVE_CONFIG` or `LITELLM_REDIS_GCP_LIVE_CONFIG` to a JSON file containing `redis` settings and `hold_seconds`. Set the waiting period from the deployment's token lifetime when checking expiry. Optional fixtures return early when unset, and configured tests run under plain `cargo test`. An unconfigured run never counts as live proof diff --git a/litellm-rust/crates/redis-connect/Cargo.toml b/litellm-rust/crates/redis-connect/Cargo.toml new file mode 100644 index 00000000000..29b3293b856 --- /dev/null +++ b/litellm-rust/crates/redis-connect/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "litellm-redis-connect" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-aws.workspace = true +litellm-auth-azure.workspace = true +litellm-auth-gcp.workspace = true +litellm-auth-types.workspace = true +litellm-error.workspace = true +litellm-secret-value.workspace = true +strum.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-types.workspace = true +litellm-core-utils.workspace = true +litellm-http.workspace = true +litellm-redis-client.workspace = true +futures-util.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +veil.workspace = true +time.workspace = true +url.workspace = true +tokio = { workspace = true, features = ["sync", "time", "macros"] } +tracing.workspace = true + +[dev-dependencies] +litellm-redis-resolve.workspace = true +litellm-http = { workspace = true, features = ["test-support"] } +litellm-redis-testing.workspace = true +rustls.workspace = true +rstest.workspace = true +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] } diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs new file mode 100644 index 00000000000..d11f6ba47ff --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -0,0 +1,152 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_circuit_breaker::{Breaker, Config as BreakerConfig}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{ + Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Namespace, RedisClient, RedisConfig, + RedisEvents, Target, Timeouts, +}; +use litellm_redis_types::{Auth, HealthCheck, Resolution, StaticAuth}; + +use crate::{ + CredentialSource, Error, RefreshingCredentials, + sources::{AzureAd, ElastiCacheIam, GcpIam}, +}; + +pub type SharedEnvironment = Arc; + +pub struct HostServices { + pub events: Option>, + pub http: litellm_http::Client, + pub environment: SharedEnvironment, + pub credential_provider: Option>, +} + +// litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only +const CLUSTER_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(25); + +/// A breaker whose transitions reach `events`, for every client that talks to one Redis. +pub(crate) fn breaker(config: BreakerConfig, events: Option>) -> Breaker { + Breaker::new( + config, + Box::new(move |transition| { + if let Some(events) = &events { + events.breaker(transition); + } + }), + ) +} + +pub fn connect( + resolved: Resolution, + host: HostServices, + breaker: Option, +) -> Result { + tracing::info!(target = ?resolved.target, namespace = ?resolved.namespace, "redis client configured"); + let credentials = credentials(resolved.auth, &host)?; + let events = host.events; + let health_check = health_check(&resolved.target, resolved.health_check); + Ok(RedisClient::new(RedisConfig { + target: resolved.target, + credentials, + namespace: resolved.namespace.map(Namespace::new).unwrap_or_default(), + timeouts: Timeouts::new( + resolved.timeout.unwrap_or(DEFAULT_TIMEOUT), + resolved.connect_timeout, + ), + breaker, + events, + health_check, + keepalive: resolved.keepalive, + })?) +} + +fn health_check(target: &Target, configured: HealthCheck) -> Option { + match (target, configured) { + (_, HealthCheck::Every(interval)) => Some(interval), + (_, HealthCheck::Disabled) => None, + (Target::Cluster(_), HealthCheck::Unset) => Some(CLUSTER_HEALTH_CHECK_INTERVAL), + (_, HealthCheck::Unset) => None, + } +} + +fn credentials(auth: Auth, host: &HostServices) -> Result { + let refreshing = |provider: RefreshingCredentials| { + Credentials::Provider(CredentialsProvider(Arc::new(provider))) + }; + Ok(match auth { + Auth::Static(StaticAuth { username, password }) => { + Credentials::Static { username, password } + } + Auth::Provider(_) => match &host.credential_provider { + Some(source) => refreshing(RefreshingCredentials::new(Arc::clone(source))), + None => return Err(Error::NoHostCredentialProvider), + }, + Auth::AzureAd(auth) => refreshing(RefreshingCredentials::new(AzureAd::new( + auth, + host.environment.clone(), + ))), + Auth::AwsIam(auth) => refreshing(RefreshingCredentials::new(ElastiCacheIam::new( + auth, + &*host.environment, + )?)), + Auth::GcpIam { service_account } => refreshing(RefreshingCredentials::new(GcpIam::new( + host.http.clone(), + service_account, + host.environment.clone(), + ))), + }) +} + +#[cfg(test)] +mod tests { + use litellm_redis_client::{Cluster, Node, Standalone}; + use rstest::rstest; + + use super::*; + + fn cluster() -> Target { + Target::Cluster(Cluster { + nodes: vec![Node { + host: "node".to_owned(), + port: 7000, + }], + tls: None, + }) + } + + fn standalone() -> Target { + Target::Standalone(Standalone { + host: "redis".to_owned(), + port: 6379, + tls: None, + db: 0, + }) + } + + #[rstest] + #[case::an_unset_cluster_takes_the_python_default( + cluster(), + HealthCheck::Unset, + Some(CLUSTER_HEALTH_CHECK_INTERVAL) + )] + #[case::zero_keeps_a_cluster_off(cluster(), HealthCheck::Disabled, None)] + #[case::an_unset_standalone_stays_off(standalone(), HealthCheck::Unset, None)] + #[case::a_configured_interval_wins( + cluster(), + HealthCheck::Every(Duration::from_secs(5)), + Some(Duration::from_secs(5)) + )] + #[case::a_negative_interval_checks_every_command( + standalone(), + HealthCheck::Every(Duration::ZERO), + Some(Duration::ZERO) + )] + fn health_checks_follow_redis_py( + #[case] target: Target, + #[case] configured: HealthCheck, + #[case] expected: Option, + ) { + assert_eq!(health_check(&target, configured), expected); + } +} diff --git a/litellm-rust/crates/redis-connect/src/error.rs b/litellm-rust/crates/redis-connect/src/error.rs new file mode 100644 index 00000000000..98343a03fb8 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/error.rs @@ -0,0 +1,102 @@ +#[derive(veil::Redact, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Connect(#[from] litellm_redis_client::ConnectError), + #[error( + "Redis setting credential_provider names a Python object, which this host cannot supply" + )] + NoHostCredentialProvider, + #[error( + "Redis credential setting {setting} must be set in the process environment, not only environment_variables" + )] + ProcessEnvironmentRequired { setting: &'static str }, + #[error("{stage} failed{}", http_status_suffix(*http_status))] + Credential { + stage: CredentialStage, + http_status: Option, + #[source] + #[redact(with = "[REDACTED]")] + cause: litellm_error::Source, + }, + #[error("credential acquisition exceeded {timeout:?}")] + CredentialTimeout { + timeout: std::time::Duration, + #[source] + cause: tokio::time::error::Elapsed, + }, + #[error("Redis credential has expired")] + ExpiredCredential, + #[error("no {provider} credential is available for Redis")] + MissingCredential { provider: &'static str }, +} + +#[derive(Clone, Copy, Debug, strum::Display, strum::VariantArray)] +pub enum CredentialStage { + #[strum(serialize = "caller identity")] + Identity, + #[strum(serialize = "Redis token signing")] + Signing, + #[strum(serialize = "Redis token exchange")] + Exchange, + #[strum(serialize = "Redis token response")] + Response, + #[strum(serialize = "Redis credential_provider")] + Provider, +} + +impl Error { + pub fn credential( + stage: CredentialStage, + ) -> impl FnOnce(E) -> Self { + move |cause| Self::Credential { + stage, + http_status: None, + cause: litellm_error::Source::new(cause), + } + } + + pub(crate) fn credential_response( + stage: CredentialStage, + status: u16, + ) -> impl FnOnce(litellm_http::Error) -> Self { + move |cause| Self::Credential { + stage, + http_status: Some(status), + cause: litellm_error::Source::new(cause), + } + } + + pub(crate) fn upstream_status(&self) -> Option { + match self { + Self::Credential { http_status, .. } => *http_status, + Self::Connect(_) + | Self::NoHostCredentialProvider + | Self::ProcessEnvironmentRequired { .. } + | Self::CredentialTimeout { .. } + | Self::ExpiredCredential + | Self::MissingCredential { .. } => None, + } + } + + pub(crate) fn kind(&self) -> &'static str { + match self { + Self::Connect(_) => "connection", + Self::NoHostCredentialProvider => "host-credential-provider", + Self::ProcessEnvironmentRequired { .. } => "process-environment-required", + Self::Credential { stage, .. } => match stage { + CredentialStage::Identity => "caller-identity", + CredentialStage::Signing => "token-signing", + CredentialStage::Exchange => "token-exchange", + CredentialStage::Response => "token-response", + CredentialStage::Provider => "credential-provider", + }, + Self::CredentialTimeout { .. } => "acquisition-timeout", + Self::ExpiredCredential => "expired-credential", + Self::MissingCredential { .. } => "missing-credential", + } + } +} + +fn http_status_suffix(status: Option) -> String { + status.map_or_else(String::new, |status| format!(" (HTTP {status})")) +} diff --git a/litellm-rust/crates/redis-connect/src/events.rs b/litellm-rust/crates/redis-connect/src/events.rs new file mode 100644 index 00000000000..a556b7d370b --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/events.rs @@ -0,0 +1,129 @@ +use std::{ + sync::Mutex, + time::{Duration, Instant}, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{CallEvent, FailureKind, Outcome, RedisEvents}; + +const DEFAULT_TIMEOUT_LOG_INTERVAL: Duration = Duration::from_secs(5); + +pub struct TracingEvents { + throttle: TimeoutLogThrottle, +} + +impl TracingEvents { + pub fn new(lookup: &dyn Lookup) -> Self { + let interval = lookup + .get("REDIS_TIMEOUT_LOG_INTERVAL") + .and_then(|seconds| seconds.parse::().ok()) + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .unwrap_or(DEFAULT_TIMEOUT_LOG_INTERVAL); + Self { + throttle: TimeoutLogThrottle::new(interval), + } + } +} + +impl RedisEvents for TracingEvents { + fn call(&self, event: &CallEvent) { + let Outcome::Failed(kind) = event.outcome else { + return; + }; + match kind { + FailureKind::BreakerOpen => { + tracing::debug!(operation = ?event.operation, "redis call rejected by the open circuit breaker") + } + FailureKind::Timeout => match self.throttle.admit(Instant::now()) { + None => { + tracing::debug!(operation = ?event.operation, duration = ?event.duration, "redis call timed out") + } + Some(0) => { + tracing::error!(operation = ?event.operation, duration = ?event.duration, "redis call timed out") + } + Some(suppressed) => tracing::error!( + operation = ?event.operation, + duration = ?event.duration, + suppressed, + "redis call timed out, more redis timeouts since the previous line were logged at debug" + ), + }, + FailureKind::Transport | FailureKind::Server | FailureKind::Decode => { + tracing::error!(operation = ?event.operation, failure = ?kind, "redis call failed") + } + } + } + + fn breaker(&self, transition: &Transition) { + match transition.to { + State::Open { .. } => tracing::warn!( + generation = transition.generation, + "redis circuit breaker opened, fast-failing redis calls" + ), + State::HalfOpen => tracing::info!( + generation = transition.generation, + "redis circuit breaker half open, probing redis" + ), + State::Closed { .. } => tracing::info!( + generation = transition.generation, + "redis circuit breaker closed, redis recovered" + ), + } + } +} + +struct TimeoutLogThrottle { + interval: Duration, + state: Mutex<(Option, u64)>, +} + +impl TimeoutLogThrottle { + fn new(interval: Duration) -> Self { + Self { + interval, + state: Mutex::new((None, 0)), + } + } + + fn admit(&self, now: Instant) -> Option { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let (last_logged_at, suppressed) = &mut *state; + if last_logged_at.is_some_and(|last| now.duration_since(last) < self.interval) { + *suppressed += 1; + return None; + } + *last_logged_at = Some(now); + Some(std::mem::take(suppressed)) + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use rstest::rstest; + + use super::TimeoutLogThrottle; + + #[rstest] + fn admits_one_line_per_interval_and_counts_the_rest() { + let throttle = TimeoutLogThrottle::new(Duration::from_secs(5)); + let start = Instant::now(); + let at = |seconds: u64| start + Duration::from_secs(seconds); + + let admitted = [ + throttle.admit(at(0)), + throttle.admit(at(1)), + throttle.admit(at(4)), + throttle.admit(at(5)), + throttle.admit(at(6)), + throttle.admit(at(20)), + ]; + + assert_eq!(admitted, [Some(0), None, None, Some(2), None, Some(1)]); + } +} diff --git a/litellm-rust/crates/redis-connect/src/hosts.rs b/litellm-rust/crates/redis-connect/src/hosts.rs new file mode 100644 index 00000000000..11e69d45aff --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/hosts.rs @@ -0,0 +1,241 @@ +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, WeakBreaker}; +use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient}; +use litellm_redis_types::{Endpoint, Resolution}; + +use crate::{Error, HostServices, breaker, connect}; + +/// The Redis servers one host talks to. Every caller on a server shares its breaker, and callers with the same +/// settings share one client, which admits through its server's breaker. Both are held weakly, so each is freed +/// with its last user and the next caller builds a fresh one. +#[derive(Default)] +pub struct RedisHosts { + breakers: Mutex, WeakBreaker)>>, + clients: Mutex>, +} + +#[derive(PartialEq)] +struct ClientKey { + resolved: Resolution, + guarded: bool, +} + +impl RedisHosts { + /// The breaker of the server `endpoint` names, built from `config` by its first caller. Without an endpoint + /// the breaker is this caller's own, which `breakers` still lists. + pub fn breaker( + &self, + endpoint: Option, + config: BreakerConfig, + events: Option>, + ) -> Breaker { + let mut breakers = live(&self.breakers, WeakBreaker::upgrade); + if let Some(shared) = breakers + .iter() + .filter(|(key, _)| key.is_some() && *key == endpoint) + .find_map(|(_, weak)| weak.upgrade()) + { + return shared; + } + let built = breaker(config, events); + breakers.push((endpoint, built.downgrade())); + built + } + + /// The client for `resolved`, connected by its first caller and guarded by its server's breaker unless `breaker` + /// is none. The lock is held while it connects, so concurrent callers with the same settings connect once. + pub fn client( + &self, + resolved: Resolution, + host: HostServices, + breaker: Option, + ) -> Result { + let breaker = breaker + .map(|config| self.breaker(resolved.target.endpoint(), config, host.events.clone())); + let key = ClientKey { + resolved, + guarded: breaker.is_some(), + }; + let mut clients = live(&self.clients, WeakRedisClient::upgrade); + if let Some(shared) = clients + .iter() + .filter(|(entry, _)| *entry == key) + .find_map(|(_, weak)| weak.upgrade()) + { + return Ok(shared); + } + let client = connect(key.resolved.clone(), host, breaker)?; + clients.push((key, client.downgrade())); + Ok(client) + } + + /// Every breaker a caller still holds, for a state gauge. + pub fn breakers(&self) -> Vec { + live(&self.breakers, WeakBreaker::upgrade) + .iter() + .filter_map(|(_, weak)| weak.upgrade()) + .collect() + } +} + +fn live( + entries: &Mutex>, + upgrade: fn(&W) -> Option, +) -> MutexGuard<'_, Vec<(K, W)>> { + let mut entries = entries.lock().unwrap_or_else(PoisonError::into_inner); + entries.retain(|(_, weak)| upgrade(weak).is_some()); + entries +} + +#[cfg(test)] +mod tests { + use std::{ + sync::Arc, + time::{Duration, Instant}, + }; + + use futures_util::FutureExt; + use litellm_circuit_breaker::{Config, Failure, State}; + use litellm_redis_types::{ + Auth, Endpoint, HealthCheck, Resolution, Standalone, StaticAuth, Target, + }; + use rstest::rstest; + + use super::RedisHosts; + use crate::HostServices; + + const CONFIG: Config = Config { + failure_threshold: 1, + recovery_timeout: Duration::from_secs(60), + timeout_min_duration: Duration::ZERO, + }; + + fn target(host: &str, db: i64) -> Target { + Target::Standalone(Standalone { + host: host.to_owned(), + port: 6379, + tls: None, + db, + }) + } + + fn endpoint(host: &str, db: i64) -> Option { + target(host, db).endpoint() + } + + fn resolved(host: &str, namespace: Option<&str>) -> Resolution { + Resolution { + target: target(host, 0), + auth: Auth::Static(StaticAuth::default()), + namespace: namespace.map(str::to_owned), + timeout: None, + connect_timeout: None, + health_check: HealthCheck::Unset, + keepalive: true, + ignored: Vec::new(), + } + } + + fn services() -> HostServices { + HostServices { + events: None, + http: litellm_http::Client::for_test(), + environment: Arc::new(|_: &str| None), + credential_provider: None, + } + } + + fn open(breaker: &litellm_circuit_breaker::Breaker) { + breaker.failed(Failure::Hard, Instant::now()); + } + + #[rstest] + fn callers_on_one_server_share_its_breaker_whatever_their_db() { + let hosts = RedisHosts::default(); + let first = hosts.breaker(endpoint("r1", 0), CONFIG, None); + let second = hosts.breaker(endpoint("r1", 1), CONFIG, None); + let other = hosts.breaker(endpoint("r2", 0), CONFIG, None); + + open(&first); + + assert!(matches!(second.state(), State::Open { .. })); + assert!(matches!(other.state(), State::Closed { .. })); + } + + #[rstest] + fn a_breaker_without_an_endpoint_is_listed_but_never_shared() { + let hosts = RedisHosts::default(); + let own = hosts.breaker(None, CONFIG, None); + let shared = hosts.breaker(endpoint("r1", 0), CONFIG, None); + + open(&own); + + assert!(matches!(shared.state(), State::Closed { .. })); + assert_eq!(hosts.breakers().len(), 2); + } + + #[rstest] + fn a_server_whose_callers_are_all_gone_gets_a_fresh_breaker() { + let hosts = RedisHosts::default(); + let first = hosts.breaker(endpoint("r1", 0), CONFIG, None); + open(&first); + drop(first); + + let again = hosts.breaker(endpoint("r1", 0), CONFIG, None); + + assert!(matches!(again.state(), State::Closed { .. })); + assert_eq!(hosts.breakers().len(), 1); + } + + #[rstest] + #[tokio::test] + async fn callers_with_the_same_settings_share_one_client() { + let hosts = RedisHosts::default(); + let first = hosts + .client(resolved("r1", Some("a")), services(), Some(CONFIG)) + .unwrap(); + + let same = hosts + .client(resolved("r1", Some("a")), services(), Some(CONFIG)) + .unwrap(); + let other_namespace = hosts + .client(resolved("r1", Some("b")), services(), Some(CONFIG)) + .unwrap(); + + assert!(first.same_client(&same)); + assert!(!first.same_client(&other_namespace)); + } + + #[rstest] + #[tokio::test] + async fn a_client_admits_through_its_servers_breaker() { + let hosts = RedisHosts::default(); + let breaker = hosts.breaker(endpoint("r1", 1), CONFIG, None); + let client = hosts + .client(resolved("r1", None), services(), Some(CONFIG)) + .unwrap(); + + open(&breaker); + + assert!(matches!( + client.ping().now_or_never(), + Some(Err(litellm_redis_client::Error::BreakerOpen)) + )); + } + + #[rstest] + #[tokio::test] + async fn a_client_without_a_breaker_is_never_shared_with_a_guarded_one() { + let hosts = RedisHosts::default(); + let guarded = hosts + .client(resolved("r1", None), services(), Some(CONFIG)) + .unwrap(); + + let unguarded = hosts + .client(resolved("r1", None), services(), None) + .unwrap(); + + assert!(!guarded.same_client(&unguarded)); + } +} diff --git a/litellm-rust/crates/redis-connect/src/lib.rs b/litellm-rust/crates/redis-connect/src/lib.rs new file mode 100644 index 00000000000..7e37f721679 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/lib.rs @@ -0,0 +1,14 @@ +mod connect; +mod error; +mod events; +mod hosts; +mod refresh; +mod sources; + +pub(crate) use connect::breaker; +pub use connect::{HostServices, SharedEnvironment, connect}; +pub use error::{CredentialStage, Error}; +pub use events::TracingEvents; +pub use hosts::RedisHosts; +pub(crate) use refresh::RefreshingCredentials; +pub use refresh::{CredentialSource, Expiring, SourceFuture}; diff --git a/litellm-rust/crates/redis-connect/src/refresh.rs b/litellm-rust/crates/redis-connect/src/refresh.rs new file mode 100644 index 00000000000..1e05d3d6289 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/refresh.rs @@ -0,0 +1,783 @@ +use std::{ + future::{Future, pending}, + pin::Pin, + sync::{Arc, Mutex}, + time::{Duration, SystemTime}, +}; + +use futures_util::{Stream, stream}; +use litellm_redis_client::{BasicAuth, RedisError, StreamingCredentialsProvider}; +use litellm_secret_value::SecretValue; +use tokio::{sync::watch, time::Instant}; + +use crate::Error; + +const TOKEN_ONLY_USERNAME: &str = "default"; +const UNKNOWN_LIFETIME: Duration = Duration::from_secs(300); +const MIN_REFRESH: Duration = Duration::from_secs(1); +const ACQUISITION_TIMEOUT: Duration = Duration::from_secs(30); +const FIRST_RETRY: Duration = Duration::from_secs(1); +const MAX_RETRY: Duration = Duration::from_secs(60); + +#[derive(Clone, PartialEq, Eq)] +pub struct Expiring { + pub username: Option, + pub password: SecretValue, + pub expires_at: Option, +} + +pub type SourceFuture<'a> = Pin> + Send + 'a>>; + +pub trait CredentialSource: Send + Sync + 'static { + fn fetch(&self) -> SourceFuture<'_>; +} + +impl CredentialSource for Arc { + fn fetch(&self) -> SourceFuture<'_> { + (**self).fetch() + } +} + +#[derive(Clone)] +enum Latest { + Pending, + Ready { + username: String, + password: SecretValue, + expires_at: Option, + }, + Unavailable(Arc), +} + +impl Latest { + fn is_same_token(&self, other: &Self) -> bool { + matches!( + (self, other), + ( + Self::Ready { username, password, .. }, + Self::Ready { username: other_username, password: other_password, .. }, + ) if username == other_username && password == other_password + ) + } + + fn expires_at(&self) -> Option { + match self { + Self::Ready { expires_at, .. } => *expires_at, + Self::Pending | Self::Unavailable(_) => None, + } + } +} + +type Start = Box; + +#[derive(Clone)] +pub struct RefreshingCredentials { + latest: watch::Receiver, + start: Arc>>, +} + +impl RefreshingCredentials { + pub fn new(source: impl CredentialSource) -> Self { + let (sender, latest) = watch::channel(Latest::Pending); + let start: Start = Box::new(move || drop(tokio::spawn(refresh(source, sender)))); + Self { + latest, + start: Arc::new(Mutex::new(Some(start))), + } + } + + fn start_once(&self) { + let start = self + .start + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(start) = start { + start(); + } + } +} + +impl StreamingCredentialsProvider for RefreshingCredentials { + fn subscribe( + &self, + ) -> Pin> + Send + 'static>> { + self.start_once(); + Box::pin(stream::unfold( + (self.latest.clone(), true), + |(mut receiver, first)| async move { + if !first { + receiver.changed().await.ok()?; + } + let latest = receiver + .wait_for(|latest| !matches!(latest, Latest::Pending)) + .await + .ok()? + .clone(); + let item = match latest { + Latest::Ready { + username, password, .. + } => Ok(BasicAuth::new(username, password.expose().to_owned())), + Latest::Unavailable(error) => Err(public_error(&error)), + Latest::Pending => unreachable!("wait_for skips Pending"), + }; + Some((item, (receiver, false))) + }, + )) + } +} + +fn public_error(error: &Error) -> RedisError { + RedisError::from(( + litellm_redis_client::RedisErrorKind::AuthenticationFailed, + "fetching Redis credentials failed", + error.to_string(), + )) +} + +struct Fresh { + latest: Latest, + refresh_in: Duration, +} + +fn fresh(credentials: Expiring, now: SystemTime) -> Result { + let remaining = credentials + .expires_at + .map(|expires_at| expires_at.duration_since(now).unwrap_or_default()); + if remaining == Some(Duration::ZERO) { + return Err(Error::ExpiredCredential); + } + Ok(Fresh { + latest: Latest::Ready { + username: credentials + .username + .unwrap_or_else(|| TOKEN_ONLY_USERNAME.to_owned()), + password: credentials.password, + expires_at: remaining.map(|remaining| Instant::now() + remaining), + }, + refresh_in: refresh_after(credentials.expires_at, now), + }) +} + +async fn until_expiry( + future: impl Future, + sender: &watch::Sender, +) -> Option { + tokio::pin!(future); + loop { + let expires_at = sender.borrow().expires_at(); + let expiration = async { + match expires_at { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => pending().await, + } + }; + tokio::select! { + biased; + () = sender.closed() => return None, + () = expiration => { + sender.send_replace(Latest::Unavailable(Arc::new(Error::ExpiredCredential))); + } + result = &mut future => return Some(result), + } + } +} + +async fn refresh(source: impl CredentialSource, sender: watch::Sender) { + let mut retry = FIRST_RETRY; + loop { + let acquisition = async { + tokio::time::timeout(ACQUISITION_TIMEOUT, source.fetch()) + .await + .map_err(|cause| Error::CredentialTimeout { + timeout: ACQUISITION_TIMEOUT, + cause, + })? + .and_then(|credentials| fresh(credentials, SystemTime::now())) + }; + let Some(result) = until_expiry(acquisition, &sender).await else { + return; + }; + let wait = match result { + Ok(Fresh { latest, refresh_in }) => { + retry = FIRST_RETRY; + sender.send_if_modified(|current| { + let changed = !current.is_same_token(&latest); + *current = latest; + changed + }); + refresh_in + } + Err(error) => { + tracing::warn!(failure = error.kind(), upstream_status = error.upstream_status(), retry_in = ?retry, "fetching Redis credentials failed"); + if !matches!(*sender.borrow(), Latest::Ready { .. }) { + sender.send_replace(Latest::Unavailable(Arc::new(error))); + } + let wait = retry; + retry = (retry * 2).min(MAX_RETRY); + wait + } + }; + if until_expiry(tokio::time::sleep(wait), &sender) + .await + .is_none() + { + return; + } + } +} + +fn refresh_after(expires_at: Option, now: SystemTime) -> Duration { + let remaining = expires_at + .map(|expires_at| expires_at.duration_since(now).unwrap_or_default()) + .unwrap_or(UNKNOWN_LIFETIME); + (remaining * 3 / 4).max(MIN_REFRESH) +} + +#[cfg(test)] +mod tests { + use litellm_redis_client::commands; + use std::{ + collections::BTreeSet, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::{Duration, SystemTime}, + }; + + use futures_util::StreamExt; + use litellm_redis_client::StreamingCredentialsProvider; + use litellm_secret_value::SecretValue; + use rstest::rstest; + + use super::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture, refresh_after}; + use litellm_redis_client::{ + Cluster, Credentials, CredentialsProvider, Namespace, Node, RedisClient, RedisConfig, + RedisUrl, Sentinel, Target, Timeouts, + }; + use litellm_redis_testing::{ + ClientConnection, RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually, + }; + + #[rstest] + #[case::three_quarters_of_the_lifetime( + Some(Duration::from_secs(900)), + Duration::from_secs(675) + )] + #[case::already_expired(Some(Duration::ZERO), Duration::from_secs(1))] + #[case::unknown_expiry(None, Duration::from_secs(225))] + fn refreshes_well_before_expiry( + #[case] lifetime: Option, + #[case] expected: Duration, + ) { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000); + + assert_eq!( + refresh_after(lifetime.map(|lifetime| now + lifetime), now), + expected + ); + } + + struct Scripted { + replies: Mutex>>, + calls: Arc, + } + + impl CredentialSource for Scripted { + fn fetch(&self) -> SourceFuture<'_> { + self.calls.fetch_add(1, Ordering::SeqCst); + let reply = self.replies.lock().unwrap().remove(0); + Box::pin(async move { + reply.map_err(|reason| { + crate::Error::credential(crate::error::CredentialStage::Identity)( + std::io::Error::other(reason), + ) + }) + }) + } + } + + fn token(password: &str, lifetime: Duration) -> Result { + Ok(Expiring { + username: None, + password: SecretValue::new(password), + expires_at: Some(SystemTime::now() + lifetime), + }) + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn pushes_each_refresh_and_keeps_the_last_token_through_failures() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + token("first", Duration::from_secs(8)), + Err("token endpoint down"), + token("second", Duration::from_secs(3600)), + ]), + calls: calls.clone(), + }); + let mut stream = provider.subscribe(); + + let first = stream.next().await.unwrap().unwrap(); + let second = stream.next().await.unwrap().unwrap(); + let late = provider.subscribe().next().await.unwrap().unwrap(); + + assert_eq!((first.username(), first.password()), ("default", "first")); + assert_eq!(second.password(), "second"); + assert_eq!(late.password(), "second"); + assert_eq!(calls.load(Ordering::SeqCst), 3); + } + + #[rstest] + fn building_outside_a_runtime_fetches_nothing_until_a_connection_subscribes() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![token("first", Duration::from_secs(3600))]), + calls: calls.clone(), + }); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let before = calls.load(Ordering::SeqCst); + let first = runtime.block_on(async { provider.subscribe().next().await.unwrap().unwrap() }); + + assert_eq!(before, 0); + assert_eq!(first.password(), "first"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn a_failure_before_any_token_fails_the_connection_instead_of_hanging() { + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + Err("no identity"), + token("late", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + + assert!(stream.next().await.unwrap().is_err()); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "late"); + } + + struct Blocking { + first: Mutex>, + started: Arc, + cancelled: Arc, + } + + struct Cancellation(Arc); + + impl Drop for Cancellation { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + impl CredentialSource for Blocking { + fn fetch(&self) -> SourceFuture<'_> { + let first = self.first.lock().unwrap().take(); + Box::pin(async move { + if let Some(first) = first { + return Ok(first); + } + let _guard = Cancellation(self.cancelled.clone()); + self.started.notify_one(); + std::future::pending().await + }) + } + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn a_refresh_that_returns_the_same_token_is_not_pushed_again() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + token("same", Duration::from_secs(8)), + token("same", Duration::from_secs(3600)), + ]), + calls: calls.clone(), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "same"); + + let pushed = tokio::time::timeout(Duration::from_secs(60), stream.next()).await; + + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert!(pushed.is_err(), "an unchanged token was pushed again"); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn expiration_is_published_while_refresh_is_stalled() { + let started = Arc::new(tokio::sync::Notify::new()); + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(Some(token("first", Duration::from_secs(4)).unwrap())), + started: started.clone(), + cancelled: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "first"); + let start = tokio::time::Instant::now(); + started.notified().await; + let expired = stream.next().await.unwrap().unwrap_err(); + assert!(start.elapsed() <= Duration::from_secs(5)); + assert!(expired.to_string().contains("expired")); + let fresh = provider.subscribe().next().await.unwrap().unwrap_err(); + assert!(fresh.to_string().contains("expired")); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn expiration_is_published_once_and_later_failures_keep_their_cause() { + let replies = std::iter::once(token("first", Duration::from_secs(8))) + .chain(std::iter::repeat_n(Err("token endpoint down"), 16)) + .collect(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(replies), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "first"); + + let failures: Vec = stream + .take(5) + .map(|item| item.unwrap_err().to_string()) + .collect() + .await; + + assert!(failures[0].contains("expired"), "{failures:?}"); + assert!( + failures[1..] + .iter() + .all(|failure| failure.contains("caller identity") && !failure.contains("expired")), + "{failures:?}" + ); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn initial_acquisition_has_a_bounded_wait() { + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(None), + started: Arc::new(tokio::sync::Notify::new()), + cancelled: Arc::new(AtomicUsize::new(0)), + }); + let error = provider.subscribe().next().await.unwrap().unwrap_err(); + assert!(error.to_string().contains("exceeded")); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn dropping_all_subscribers_cancels_inflight_acquisition() { + let started = Arc::new(tokio::sync::Notify::new()); + let cancelled = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(None), + started: started.clone(), + cancelled: cancelled.clone(), + }); + let stream = provider.subscribe(); + started.notified().await; + drop(stream); + drop(provider); + tokio::task::yield_now().await; + assert_eq!(cancelled.load(Ordering::SeqCst), 1); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn an_expired_fetch_is_rejected_and_a_later_valid_token_recovers() { + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + Ok(Expiring { + username: None, + password: SecretValue::new("expired"), + expires_at: Some(SystemTime::UNIX_EPOCH), + }), + token("recovered", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert!( + stream + .next() + .await + .unwrap() + .unwrap_err() + .to_string() + .contains("expired") + ); + assert_eq!( + stream.next().await.unwrap().unwrap().password(), + "recovered" + ); + } + + #[rstest] + fn credential_failure_keeps_its_typed_cause_but_hides_it_at_the_redis_boundary() { + let error = crate::Error::credential(crate::error::CredentialStage::Identity)( + std::io::Error::new(std::io::ErrorKind::PermissionDenied, "credential-value"), + ); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .kind(), + std::io::ErrorKind::PermissionDenied + ); + let debug = format!("{error:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("[REDACTED]")); + let public = super::public_error(&error).to_string(); + assert!(public.contains("caller identity")); + assert!(!public.contains("credential-value")); + } + + #[rstest] + fn credential_http_failures_expose_status_without_exposing_the_response_body() { + let error = crate::Error::credential_response(crate::CredentialStage::Exchange, 503)( + litellm_http::Error::upstream(503, "credential-value".into()), + ); + let public = super::public_error(&error).to_string(); + assert!(public.contains("HTTP 503")); + assert!(!public.contains("credential-value")); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .upstream_response() + .unwrap() + .0, + 503 + ); + } + + fn endpoint(deployment: &RedisDeployment) -> Target { + match deployment { + RedisDeployment::Standalone(server) => Target::Url(RedisUrl { + url: SecretValue::new(server.url()), + tls: None, + db: None, + username: None, + password: None, + }), + RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster { + nodes: cluster + .nodes() + .map(|(host, port)| Node { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + } + } + + fn user(username: &str, password: &str, lifetime: Duration) -> Result { + Ok(Expiring { + username: Some(username.to_owned()), + password: SecretValue::new(password), + expires_at: Some(SystemTime::now() + lifetime), + }) + } + + fn connection_users(connections: &[ClientConnection]) -> BTreeSet<&str> { + connections + .iter() + .map(|connection| connection.user.as_str()) + .collect() + } + + fn connection_ids(connections: &[ClientConnection]) -> BTreeSet<(u16, u64)> { + connections + .iter() + .map(|connection| (connection.node, connection.id)) + .collect() + } + + #[rstest] + #[tokio::test] + async fn a_refreshed_token_reauthenticates_live_connections_before_the_old_one_expires( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + user("alice", "first", Duration::from_secs(2)), + user("bob", "second", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))), + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = server.clients().await.unwrap(); + + let after = eventually(async || { + let connections = server.clients().await.ok()?; + (!connections.is_empty() && connection_users(&connections) == BTreeSet::from(["bob"])) + .then_some(connections) + }) + .await; + server + .cli_on_every_node(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + + assert!(!before.is_empty()); + assert_eq!(connection_users(&before), BTreeSet::from(["alice"])); + assert_eq!(connection_ids(&after), connection_ids(&before)); + assert_eq!( + client + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + server + .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let value = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + let reconnected = server.clients().await.unwrap(); + + assert_eq!(value, "1"); + assert!(!reconnected.is_empty()); + assert_eq!(connection_users(&reconnected), BTreeSet::from(["bob"])); + assert!(connection_ids(&reconnected).is_disjoint(&connection_ids(&after))); + } + + /// Cloud credentials rotate while Sentinel fails over: the refreshed user must be what the + /// client presents to the master Sentinel promotes, not the one it connected with first. + #[rstest] + #[tokio::test] + async fn a_refreshed_user_reaches_the_master_sentinel_promotes() { + let sentinel = RedisSentinel::start().await.unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + user("alice", "first", Duration::from_secs(2)), + user("bob", "second", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let (host, port) = sentinel.sentinel(); + let client = RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))), + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(1), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("before", "1", Duration::from_secs(60))) + .await + .unwrap(); + let master = sentinel.master_port().await.unwrap(); + let app_users = |connections: &[ClientConnection]| -> BTreeSet { + connections + .iter() + .filter(|connection| connection.user != "default") + .map(|connection| connection.user.clone()) + .collect() + }; + eventually(async || { + let connections = sentinel.clients_on(master).await.ok()?; + (app_users(&connections) == BTreeSet::from(["bob".to_owned()])).then_some(()) + }) + .await; + sentinel + .cli_on_data_nodes(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(60))) + .await + .ok() + }) + .await; + let on_promoted = sentinel.clients_on(promoted).await.unwrap(); + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!( + app_users(&on_promoted), + BTreeSet::from(["bob".to_owned()]), + "{on_promoted:?}" + ); + assert_eq!( + client + .send(commands::get::("after")) + .await + .unwrap() + .as_deref(), + Some("2") + ); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/aws.rs b/litellm-rust/crates/redis-connect/src/sources/aws.rs new file mode 100644 index 00000000000..1e87bdd1b12 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/aws.rs @@ -0,0 +1,188 @@ +use std::time::{Duration, SystemTime}; + +use litellm_auth_aws::{AwsAuthService, SigningScope, presign_url}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_types::AwsIamAuth; + +use crate::Error; +use crate::error::CredentialStage; +use litellm_secret_value::SecretValue; + +use crate::{CredentialSource, Expiring, SourceFuture}; + +const SERVICE: &str = "elasticache"; +const TOKEN_LIFETIME: Duration = Duration::from_secs(900); + +const SDK_IDENTITY_ENVIRONMENT: &[&str] = &[ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_SHARED_CREDENTIALS_FILE", + "AWS_CONFIG_FILE", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_ARN", + "AWS_ROLE_SESSION_NAME", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_AUTHORIZATION_TOKEN", + "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", + "AWS_EC2_METADATA_DISABLED", + "AWS_EC2_METADATA_SERVICE_ENDPOINT", + "AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE", + "AWS_ENDPOINT_URL", + "AWS_ENDPOINT_URL_STS", +]; + +pub(crate) struct ElastiCacheIam { + service: AwsAuthService, + auth: AwsIamAuth, +} + +impl ElastiCacheIam { + pub(crate) fn new(auth: AwsIamAuth, environment: &dyn Lookup) -> Result { + validate_environment(environment, &|name: &str| std::env::var(name).ok())?; + Ok(Self { + service: AwsAuthService::default(), + auth, + }) + } +} + +fn validate_environment(environment: &dyn Lookup, process: &dyn Lookup) -> Result<(), Error> { + match SDK_IDENTITY_ENVIRONMENT.iter().find(|name| { + environment.get(name).is_some_and(|value| { + Some(SecretValue::new(value)) != process.get(name).map(SecretValue::new) + }) + }) { + Some(setting) => Err(Error::ProcessEnvironmentRequired { setting }), + None => Ok(()), + } +} + +pub(crate) fn connect_url(auth: &AwsIamAuth) -> String { + let query = [ + Some(("Action", "connect")), + Some(("User", auth.user_name.as_str())), + auth.serverless + .then_some(("ResourceType", "ServerlessCache")), + ]; + let mut url = String::from("https://"); + url.push_str(&auth.cache_name.to_lowercase()); + url.push_str("/?"); + url.push_str( + &url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(query.into_iter().flatten()) + .finish(), + ); + url +} + +impl CredentialSource for ElastiCacheIam { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let credentials = self + .service + .default_chain_credentials(&self.auth.region) + .await + .map_err(Error::credential(CredentialStage::Identity))?; + let now = SystemTime::now(); + let signed = presign_url( + &connect_url(&self.auth), + TOKEN_LIFETIME, + SigningScope { + region: &self.auth.region, + service: SERVICE, + credentials: &credentials, + time: now, + }, + ) + .map_err(Error::credential(CredentialStage::Signing))?; + let token_expiry = now + TOKEN_LIFETIME; + Ok(Expiring { + username: Some(self.auth.user_name.clone()), + password: SecretValue::new(signed.trim_start_matches("https://")), + expires_at: Some( + credentials + .expiry() + .map_or(token_expiry, |expiry| expiry.min(token_expiry)), + ), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use litellm_redis_types::AwsIamAuth; + use rstest::rstest; + + use super::{connect_url, validate_environment}; + + #[rstest] + #[case::profile("AWS_PROFILE")] + #[case::role("AWS_ROLE_ARN")] + #[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")] + #[case::source_key("AWS_ACCESS_KEY_ID")] + #[case::source_secret("AWS_SECRET_ACCESS_KEY")] + #[case::session("AWS_SESSION_TOKEN")] + #[case::profile_path("AWS_SHARED_CREDENTIALS_FILE")] + #[case::container("AWS_CONTAINER_CREDENTIALS_FULL_URI")] + fn configured_identity_overrides_cannot_silently_use_an_ambient_identity( + #[case] setting: &'static str, + ) { + let configured = |name: &str| (name == setting).then(|| "credential-value".to_owned()); + let process = |name: &str| (name == setting).then(|| "ambient".to_owned()); + let error = validate_environment(&configured, &process).unwrap_err(); + + assert!(matches!( + error, + crate::Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting + )); + assert!(!error.to_string().contains("credential-value")); + assert!(!format!("{error:?}").contains("credential-value")); + } + + #[rstest] + #[case::process_environment("AWS_PROFILE", Some("profile"), Some("profile"))] + #[case::ambient_only("AWS_PROFILE", None, Some("profile"))] + #[case::unconfigured("AWS_PROFILE", None, None)] + #[case::unrelated_bedrock_setting("AWS_ROLE_NAME", Some("bedrock"), None)] + fn unchanged_or_unrelated_settings_preserve_the_sdk_chain( + #[case] setting: &str, + #[case] configured: Option<&str>, + #[case] ambient: Option<&str>, + ) { + let configured = |name: &str| { + (name == setting) + .then_some(configured) + .flatten() + .map(str::to_owned) + }; + let ambient = |name: &str| { + (name == setting) + .then_some(ambient) + .flatten() + .map(str::to_owned) + }; + assert!(validate_environment(&configured, &ambient).is_ok()); + } + + #[rstest] + #[case::provisioned(false, "https://cache-1/?Action=connect&User=iam+user%2F1")] + #[case::serverless( + true, + "https://cache-1/?Action=connect&User=iam+user%2F1&ResourceType=ServerlessCache" + )] + fn the_connect_url_matches_python_urlencode(#[case] serverless: bool, #[case] expected: &str) { + let auth = AwsIamAuth { + user_name: "iam user/1".to_owned(), + cache_name: "Cache-1".to_owned(), + region: "us-east-1".to_owned(), + serverless, + }; + + assert_eq!(connect_url(&auth), expected); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/azure.rs b/litellm-rust/crates/redis-connect/src/sources/azure.rs new file mode 100644 index 00000000000..6dc72e6ce72 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/azure.rs @@ -0,0 +1,64 @@ +use litellm_auth_azure::{AzureAuthService, AzureIdentity}; +use litellm_auth_types::ResolvedCredential; +use litellm_redis_types::AzureAdAuth; + +use crate::Error; +use crate::error::CredentialStage; + +use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture}; + +const AZURE_REDIS_SCOPE: &str = "https://redis.azure.com/.default"; + +pub(crate) struct AzureAd { + service: AzureAuthService, + identity: AzureIdentity, + username: Option, + lookup: SharedEnvironment, +} + +impl AzureAd { + pub(crate) fn new(auth: AzureAdAuth, lookup: SharedEnvironment) -> Self { + // _build_azure_credential in litellm/_redis.py picks the credential class the same way + let identity = match (auth.client_id, auth.tenant_id, auth.client_secret) { + (Some(client_id), Some(tenant_id), Some(client_secret)) => { + AzureIdentity::ClientSecret { + tenant_id, + client_id, + client_secret, + } + } + (Some(client_id), _, _) => AzureIdentity::ManagedIdentity { client_id }, + (None, _, _) => AzureIdentity::Default, + }; + Self { + service: AzureAuthService::default(), + identity, + username: auth.username, + lookup, + } + } +} + +impl CredentialSource for AzureAd { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let credential = self + .service + .identity_token(&self.identity, AZURE_REDIS_SCOPE, &|name| { + self.lookup.get(name) + }) + .await + .map_err(Error::credential(CredentialStage::Identity))? + .ok_or(Error::MissingCredential { provider: "Azure" })?; + let (token, expires_at) = match credential { + ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on), + ResolvedCredential::Static(token) => (token, None), + }; + Ok(Expiring { + username: self.username.clone(), + password: token, + expires_at, + }) + }) + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/gcp.rs b/litellm-rust/crates/redis-connect/src/sources/gcp.rs new file mode 100644 index 00000000000..9fb238709a1 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/gcp.rs @@ -0,0 +1,152 @@ +use std::time::SystemTime; + +use litellm_auth_gcp::{CredentialSource as GcpCredentialSource, GcpAuth}; +use serde::Deserialize; + +use crate::Error; +use crate::error::CredentialStage; +use litellm_secret_value::SecretValue; + +use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture}; + +const IAM_CREDENTIALS: &str = "https://iamcredentials.googleapis.com/v1"; +const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +pub(crate) struct GcpIam { + http: litellm_http::Client, + identity: GcpAuth, + service_account: String, + lookup: SharedEnvironment, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct AccessToken { + access_token: String, + expire_time: String, +} + +impl GcpIam { + pub(crate) fn new( + http: litellm_http::Client, + service_account: String, + lookup: SharedEnvironment, + ) -> Self { + Self { + http, + identity: GcpAuth::default(), + service_account, + lookup, + } + } +} + +fn token_url(service_account: &str) -> String { + match service_account.starts_with("projects/") { + true => format!("{IAM_CREDENTIALS}/{service_account}:generateAccessToken"), + false => format!( + "{IAM_CREDENTIALS}/projects/-/serviceAccounts/{service_account}:generateAccessToken" + ), + } +} + +impl CredentialSource for GcpIam { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let lookup = |name: &str| self.lookup.get(name); + let caller = self + .identity + .access_token(GcpCredentialSource::application_default(&lookup)) + .await + .map_err(Error::credential(CredentialStage::Identity))?; + let response = self + .http + .post(token_url(&self.service_account)) + .bearer_auth(caller.expose()) + .json(&serde_json::json!({ "scope": [CLOUD_PLATFORM_SCOPE] })) + .send() + .await + .map_err(Error::credential(CredentialStage::Exchange))?; + let status = response.status().as_u16(); + let reply: AccessToken = response + .error_for_status() + .map_err(litellm_http::Error::from) + .map_err(Error::credential_response( + CredentialStage::Exchange, + status, + ))? + .json() + .await + .map_err(litellm_http::Error::from) + .map_err(Error::credential(CredentialStage::Response))?; + let expires_at = time::OffsetDateTime::parse( + &reply.expire_time, + &time::format_description::well_known::Rfc3339, + ) + .map(SystemTime::from) + .map_err(Error::credential(CredentialStage::Response))?; + Ok(Expiring { + username: None, + password: SecretValue::new(reply.access_token), + expires_at: Some(expires_at), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::{GcpIam, token_url}; + use crate::CredentialSource; + + #[rstest] + #[case::bare_email( + "redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + #[case::resource_name( + "projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + #[case::project_scoped_resource_name( + "projects/project/serviceAccounts/redis@project.iam.gserviceaccount.com", + "https://iamcredentials.googleapis.com/v1/projects/project/serviceAccounts/redis@project.iam.gserviceaccount.com:generateAccessToken" + )] + fn a_bare_service_account_email_is_addressed_as_a_resource_name( + #[case] service_account: &str, + #[case] expected: &str, + ) { + assert_eq!(token_url(service_account), expected); + } + + #[rstest] + #[tokio::test] + async fn the_caller_identity_comes_from_application_default_credentials_not_vertex_settings() { + let environment = |name: &str| { + match name { + "VERTEXAI_CREDENTIALS" => { + Some(r#"{"type":"service_account","project_id":"vertex"}"#) + } + "GOOGLE_APPLICATION_CREDENTIALS" => Some("/nonexistent/redis-adc.json"), + _ => None, + } + .map(str::to_owned) + }; + let source = GcpIam::new( + litellm_http::Client::for_test(), + "projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com".to_owned(), + std::sync::Arc::new(environment), + ); + + let Err(error) = source.fetch().await else { + panic!("a missing credentials file must fail"); + }; + + assert_eq!( + litellm_error::find::(&error).map(std::io::Error::kind), + Some(std::io::ErrorKind::NotFound) + ); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/mod.rs b/litellm-rust/crates/redis-connect/src/sources/mod.rs new file mode 100644 index 00000000000..988ad3e1d23 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/mod.rs @@ -0,0 +1,7 @@ +mod aws; +mod azure; +mod gcp; + +pub(crate) use aws::ElastiCacheIam; +pub(crate) use azure::AzureAd; +pub(crate) use gcp::GcpIam; diff --git a/litellm-rust/crates/redis-connect/tests/connect.rs b/litellm-rust/crates/redis-connect/tests/connect.rs new file mode 100644 index 00000000000..b1d541f4a5e --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/connect.rs @@ -0,0 +1,230 @@ +use litellm_redis_client::commands; +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use litellm_redis_connect::{ + CredentialSource, Error, Expiring, HostServices, SourceFuture, connect, +}; +use litellm_redis_testing::{ + RedisDeployment, RedisSentinel, RedisServer, RedisTlsServer, SERVICE_NAME, Topology, eventually, +}; +use litellm_redis_types::{ProviderHandle, Resolution}; +use litellm_secret_value::SecretValue; +use rstest::rstest; + +// a host hands its credential_provider object over as a handle, so these tests do the same +fn resolved(kwargs: serde_json::Value) -> Resolution { + let provider = kwargs + .get("credential_provider") + .map(|_| ProviderHandle::new(1)); + let kwargs = serde_json::from_value::>(kwargs).unwrap(); + litellm_redis_resolve::interpret(&litellm_redis_resolve::kwargs(kwargs), provider, None) + .unwrap() +} + +fn services() -> HostServices { + HostServices { + events: None, + http: litellm_http::Client::for_test(), + environment: Arc::new(|_: &str| None), + credential_provider: None, + } +} + +struct HostPassword(&'static str); + +impl CredentialSource for HostPassword { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + Ok(Expiring { + username: None, + password: SecretValue::new(self.0), + expires_at: None, + }) + }) + } +} + +fn port(url: &str) -> u16 { + url.rsplit(':').next().unwrap().parse().unwrap() +} + +#[rstest] +#[case::trusted_ca(true, true)] +#[case::system_roots_only(false, false)] +#[tokio::test] +async fn python_tls_kwargs_reach_the_server(#[case] with_ca: bool, #[case] connects: bool) { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + let server = RedisTlsServer::start(false).await.unwrap(); + let kwargs = serde_json::json!({ + "host": server.host(), + "port": server.port(), + "ssl": true, + "ssl_ca_certs": with_ca.then(|| server.files().ca.display().to_string()), + "socket_timeout": 2, + }); + + let pinged = connect(resolved(kwargs), services(), None) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} + +#[rstest] +#[case::right_password("right", true)] +#[case::wrong_password("wrong", false)] +#[tokio::test] +async fn static_credentials_authenticate( + #[case] password: &str, + #[case] connects: bool, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let kwargs = match &server { + RedisDeployment::Standalone(server) => serde_json::json!({ + "host": "127.0.0.1", + "port": port(server.url()), + "password": password, + "socket_timeout": 2, + }), + RedisDeployment::Cluster(cluster) => serde_json::json!({ + "startup_nodes": cluster.nodes().map(|(host, port)| serde_json::json!({ + "host": host, + "port": port, + })).collect::>(), + "password": password, + "socket_timeout": 2, + }), + }; + + let pinged = connect(resolved(kwargs), services(), None) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} + +#[rstest] +#[tokio::test] +async fn python_sentinel_kwargs_follow_the_master_through_a_failover() { + let sentinel = RedisSentinel::start().await.unwrap(); + let (host, sentinel_port) = sentinel.sentinel(); + let kwargs = serde_json::json!({ + "sentinel_nodes": [[host, sentinel_port]], + "service_name": SERVICE_NAME, + "socket_timeout": 1, + }); + let client = connect(resolved(kwargs), services(), None).unwrap(); + client + .send(commands::set("before", "1", Duration::from_secs(600))) + .await + .unwrap(); + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(600))) + .await + .ok() + }) + .await; + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!(sentinel.master_port().await.unwrap(), promoted); + assert_eq!( + client + .send(commands::get::("after")) + .await + .unwrap() + .as_deref(), + Some("2") + ); +} + +#[rstest] +fn a_python_credential_provider_without_a_host_source_fails_at_boot() { + let error = connect( + resolved(serde_json::json!({"host": "redis", "credential_provider": true})), + services(), + None, + ) + .err() + .unwrap(); + + assert!( + matches!(error, Error::NoHostCredentialProvider), + "{error:?}" + ); +} + +#[rstest] +#[case::profile("AWS_PROFILE")] +#[case::role("AWS_ROLE_ARN")] +#[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")] +fn config_only_aws_identities_fail_before_connecting(#[case] setting: &'static str) { + let configured = format!("{}-override", std::env::var(setting).unwrap_or_default()); + let host = HostServices { + environment: Arc::new(move |name: &str| (name == setting).then(|| configured.clone())), + ..services() + }; + let error = connect( + resolved(serde_json::json!({ + "host": "redis", "ssl": true, "aws_iam_auth": true, + "aws_iam_user_name": "user", "aws_iam_cache_name": "cache", + "aws_iam_region": "us-east-1", + })), + host, + None, + ) + .err() + .expect("must not select ambient credentials"); + + assert!(matches!( + error, + Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting + )); +} + +#[rstest] +#[case::right_password("right", true)] +#[case::wrong_password("wrong", false)] +#[tokio::test] +async fn a_python_credential_provider_authenticates_through_the_host_source( + #[case] password: &'static str, + #[case] connects: bool, +) { + let server = RedisServer::start().await.unwrap(); + server + .cli(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let kwargs = serde_json::json!({ + "host": "127.0.0.1", + "port": port(server.url()), + "password": "ignored", + "credential_provider": true, + "socket_timeout": 2, + }); + + let pinged = connect( + resolved(kwargs), + HostServices { + credential_provider: Some(Arc::new(HostPassword(password))), + ..services() + }, + None, + ) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} diff --git a/litellm-rust/crates/redis-connect/tests/live/aws.rs b/litellm-rust/crates/redis-connect/tests/live/aws.rs new file mode 100644 index 00000000000..da5ec323392 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/aws.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_AWS_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/azure.rs b/litellm-rust/crates/redis-connect/tests/live/azure.rs new file mode 100644 index 00000000000..a7f678edef4 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/azure.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_AZURE_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/gcp.rs b/litellm-rust/crates/redis-connect/tests/live/gcp.rs new file mode 100644 index 00000000000..0965a7a9e9a --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/gcp.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_GCP_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/main.rs b/litellm-rust/crates/redis-connect/tests/live/main.rs new file mode 100644 index 00000000000..90e7637530e --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/main.rs @@ -0,0 +1,111 @@ +mod aws; +mod azure; +mod gcp; + +use std::{collections::BTreeMap, path::PathBuf, sync::Arc, time::Duration}; + +use litellm_http::{ + ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, + media::PublicDnsResolver, +}; +use litellm_redis_client::commands; +use litellm_redis_connect::{HostServices, connect}; +use serde::Deserialize; + +#[derive(Deserialize)] +struct LiveConfig { + redis: BTreeMap, + hold_seconds: u64, +} + +fn config(variable: &str) -> Option { + let path = PathBuf::from(std::env::var_os(variable)?); + Some( + serde_json::from_slice(&std::fs::read(path).expect("read live Redis config")) + .expect("parse live Redis config"), + ) +} + +async fn roundtrip(config: LiveConfig) { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + let lookup = |name: &str| std::env::var(name).ok(); + let resolved = litellm_redis_resolve::interpret( + &litellm_redis_resolve::kwargs(config.redis), + None, + lookup("REDIS_USERNAME").as_deref(), + ) + .expect("resolve live Redis config"); + let http_settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&lookup)]); + let pool = HttpClientPool::new(Arc::new(PublicDnsResolver)); + let http = pool + .client( + &Resolution::from(&http_settings).config, + ClientVariant::Provider, + ) + .expect("build production HTTP client"); + let services = || HostServices { + events: None, + http: http.clone(), + environment: Arc::new(|name: &str| std::env::var(name).ok()), + credential_provider: None, + }; + let client = connect(resolved.clone(), services(), None).expect("build live Redis client"); + let marker = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let key = format!("litellm:redis-connect:live:{}:{marker}", std::process::id()); + let ttl = Duration::from_secs(config.hold_seconds + 120); + + client.ping().await.expect("authenticate and ping"); + client + .send(commands::set(&key, "before-refresh", ttl)) + .await + .expect("set before refresh"); + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + + tokio::time::sleep(Duration::from_secs(config.hold_seconds)).await; + + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + let fresh = connect(resolved, services(), None).expect("build a fresh live Redis client"); + assert_eq!( + fresh + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + fresh + .send(commands::set(&key, "after-refresh", ttl)) + .await + .expect("set after refresh"); + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("after-refresh") + ); + client + .send(commands::del(&[&key])) + .await + .expect("delete live test key"); +} diff --git a/litellm-rust/crates/redis-resolve/AGENTS.md b/litellm-rust/crates/redis-resolve/AGENTS.md new file mode 100644 index 00000000000..b7ce8278975 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/AGENTS.md @@ -0,0 +1,31 @@ +## Interpreting merged kwargs + +- The boundary is merged redis-py kwargs in, `litellm_redis_types::Resolution` out. `interpret` takes the dict a host already merged the way `_merge_redis_kwargs` in `litellm/_redis.py` does, with its `REDIS_*` values and `get_secret` fallbacks filled, and stands in for redis-py's constructors and `from_url`. It never reads the environment or a secret, so both hosts share it. Every `Entry` keeps the source a warning or error names, which `kwargs` sets to the kwarg name and the gateway merge sets to `REDIS_` or the fallback name +- The `Resolution` holds the target, the auth and the timeouts Rust models, and every setting the Rust client ignores as an `Ignored` with its source and consequence, which the host logs. A host passes its `credential_provider` object as a `ProviderHandle`, and a `credential_provider` key without one is `Error::ProviderWithoutHandle`. Python's Azure AD auth reads `REDIS_USERNAME` from `os.environ` instead of the merged kwargs, so it is a required argument of `interpret` that each host reads from its own process environment, the gateway through its `Lookup` and python-bridge through Python's `os.environ`. An empty one is unset, like `or None` +- `options.rs` reads the kwargs into the typed `Options` and collects the ignored ones. `endpoint.rs`, `auth.rs` and `tls.rs` pick the target, the credentials and TLS from `Options` alone +- Rust never introspects redis-py. Every Redis kwarg redis-py or `litellm/_redis.py` accepts is one row in `SETTINGS` in `src/keys.rs`, and that row is the whole answer for the kwarg: its `Handling` (`Read`, `Warn` with the consequence the warning names, or `Reject`), whether a url query carries it (`FromUrl`), whether `_get_redis_env_kwarg_mapping` reads it as `REDIS_`, and the names `_merge_redis_kwargs` falls back to. The last two columns are only read by the gateway merge, and they stay on the row so there is one table. A `Read` row is a field of `Options`, which an inline test pins. A kwarg with no row is not a Redis setting, so it is dropped without a warning +- `redis-py-kwargs.json` holds the Python facts each row answers: which redis-py versions name the kwarg in which signature, and whether Python reads its `REDIS_*` name under the locked redis-py. Never edit it by hand. `redis_py_kwargs.py` next to it regenerates it, and its docstring says how to run it. A kwarg litellm reads that no signature or `REDIS_*` name covers goes in its `LITELLM_ONLY`, and a new redis-py version in its `VERSIONS` and the `redis-compat` CI matrix together. `test_redis_py_kwargs_catalog_matches_the_installed_redis_py` in `tests/unit/test_redis.py` checks it against each version the `redis-compat` CI matrix installs, and `tests/catalog.rs` fails until `SETTINGS` has exactly one row per kwarg in it with the same `environment`. A redis-py bump that adds a kwarg therefore cannot land without a row deciding how Rust handles it +- The contract is `resolves_like_python_redis` in `tests/interpret.rs` and its mirror `test_redis_settings_resolve_like_the_rust_resolver_expects` in `tests/unit/test_redis.py`. Each case holds the merged dict, the sources the merge gave it, the `REDIS_USERNAME` the host passes, and the expected outcome written next to it. `Resolves` is the view both sides build. `Refuses` is a config Rust refuses and Python connects with, a disagreement on purpose. `Fails` is a config redis-py 5.3.1 already raises on, building the client or its first connection, so Rust failing changes no deployment that works today. Its Rust side names the class `Error::raised_building_the_python_client` answers, `None` when Python only raises once a command connects, and its Python mirror `_Fails` asserts the class and message redis-py raises. The Python side asserts that `_merge_redis_kwargs` returns that merged dict and reads the effective redis-py 5.3.1 connection objects without connecting, and `merges_like_python_redis` in `gateway/src/redis/surface_tests.rs` pins the same merged dicts from the Rust merge. A new rule is a new case in every file, never a generated fixture. Python and Rust pods share one Redis during the switch, so a change here is a change to a wire contract +- A url carries its own TLS. redis-py's url allowlist is the plain `Connection` signature, so every `ssl_*` setting outside the url is dropped, and a credential provider strips the url's userinfo. A `rediss://` url's `ssl_cert_reqs`, `ssl_ca_certs`, `ssl_ca_data`, `ssl_ca_path`, `ssl_certfile`, `ssl_keyfile`, `ssl_password` and `ssl_check_hostname` query parameters are read like `parse_url` reads them +- TLS options and `db` are only read for the endpoint that uses them, so a bad `ssl_cert_reqs` never fails a url target. A cluster with a non-zero `db`, and Sentinel with a `host` or `port`, are errors because redis-py raises building those clients +- An empty value (`""`, `[]`, `{}`) is the same as an unset one, like Python's `x or ...`. A null is unset too, except `ssl_cert_reqs`, where redis-py reads `None` as `CERT_NONE`, so a null skips certificate verification +- `ssl_cert_reqs` takes `none`, `optional` and `required` case-sensitively like redis-py's lookup, so `None` (what `REDIS_SSL_CERT_REQS=None` gives) is `Error::Invalid`, and an `ssl.VerifyMode` int, where 0 skips verification and 1 or 2 verifies. `optional` verifies, because rustls has no optional mode +- Precedence matches Python. The target is cluster, then URL, then Sentinel (`sentinel_nodes` plus `service_name`), then host. Auth is a `credential_provider` handle, then GCP, then Azure, then AWS, and the first one wins with a warning +- Numbers are read with `litellm_python_compat::number`, so whitespace and digit underscores parse like Python's `int()` and `float()`. A value they cannot parse (`port: abc`, `socket_timeout: slow`) is dropped with a warning and the default is used, like `_coerce_redis_kwargs_types`. A number that parses but is out of range (negative seconds, a port above 65535) fails with `Error::Invalid` +- Flags (`ssl`, `ssl_check_hostname`, `aws_iam_auth`, `aws_iam_serverless`) are read one way on every surface. `true`, `1` and `yes` are true and `false`, `0`, `no`, `off`, `f` and `n` are false, the text redis-py's `_str_to_bool` and the proxy's pydantic models agree on. Any other text is `Error::Invalid`, so a TLS setting is never guessed +- The resolver fails on a setting it cannot honor only when that setting is a security check. A `Warn` row is dropped and returned as an `Ignored` naming its source (`ssl_ciphers` or `REDIS_SSL_CIPHERS`) and the consequence on its row, because a Python deployment that worked must keep working when its stores move to Rust. Only a value that is malformed (`Error::Invalid` or `Error::UrlQuery`), a target that is missing or a `Reject` row that asks for its check fails. Promoting a `Warn` row to `Read` is then a pure improvement with no behavior cliff, while moving one to `Reject` can fail a config that connects today + - A `Warn` row with a `met` check is dropped silently when its value asks for nothing more than Rust does, so `ssl_min_version` of TLS 1.2 or lower never warns + - A `Reject` row is a security check rustls cannot perform that Python's locked async client already refuses with a `TypeError`, so failing closed breaks no proxy that runs today. `tests/catalog.rs` fails if the locked redis-py's async, cluster or connection signature accepts a `Reject` row, so a redis-py bump forces the decision again. A row's `Asked` says which values ask for the check, from the kwargs, `REDIS_*` or a url query, and those are `Error::Unenforceable` + - `ssl_validate_ocsp` and `ssl_validate_ocsp_stapled` ask when true, because rustls has no client-side OCSP. False asks for nothing and is dropped without a warning + - `ssl_include_verify_flags` and `ssl_exclude_verify_flags` ask when set, because rustls takes no X.509 verify flags such as CRL checking. An empty list asks for nothing + - `ssl_ca_data` and `ssl_ca_path` are read into `TlsOptions` and trusted next to `ssl_ca_certs`. `ssl_password` is read into `TlsOptions::key_password`, which `redis-client` uses to decrypt a PKCS#8 key like the sync client's `load_cert_chain(password=...)`. redis-py 5.3.1's async client refuses it, so this is a superset of the path the proxy uses, like `ssl_ca_path` + - `unix_socket_path` resolves to `Target::Unix`, which like redis-py's `UnixDomainSocketConnection` replaces host, port and TLS. Python still requires `host` next to it, so Rust does too + - A url's `?db=`, `?username=`, `?password=`, `?socket_timeout=`, `?socket_connect_timeout=` and `?health_check_interval=` are read like `parse_url` reads them, so `?db=` wins over the path, the userinfo wins over `?username=` and `?password=`, and the timeouts and health check win over the kwargs like `from_url`. A repeated parameter takes its first value like `parse_qs`. A number `int()` or `float()` cannot read is `Error::UrlQuery`, as `parse_url` raises `ValueError`. Any other query parameter, including one with no row, is ignored naming `url ?`, because redis-py turns each one into a client kwarg. Which ones a url carries is the `FromUrl` column + - `socket_keepalive` is a flag that defaults to true on every target, like `_get_redis_cluster_kwargs` defaults it for a cluster. Python leaves it off for a standalone or Sentinel client, so an unset value there is a difference on purpose, because without keepalive a load balancer drops an idle connection silently. False turns it off. `socket_keepalive_options` is warned, and `redis-client` probes after its own idle time instead + - `health_check_interval` is read as an `int` of seconds like `_coerce_redis_kwargs_types` into `HealthCheck`. Zero is `Disabled` and stays off on every target, a negative interval checks before every command like redis-py, and only `Unset` takes a default. `redis-connect` gives an unset cluster Python's 25s default and `redis-client` pings every node before the next command once the client sat idle that long, outside the breaker and the call events +- `Error::raised_building_the_python_client` is how python-bridge treats a config Rust cannot use. Where Python raises building its client, the bridge raises the same class from `NativeRedisClient.connect`. Where Python builds the client and only fails once a command connects, which `RedisCache` catches and logs on every call, the stores stay on memory and the bridge logs an error +- Cases where Rust and Python disagree on purpose are cases in both contract tests with a different expected outcome on each side + - `ssl_ca_path`: redis-py 5.3.1's async `Redis` has no such parameter and raises `TypeError` building the client, while Rust trusts the directory like `load_verify_locations(capath=...)` + - `on`, `t` and `y` for a flag are `Error::Invalid` on every surface. pydantic reads them as true for `coordination_redis` and `_str_to_bool` reads them as false everywhere else, so Rust refuses to pick one. Other unrecognized text is false in Python and an error here + - `ssl_cert_reqs` as an `ssl.VerifyMode` int: redis-py 5.3.1's sync client takes it, while its async `RedisSSLContext` leaves `cert_reqs` unset and raises `AttributeError` once a command connects. Rust reads it like the sync client + - Empty `sentinel_nodes` with no `REDIS_SENTINEL_NODES` is no target at all, so it fails as "either host or url must be specified". redis-py keeps the empty list and fails building the Sentinel client diff --git a/litellm-rust/crates/redis-resolve/Cargo.toml b/litellm-rust/crates/redis-resolve/Cargo.toml new file mode 100644 index 00000000000..b3d3fc01e44 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-redis-resolve" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true +litellm-secret-value.workspace = true +serde_json.workspace = true +strum.workspace = true +thiserror.workspace = true +tracing.workspace = true +url.workspace = true + +[dev-dependencies] +rstest.workspace = true +serde.workspace = true diff --git a/litellm-rust/crates/redis-resolve/redis-py-kwargs.json b/litellm-rust/crates/redis-resolve/redis-py-kwargs.json new file mode 100644 index 00000000000..22c4495818b --- /dev/null +++ b/litellm-rust/crates/redis-resolve/redis-py-kwargs.json @@ -0,0 +1,118 @@ +{ + "meta": { + "description": "What redis-py and litellm/_redis.py accept for every Redis kwarg the Rust resolver can be handed. SETTINGS in src/keys.rs has exactly one row per arg here", + "redis_py_versions": [ + "5.3.1", + "6.4.0", + "7.4.1", + "8.0.1" + ], + "python": "3.12", + "method": "python redis_py_kwargs.py next to this file, see its docstring. It reads inspect.getfullargspec(inspect.unwrap(cls.__init__)) args and kwonlyargs, minus self", + "signatures": { + "sync": "redis.Redis.__init__. _get_redis_kwargs uses it as the config allowlist and the REDIS_ environment mapping", + "async": "redis.asyncio.Redis.__init__, which the proxy constructs. A kwarg in sync but not async raises TypeError when the client is built", + "connection": "redis.Connection or redis.asyncio.Connection MRO, the url allowlist from _get_redis_url_kwargs", + "cluster": "redis.asyncio.RedisCluster.__init__, the allowlist from _get_redis_cluster_kwargs" + }, + "environment": "true when _get_redis_env_kwarg_mapping reads the arg as REDIS_ under the locked redis-py", + "redis_py": "each redis-py version whose signatures name the arg, with which ones. An empty object means only litellm reads it", + "checked_by": "tests/unit/test_redis.py checks the installed redis-py's column and the REDIS_* names against this file, and tests/catalog.rs checks SETTINGS against it", + "redis_py_locked": "5.3.1" + }, + "args": { + "address_remap": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "auto_close_connection_pool": {"environment": false, "redis_py": {"5.3.1": ["async"], "6.4.0": ["async"], "7.4.1": ["async"], "8.0.1": ["async"]}}, + "aws_iam_auth": {"environment": true, "redis_py": {}}, + "aws_iam_cache_name": {"environment": true, "redis_py": {}}, + "aws_iam_region": {"environment": true, "redis_py": {}}, + "aws_iam_serverless": {"environment": true, "redis_py": {}}, + "aws_iam_user_name": {"environment": true, "redis_py": {}}, + "azure_client_id": {"environment": true, "redis_py": {}}, + "azure_client_secret": {"environment": true, "redis_py": {}}, + "azure_redis_ad_token": {"environment": true, "redis_py": {}}, + "azure_tenant_id": {"environment": true, "redis_py": {}}, + "cache": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "cache_config": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "charset": {"environment": true, "redis_py": {"5.3.1": ["sync"]}}, + "client_name": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "cluster_error_retry_attempts": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "command_packer": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "connection_class": {"environment": false, "redis_py": {}}, + "connection_error_retry_attempts": {"environment": false, "redis_py": {"5.3.1": ["cluster"]}}, + "connection_pool": {"environment": false, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "credential_provider": {"environment": false, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "db": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "decode_responses": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "driver_info": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "dynamic_startup_nodes": {"environment": false, "redis_py": {"6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "encoder_class": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "encoding": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "encoding_errors": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "errors": {"environment": true, "redis_py": {"5.3.1": ["sync"]}}, + "event_dispatcher": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "gcp_service_account": {"environment": true, "redis_py": {}}, + "gcp_ssl_ca_certs": {"environment": true, "redis_py": {}}, + "health_check_interval": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "host": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "legacy_responses": {"environment": false, "redis_py": {"8.0.1": ["sync", "async", "connection", "cluster"]}}, + "lib_name": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "lib_version": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "load_balancing_strategy": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "maint_notifications_config": {"environment": false, "redis_py": {"7.4.1": ["sync", "connection"], "8.0.1": ["sync", "connection"]}}, + "maint_notifications_pool_handler": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "maintenance_notification_hash": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "maintenance_state": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "max_connections": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "namespace": {"environment": false, "redis_py": {}}, + "orig_host_address": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "orig_socket_connect_timeout": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "orig_socket_timeout": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "oss_cluster_maint_notifications_handler": {"environment": false, "redis_py": {"7.4.1": ["sync", "connection"], "8.0.1": ["sync", "connection"]}}, + "parser": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "parser_class": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "password": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "path": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "policy_resolver": {"environment": false, "redis_py": {"7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "port": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "protocol": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "read_from_replicas": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "redis_connect_func": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection"], "6.4.0": ["sync", "async", "connection"], "7.4.1": ["sync", "async", "connection"], "8.0.1": ["sync", "async", "connection"]}}, + "reinitialize_steps": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "require_full_coverage": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "retry": {"environment": false, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "retry_on_error": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "retry_on_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection"], "6.4.0": ["sync", "async", "connection"], "7.4.1": ["sync", "async", "connection"], "8.0.1": ["sync", "async", "connection"]}}, + "sentinel_nodes": {"environment": false, "redis_py": {}}, + "sentinel_password": {"environment": false, "redis_py": {}}, + "service_name": {"environment": false, "redis_py": {}}, + "single_connection_client": {"environment": true, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "socket_connect_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_keepalive": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_keepalive_options": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_read_size": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_type": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "ssl": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_certs": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_data": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_path": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "ssl_cert_reqs": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_certfile": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_check_hostname": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ciphers": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_exclude_verify_flags": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_include_verify_flags": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_keyfile": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_min_version": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ocsp_context": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_ocsp_expected_cert": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_password": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "ssl_validate_ocsp": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_validate_ocsp_stapled": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "startup_nodes": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "unix_socket_path": {"environment": true, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "url": {"environment": true, "redis_py": {}}, + "username": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}} + } +} diff --git a/litellm-rust/crates/redis-resolve/redis_py_kwargs.py b/litellm-rust/crates/redis-resolve/redis_py_kwargs.py new file mode 100644 index 00000000000..9459cfd3761 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/redis_py_kwargs.py @@ -0,0 +1,110 @@ +"""Regenerates redis-py-kwargs.json, the Python facts SETTINGS in src/keys.rs answers. + +Run it from the repository root in the project environment, which has the locked redis-py and litellm: + + uv run python litellm-rust/crates/redis-resolve/redis_py_kwargs.py + +It reads every other redis-py version's signatures in a throwaway environment through +`uv run --no-project --with redis==`, so nothing else is installed. Then update SETTINGS until +tests/catalog.rs passes. +""" + +import inspect +import json +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +VERSIONS: Final = ("5.3.1", "6.4.0", "7.4.1", "8.0.1") +PYTHON: Final = "3.12" +CATALOG: Final = Path(__file__).with_name("redis-py-kwargs.json") +# Keys litellm/_redis.py or its callers read that neither a redis-py signature nor the REDIS_* mapping names +LITELLM_ONLY: Final = frozenset( + {"connection_class", "namespace", "sentinel_nodes", "sentinel_password", "service_name"} +) + +META: Final = { + "description": "What redis-py and litellm/_redis.py accept for every Redis kwarg the Rust resolver can be handed. SETTINGS in src/keys.rs has exactly one row per arg here", + "redis_py_versions": list(VERSIONS), + "python": PYTHON, + "method": "python redis_py_kwargs.py next to this file, see its docstring. It reads inspect.getfullargspec(inspect.unwrap(cls.__init__)) args and kwonlyargs, minus self", + "signatures": { + "sync": "redis.Redis.__init__. _get_redis_kwargs uses it as the config allowlist and the REDIS_ environment mapping", + "async": "redis.asyncio.Redis.__init__, which the proxy constructs. A kwarg in sync but not async raises TypeError when the client is built", + "connection": "redis.Connection or redis.asyncio.Connection MRO, the url allowlist from _get_redis_url_kwargs", + "cluster": "redis.asyncio.RedisCluster.__init__, the allowlist from _get_redis_cluster_kwargs", + }, + "environment": "true when _get_redis_env_kwarg_mapping reads the arg as REDIS_ under the locked redis-py", + "redis_py": "each redis-py version whose signatures name the arg, with which ones. An empty object means only litellm reads it", + "checked_by": "tests/unit/test_redis.py checks the installed redis-py's column and the REDIS_* names against this file, and tests/catalog.rs checks SETTINGS against it", +} + + +def _named_args(*classes: type) -> frozenset[str]: + return frozenset( + name + for cls in classes + for spec in (inspect.getfullargspec(inspect.unwrap(cls.__init__)),) + for name in spec.args + spec.kwonlyargs + if name != "self" + ) + + +def arg_signatures() -> dict[str, list[str]]: + """Each kwarg the installed redis-py names, with the signatures that name it.""" + import redis + import redis.asyncio as async_redis + + connection_classes: Final = (*redis.Connection.__mro__, *async_redis.Connection.__mro__) + signatures: Final = { + "sync": _named_args(redis.Redis), + "async": _named_args(async_redis.Redis), + "connection": _named_args(*(cls for cls in connection_classes if cls is not object)), + "cluster": _named_args(async_redis.RedisCluster), + } + names: Final = sorted(frozenset().union(*signatures.values())) + return {name: [kind for kind, named in signatures.items() if name in named] for name in names} + + +def _arg_signatures_of(version: str) -> Mapping[str, list[str]]: + command: Final = ("uv", "run", "--no-project", "--python", PYTHON, "--with", f"redis=={version}") + script: Final = ("python", "-I", __file__, "--arg-signatures") + return json.loads(subprocess.run((*command, *script), check=True, capture_output=True, text=True).stdout) + + +def catalog() -> dict[str, object]: + import redis + + from litellm._redis import _get_redis_env_kwarg_mapping + + environment: Final = frozenset(_get_redis_env_kwarg_mapping().values()) + by_version: Final = {version: _arg_signatures_of(version) for version in VERSIONS} + names: Final = sorted(LITELLM_ONLY | environment | frozenset().union(*by_version.values())) + return { + "meta": {**META, "redis_py_locked": redis.__version__}, + "args": { + name: { + "environment": name in environment, + "redis_py": {version: args[name] for version, args in by_version.items() if name in args}, + } + for name in names + }, + } + + +# one line per arg, so a version bump reads as a diff of the args it touched +def render(catalog: Mapping[str, object]) -> str: + meta: Final = json.dumps(catalog["meta"], indent=2).replace("\n", "\n ") + args: Final = catalog["args"] + assert isinstance(args, dict) + rows: Final = ",\n".join(f" {json.dumps(name)}: {json.dumps(arg)}" for name, arg in args.items()) + return f'{{\n "meta": {meta},\n "args": {{\n{rows}\n }}\n}}\n' + + +if __name__ == "__main__": + if sys.argv[1:] == ["--arg-signatures"]: + sys.stdout.write(json.dumps(arg_signatures())) + else: + CATALOG.write_text(render(catalog())) diff --git a/litellm-rust/crates/redis-resolve/src/auth.rs b/litellm-rust/crates/redis-resolve/src/auth.rs new file mode 100644 index 00000000000..620f1a2e11b --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/auth.rs @@ -0,0 +1,95 @@ +use litellm_redis_types::{Auth, AwsIamAuth, AzureAdAuth, ProviderHandle, StaticAuth}; + +use crate::{ + Error, + options::{self, AwsOptions, Options, flag, secret, text}, +}; + +pub(super) fn redis_auth( + options: &Options, + provider: Option, + redis_username: Option<&str>, +) -> Result { + match (provider, &options.credential_provider) { + (Some(handle), _) => return Ok(Auth::Provider(handle)), + (None, Some(_)) => return Err(Error::ProviderWithoutHandle), + (None, None) => {} + } + let gcp = text(options.gcp.service_account.as_ref()); + let azure = text(options.azure.ad_token.as_ref()) + .is_some_and(|enabled| enabled.eq_ignore_ascii_case("true")); + let aws = flag(options.aws.auth.as_ref())?; + match (gcp, azure, aws) { + (Some(service_account), azure, aws) => { + if azure { + tracing::warn!( + "Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. Using GCP IAM. Remove one to avoid misconfiguration." + ); + } + if aws { + tracing::warn!( + "Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured for Redis. Using GCP IAM. Remove one to avoid misconfiguration." + ); + } + Ok(Auth::GcpIam { service_account }) + } + (None, true, aws) => { + if aws { + tracing::warn!( + "Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured for Redis. Using Azure AD. Remove one to avoid misconfiguration." + ); + } + Ok(Auth::AzureAd(AzureAdAuth { + client_id: text(options.azure.client_id.as_ref()), + tenant_id: text(options.azure.tenant_id.as_ref()), + client_secret: secret(options.azure.client_secret.as_ref()), + username: redis_username + .filter(|username| !username.is_empty()) + .map(str::to_owned), + })) + } + (None, false, true) => aws_iam(&options.aws, uses_tls(options)?).map(Auth::AwsIam), + (None, false, false) => Ok(Auth::Static(StaticAuth { + username: secret(options.username.as_ref()), + password: secret(options.password.as_ref()), + })), + } +} + +// _uses_tls: a url's scheme decides unless cluster nodes are set, otherwise the ssl kwarg does +fn uses_tls(options: &Options) -> Result { + match (&options.startup_nodes, text(options.url.as_ref())) { + (None, Some(url)) => Ok(url + .split_once("://") + .is_some_and(|(scheme, _)| scheme.eq_ignore_ascii_case("rediss"))), + _ => flag(options.ssl.as_ref()), + } +} + +fn aws_iam(aws: &AwsOptions, uses_tls: bool) -> Result { + if !uses_tls { + return Err(Error::AwsIamWithoutTls); + } + match ( + text(aws.user_name.as_ref()), + text(aws.cache_name.as_ref()), + text(aws.region.as_ref()), + ) { + (Some(user_name), Some(cache_name), Some(region)) => Ok(AwsIamAuth { + user_name, + cache_name, + region, + serverless: options::flag(aws.serverless.as_ref())?, + }), + (user_name, cache_name, region) => Err(Error::AwsIamIncomplete { + missing: [ + ("aws_iam_user_name", user_name.is_none()), + ("aws_iam_cache_name", cache_name.is_none()), + ("aws_iam_region", region.is_none()), + ] + .into_iter() + .filter_map(|(name, missing)| missing.then_some(name)) + .collect(), + }), + } +} diff --git a/litellm-rust/crates/redis-resolve/src/endpoint.rs b/litellm-rust/crates/redis-resolve/src/endpoint.rs new file mode 100644 index 00000000000..e3722b1b7cc --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/endpoint.rs @@ -0,0 +1,180 @@ +use litellm_redis_types::{ + Auth, Cluster, Node, RedisUrl, Sentinel, Standalone, StaticAuth, Target, TlsOptions, UnixSocket, +}; +use litellm_secret_value::SecretValue; + +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, + options::{Options, as_text, flag, integer, invalid, secret, text}, + tls::{self, QuerySettings}, +}; + +/// The target, first match wins like `get_redis_client`: cluster, url, Sentinel, then host. A +/// url also returns the settings its query carries. +pub(super) fn target(options: &Options, auth: &Auth) -> Result<(Target, QuerySettings), Error> { + if let Some(nodes) = cluster_nodes(options.startup_nodes.as_ref())? { + reject_cluster_db(options.db.as_ref())?; + let cluster = Cluster { + nodes, + tls: tls_for(options, auth)?, + }; + return Ok((Target::Cluster(cluster), QuerySettings::default())); + } + if let Some(url) = text(options.url.as_ref()) { + let query = tls::url_query(&url)?; + let url = RedisUrl { + url: SecretValue::new(match auth { + Auth::Static(_) => url, + _ => without_userinfo(&url), + }), + tls: query.tls, + db: query.db, + username: query.username, + password: query.password, + }; + return Ok((Target::Url(url), query.settings)); + } + if let Some(sentinels) = options.sentinel.nodes.as_ref() { + let sentinel = sentinel(options, auth, sentinels)?; + return Ok((Target::Sentinel(sentinel), QuerySettings::default())); + } + let host = text(options.host.as_ref()).ok_or(Error::NoTarget)?; + let target = match text(options.unix_socket_path.as_ref()) { + Some(path) => Target::Unix(UnixSocket { + path: path.into(), + db: db(options)?, + }), + None => Target::Standalone(Standalone { + host, + port: integer(options.port.as_ref())?.unwrap_or(6379), + tls: tls_for(options, auth)?, + db: db(options)?, + }), + }; + Ok((target, QuerySettings::default())) +} + +fn tls_for(options: &Options, auth: &Auth) -> Result, Error> { + if !flag(options.ssl.as_ref())? { + return Ok(None); + } + let gcp_ca_certs = match auth { + Auth::GcpIam { .. } => options.gcp.ssl_ca_certs.as_ref(), + _ => None, + }; + tls::tls_options(&options.tls, gcp_ca_certs).map(Some) +} + +fn db(options: &Options) -> Result { + integer(options.db.as_ref()).map(|db| db.unwrap_or(0)) +} + +fn sentinel(options: &Options, auth: &Auth, sentinels: &Entry) -> Result { + let monitors = sentinel_nodes(&sentinels.value) + .ok_or_else(|| invalid(&sentinels.source, Expected::SentinelNodes))?; + let service_name = + text(options.sentinel.service_name.as_ref()).ok_or(Error::NoSentinelService)?; + // redis.asyncio.Sentinel passes the connection kwargs to each monitor's Redis(host=..., port=...), so a + // host or port setting raises TypeError when litellm builds the client + if let Some(entry) = options.host.as_ref().or(options.port.as_ref()) { + return Err(Error::SentinelWithHost { + setting: entry.source.clone(), + }); + } + Ok(Sentinel { + monitors, + service_name, + monitor_auth: monitor_auth(options, auth), + tls: tls_for(options, auth)?, + db: db(options)?, + }) +} + +// _sentinel_auth_kwargs gives the monitors the data node's connection kwargs with sentinel_password as the +// password, so a static username reaches them too +fn monitor_auth(options: &Options, auth: &Auth) -> Option { + let password = secret(options.sentinel.password.as_ref())?; + let username = match auth { + Auth::Static(StaticAuth { username, .. }) => username.clone(), + _ => None, + }; + Some(StaticAuth { + username, + password: Some(password), + }) +} + +fn cluster_nodes(entry: Option<&Entry>) -> Result>, Error> { + let Some(entry) = entry else { + return Ok(None); + }; + let nodes = nodes(&entry.value, |node| { + node_address( + node.get("host").and_then(as_text), + node.get("port").and_then(as_text), + ) + }) + .ok_or_else(|| invalid(&entry.source, Expected::StartupNodes))?; + Ok((!nodes.is_empty()).then_some(nodes)) +} + +fn sentinel_nodes(value: &Value) -> Option> { + let nodes = nodes(value, |pair| match pair { + Value::Array(pair) if pair.len() == 2 => node_address(as_text(&pair[0]), as_text(&pair[1])), + _ => None, + })?; + (!nodes.is_empty()).then_some(nodes) +} + +fn nodes(value: &Value, parse: impl Fn(&Value) -> Option) -> Option> { + match value { + Value::String(json) => serde_json::from_str::>(json) + .ok()? + .iter() + .map(parse) + .collect(), + Value::Array(nodes) => nodes.iter().map(parse).collect(), + _ => None, + } +} + +fn node_address(host: Option, port: Option) -> Option { + Some(Node { + host: host?, + port: port?.parse().ok()?, + }) +} + +// redis-py's RedisCluster raises unless db is 0 or None +fn reject_cluster_db(db: Option<&Entry>) -> Result<(), Error> { + let Some(entry) = db else { + return Ok(()); + }; + let is_zero = match &entry.value { + Value::String(_) => matches!(integer::(Some(entry)), Ok(None | Some(0))), + Value::Bool(flag) => !flag, + Value::Number(number) => number.as_f64() == Some(0.0), + _ => false, + }; + match is_zero { + true => Ok(()), + false => Err(Error::DbInCluster { + setting: entry.source.clone(), + }), + } +} + +// _get_redis_url_kwargs keeps only plain Connection kwargs, so redis-py drops every ssl_* setting for a url +// and a credential provider replaces the url's userinfo, which _url_without_userinfo strips +fn without_userinfo(url: &str) -> String { + let Some((scheme, rest)) = url.split_once("://") else { + return url.to_owned(); + }; + let (authority, tail) = rest.split_at(rest.find(['/', '?', '#']).unwrap_or(rest.len())); + let host = authority + .rsplit_once('@') + .map_or(authority, |(_, host)| host); + format!("{scheme}://{host}{tail}") +} diff --git a/litellm-rust/crates/redis-resolve/src/error.rs b/litellm-rust/crates/redis-resolve/src/error.rs new file mode 100644 index 00000000000..b13b65edfb5 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/error.rs @@ -0,0 +1,80 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid Redis setting {setting}: expected {expected}")] + Invalid { setting: String, expected: Expected }, + #[error("invalid Redis setting {setting}: redis-py's parse_url cannot read it as a number")] + UrlQuery { setting: String }, + #[error("either host or url must be specified for redis")] + NoTarget, + #[error("service_name is required with sentinel_nodes")] + NoSentinelService, + #[error("{setting} cannot be combined with sentinel_nodes, which redis-py rejects at startup")] + SentinelWithHost { setting: String }, + #[error("{setting} must be 0 in cluster mode")] + DbInCluster { setting: String }, + #[error("credential_provider is a host object, and the host passed no handle for it")] + ProviderWithoutHandle, + #[error( + "{setting} asks for {check}, which rustls cannot perform, so Redis is not connected without them" + )] + Unenforceable { + setting: String, + check: &'static str, + }, + #[error("AWS ElastiCache IAM Redis authentication requires TLS")] + AwsIamWithoutTls, + #[error("AWS ElastiCache IAM Redis authentication requires: {}", missing.join(", "))] + AwsIamIncomplete { missing: Vec<&'static str> }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)] +pub enum Expected { + #[strum(serialize = "an integer")] + Integer, + #[strum(serialize = "a boolean")] + Boolean, + #[strum(serialize = "a non-negative number of seconds")] + Seconds, + #[strum(serialize = "none, optional, required, null or an ssl.VerifyMode value")] + CertRequirement, + #[strum(serialize = "a list of objects with host and port")] + StartupNodes, + #[strum(serialize = "a non-empty list of [host, port] pairs")] + SentinelNodes, +} + +/// The exception class redis-py 5.3.1 raises building its client. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PythonException { + ValueError, + TypeError, + RedisClusterException, +} + +impl Error { + /// What `_get_redis_client_logic` and redis-py's async client constructors raise on the same kwargs, or `None` + /// when Python builds the client and fails only once a command opens a connection, or never fails. + pub fn raised_building_the_python_client(&self) -> Option { + match self { + Self::Invalid { expected, .. } => match expected { + Expected::StartupNodes | Expected::SentinelNodes => { + Some(PythonException::ValueError) + } + Expected::Integer + | Expected::Boolean + | Expected::Seconds + | Expected::CertRequirement => None, + }, + Self::UrlQuery { .. } + | Self::NoTarget + | Self::NoSentinelService + | Self::AwsIamWithoutTls + | Self::AwsIamIncomplete { .. } => Some(PythonException::ValueError), + Self::SentinelWithHost { .. } | Self::Unenforceable { .. } => { + Some(PythonException::TypeError) + } + Self::DbInCluster { .. } => Some(PythonException::RedisClusterException), + Self::ProviderWithoutHandle => None, + } + } +} diff --git a/litellm-rust/crates/redis-resolve/src/keys.rs b/litellm-rust/crates/redis-resolve/src/keys.rs new file mode 100644 index 00000000000..c4897701d6f --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/keys.rs @@ -0,0 +1,304 @@ +use litellm_redis_types::Ignored; +use serde_json::Value; + +/// How the Rust client honors a redis-py kwarg. +#[derive(Clone, Copy, Debug)] +pub enum Handling { + /// Modeled by the resolver into a field of `Options`. + Read, + /// Dropped and returned as an `Ignored` naming `consequence`, unless `met` says the value asks for nothing the + /// Rust client does not already do. + Warn { + consequence: &'static str, + met: Option bool>, + }, + /// A security check the Rust client cannot perform and Python's async client already refuses, so asking for it + /// fails instead of connecting without it. + Reject { check: &'static str, asked: Asked }, +} + +/// Which values of a `Reject` row ask for its check. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Asked { + /// A flag read as true. False asks for nothing. + WhenTrue, + /// Any present value, like a list of flags. An empty one asks for nothing. + WhenSet, +} + +/// Whether redis-py's `parse_url` reads the setting from a url's query and the Rust resolver does too. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FromUrl { + No, + Any, + Rediss, +} + +#[derive(Clone, Copy, Debug)] +pub struct Setting { + pub kwarg: &'static str, + pub handling: Handling, + pub from_url: FromUrl, + /// `_get_redis_env_kwarg_mapping` reads it as `REDIS_`. + pub environment: bool, + /// The names `_merge_redis_kwargs` reads with `get_secret` when the merged kwargs lack the setting. + pub fallback: &'static [&'static str], +} + +impl Setting { + const fn new(kwarg: &'static str, handling: Handling) -> Self { + Self { + kwarg, + handling, + from_url: FromUrl::No, + environment: false, + fallback: &[], + } + } + + const fn env(self) -> Self { + Self { + environment: true, + ..self + } + } + + const fn in_url(self, from_url: FromUrl) -> Self { + Self { from_url, ..self } + } + + const fn or_secrets(self, fallback: &'static [&'static str]) -> Self { + Self { fallback, ..self } + } + + pub fn environment_name(&self) -> Option { + self.environment + .then(|| format!("REDIS_{}", self.kwarg.to_ascii_uppercase())) + } + + pub fn is_read(&self) -> bool { + matches!(self.handling, Handling::Read) + } + + /// The warning a set value earns, or `None` when the setting is read, refused or already met. + pub(crate) fn warning(&self, value: &Value) -> Option<&'static str> { + match self.handling { + Handling::Warn { consequence, met } if !met.is_some_and(|met| met(value)) => { + Some(consequence) + } + _ => None, + } + } +} + +const fn read(kwarg: &'static str) -> Setting { + Setting::new(kwarg, Handling::Read) +} + +const fn warn(kwarg: &'static str, consequence: &'static str) -> Setting { + Setting::new( + kwarg, + Handling::Warn { + consequence, + met: None, + }, + ) +} + +const fn warn_unless( + kwarg: &'static str, + consequence: &'static str, + met: fn(&Value) -> bool, +) -> Setting { + Setting::new( + kwarg, + Handling::Warn { + consequence, + met: Some(met), + }, + ) +} + +const fn reject(kwarg: &'static str, check: &'static str, asked: Asked) -> Setting { + Setting::new(kwarg, Handling::Reject { check, asked }) +} + +const NO_EQUIVALENT: &str = "the Rust client has no equivalent"; +const OWN_STACK: &str = "the Rust client uses its own pool, retries and events"; +const OWN_CONNECTION: &str = "it tunes redis-py's own connection, and the Rust client has its own"; +const OWN_CLUSTER: &str = "the Rust cluster client discovers, routes and retries on its own"; +const MAINTENANCE: &str = "the Rust client does not follow server maintenance notifications"; +const RUSTLS: &str = "rustls negotiates TLS 1.2+ with its own ECDHE AEAD suites"; +const REVOCATION: &str = "OCSP revocation checks"; +const OCSP: &str = "it only configures ssl_validate_ocsp, which the Rust client refuses when set"; +const VERIFY_FLAGS: &str = "X.509 verify flags such as CRL checking"; + +/// Every Redis kwarg redis-py or `litellm/_redis.py` accepts, one row each. `redis-py-kwargs.json` records which +/// redis-py versions and `REDIS_*` names accept them, and `tests/catalog.rs` keeps the two in step. +pub const SETTINGS: &[Setting] = &[ + warn("address_remap", OWN_CLUSTER), + warn("auto_close_connection_pool", OWN_STACK), + read("aws_iam_auth").env(), + read("aws_iam_cache_name").env(), + read("aws_iam_region") + .env() + .or_secrets(&["AWS_REGION", "AWS_DEFAULT_REGION"]), + read("aws_iam_serverless").env(), + read("aws_iam_user_name").env(), + read("azure_client_id") + .env() + .or_secrets(&["AZURE_CLIENT_ID"]), + read("azure_client_secret") + .env() + .or_secrets(&["AZURE_CLIENT_SECRET"]), + read("azure_redis_ad_token") + .env() + .or_secrets(&["REDIS_AZURE_AD_TOKEN"]), + read("azure_tenant_id") + .env() + .or_secrets(&["AZURE_TENANT_ID"]), + warn("cache", NO_EQUIVALENT).env(), + warn("cache_config", NO_EQUIVALENT).env(), + warn("charset", NO_EQUIVALENT).env(), + warn("client_name", NO_EQUIVALENT).env(), + warn("cluster_error_retry_attempts", OWN_CLUSTER), + warn("command_packer", OWN_CONNECTION), + warn("connection_class", OWN_STACK), + warn("connection_error_retry_attempts", OWN_CLUSTER), + warn("connection_pool", OWN_STACK), + read("credential_provider"), + read("db").env().in_url(FromUrl::Any), + warn("decode_responses", NO_EQUIVALENT).env(), + warn("driver_info", NO_EQUIVALENT), + warn("dynamic_startup_nodes", OWN_CLUSTER), + warn("encoder_class", OWN_CONNECTION), + warn("encoding", NO_EQUIVALENT).env(), + warn("encoding_errors", NO_EQUIVALENT).env(), + warn("errors", NO_EQUIVALENT).env(), + warn("event_dispatcher", OWN_STACK).env(), + read("gcp_service_account") + .env() + .or_secrets(&["REDIS_GCP_SERVICE_ACCOUNT"]), + read("gcp_ssl_ca_certs") + .env() + .or_secrets(&["REDIS_GCP_SSL_CA_CERTS"]), + read("health_check_interval").env().in_url(FromUrl::Any), + read("host").env(), + warn("legacy_responses", NO_EQUIVALENT), + warn("lib_name", NO_EQUIVALENT).env(), + warn("lib_version", NO_EQUIVALENT).env(), + warn("load_balancing_strategy", OWN_CLUSTER), + warn("maint_notifications_config", MAINTENANCE), + warn("maint_notifications_pool_handler", MAINTENANCE), + warn("maintenance_notification_hash", MAINTENANCE), + warn("maintenance_state", MAINTENANCE), + warn("max_connections", "the Rust client has no connection cap").env(), + read("namespace"), + warn("orig_host_address", MAINTENANCE), + warn("orig_socket_connect_timeout", MAINTENANCE), + warn("orig_socket_timeout", MAINTENANCE), + warn("oss_cluster_maint_notifications_handler", MAINTENANCE), + warn("parser", OWN_CONNECTION), + warn("parser_class", OWN_CONNECTION), + read("password").env().in_url(FromUrl::Any), + warn("path", OWN_CLUSTER), + warn("policy_resolver", OWN_CLUSTER), + read("port").env(), + warn("protocol", NO_EQUIVALENT).env(), + warn("read_from_replicas", OWN_CLUSTER), + warn( + "redis_connect_func", + "the Python object is not callable from Rust, so static username and password are used", + ) + .env(), + warn("reinitialize_steps", OWN_CLUSTER), + warn("require_full_coverage", OWN_CLUSTER), + warn("retry", OWN_STACK), + warn("retry_on_error", NO_EQUIVALENT).env(), + warn("retry_on_timeout", NO_EQUIVALENT).env(), + read("sentinel_nodes").or_secrets(&["REDIS_SENTINEL_NODES"]), + read("sentinel_password").or_secrets(&["REDIS_SENTINEL_PASSWORD"]), + read("service_name").or_secrets(&["REDIS_SERVICE_NAME"]), + warn("single_connection_client", NO_EQUIVALENT).env(), + read("socket_connect_timeout").env().in_url(FromUrl::Any), + read("socket_keepalive").env(), + warn("socket_keepalive_options", NO_EQUIVALENT).env(), + warn("socket_read_size", OWN_CONNECTION), + read("socket_timeout").env().in_url(FromUrl::Any), + warn("socket_type", OWN_CONNECTION), + read("ssl").env(), + read("ssl_ca_certs").env().in_url(FromUrl::Rediss), + read("ssl_ca_data").env().in_url(FromUrl::Rediss), + read("ssl_ca_path").env().in_url(FromUrl::Rediss), + read("ssl_cert_reqs").env().in_url(FromUrl::Rediss), + read("ssl_certfile").env().in_url(FromUrl::Rediss), + read("ssl_check_hostname").env().in_url(FromUrl::Rediss), + warn("ssl_ciphers", RUSTLS).env(), + reject("ssl_exclude_verify_flags", VERIFY_FLAGS, Asked::WhenSet), + reject("ssl_include_verify_flags", VERIFY_FLAGS, Asked::WhenSet), + read("ssl_keyfile").env().in_url(FromUrl::Rediss), + warn_unless("ssl_min_version", RUSTLS, is_met_by_rustls).env(), + warn("ssl_ocsp_context", OCSP).env(), + warn("ssl_ocsp_expected_cert", OCSP).env(), + read("ssl_password").env().in_url(FromUrl::Rediss), + reject("ssl_validate_ocsp", REVOCATION, Asked::WhenTrue).env(), + reject("ssl_validate_ocsp_stapled", REVOCATION, Asked::WhenTrue).env(), + read("startup_nodes").or_secrets(&["REDIS_CLUSTER_NODES"]), + read("unix_socket_path").env(), + read("url").env(), + read("username").env().in_url(FromUrl::Any), +]; + +// rustls only speaks TLS 1.2 and 1.3, so a minimum of 1.2 or lower is already met. redis-py takes an +// ssl.TLSVersion, which reaches Rust as its name or its value +fn is_met_by_rustls(value: &Value) -> bool { + match value { + Value::Number(number) => matches!(number.as_i64(), Some(-2 | 769 | 770 | 771)), + Value::String(text) => matches!( + text.strip_prefix("TLSVersion.") + .unwrap_or(text) + .to_ascii_uppercase() + .as_str(), + "MINIMUM_SUPPORTED" | "TLSV1" | "TLSV1_1" | "TLSV1_2" + ), + _ => false, + } +} + +pub fn setting(kwarg: &str) -> Option<&'static Setting> { + SETTINGS.iter().find(|setting| setting.kwarg == kwarg) +} + +/// The check a `Reject` row stands for and which values ask for it. +pub(crate) fn refusal(kwarg: &str) -> Option<(&'static str, Asked)> { + match setting(kwarg)?.handling { + Handling::Reject { check, asked } => Some((check, asked)), + Handling::Read | Handling::Warn { .. } => None, + } +} + +/// A warned setting the kwargs set, with what dropping it means. A kwarg with no row is not a Redis setting, so it +/// is dropped without a warning. +pub(crate) fn ignored(source: &str, kwarg: &str, value: &Value) -> Option { + let consequence = setting(kwarg)?.warning(value)?; + Some(Ignored { + setting: source.to_owned(), + consequence, + }) +} + +/// A url query parameter the resolver does not read. redis-py turns every one into a client kwarg, so even one with +/// no row is reported. +pub(crate) fn ignored_query(source: &str, name: &str, value: &Value) -> Option { + let consequence = match setting(name) { + None => NO_EQUIVALENT, + Some(setting) if setting.is_read() => { + "the Rust client does not read it from this url's query" + } + Some(setting) => setting.warning(value)?, + }; + Some(Ignored { + setting: source.to_owned(), + consequence, + }) +} diff --git a/litellm-rust/crates/redis-resolve/src/lib.rs b/litellm-rust/crates/redis-resolve/src/lib.rs new file mode 100644 index 00000000000..b7aa5d18f64 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/lib.rs @@ -0,0 +1,95 @@ +use std::{collections::BTreeMap, time::Duration}; + +use litellm_redis_types::{Auth, HealthCheck, ProviderHandle, Resolution, StaticAuth, Target}; +use serde_json::Value; + +mod auth; +mod endpoint; +mod error; +mod keys; +mod options; +mod tls; + +pub use error::{Error, Expected, PythonException}; +pub use keys::{Asked, FromUrl, Handling, SETTINGS, Setting, setting}; +pub use options::is_present; + +use options::Options; + +/// One merged redis-py kwarg and the name a warning or error reports it under. +#[derive(Clone, Debug, PartialEq)] +pub struct Entry { + pub value: Value, + pub source: String, +} + +pub type Kwargs = BTreeMap; + +/// Kwargs whose every entry is reported under its own kwarg name. +pub fn kwargs(fields: impl IntoIterator) -> Kwargs { + fields + .into_iter() + .map(|(kwarg, value)| { + let entry = Entry { + value, + source: kwarg.clone(), + }; + (kwarg, entry) + }) + .collect() +} + +/// What redis-py's constructors make of kwargs a host already merged the way `_get_redis_client_logic` does. A +/// host passes its `credential_provider` object as `provider`, and `REDIS_USERNAME` as its process environment +/// holds it, which Python's Azure AD auth reads from `os.environ` instead of the merged kwargs. +pub fn interpret( + kwargs: &Kwargs, + provider: Option, + redis_username: Option<&str>, +) -> Result { + let (options, ignored) = Options::read(kwargs)?; + let auth = auth::redis_auth(&options, provider, redis_username)?; + let (target, query) = endpoint::target(&options, &auth)?; + let auth = match (&target, auth) { + (Target::Url(_), Auth::Static(StaticAuth { username, .. })) => Auth::Static(StaticAuth { + username, + password: None, + }), + (_, auth) => auth, + }; + let timeout = query.timeout.map_or_else( + || options::seconds(options.socket_timeout.as_ref()), + |timeout| Ok(Some(timeout)), + )?; + let connect_timeout = query.connect_timeout.map_or_else( + || options::seconds(options.socket_connect_timeout.as_ref()), + |timeout| Ok(Some(timeout)), + )?; + let health_check_interval = query.health_check_interval.map_or_else( + || options::integer::(options.health_check_interval.as_ref()), + |interval| Ok(Some(interval)), + )?; + Ok(Resolution { + target, + auth, + namespace: options::text(options.namespace.as_ref()), + timeout, + connect_timeout, + health_check: health_check(health_check_interval), + keepalive: options + .socket_keepalive + .as_ref() + .map_or(Ok(true), |entry| options::flag(Some(entry)))?, + ignored: ignored.into_iter().chain(query.ignored).collect(), + }) +} + +// redis-py checks health only when the interval is truthy, so zero turns it off and a negative interval checks +// before every command +fn health_check(seconds: Option) -> HealthCheck { + match seconds { + None => HealthCheck::Unset, + Some(0) => HealthCheck::Disabled, + Some(seconds) => HealthCheck::Every(Duration::from_secs(seconds.try_into().unwrap_or(0))), + } +} diff --git a/litellm-rust/crates/redis-resolve/src/options.rs b/litellm-rust/crates/redis-resolve/src/options.rs new file mode 100644 index 00000000000..f2539e46ca5 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/options.rs @@ -0,0 +1,281 @@ +use std::time::Duration; + +use litellm_python_compat::{BigInt, number}; +use litellm_redis_types::Ignored; +use litellm_secret_value::SecretValue; +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, Kwargs, + keys::{self, Asked}, +}; + +/// Every setting the resolver reads. A setting is `None` when it is unset, null or empty, except that a null +/// `tls.cert_reqs` is kept because redis-py reads `ssl_cert_reqs=None` as `CERT_NONE`. +pub(super) struct Options { + pub(super) url: Option, + pub(super) host: Option, + pub(super) port: Option, + pub(super) db: Option, + pub(super) unix_socket_path: Option, + pub(super) startup_nodes: Option, + pub(super) sentinel: SentinelOptions, + pub(super) username: Option, + pub(super) password: Option, + pub(super) credential_provider: Option, + pub(super) gcp: GcpOptions, + pub(super) azure: AzureOptions, + pub(super) aws: AwsOptions, + pub(super) ssl: Option, + pub(super) tls: TlsInput, + pub(super) namespace: Option, + pub(super) socket_timeout: Option, + pub(super) socket_connect_timeout: Option, + pub(super) health_check_interval: Option, + pub(super) socket_keepalive: Option, +} + +pub(super) struct SentinelOptions { + pub(super) nodes: Option, + pub(super) service_name: Option, + pub(super) password: Option, +} + +pub(super) struct GcpOptions { + pub(super) service_account: Option, + pub(super) ssl_ca_certs: Option, +} + +pub(super) struct AzureOptions { + pub(super) ad_token: Option, + pub(super) client_id: Option, + pub(super) tenant_id: Option, + pub(super) client_secret: Option, +} + +pub(super) struct AwsOptions { + pub(super) auth: Option, + pub(super) user_name: Option, + pub(super) cache_name: Option, + pub(super) region: Option, + pub(super) serverless: Option, +} + +/// The TLS settings, read from the kwargs or from a `rediss://` url's query. +pub(super) struct TlsInput { + pub(super) cert_reqs: Option, + pub(super) ca_certs: Option, + pub(super) ca_data: Option, + pub(super) ca_path: Option, + pub(super) certfile: Option, + pub(super) keyfile: Option, + pub(super) key_password: Option, + pub(super) check_hostname: Option, +} + +/// The settings the kwargs set that the Rust client drops, once none of them asks for a check it cannot perform. +fn ignored(kwargs: &Kwargs) -> Result, Error> { + refuse_unenforceable(kwargs.iter())?; + Ok(kwargs + .iter() + .filter_map(|(key, entry)| keys::ignored(&entry.source, key, &entry.value)) + .collect()) +} + +impl Options { + /// The typed settings and the ones the Rust client drops. + pub(super) fn read(kwargs: &Kwargs) -> Result<(Self, Vec), Error> { + let ignored = ignored(kwargs)?; + let options = Self::from_lookup(|key| kwargs.get(key).cloned()); + Ok((options, ignored)) + } + + fn from_lookup(lookup: impl Fn(&'static str) -> Option) -> Self { + let get = |key| lookup(key).filter(|entry| is_present(&entry.value)); + Self { + url: get("url"), + host: get("host"), + port: get("port"), + db: get("db"), + unix_socket_path: get("unix_socket_path"), + startup_nodes: get("startup_nodes"), + sentinel: SentinelOptions { + nodes: get("sentinel_nodes"), + service_name: get("service_name"), + password: get("sentinel_password"), + }, + username: get("username"), + password: get("password"), + credential_provider: get("credential_provider"), + gcp: GcpOptions { + service_account: get("gcp_service_account"), + ssl_ca_certs: get("gcp_ssl_ca_certs"), + }, + azure: AzureOptions { + ad_token: get("azure_redis_ad_token"), + client_id: get("azure_client_id"), + tenant_id: get("azure_tenant_id"), + client_secret: get("azure_client_secret"), + }, + aws: AwsOptions { + auth: get("aws_iam_auth"), + user_name: get("aws_iam_user_name"), + cache_name: get("aws_iam_cache_name"), + region: get("aws_iam_region"), + serverless: get("aws_iam_serverless"), + }, + ssl: get("ssl"), + tls: TlsInput { + cert_reqs: lookup("ssl_cert_reqs") + .filter(|entry| entry.value.is_null() || is_present(&entry.value)), + ca_certs: get("ssl_ca_certs"), + ca_data: get("ssl_ca_data"), + ca_path: get("ssl_ca_path"), + certfile: get("ssl_certfile"), + keyfile: get("ssl_keyfile"), + key_password: get("ssl_password"), + check_hostname: get("ssl_check_hostname"), + }, + namespace: get("namespace"), + socket_timeout: get("socket_timeout"), + socket_connect_timeout: get("socket_connect_timeout"), + health_check_interval: get("health_check_interval"), + socket_keepalive: get("socket_keepalive"), + } + } +} + +/// A `Reject` row whose value asks for its check fails, because Rust cannot perform it. +pub(super) fn refuse_unenforceable<'a>( + entries: impl IntoIterator, +) -> Result<(), Error> { + entries + .into_iter() + .filter_map(|(key, entry)| Some((keys::refusal(key)?, entry))) + .try_for_each(|((check, asked), entry)| { + let asks = match asked { + Asked::WhenTrue => flag(Some(entry))?, + Asked::WhenSet => is_present(&entry.value), + }; + match asks { + true => Err(Error::Unenforceable { + setting: entry.source.clone(), + check, + }), + false => Ok(()), + } + }) +} + +// Python reads most of these with `x or fallback`, so an empty value counts as unset everywhere +pub fn is_present(value: &Value) -> bool { + match value { + Value::Null => false, + Value::String(text) => !text.is_empty(), + Value::Array(items) => !items.is_empty(), + Value::Object(entries) => !entries.is_empty(), + Value::Bool(_) | Value::Number(_) => true, + } +} + +pub(super) fn invalid(setting: &str, expected: Expected) -> Error { + Error::Invalid { + setting: setting.to_owned(), + expected, + } +} + +pub(super) fn as_text(value: &Value) -> Option { + match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} + +pub(super) fn text(entry: Option<&Entry>) -> Option { + entry.and_then(|entry| as_text(&entry.value)) +} + +pub(super) fn secret(entry: Option<&Entry>) -> Option { + text(entry).map(SecretValue::new) +} + +/// A flag every Python surface reads the same way. `_str_to_bool` takes only `true`, `1` and `yes` +/// as true while the proxy's pydantic models also take `on`, `t` and `y`, so text they disagree on, +/// or text neither accepts, is an error instead of a guess. +pub(super) fn flag(entry: Option<&Entry>) -> Result { + let Some(entry) = entry else { + return Ok(false); + }; + match &entry.value { + Value::Bool(flag) => Ok(*flag), + Value::Number(number) => Ok(number.as_f64().is_some_and(|number| number != 0.0)), + Value::String(text) => match text.to_ascii_lowercase().as_str() { + "true" | "1" | "yes" => Ok(true), + "false" | "0" | "no" | "off" | "f" | "n" => Ok(false), + _ => Err(invalid(&entry.source, Expected::Boolean)), + }, + _ => Err(invalid(&entry.source, Expected::Boolean)), + } +} + +pub(super) fn integer>(entry: Option<&Entry>) -> Result, Error> { + coerced(entry, Expected::Integer, number::int, |number| { + T::try_from(number).ok() + }) +} + +pub(super) fn seconds(entry: Option<&Entry>) -> Result, Error> { + coerced(entry, Expected::Seconds, number::float, |seconds| { + Duration::try_from_secs_f64(seconds).ok() + }) +} + +fn coerced( + entry: Option<&Entry>, + expected: Expected, + parse: impl Fn(&str) -> Option, + convert: impl Fn(N) -> Option, +) -> Result, Error> { + let Some((entry, text)) = entry.and_then(|entry| Some((entry, as_text(&entry.value)?))) else { + return Ok(None); + }; + match parse(&text) { + Some(number) => convert(number) + .map(Some) + .ok_or_else(|| invalid(&entry.source, expected)), + None => { + tracing::warn!(setting = %entry.source, "Redis setting is not a number, so the default is used like redis-py"); + Ok(None) + } + } +} + +#[cfg(test)] +mod tests { + use std::{cell::RefCell, collections::BTreeSet}; + + use rstest::rstest; + + use super::Options; + use crate::SETTINGS; + + // a Read row that Options never asks for would be neither modeled nor passed through nor warned about + #[rstest] + fn options_ask_for_exactly_the_read_rows() { + let asked = RefCell::new(BTreeSet::new()); + let _ = Options::from_lookup(|kwarg| { + asked.borrow_mut().insert(kwarg); + None + }); + + let read: BTreeSet<&str> = SETTINGS + .iter() + .filter(|setting| setting.is_read()) + .map(|setting| setting.kwarg) + .collect(); + assert_eq!(asked.into_inner(), read); + } +} diff --git a/litellm-rust/crates/redis-resolve/src/tls.rs b/litellm-rust/crates/redis-resolve/src/tls.rs new file mode 100644 index 00000000000..a46a6ccc650 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/tls.rs @@ -0,0 +1,189 @@ +use std::time::Duration; + +use litellm_python_compat::number; +use litellm_redis_types::{Ignored, TlsOptions}; +use litellm_secret_value::SecretValue; + +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, + keys::{self, FromUrl}, + options::{self, TlsInput, flag, invalid, secret, text}, +}; + +/// `gcp_ca_certs` replaces `ssl_ca_certs` when GCP IAM is on, as `_get_redis_client_logic` does. +pub(super) fn tls_options( + input: &TlsInput, + gcp_ca_certs: Option<&Entry>, +) -> Result { + Ok(TlsOptions { + ca_certs: text(gcp_ca_certs.or(input.ca_certs.as_ref())).map(Into::into), + ca_data: text(input.ca_data.as_ref()), + ca_path: text(input.ca_path.as_ref()).map(Into::into), + certfile: text(input.certfile.as_ref()).map(Into::into), + keyfile: text(input.keyfile.as_ref()).map(Into::into), + key_password: secret(input.key_password.as_ref()), + verify_certificate: verify_certificate(input.cert_reqs.as_ref())?, + check_hostname: flag(input.check_hostname.as_ref())?, + }) +} + +// RedisSSLContext looks text up case-sensitively and takes an ssl.VerifyMode as is +fn verify_certificate(cert_reqs: Option<&Entry>) -> Result { + let Some(entry) = cert_reqs else { + return Ok(true); + }; + match &entry.value { + // Explicit None intentionally means CERT_NONE, matching redis-py rather than the unset default. + Value::Null => Ok(false), + Value::String(text) if text == "none" => Ok(false), + Value::String(text) if text == "optional" || text == "required" => Ok(true), + Value::Number(mode) => match mode.as_i64() { + Some(0) => Ok(false), + Some(1 | 2) => Ok(true), + _ => Err(invalid(&entry.source, Expected::CertRequirement)), + }, + _ => Err(invalid(&entry.source, Expected::CertRequirement)), + } +} + +/// What redis-py's `parse_url` reads from a url's query on top of the url itself. +pub(super) struct UrlQuery { + pub(super) tls: Option, + pub(super) db: Option, + pub(super) username: Option, + pub(super) password: Option, + pub(super) settings: QuerySettings, +} + +/// The connection settings a url's query sets, which win over the kwargs because `from_url` applies them last. +#[derive(Default)] +pub(super) struct QuerySettings { + pub(super) timeout: Option, + pub(super) connect_timeout: Option, + pub(super) health_check_interval: Option, + pub(super) ignored: Vec, +} + +pub(super) fn url_query(url: &str) -> Result { + let Ok(url) = url::Url::parse(url) else { + return Ok(UrlQuery { + tls: None, + db: None, + username: None, + password: None, + settings: QuerySettings::default(), + }); + }; + let tls_scheme = url.scheme().eq_ignore_ascii_case("rediss"); + let pairs: Vec<(String, Entry)> = url + .query_pairs() + .filter(|(_, value)| !value.is_empty()) + .map(|(name, value)| { + ( + name.to_string(), + Entry { + value: Value::String(value.into_owned()), + source: format!("url ?{name}"), + }, + ) + }) + .collect(); + // parse_qs keeps every value of a repeated parameter and parse_url reads the first + let query: Vec<(String, Entry)> = pairs + .iter() + .enumerate() + .filter(|(index, (name, _))| !pairs[..*index].iter().any(|(earlier, _)| earlier == name)) + .map(|(_, pair)| pair.clone()) + .collect(); + options::refuse_unenforceable(query.iter().map(|(name, entry)| (name, entry)))?; + let (query, ignored): (Vec<_>, Vec<_>) = query.into_iter().partition(|(name, _)| { + keys::setting(name).is_some_and(|setting| match setting.from_url { + FromUrl::Any => true, + FromUrl::Rediss => tls_scheme, + FromUrl::No => false, + }) + }); + let get = |key: &str| { + query + .iter() + .find(|(name, _)| name == key) + .map(|(_, entry)| entry.clone()) + }; + let seconds = |key: &str| { + query_number( + get(key).as_ref(), + number::float, + Expected::Seconds, + |seconds| Duration::try_from_secs_f64(seconds).ok(), + ) + }; + let integer = |key: &str| { + query_number( + get(key).as_ref(), + number::int, + Expected::Integer, + |number| i64::try_from(number).ok(), + ) + }; + let db = integer("db")?; + let settings = QuerySettings { + timeout: seconds("socket_timeout")?, + connect_timeout: seconds("socket_connect_timeout")?, + health_check_interval: integer("health_check_interval")?, + ignored: ignored + .iter() + .filter_map(|(name, entry)| keys::ignored_query(&entry.source, name, &entry.value)) + .collect(), + }; + let tls = match tls_scheme { + true => Some(TlsOptions { + ca_certs: text(get("ssl_ca_certs").as_ref()).map(Into::into), + ca_data: text(get("ssl_ca_data").as_ref()), + ca_path: text(get("ssl_ca_path").as_ref()).map(Into::into), + certfile: text(get("ssl_certfile").as_ref()).map(Into::into), + keyfile: text(get("ssl_keyfile").as_ref()).map(Into::into), + key_password: secret(get("ssl_password").as_ref()), + verify_certificate: verify_certificate(get("ssl_cert_reqs").as_ref())?, + check_hostname: text(get("ssl_check_hostname").as_ref()) + .is_some_and(|text| url_bool(&text)), + }), + false => None, + }; + Ok(UrlQuery { + tls, + db, + username: secret(get("username").as_ref()), + password: secret(get("password").as_ref()), + settings, + }) +} + +// parse_url reads these with int() or float() and raises ValueError building the client on anything else +fn query_number( + entry: Option<&Entry>, + parse: impl Fn(&str) -> Option, + expected: Expected, + convert: impl Fn(N) -> Option, +) -> Result, Error> { + let Some(entry) = entry else { + return Ok(None); + }; + let number = text(Some(entry)) + .and_then(|text| parse(&text)) + .ok_or_else(|| Error::UrlQuery { + setting: entry.source.clone(), + })?; + convert(number) + .map(Some) + .ok_or_else(|| invalid(&entry.source, expected)) +} + +// redis-py's to_bool: only 0, F, FALSE, N and NO (any case) are false +fn url_bool(text: &str) -> bool { + !matches!( + text.to_ascii_uppercase().as_str(), + "" | "0" | "F" | "FALSE" | "N" | "NO" + ) +} diff --git a/litellm-rust/crates/redis-resolve/tests/catalog.rs b/litellm-rust/crates/redis-resolve/tests/catalog.rs new file mode 100644 index 00000000000..b26ad0eebfb --- /dev/null +++ b/litellm-rust/crates/redis-resolve/tests/catalog.rs @@ -0,0 +1,129 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_redis_resolve::{Asked, Error, Handling, SETTINGS, Setting, interpret, kwargs}; +use rstest::rstest; +use serde::Deserialize; +use serde_json::json; + +#[derive(Deserialize)] +struct Catalog { + meta: Meta, + args: BTreeMap, +} + +#[derive(Deserialize)] +struct Meta { + redis_py_locked: String, +} + +#[derive(Deserialize)] +struct Arg { + environment: bool, + redis_py: BTreeMap>, +} + +// the signatures of the clients the proxy builds, where an unknown kwarg raises TypeError +const PROXY_SIGNATURES: &[&str] = &["async", "cluster", "connection"]; + +fn rejected() -> impl Iterator { + SETTINGS + .iter() + .filter_map(|setting| match setting.handling { + Handling::Reject { asked, .. } => Some((setting, asked)), + Handling::Read | Handling::Warn { .. } => None, + }) +} + +fn catalog() -> Catalog { + serde_json::from_str(include_str!("../redis-py-kwargs.json")).unwrap() +} + +fn host_with(kwarg: &str, value: serde_json::Value) -> litellm_redis_resolve::Kwargs { + kwargs([ + ("host".to_owned(), json!("redis")), + (kwarg.to_owned(), value), + ]) +} + +#[rstest] +fn every_catalogued_kwarg_has_exactly_one_settings_row() { + let rows: Vec<&str> = SETTINGS.iter().map(|setting| setting.kwarg).collect(); + let unique: BTreeSet<&str> = rows.iter().copied().collect(); + assert_eq!(rows.len(), unique.len(), "a kwarg has two SETTINGS rows"); + + let catalog = catalog(); + let catalogued: BTreeSet<&str> = catalog.args.keys().map(String::as_str).collect(); + assert_eq!(unique, catalogued); +} + +#[rstest] +fn settings_read_redis_environment_names_exactly_where_python_does() { + let catalog = catalog(); + + let differing: Vec<&str> = SETTINGS + .iter() + .filter(|setting| { + catalog + .args + .get(setting.kwarg) + .is_some_and(|arg| arg.environment != setting.environment) + }) + .map(|setting| setting.kwarg) + .collect(); + assert!(differing.is_empty(), "{differing:?}"); +} + +#[rstest] +fn every_warned_setting_is_reported() { + let warned = SETTINGS + .iter() + .filter(|setting| matches!(setting.handling, Handling::Warn { met: None, .. })); + + for setting in warned { + let resolved = interpret(&host_with(setting.kwarg, json!("set")), None, None).unwrap(); + + let reported: Vec<&str> = resolved + .ignored + .iter() + .map(|ignored| ignored.setting.as_str()) + .collect(); + assert_eq!(reported, vec![setting.kwarg]); + } +} + +#[rstest] +fn every_rejected_setting_fails_when_asked_for_and_not_otherwise() { + for (setting, asked) in rejected() { + let (asking, not_asking) = match asked { + Asked::WhenTrue => (json!(true), json!(false)), + Asked::WhenSet => (json!(["VERIFY_CRL_CHECK_LEAF"]), json!([])), + }; + + let error = interpret(&host_with(setting.kwarg, asking), None, None).unwrap_err(); + assert!( + matches!(&error, Error::Unenforceable { setting: source, .. } if source == setting.kwarg), + "{error:?}" + ); + assert!(interpret(&host_with(setting.kwarg, not_asking), None, None).is_ok()); + } +} + +#[rstest] +fn python_already_refuses_every_rejected_setting() { + let catalog = catalog(); + + let accepted: Vec<&str> = rejected() + .map(|(setting, _)| setting.kwarg) + .filter(|kwarg| { + catalog.args[*kwarg] + .redis_py + .get(&catalog.meta.redis_py_locked) + .is_some_and(|signatures| { + signatures + .iter() + .any(|signature| PROXY_SIGNATURES.contains(&signature.as_str())) + }) + }) + .collect(); + assert!(accepted.is_empty(), "{accepted:?}"); +} diff --git a/litellm-rust/crates/redis-resolve/tests/interpret.rs b/litellm-rust/crates/redis-resolve/tests/interpret.rs new file mode 100644 index 00000000000..fd59ec2acb1 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/tests/interpret.rs @@ -0,0 +1,790 @@ +use std::collections::BTreeMap; + +use litellm_redis_resolve::{Entry, Error, Expected, Kwargs, PythonException, interpret, kwargs}; +use litellm_redis_types::{ + Auth, HealthCheck, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions, +}; +use rstest::rstest; +use serde_json::json; + +use Outcome::{Fails, Refuses, Resolves}; +use PythonException::{RedisClusterException, TypeError, ValueError}; + +const PROVIDER: ProviderHandle = ProviderHandle::new(7); + +fn merged(fields: serde_json::Value, sources: &[(&str, &str)]) -> Kwargs { + let fields = serde_json::from_value::>(fields).unwrap(); + kwargs(fields) + .into_iter() + .map(|(kwarg, entry)| { + let source = sources + .iter() + .find(|(named, _)| *named == kwarg) + .map_or(entry.source, |(_, source)| (*source).to_owned()); + (kwarg, Entry { source, ..entry }) + }) + .collect() +} + +// a host hands its credential_provider object over as a handle, so the cases do the same +fn interpret_merged( + fields: serde_json::Value, + sources: &[(&str, &str)], + redis_username: Option<&str>, +) -> Result { + let has_provider = fields + .get("credential_provider") + .is_some_and(|provider| !provider.is_null()); + interpret( + &merged(fields, sources), + has_provider.then_some(PROVIDER), + redis_username, + ) +} + +fn interpret_kwargs(fields: serde_json::Value) -> Result { + interpret_merged(fields, &[], None) +} + +/// What a contract case expects of the resolver, next to what its Python mirror expects of redis-py. +enum Outcome { + Resolves(serde_json::Value), + /// Rust refuses a config Python connects with, a disagreement on purpose. + Refuses(&'static str), + /// redis-py 5.3.1 already raises on the config, building the client or its first connection, so failing + /// here breaks no deployment that works today. The class is what it raises building the client, and `None` + /// when it only raises once a command connects. + Fails(&'static str, Option), +} + +#[rstest] +fn a_provider_handle_beats_every_cloud_auth_and_static_credentials() { + let resolved = interpret_kwargs(json!({ + "host": "redis", "password": "pass", "credential_provider": "host-object", + "gcp_service_account": "account", "aws_iam_auth": true, "azure_redis_ad_token": "true", + })) + .unwrap(); + + assert_eq!(resolved.auth, Auth::Provider(PROVIDER)); +} + +#[rstest] +fn a_rediss_url_reads_ssl_password_from_its_query() { + let resolved = + interpret_kwargs(json!({"url": "rediss://cache:6380/0?ssl_password=query-key-secret"})) + .unwrap(); + + let Target::Url(parsed) = resolved.target else { + panic!("expected a url target, got {:?}", resolved.target); + }; + let key_password = parsed.tls.and_then(|tls| tls.key_password); + assert_eq!(expose(key_password.as_ref()), Some("query-key-secret")); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +#[case::read_on_a_plain_url( + "redis://cache:6379/1?db=3&username=query-user&password=query-secret", + Some(3), + Some("query-user"), + Some("query-secret") +)] +#[case::absent("rediss://cache:6380/1", None, None, None)] +fn a_url_query_carries_db_and_credentials_like_parse_url( + #[case] url: &str, + #[case] db: Option, + #[case] username: Option<&str>, + #[case] password: Option<&str>, +) { + let resolved = interpret_kwargs(json!({"url": url})).unwrap(); + + let Target::Url(parsed) = resolved.target else { + panic!("expected a url target, got {:?}", resolved.target); + }; + assert_eq!(parsed.db, db); + assert_eq!(expose(parsed.username.as_ref()), username); + assert_eq!(expose(parsed.password.as_ref()), password); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +fn a_url_query_db_that_is_not_an_integer_fails_like_parse_url() { + let error = interpret_kwargs(json!({"url": "redis://cache:6379?db=first"})).unwrap_err(); + + assert!(matches!(error, Error::UrlQuery { setting } if setting == "url ?db")); +} + +#[rstest] +#[case::tls_1_2_is_already_met(json!("TLSv1_2"), false)] +#[case::the_enum_name(json!("TLSVersion.TLSv1_2"), false)] +#[case::the_enum_value(json!(771), false)] +#[case::tls_1_3_is_not_enforced(json!("TLSv1_3"), true)] +fn ssl_min_version_warns_only_when_rustls_cannot_meet_it( + #[case] version: serde_json::Value, + #[case] warns: bool, +) { + let resolved = + interpret_kwargs(json!({"host": "redis", "ssl": true, "ssl_min_version": version})) + .unwrap(); + + assert_eq!( + !resolved.ignored.is_empty(), + warns, + "{:?}", + resolved.ignored + ); +} + +#[rstest] +#[case::unset(json!({"host": "redis"}), true)] +#[case::empty_like_unset(json!({"host": "redis", "socket_keepalive": ""}), true)] +#[case::enabled(json!({"host": "redis", "socket_keepalive": "yes"}), true)] +#[case::disabled(json!({"host": "redis", "socket_keepalive": false}), false)] +#[case::disabled_as_text(json!({"host": "redis", "socket_keepalive": "0"}), false)] +#[case::disabled_on_a_cluster(json!({"startup_nodes": [{"host": "node-a", "port": 7000}], "socket_keepalive": "false"}), false)] +fn keepalive_stays_on_unless_socket_keepalive_turns_it_off( + #[case] kwargs: serde_json::Value, + #[case] keepalive: bool, +) { + let resolved = interpret_kwargs(kwargs).unwrap(); + + assert_eq!(resolved.keepalive, keepalive); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +fn a_socket_keepalive_no_python_surface_agrees_on_is_invalid() { + let error = interpret_kwargs(json!({"host": "redis", "socket_keepalive": "on"})).unwrap_err(); + + assert!( + matches!(&error, Error::Invalid { setting, expected: Expected::Boolean } if setting == "socket_keepalive"), + "{error:?}" + ); +} + +fn tls_view(tls: Option<&TlsOptions>) -> serde_json::Value { + tls.map_or(serde_json::Value::Null, |tls| { + serde_json::json!({ + "ca_certs": tls.ca_certs, + "ca_data": tls.ca_data, + "ca_path": tls.ca_path, + "certfile": tls.certfile, + "keyfile": tls.keyfile, + "verify_certificate": tls.verify_certificate, + "check_hostname": tls.check_hostname, + }) + }) +} + +fn nodes_view(nodes: &[Node]) -> Vec { + nodes + .iter() + .map(|node| format!("{}:{}", node.host, node.port)) + .collect() +} + +fn expose(value: Option<&litellm_secret_value::SecretValue>) -> Option<&str> { + value.map(litellm_secret_value::SecretValue::expose) +} + +fn view(resolved: &Resolution) -> serde_json::Value { + let target = match &resolved.target { + Target::Unix(socket) => serde_json::json!({ + "mode": "unix", "path": socket.path, "tls": null, + }), + Target::Standalone(host) => serde_json::json!({ + "mode": "host", + "host": host.host, + "port": host.port, + "db": host.db, + "tls": tls_view(host.tls.as_ref()), + }), + Target::Url(url) => serde_json::json!({ + "mode": "url", + "url": url.url.expose(), + "tls": tls_view(url.tls.as_ref()), + }), + Target::Cluster(cluster) => serde_json::json!({ + "mode": "cluster", + "nodes": nodes_view(&cluster.nodes), + "tls": tls_view(cluster.tls.as_ref()), + }), + Target::Sentinel(sentinel) => serde_json::json!({ + "mode": "sentinel", + "nodes": nodes_view(&sentinel.monitors), + "service_name": sentinel.service_name, + "sentinel_password": expose(sentinel.monitor_auth.as_ref().and_then(|auth| auth.password.as_ref())), + "db": sentinel.db, + "tls": tls_view(sentinel.tls.as_ref()), + }), + }; + let auth = match &resolved.auth { + Auth::Static(StaticAuth { username, password }) => serde_json::json!({ + "kind": "static", + "username": expose(username.as_ref()), + "password": expose(password.as_ref()), + }), + Auth::Provider(_) => serde_json::json!({"kind": "host"}), + Auth::GcpIam { service_account } => { + serde_json::json!({"kind": "gcp_iam", "service_account": service_account}) + } + Auth::AzureAd(azure) => { + serde_json::json!({"kind": "azure_ad", "username": azure.username}) + } + Auth::AwsIam(aws) => serde_json::json!({ + "kind": "aws_iam", + "user_name": aws.user_name, + "cache_name": aws.cache_name, + "region": aws.region, + "serverless": aws.serverless, + }), + }; + let serde_json::Value::Object(target) = target else { + unreachable!("every target view is an object") + }; + let rest = serde_json::json!({ + "auth": auth, + "socket_timeout": resolved.timeout.map_or(5.0, |timeout| timeout.as_secs_f64()), + "socket_connect_timeout": resolved.connect_timeout.or(resolved.timeout).map_or(5.0, |timeout| timeout.as_secs_f64()), + "health_check_interval": match (&resolved.target, resolved.health_check) { + (_, HealthCheck::Every(interval)) => Some(interval.as_secs_f64()), + (Target::Cluster(_), HealthCheck::Unset) => Some(25.0), + (_, HealthCheck::Unset | HealthCheck::Disabled) => None, + }, + }); + let serde_json::Value::Object(rest) = rest else { + unreachable!("the view is an object") + }; + serde_json::Value::Object(target.into_iter().chain(rest).collect()) +} + +fn expected_view(fields: serde_json::Value) -> serde_json::Value { + let serde_json::Value::Object(fields) = fields else { + unreachable!("expected views are objects") + }; + let defaults = json!({ + "auth": {"kind": "static", "username": null, "password": null}, + "tls": null, + "socket_timeout": 5.0, + "socket_connect_timeout": 5.0, + "health_check_interval": null, + }); + let serde_json::Value::Object(defaults) = defaults else { + unreachable!("the defaults are an object") + }; + serde_json::Value::Object(defaults.into_iter().chain(fields).collect()) +} + +#[rstest] +#[case::host_defaults( + json!({"host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::empty_cluster_nodes_use_the_environment( + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::empty_sentinel_nodes_use_the_environment( + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))), +)] +#[case::empty_service_name_uses_the_environment( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[("service_name", "REDIS_SERVICE_NAME")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))), +)] +#[case::plaintext_ignores_tls_requirements( + json!({"host": "redis", "ssl_cert_reqs": "strict"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::an_ambiguous_flag_is_refused( + json!({"host": "redis", "ssl": "on"}), + &[], + None, + Refuses("invalid Redis setting ssl: expected a boolean"), +)] +#[case::empty_sentinel_nodes_without_the_environment_name_no_target( + json!({"sentinel_nodes": [], "service_name": "primary"}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +#[case::ssl_ca_path_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_ca_path": "/ca"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/ca", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::ssl_validate_ocsp_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), + &[], + None, + Fails("ssl_validate_ocsp asks for OCSP revocation checks, which rustls cannot perform, so Redis is not connected without them", Some(TypeError)), +)] +#[case::ssl_include_verify_flags_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_include_verify_flags": ["VERIFY_CRL_CHECK_LEAF"]}), + &[], + None, + Fails("ssl_include_verify_flags asks for X.509 verify flags such as CRL checking, which rustls cannot perform, so Redis is not connected without them", Some(TypeError)), +)] +#[case::ssl_password_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_password": "key-secret"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::host_with_every_basic_setting( + json!({"db": 2, "host": "redis", "password": "pass", "port": 6380, "socket_timeout": 1.5, "ssl": true, "username": "user"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 2, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": "pass"}, "socket_timeout": 1.5, "socket_connect_timeout": 1.5}))), +)] +#[case::environment_host_port_and_password( + json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}), + &[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "environment", "port": 6390, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": null, "password": "environment-secret"}}))), +)] +#[case::string_numbers_are_coerced( + json!({"db": "3", "host": "redis", "port": "6380", "socket_timeout": "2.5"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 3, "socket_timeout": 2.5, "socket_connect_timeout": 2.5}))), +)] +#[case::padded_and_underscored_numbers_parse_like_python_int_and_float( + json!({"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}), + &[], + None, + Resolves(expected_view(json!({"db": 2, "host": "redis", "mode": "host", "port": 6380, "socket_timeout": 10.5, "socket_connect_timeout": 10.5}))), +)] +#[case::padded_environment_numbers_parse_like_python_int( + json!({"port": "6_380", "host": "redis", "db": "3\n"}), + &[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")], + None, + Resolves(expected_view(json!({"db": 3, "host": "redis", "mode": "host", "port": 6380}))), +)] +#[case::socket_timeout_from_the_environment( + json!({"socket_timeout": "0.75", "host": "redis"}), + &[("socket_timeout", "REDIS_SOCKET_TIMEOUT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 0.75, "socket_connect_timeout": 0.75}))), +)] +#[case::environment_url( + json!({"url": "redis://environment:6379/1"}), + &[("url", "REDIS_URL")], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://environment:6379/1"}))), +)] +#[case::environment_url_beats_the_environment_host( + json!({"host": "ignored", "url": "redis://from-url:6379"}), + &[("host", "REDIS_HOST"), ("url", "REDIS_URL")], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://from-url:6379"}))), +)] +#[case::an_explicit_host_drops_the_environment_url( + json!({"host": "configured"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "configured", "port": 6379, "db": 0}))), +)] +#[case::a_url_drops_host_db_and_password_and_keeps_the_username( + json!({"db": 4, "host": "dropped", "password": "dropped", "url": "rediss://cache:6380", "username": "user"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": null}}))), +)] +#[case::a_url_ignores_tls_settings_outside_the_url( + json!({"ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true, "url": "rediss://cache:6380"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::startup_nodes_beat_a_url( + json!({"password": "cluster-secret", "startup_nodes": [{"host": "node-a", "port": 7000}], "url": "redis://single:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "auth": {"kind": "static", "username": null, "password": "cluster-secret"}, "health_check_interval": 25.0}))), +)] +#[case::environment_cluster_nodes( + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"], "health_check_interval": 25.0}))), +)] +#[case::cluster_tls_options( + json!({"ssl": "true", "ssl_ca_certs": "/ca.pem", "ssl_check_hostname": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": true}, "health_check_interval": 25.0}))), +)] +#[case::sentinel_from_settings( + json!({"password": "data-secret", "sentinel_nodes": [["sentinel-a", 26379], ["sentinel-b", 26380]], "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379", "sentinel-b:26380"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0, "auth": {"kind": "static", "username": null, "password": "data-secret"}}))), +)] +#[case::sentinel_from_the_environment( + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0}))), +)] +#[case::a_host_breaks_sentinel( + json!({"host": "ignored", "sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[], + None, + Fails("host cannot be combined with sentinel_nodes, which redis-py rejects at startup", Some(TypeError)), +)] +#[case::a_url_beats_sentinels( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary", "url": "redis://single:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://single:6379"}))), +)] +#[case::an_explicit_credential_provider_beats_cloud_auth( + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "host"}}))), +)] +#[case::gcp_beats_azure_and_aws( + json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "gcp_iam", "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}}))), +)] +#[case::gcp_ca_certs_replace_the_configured_ones_under_ssl( + json!({"gcp_service_account": "account", "gcp_ssl_ca_certs": "/gcp-ca.pem", "host": "redis", "ssl": true, "ssl_ca_certs": "/configured-ca.pem"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/gcp-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "gcp_iam", "service_account": "account"}}))), +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[], + Some("environment-user"), + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": "environment-user"}}))), +)] +#[case::an_empty_environment_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}), + &[("username", "REDIS_USERNAME")], + Some(""), + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))), +)] +#[case::azure_without_a_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))), +)] +#[case::azure_needs_the_literal_true( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": "1", "azure_tenant_id": "tenant", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::aws_iam_over_tls_with_the_default_region( + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}), + &[("aws_iam_region", "AWS_DEFAULT_REGION")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "us-east-1", "serverless": false}}))), +)] +#[case::aws_region_beats_the_default_region( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}), + &[("aws_iam_region", "AWS_REGION")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-central-1", "serverless": false}}))), +)] +#[case::aws_iam_serverless_over_a_rediss_url_drops_its_userinfo( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "eu-west-1", "aws_iam_serverless": true, "aws_iam_user_name": "iam-user", "url": "rediss://user:pass@cache-1:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache-1:6379", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-west-1", "serverless": true}}))), +)] +#[case::aws_iam_without_tls_fails( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "us-east-1", "aws_iam_user_name": "iam-user", "host": "cache-1"}), + &[], + None, + Fails("AWS ElastiCache IAM Redis authentication requires TLS", Some(ValueError)), +)] +#[case::aws_iam_without_a_region_fails( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}), + &[], + None, + Fails("AWS ElastiCache IAM Redis authentication requires: aws_iam_region", Some(ValueError)), +)] +#[case::tls_options_from_settings( + json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))), +)] +#[case::a_null_cert_requirement_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": false}}))), +)] +#[case::a_verify_mode_of_none_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": 0}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": false}}))), +)] +#[case::a_verify_mode_of_optional_verifies( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": 1}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::a_cert_requirement_spelled_like_python_none_fails( + json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], + None, + Fails("invalid Redis setting REDIS_SSL_CERT_REQS: expected none, optional, required, null or an ssl.VerifyMode value", None), +)] +#[case::tls_options_from_the_environment( + json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/env-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::connect_timeout_is_separate_from_the_socket_timeout( + json!({"host": "redis", "socket_connect_timeout": 0.25, "socket_timeout": 2.5}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 2.5, "socket_connect_timeout": 0.25}))), +)] +#[case::environment_references_resolve( + json!({"host": "referenced", "password": "referenced-secret"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "referenced", "port": 6379, "db": 0, "auth": {"kind": "static", "username": null, "password": "referenced-secret"}}))), +)] +#[case::no_target_fails( + json!({"port": 6379}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +#[case::client_certificates( + json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": "/client.key", "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::a_rediss_url_takes_tls_options_from_its_query( + json!({"url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes", "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))), +)] +#[case::a_url_query_carries_socket_timeout( + json!({"url": "redis://cache:6379/0?socket_timeout=2"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2", "socket_timeout": 2.0, "socket_connect_timeout": 2.0}))), +)] +#[case::url_query_settings_win_over_the_kwargs( + json!({"health_check_interval": 4, "socket_connect_timeout": 7, "socket_timeout": 1, "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3", "socket_timeout": 3.0, "socket_connect_timeout": 0.5, "health_check_interval": 9.0}))), +)] +#[case::a_repeated_url_query_parameter_takes_its_first_value( + json!({"url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9", "socket_timeout": 2.0, "socket_connect_timeout": 2.0}))), +)] +#[case::a_url_query_number_parse_url_cannot_read( + json!({"url": "redis://cache:6379/0?socket_timeout=slow"}), + &[], + None, + Fails("invalid Redis setting url ?socket_timeout: redis-py's parse_url cannot read it as a number", Some(ValueError)), +)] +#[case::tls_settings_outside_a_url_never_fail_it( + json!({"db": "99999999999999999999", "ssl_cert_reqs": "strict", "url": "rediss://cache:6380"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::db_is_not_read_for_a_cluster( + json!({"db": 99_999_999_999_i64, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_true( + json!({"db": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_false( + json!({"db": false, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::cluster_db_fraction( + json!({"db": 0.5, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_float_zero( + json!({"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::cluster_db_nonempty_list( + json!({"db": [1], "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_empty_list( + json!({"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::ssl_ciphers( + json!({"host": "redis", "ssl": true, "ssl_ciphers": "HIGH"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::retry_on_timeout( + json!({"host": "redis", "retry_on_timeout": true}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_port( + json!({"host": "redis", "port": "abc"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_db( + json!({"db": "first", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::socket_timeout_not_a_number( + json!({"host": "redis", "socket_timeout": "slow"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_port_from_the_environment( + json!({"port": "abc", "host": "redis"}), + &[("host", "REDIS_HOST"), ("port", "REDIS_PORT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::protocol_3( + json!({"host": "redis", "protocol": 3}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::retry_on_error_from_the_environment( + json!({"retry_on_error": true, "host": "redis"}), + &[("retry_on_error", "REDIS_RETRY_ON_ERROR")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::unix_socket_path( + json!({"host": "redis", "unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Resolves(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))), +)] +#[case::unix_socket_path_keeps_tls_settings_out_of_the_target( + json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Resolves(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))), +)] +#[case::ssl_ca_data( + json!({"host": "redis", "ssl": true, "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::ssl_ca_path_from_the_environment( + json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/etc/ssl/private-cas", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::health_check_interval_is_whole_seconds( + json!({"health_check_interval": "7", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "health_check_interval": 7.0}))), +)] +#[case::health_check_interval_zero_is_off( + json!({"health_check_interval": 0, "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::health_check_interval_zero_keeps_a_cluster_off( + json!({"health_check_interval": 0, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))), +)] +#[case::a_cluster_rejects_a_nonzero_db( + json!({"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::unix_socket_path_without_a_host_fails_like_python( + json!({"unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +fn resolves_like_python_redis( + #[case] merged: serde_json::Value, + #[case] sources: &[(&str, &str)], + #[case] redis_username: Option<&str>, + #[case] expected: Outcome, +) { + let outcome = interpret_merged(merged, sources, redis_username); + let rendered = outcome + .as_ref() + .map(view) + .map_err(|error| error.to_string()); + + match expected { + Resolves(expected) => assert_eq!(rendered, Ok(expected)), + Refuses(message) => assert_eq!(rendered, Err(message.to_owned())), + Fails(message, raised) => { + assert_eq!(rendered, Err(message.to_owned())); + assert_eq!( + outcome + .err() + .and_then(|error| error.raised_building_the_python_client()), + raised + ); + } + } +} diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md new file mode 100644 index 00000000000..963f7ff168f --- /dev/null +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -0,0 +1,13 @@ +- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`) +- An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on, or `None` when it returns a settlement the caller waits on, like `SpendCounterStore::increment_many_post_call` +- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s +- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec +- A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` +- Every family works with no Redis, answering from memory. `attach` adds the Redis tier at most once and keeps memory, as `DualCache.attach_redis_cache` does, so a host builds the store at startup and never rebuilds it to connect +- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not +- Key suffixes and wire spellings derive from strum on the family's enums. A value the family stores in Redis implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`. An argument or reply is a typed record, never a positional tuple: a frozen Python dataclass read in `python-bridge` by a `#[derive(FromPyObject)]` struct with `#[pyo3(attribute)]` fields and converted into the family's type with `From`, and a reply returned as a frozen `#[pyclass]`. The bridge reads an enum through its strum spelling, never a hand-written match +- Wire contracts are hand-written cases with expected values, inlined as the same literals in this crate's test and in `tests/unit/redis/stores/`, checked here against real Redis and there against Python's own codec. Change both sides together +- `rate_limits.rs` keeps the v3 limiter's Lua byte for byte (`rate_limits/lua.rs`) so pods on either stack share one script SHA; the admission decisions stay in the Python limiter +- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. Each op exists once and takes a `Send` (`Now`, `PreCall`, `PostCall`) picked by the caller, and the store calls `send.batch(&client)`. A store never encodes a phase in a method name. A post-call method exists only where a follow-up runs after the write settles (`token_increment_post_call`), as a thin wrapper over the shared op. `Send::Now` is an alone batch, so a script call outside a request goes out alone, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here +- Every store sends the commands main's `RedisCache` sends for that operation: `MGET` where Python reads several keys at once, `SET ... EX`, and `INCRBYFLOAT` with `EXPIRE` or `TTL`. A `MONITOR` diff against main is a parity check, and any difference is a defect +- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml new file mode 100644 index 00000000000..1dd61d76d3c --- /dev/null +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "litellm-redis-store" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-python-compat.workspace = true +litellm-redis-client.workspace = true +strum.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tracing.workspace = true + +[dev-dependencies] +litellm-circuit-breaker.workspace = true +litellm-redis-testing.workspace = true +litellm-secret-value.workspace = true +litellm-tracing.workspace = true +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/redis-store/src/auth_cache/keys.rs b/litellm-rust/crates/redis-store/src/auth_cache/keys.rs new file mode 100644 index 00000000000..feb87a40b00 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/keys.rs @@ -0,0 +1,8 @@ +/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every +/// other object under a prefixed key. +pub fn is_key_object(key: &str) -> bool { + key.len() == 64 + && key + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/mod.rs b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs new file mode 100644 index 00000000000..b5a83b5aec6 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/mod.rs @@ -0,0 +1,258 @@ +mod keys; +mod types; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_python_compat::{Value, json}; +use litellm_redis_client::{Declared, RedisBatch, commands}; + +pub use self::keys::is_key_object; +pub use self::types::Settings; +use crate::{ + Error, Number, Pending, RequestRedis, Send, + dual_cache::DualCache, + memory::{Clock, lock, system_clock}, + remote::{Remote, RemoteSlot}, + request::ready, +}; + +/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it. +const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"; + +/// The proxy's auth object cache with `enable_redis_auth_cache` on: `UserApiKeyCache`'s two +/// `DualCache` partitions, so churn in management objects cannot evict user-key objects, in front of +/// one Redis shared by every pod. Keys are the logical cache keys `UserApiKeyCache` callers already +/// use; values are what `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`. +pub struct AuthCache { + key_objects: DualCache, + objects: DualCache, + redis: Arc, + settings: Mutex, +} + +impl AuthCache { + pub fn new(remote: Option, settings: Settings) -> Self { + Self::with_clock(remote, settings, system_clock()) + } + + pub fn with_clock(remote: Option, settings: Settings, now: Clock) -> Self { + let redis = Arc::new(RemoteSlot::new(remote)); + let partition = || DualCache::new(redis.clone(), settings.memory(), now.clone()); + Self { + key_objects: partition(), + objects: partition(), + redis, + settings: Mutex::new(settings), + } + } + + pub fn has_redis(&self) -> bool { + self.redis.get().is_some() + } + + pub fn settings(&self) -> Settings { + *lock(&self.settings) + } + + /// `UserApiKeyCache.update_cache_ttl` and `update_in_memory_max_size` on both partitions, which + /// keep what they hold. + pub fn reconfigure(&self, settings: Settings) { + *lock(&self.settings) = settings; + self.key_objects.reconfigure(settings.memory()); + self.objects.reconfigure(settings.memory()); + } + + pub fn get_local(&self, key: &str) -> Result, Error> { + Ok(self.partition(key).memory().get(key)?) + } + + /// A memory write with `ttl`, or `default_in_memory_ttl` when it names none. + pub fn set_local(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { + let partition = self.partition(key); + Ok(partition + .memory() + .set(key, value, ttl.or(partition.default_in_memory_ttl()))?) + } + + pub fn delete_local(&self, key: &str) -> Result<(), Error> { + Ok(self.partition(key).memory().delete(key)?) + } + + pub fn flush_local(&self) -> Result<(), Error> { + self.key_objects.memory().flush()?; + Ok(self.objects.memory().flush()?) + } + + /// Entries held by the management-object and the user-key partitions. + pub fn len_local(&self) -> Result<(usize, usize), Error> { + Ok(( + self.objects.memory().len()?, + self.key_objects.memory().len()?, + )) + } + + /// `DualCache.async_get_cache`. A key `request` already read as missing from Redis is not read + /// again. + pub async fn get( + &self, + request: Option<&RequestRedis>, + key: &str, + ) -> Result, Error> { + self.partition(key).get(request, key, None).await + } + + /// `DualCache.async_set_cache`, writing `json.dumps(value)`. + pub async fn set(&self, key: &str, value: Value, ttl: Option) -> Result<(), Error> { + self.partition(key).set(key, value, ttl).await + } + + /// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`. + pub async fn set_str( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), Error> { + self.partition(key).set_str(key, value, ttl).await + } + + /// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks. + pub fn declare_set( + &self, + send: Send<'_>, + key: &str, + value: Value, + ttl: Option, + ) -> Result, Error> { + self.partition(key).declare_set(send, key, value, ttl) + } + + /// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks. + pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { + self.partition(key).declare_delete(send, key) + } + + /// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then + /// the rest in another, each with one TTL. + pub async fn set_many( + &self, + entries: Vec<(String, Value)>, + ttl: Option, + ) -> Result<(), Error> { + let (key_objects, others): (Vec<_>, Vec<_>) = + entries.into_iter().partition(|(key, _)| is_key_object(key)); + self.key_objects.set_pipeline(key_objects, ttl).await?; + self.objects.set_pipeline(others, ttl).await + } + + /// The auth prefetch write-back: memory first, an entry without a TTL taking + /// `default_in_memory_ttl`, then `RedisCache.async_set_cache_pipeline_with_ttls`, which encodes + /// every value before it sends anything and gives an entry without a TTL the connection's + /// default. On a request batch the `SET`s are declared and nothing is left to wait for. + /// [`Send::Now`] is one pipeline of its own rather than one round trip per entry. + pub fn write_back( + &self, + send: Send<'_>, + entries: Vec<(String, Value, Option)>, + ) -> Result, Error> { + let encoded = self + .redis + .get() + .map(|_| { + entries + .iter() + .map(|(_, value, _)| json::dumps(value)) + .collect::, _>>() + }) + .transpose(); + let writes = entries + .into_iter() + .map(|(key, value, ttl)| { + self.set_local(&key, value, ttl)?; + Ok((key, ttl)) + }) + .collect::, Error>>()?; + let (Some(remote), Some(encoded)) = (self.redis.get(), encoded?) else { + return Ok(ready(())); + }; + let sets = writes.iter().zip(encoded).map(|((key, ttl), value)| { + commands::set(key, value, ttl.unwrap_or(remote.default_ttl)) + }); + let Send::Now = send else { + let batch = send.batch(&remote.client); + sets.for_each(|set| drop(batch.add(set))); + return Ok(ready(())); + }; + let pipeline = RedisBatch::pipelined(remote.client.clone()); + let written: Vec<_> = sets.map(|set| pipeline.add(set)).collect(); + Ok(Box::pin(async move { + Declared::all(written) + .await + .into_iter() + .try_for_each(|reply| Ok(reply?)) + })) + } + + /// `UserApiKeyCache.async_delete_cache_keys`: the key-object partition and the rest as two + /// concurrent `DualCache.async_delete_cache_keys`. Both run even when one fails. + pub async fn delete(&self, keys: &[String]) -> Result<(), Error> { + let (key_objects, others): (Vec, Vec) = + keys.iter().cloned().partition(|key| is_key_object(key)); + let (key_objects, others) = tokio::join!( + self.key_objects.delete_keys(&key_objects), + self.objects.delete_keys(&others) + ); + key_objects.and(others) + } + + /// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on + /// the batch `send` picks, noting the keys Redis lacks as the request's misses. `None` without + /// Redis. Any failed read or undecodable value fails the whole read, as a failed `MGET` does. + pub fn get_many(&self, send: Send<'_>, keys: &[String]) -> Option>>> { + let client = self.redis.client()?; + let misses = send.misses(client); + let read = send.batch(client).add(commands::mget::>( + &keys.iter().map(String::as_str).collect::>(), + )); + let keys = keys.to_vec(); + Some(Box::pin(async move { + let raw = read.resolve().await?; + if let Some(misses) = misses { + misses.note( + keys.into_iter() + .zip(&raw) + .filter(|(_, raw)| raw.is_none()) + .map(|(key, _)| key), + ); + } + raw.into_iter().map(|raw| Ok(types::decode(raw)?)).collect() + })) + } + + /// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds. + /// `None` without Redis or for an absent entry. + pub async fn increment_if_present( + &self, + key: &str, + amount: Number, + ) -> Result, Error> { + let Some(client) = self.redis.client() else { + return Ok(None); + }; + let amount = litellm_redis_client::Amount::from(amount); + Ok(client + .send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount)) + .await?) + } + + fn partition(&self, key: &str) -> &DualCache { + if is_key_object(key) { + &self.key_objects + } else { + &self.objects + } + } +} diff --git a/litellm-rust/crates/redis-store/src/auth_cache/types.rs b/litellm-rust/crates/redis-store/src/auth_cache/types.rs new file mode 100644 index 00000000000..e933ac19af9 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache/types.rs @@ -0,0 +1,35 @@ +use std::time::Duration; + +use litellm_python_compat::Value; + +use crate::{DecodeError, codec::python_value, dual_cache::MemorySettings}; + +/// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Settings { + pub memory_entries: usize, + pub memory_ttl: Option, +} + +impl Default for Settings { + /// `InMemoryCache`'s 200 entries. + fn default() -> Self { + Self { + memory_entries: 200, + memory_ttl: None, + } + } +} + +impl Settings { + pub(super) fn memory(self) -> MemorySettings { + MemorySettings { + entries: Some(self.memory_entries), + ttl: self.memory_ttl, + } + } +} + +pub(super) fn decode(raw: Option>) -> Result, DecodeError> { + raw.as_deref().map(python_value).transpose() +} diff --git a/litellm-rust/crates/redis-store/src/codec.rs b/litellm-rust/crates/redis-store/src/codec.rs new file mode 100644 index 00000000000..21cde7d0b81 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/codec.rs @@ -0,0 +1,74 @@ +use litellm_python_compat::{Value, json, literal, repr::to_str}; + +use crate::{Error, FromValue, IntoValue, Number}; + +#[derive(Debug, thiserror::Error)] +pub enum DecodeError { + #[error("Redis value is not UTF-8")] + NotUtf8, + #[error("Redis value is neither JSON nor a Python literal: {0}")] + Unparsed(#[source] litellm_python_compat::Error), + #[error("Redis value has the wrong shape: {0}")] + Shape(&'static str), +} + +/// `RedisCache._get_cache_logic`: `json.loads`, falling back to `ast.literal_eval`. +pub(crate) fn python_value(raw: &[u8]) -> Result { + let text = std::str::from_utf8(raw).map_err(|_| DecodeError::NotUtf8)?; + json::loads(text).or_else(|_| literal::literal_eval(text).map_err(DecodeError::Unparsed)) +} + +/// A cached value as `RedisCache._get_cache_logic` reads it. +pub(crate) fn decode(raw: &[u8]) -> Result { + V::from_value(&python_value(raw)?) +} + +/// `json.dumps(value)`, what `RedisCache.async_set_cache` writes. +pub(crate) fn dumps(value: V) -> Result { + Ok(json::dumps(&value.into_value())?) +} + +/// `str(value)`, what the sync `RedisCache.set_cache` writes. +pub(crate) fn str(value: V) -> String { + to_str(&value.into_value()) +} + +/// A number written by Python, `shape` naming what it was meant to be when it is not one. +pub(crate) fn number(raw: &[u8], shape: &'static str) -> Result { + Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape(shape)) +} + +#[cfg(test)] +mod tests { + use litellm_python_compat::{Value, json, repr::repr}; + use rstest::rstest; + + use super::{number, python_value}; + use crate::Number; + + // Expected values are the repr of main's RedisCache._get_cache_logic on these bytes + #[rstest] + #[case::nan(b"NaN", "nan")] + #[case::infinity(b"Infinity", "inf")] + #[case::negative_infinity(b"-Infinity", "-inf")] + #[case::nan_inside_a_document(br#"{"spend": NaN}"#, "{'spend': nan}")] + #[case::python_literal_fallback(b"{'a': (1, 2)}", "{'a': (1, 2)}")] + fn reads_what_python_reads(#[case] raw: &[u8], #[case] expected: &str) { + assert_eq!(repr(&python_value(raw).unwrap()), expected); + } + + #[rstest] + #[case::nan(f64::NAN)] + #[case::infinity(f64::INFINITY)] + #[case::negative_infinity(f64::NEG_INFINITY)] + fn a_non_finite_number_reads_back_as_written(#[case] value: f64) { + let written = json::dumps(&Value::Float(value)).unwrap(); + + let read = number(written.as_bytes(), "a number").unwrap(); + + assert_eq!( + repr(&read.to_value()), + repr(&Number::Float(value).to_value()) + ); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs b/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs new file mode 100644 index 00000000000..71dfda49233 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs @@ -0,0 +1,242 @@ +use std::{ + collections::{HashMap, VecDeque}, + sync::Mutex, + time::Duration, +}; + +use litellm_redis_client::{Declared, Error as RedisError, commands}; + +use super::{DualCache, MemorySettings}; +use crate::{Error, FromValue, Send, codec, memory::lock, request::Misses}; + +/// A batch-reading family's `DualCache` settings, as the host reads them from `litellm` globals. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TierPolicy { + /// Python's `redis_batch_cache_expiry`: a key memory misses is read from Redis at most this often. + pub refresh_after: Duration, + /// Python's `default_max_redis_batch_cache_size`: how many keys the refresh throttle tracks. + pub tracked_reads: usize, + /// `litellm.default_in_memory_ttl`, which a value read from Redis keeps in memory. + pub default_in_memory_ttl: Option, +} + +impl TierPolicy { + pub(crate) fn memory(&self) -> MemorySettings { + MemorySettings { + entries: None, + ttl: self.default_in_memory_ttl, + } + } +} + +/// `DualCache.last_redis_batch_access_time` with its expiry, which only `async_batch_get_cache` reads. +pub(crate) struct BatchReads { + refresh_after: Duration, + recent: Mutex, +} + +impl BatchReads { + pub(super) fn new(policy: TierPolicy) -> Self { + Self { + refresh_after: policy.refresh_after, + recent: Mutex::new(RecentReads::new(policy.tracked_reads)), + } + } +} + +impl DualCache { + /// `DualCache.async_batch_get_cache`: memory first, then one `MGET` for the misses whose refresh + /// is due, written back to memory. A failed Redis read serves memory, as `RedisCache` swallows + /// it, and only a breaker refusal releases the refresh it reserved. + pub(crate) async fn batch_get(&self, keys: &[String]) -> Result>, Error> { + let declared = self.declare_batch_get(Send::Now, keys)?; + self.resolve_batch_get(declared).await + } + + /// The first half of `batch_get`: reserves the due refreshes and declares their `MGET` on the + /// batch `send` picks, without sending anything. + pub(crate) fn declare_batch_get( + &self, + send: Send<'_>, + keys: &[String], + ) -> Result { + let fetch = match self.redis.client() { + None => None, + Some(client) => { + let reservation = self.reserve(keys, &self.read_memory(keys)?); + (!reservation.keys.is_empty()).then(|| { + let keys: Vec<&str> = reservation.keys.iter().map(String::as_str).collect(); + let replies = send.batch(client).add(commands::mget::>(&keys)); + Fetch { + reservation, + replies, + misses: send.misses(client), + } + }) + } + }; + Ok(DeclaredRead { + keys: keys.to_vec(), + fetch, + }) + } + + /// The second half of `batch_get`. Memory is read again here, so a value written to memory + /// since the declaration wins over what Redis answered. The keys Redis lacks are noted as the + /// request's misses. An undecodable value serves memory for the whole read, where `main`'s + /// `RedisCache` returns the raw bytes and `DualCache` backfills memory with them. + pub(crate) async fn resolve_batch_get( + &self, + declared: DeclaredRead, + ) -> Result>, Error> { + let DeclaredRead { keys, fetch } = declared; + let Some(Fetch { + reservation, + replies, + misses, + }) = fetch + else { + return self.read_memory(&keys); + }; + let fetched = replies.resolve().await; + let local = self.read_memory(&keys)?; + let fetched = match fetched { + Ok(fetched) => fetched, + Err(RedisError::BreakerOpen) => { + self.release(reservation); + return Ok(local); + } + Err(error) => { + crate::failure_log::redis_failure!( + &error, + "Redis read failed, serving the memory tier" + ); + return Ok(local); + } + }; + if let Some(misses) = misses { + misses.note( + reservation + .keys + .iter() + .zip(&fetched) + .filter(|(_, raw)| raw.is_none()) + .map(|(key, _)| key.clone()), + ); + } + let decoded = match fetched + .iter() + .map(|raw| raw.as_deref().map(codec::decode).transpose()) + .collect::, _>>() + { + Ok(decoded) => decoded, + Err(error) => { + tracing::warn!(%error, "a Redis value did not decode, serving the memory tier"); + return Ok(local); + } + }; + let read = reservation + .keys + .into_iter() + .zip(decoded) + .collect::>(); + keys.iter() + .zip(local) + .map(|(key, local)| match local { + Some(value) => Ok(Some(value)), + None => self.backfill(key, read.get(key).cloned().flatten(), None), + }) + .collect() + } + + fn read_memory(&self, keys: &[String]) -> Result>, Error> { + keys.iter().map(|key| Ok(self.memory.get(key)?)).collect() + } + + fn reserve(&self, keys: &[String], local: &[Option]) -> Reservation { + let now = self.now(); + let refresh_after = self.reads.refresh_after; + let mut recent = lock(&self.reads.recent); + let due = keys + .iter() + .zip(local) + .filter(|(_, value)| value.is_none()) + .filter_map(|(key, _)| { + let previous = recent.get(key); + let due = previous.is_none_or(|read| now.saturating_sub(read) >= refresh_after); + due.then(|| { + recent.insert(key.clone(), now); + (key.clone(), previous) + }) + }) + .collect::>(); + let (keys, previous) = due.into_iter().unzip(); + Reservation { keys, previous } + } + + fn release(&self, reservation: Reservation) { + let mut recent = lock(&self.reads.recent); + for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) { + match previous { + Some(read) => recent.insert(key, read), + None => recent.remove(&key), + } + } + } +} + +/// The `MGET` `declare_batch_get` put on a batch, with the refreshes it reserved. +pub(crate) struct DeclaredRead { + keys: Vec, + fetch: Option, +} + +struct Fetch { + reservation: Reservation, + replies: Declared>>>, + misses: Option, +} + +struct Reservation { + keys: Vec, + previous: Vec>, +} + +/// Python's `LimitedSizeOrderedDict`: insertion-ordered, and any write at capacity drops the +/// oldest entry first, even when the written key is already present. +struct RecentReads { + capacity: usize, + order: VecDeque, + times: HashMap, +} + +impl RecentReads { + fn new(capacity: usize) -> Self { + Self { + capacity, + order: VecDeque::new(), + times: HashMap::new(), + } + } + + fn get(&self, key: &str) -> Option { + self.times.get(key).copied() + } + + fn insert(&mut self, key: String, time: Duration) { + if self.times.len() >= self.capacity + && let Some(oldest) = self.order.pop_front() + { + self.times.remove(&oldest); + } + if self.times.insert(key.clone(), time).is_none() { + self.order.push_back(key); + } + } + + fn remove(&mut self, key: &str) { + if self.times.remove(key).is_some() { + self.order.retain(|tracked| tracked != key); + } + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/increment.rs b/litellm-rust/crates/redis-store/src/dual_cache/increment.rs new file mode 100644 index 00000000000..705bc71b2c1 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/increment.rs @@ -0,0 +1,129 @@ +use std::time::Duration; + +use futures_util::future::join_all; +use litellm_redis_client::{Declared, Expiry, RedisBatch, commands}; + +use super::DualCache; +use crate::{Error, Number, RequestRedis, Send, failure_log::redis_failure, memory::Counter}; + +const INCREMENT_FAILED: &str = + "Redis async_increment_cache failed, falling back to in-memory result"; +const PIPELINE_FAILED: &str = + "Redis async_increment_cache_pipeline failed, falling back to in-memory result"; + +/// A `RedisPipelineIncrementOperation`. +pub(crate) struct PipelineIncrement { + pub(crate) key: String, + pub(crate) amount: Number, + pub(crate) ttl: Option, +} + +impl DualCache { + /// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL, else + /// setting `ttl` or Redis's default. `None` when memory refuses the increment, which then never + /// reaches Redis, and the memory count when Redis fails. Both only log, as there. + pub(crate) async fn increment( + &self, + key: &str, + amount: Number, + ttl: Option, + ) -> Option { + let local = self + .memory + .increment(key, amount, ttl) + .map_err(|error| redis_failure!(&Error::from(error), INCREMENT_FAILED)) + .ok()?; + let Some(remote) = self.redis.get() else { + return Some(local); + }; + let expiry = Expiry::KeepExisting(ttl.unwrap_or(remote.default_ttl)); + let increment = commands::incr_by_float(key, amount.into(), Some(expiry)); + match remote.client.send(increment).await { + Ok(total) => Some(Number::Float(total)), + Err(error) => { + redis_failure!(&error, INCREMENT_FAILED); + Some(local) + } + } + } + + /// `DualCache.async_increment_cache_pipeline`: memory, then one pipeline of `INCRBYFLOAT` with + /// `EXPIRE` per increment. `None` when memory refuses an increment, which leaves the ones before + /// it applied and sends nothing, and the memory totals when Redis fails. + pub(crate) async fn increment_pipeline( + &self, + increments: &[PipelineIncrement], + ) -> Option> { + let local = increments + .iter() + .map(|increment| { + self.memory + .increment(&increment.key, increment.amount, increment.ttl) + }) + .collect::, _>>() + .map_err(|error| redis_failure!(&Error::from(error), PIPELINE_FAILED)) + .ok()?; + let Some(client) = self.redis.client() else { + return Some(local); + }; + let pipeline = RedisBatch::pipelined(client.clone()); + let totals = increments + .iter() + .map(|increment| { + pipeline.add(commands::incr_by_float( + &increment.key, + increment.amount.into(), + increment.ttl.map(Expiry::Refresh), + )) + }) + .collect::>(); + match Declared::all(totals) + .await + .into_iter() + .collect::, _>>() + { + Ok(totals) => Some(totals.into_iter().map(Number::Float).collect()), + Err(error) => { + redis_failure!(&error, PIPELINE_FAILED); + Some(local) + } + } + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, stopping at the first + /// increment memory refuses and only logging it, then each increment on the request's post-call + /// batch, whose failure is logged once it settles. + pub(crate) fn increment_pipeline_post_call( + &self, + request: &RequestRedis, + increments: &[PipelineIncrement], + ) { + if let Err(error) = increments.iter().try_for_each(|increment| { + self.memory + .increment(&increment.key, increment.amount, increment.ttl) + .map(drop) + }) { + tracing::warn!(%error, "in-memory increment failed"); + } + let Some(batch) = self.redis.batch(Send::PostCall(request)) else { + return; + }; + let declared = increments + .iter() + .map(|increment| { + batch.add(commands::incr_by_float( + &increment.key, + increment.amount.into(), + increment.ttl.map(Expiry::KeepExisting), + )) + }) + .collect::>(); + request.spawn(async move { + join_all(declared.into_iter().map(Declared::settled)) + .await + .into_iter() + .filter_map(Result::err) + .for_each(|error| redis_failure!(&error, "deferred Redis increment failed")); + }); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/mod.rs b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs new file mode 100644 index 00000000000..b2c277a118e --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs @@ -0,0 +1,279 @@ +mod batch_read; +mod increment; +mod read; +mod write; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +pub use self::batch_read::TierPolicy; +pub(crate) use self::{batch_read::BatchReads, increment::PipelineIncrement}; +use crate::{ + memory::{Clock, Memory, lock}, + remote::RemoteSlot, +}; + +/// The memory tier's `InMemoryCache` size and the `default_in_memory_ttl` a write or a Redis +/// backfill takes when it names no TTL. `None` keeps `InMemoryCache`'s own default. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct MemorySettings { + pub(crate) entries: Option, + pub(crate) ttl: Option, +} + +/// Python's `DualCache`: a per-worker `InMemoryCache` in front of the Redis every pod shares. Each +/// method is the `DualCache` method it names. A family whose Python reached past `DualCache` to its +/// `in_memory_cache` or `redis_cache` uses `memory` and `redis` the same way. +pub(crate) struct DualCache { + memory: Memory, + redis: Arc, + default_in_memory_ttl: Mutex>, + now: Clock, + reads: R, +} + +impl DualCache { + /// `redis` may be shared with the family's other caches, as one `RedisCache` is attached to each. + pub(crate) fn new(redis: Arc, memory: MemorySettings, now: Clock) -> Self { + Self { + memory: Memory::new(memory.entries, &now), + redis, + default_in_memory_ttl: Mutex::new(memory.ttl), + now, + reads: (), + } + } + + /// A cache that batch-reads, with the `async_batch_get_cache` throttle `policy` configures. + pub(crate) fn batch_reading( + redis: Arc, + policy: TierPolicy, + now: Clock, + ) -> DualCache { + let Self { + memory, + redis, + default_in_memory_ttl, + now, + reads: (), + } = Self::new(redis, policy.memory(), now); + DualCache { + memory, + redis, + default_in_memory_ttl, + now, + reads: BatchReads::new(policy), + } + } +} + +impl DualCache { + pub(crate) fn memory(&self) -> &Memory { + &self.memory + } + + pub(crate) fn redis(&self) -> &RemoteSlot { + &self.redis + } + + pub(crate) fn now(&self) -> Duration { + (self.now)() + } + + pub(crate) fn default_in_memory_ttl(&self) -> Option { + *lock(&self.default_in_memory_ttl) + } + + /// `DualCache.update_cache_ttl` and `update_in_memory_max_size` together: both apply in place, + /// and the entries memory already holds stay. + pub(crate) fn reconfigure(&self, memory: MemorySettings) { + *lock(&self.default_in_memory_ttl) = memory.ttl; + self.memory.update_max_size(memory.entries); + } +} + +#[cfg(test)] +mod tests { + use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, + }; + + use rstest::rstest; + use tokio::runtime::Handle; + + use super::{DualCache, MemorySettings, PipelineIncrement}; + use crate::{Clock, Number, RequestRedis, memory::Counter, remote::RemoteSlot}; + + #[derive(Clone, Debug, PartialEq)] + enum Entry { + Number(Number), + Text, + } + + impl From for Entry { + fn from(number: Number) -> Self { + Self::Number(number) + } + } + + impl Counter for Entry { + fn as_number(&self) -> Option { + match self { + Self::Number(number) => Some(*number), + Self::Text => None, + } + } + } + + fn clock() -> (Clock, Arc) { + let seconds = Arc::new(AtomicU64::new(1_000)); + let read = seconds.clone(); + let clock: Clock = Arc::new(move || Duration::from_secs(read.load(Ordering::SeqCst))); + (clock, seconds) + } + + fn without_redis( + default_in_memory_ttl: Option, + now: Clock, + ) -> DualCache { + DualCache::new( + Arc::new(RemoteSlot::new(None)), + MemorySettings { + entries: None, + ttl: default_in_memory_ttl, + }, + now, + ) + } + + fn increment(key: &str, amount: Number) -> PipelineIncrement { + PipelineIncrement { + key: key.to_owned(), + amount, + ttl: None, + } + } + + // InMemoryCache's own default TTL is 600 s, which an entry outlives here only when the cache + // configured none + #[rstest] + #[case::default_in_memory_ttl(Some(Duration::from_secs(5)), None)] + #[case::in_memory_cache_default(None, Some(Number::Int(1)))] + #[tokio::test] + async fn a_set_without_a_ttl_takes_default_in_memory_ttl( + #[case] default_in_memory_ttl: Option, + #[case] after_six_seconds: Option, + ) { + let (now, seconds) = clock(); + let cache = without_redis(default_in_memory_ttl, now); + + cache.set("k", Number::Int(1), None).await.unwrap(); + seconds.fetch_add(6, Ordering::SeqCst); + + assert_eq!(cache.memory().get("k").unwrap(), after_six_seconds); + } + + #[rstest] + #[case::absent_starts_at_int_zero(None, Number::Int(2), Number::Int(2))] + #[case::int_plus_int_stays_int(Some(Number::Int(1)), Number::Int(2), Number::Int(3))] + #[case::a_float_makes_a_float(Some(Number::Int(1)), Number::Float(0.5), Number::Float(1.5))] + fn a_local_increment_adds_as_python_adds( + #[case] stored: Option, + #[case] amount: Number, + #[case] expected: Number, + ) { + let (now, _) = clock(); + let cache = without_redis::(None, now); + if let Some(stored) = stored { + cache.memory().set("k", stored, None).unwrap(); + } + + let total = cache.memory().increment("k", amount, None).unwrap(); + + assert_eq!(total, expected); + assert_eq!(cache.memory().get("k").unwrap(), Some(expected)); + } + + #[test] + fn a_local_increment_of_a_non_number_fails_and_keeps_the_entry() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + cache.memory().set("k", Entry::Text, None).unwrap(); + + let total = cache.memory().increment("k", Number::Int(1), None); + + assert!(total.is_err()); + assert_eq!(cache.memory().get("k").unwrap(), Some(Entry::Text)); + } + + // async_increment_cache and async_increment_cache_pipeline catch the in-memory failure with the + // Redis one, so the caller gets None and nothing is sent + #[tokio::test] + async fn an_increment_memory_refuses_answers_none() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + cache.memory().set("text", Entry::Text, None).unwrap(); + + let single = cache.increment("text", Number::Int(1), None).await; + let pipelined = cache + .increment_pipeline(&[ + increment("before", Number::Int(1)), + increment("text", Number::Int(1)), + ]) + .await; + + assert_eq!(single, None); + assert_eq!(pipelined, None); + assert_eq!(cache.memory().get("text").unwrap(), Some(Entry::Text)); + } + + // update_cache_ttl and update_in_memory_max_size change the settings in place, so what memory + // already holds survives and the next write takes the new TTL + #[tokio::test] + async fn reconfigure_keeps_entries_and_applies_the_new_ttl() { + let (now, seconds) = clock(); + let cache = without_redis(None, now); + cache.set("kept", Number::Int(1), None).await.unwrap(); + + cache.reconfigure(MemorySettings { + entries: Some(10), + ttl: Some(Duration::from_secs(5)), + }); + cache.set("short", Number::Int(2), None).await.unwrap(); + seconds.fetch_add(6, Ordering::SeqCst); + + assert_eq!(cache.memory().get("kept").unwrap(), Some(Number::Int(1))); + assert_eq!(cache.memory().get("short").unwrap(), None); + } + + // async_increment_cache_pipeline_post_call: InMemoryCache.async_increment_pipeline raises at the + // first entry it cannot add to, so later increments never reach memory, and DualCache only logs + #[tokio::test] + async fn a_post_call_increment_stops_at_the_first_entry_memory_refuses() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + cache.memory().set("text", Entry::Text, None).unwrap(); + + cache.increment_pipeline_post_call( + &request, + &[ + increment("before", Number::Int(1)), + increment("text", Number::Int(1)), + increment("after", Number::Int(1)), + ], + ); + + assert_eq!( + cache.memory().get("before").unwrap(), + Some(Entry::Number(Number::Int(1))) + ); + assert_eq!(cache.memory().get("after").unwrap(), None); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/read.rs b/litellm-rust/crates/redis-store/src/dual_cache/read.rs new file mode 100644 index 00000000000..14156483c6f --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/read.rs @@ -0,0 +1,55 @@ +use std::time::Duration; + +use litellm_redis_client::commands; + +use super::DualCache; +use crate::{Error, FromValue, RequestRedis, codec, failure_log::redis_failure}; + +impl DualCache { + /// `DualCache.async_get_cache`: memory, else `None` for a key `request` already read as missing + /// from Redis, else a `GET` written back to memory for `ttl` or `default_in_memory_ttl`. A failed + /// or undecodable Redis read is a miss, as there. + pub(crate) async fn get( + &self, + request: Option<&RequestRedis>, + key: &str, + ttl: Option, + ) -> Result, Error> { + if let Some(value) = self.memory.get(key)? { + return Ok(Some(value)); + } + let Some(client) = self.redis.client() else { + return Ok(None); + }; + if request.is_some_and(|request| request.misses(client).contains(key)) { + return Ok(None); + } + let value = client + .send(commands::get::>(key)) + .await + .map_err(Error::from) + .and_then(|raw| Ok(raw.as_deref().map(codec::decode).transpose()?)) + .unwrap_or_else(|error| { + redis_failure!(&error, "LiteLLM Cache: exception in async_get_cache"); + None + }); + self.backfill(key, value, ttl) + } +} + +impl DualCache { + /// `_backfill_kwargs`: a value read from Redis stays in memory for `ttl` or + /// `default_in_memory_ttl`. + pub(super) fn backfill( + &self, + key: &str, + value: Option, + ttl: Option, + ) -> Result, Error> { + if let Some(value) = &value { + self.memory + .set(key, value.clone(), ttl.or(self.default_in_memory_ttl()))?; + } + Ok(value) + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/write.rs b/litellm-rust/crates/redis-store/src/dual_cache/write.rs new file mode 100644 index 00000000000..6fcc6f32741 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/write.rs @@ -0,0 +1,178 @@ +use std::time::Duration; + +use litellm_redis_client::{Declared, RedisBatch, commands}; + +use super::DualCache; +use crate::{ + Error, IntoValue, Pending, Send, codec, + failure_log::redis_failure, + request::{pending, ready}, +}; + +/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `async_delete_cache_keys` deletes at once. +const DELETE_CHUNK: usize = 1000; + +impl DualCache { + /// `DualCache.async_set_cache`: memory, then a `SET` of `json.dumps(value)`. A Redis failure, or + /// a value `json.dumps` refuses, only logs. + pub(crate) async fn set( + &self, + key: &str, + value: V, + ttl: Option, + ) -> Result<(), Error> { + self.set_now(key, value, ttl, codec::dumps).await + } + + /// `DualCache.set_cache`, whose sync `RedisCache.set_cache` writes `str(value)`. + pub(crate) async fn set_str( + &self, + key: &str, + value: V, + ttl: Option, + ) -> Result<(), Error> { + self.set_now(key, value, ttl, |value| Ok(codec::str(value))) + .await + } + + /// `DualCache._set_on_batch`: memory now, the `SET` on the batch `send` picks. Memory and Redis + /// both take `ttl`, else `default_in_memory_ttl`, and Redis falls back to its own default. + pub(crate) fn declare_set( + &self, + send: Send<'_>, + key: &str, + value: V, + ttl: Option, + ) -> Result, Error> { + self.declare_encoded(send, key, value, ttl, codec::dumps) + } + + /// `DualCache.async_set_cache_pipeline`: memory, then one pipeline of `SET`s. A value + /// `json.dumps` refuses aborts the whole pipeline, and any failure only logs, as `RedisCache` + /// does. + pub(crate) async fn set_pipeline( + &self, + entries: Vec<(String, V)>, + ttl: Option, + ) -> Result<(), Error> { + let ttl = ttl.or(self.default_in_memory_ttl()); + let encoded = self.redis.get().map(|_| { + entries + .iter() + .map(|(_, value)| codec::dumps(value.clone())) + .collect::, Error>>() + }); + let keys = entries + .into_iter() + .map(|(key, value)| { + self.memory.set(&key, value, ttl)?; + Ok(key) + }) + .collect::, Error>>()?; + let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else { + return Ok(()); + }; + let failure = match encoded { + Err(error) => Some(error), + Ok(values) => { + let pipeline = RedisBatch::pipelined(remote.client.clone()); + let ttl = ttl.unwrap_or(remote.default_ttl); + let written = keys + .iter() + .zip(values) + .map(|(key, value)| pipeline.add(commands::set(key, value, ttl))) + .collect::>(); + Declared::all(written) + .await + .into_iter() + .find_map(Result::err) + .map(Error::from) + } + }; + if let Some(error) = failure { + redis_failure!( + &error, + "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS", + ); + } + Ok(()) + } + + async fn set_now( + &self, + key: &str, + value: V, + ttl: Option, + encode: impl FnOnce(V) -> Result, + ) -> Result<(), Error> { + if let Err(error) = self + .declare_encoded(Send::Now, key, value, ttl, encode)? + .await + { + redis_failure!(&error, "LiteLLM Cache: exception in async add_cache"); + } + Ok(()) + } + + fn declare_encoded( + &self, + send: Send<'_>, + key: &str, + value: V, + ttl: Option, + encode: impl FnOnce(V) -> Result, + ) -> Result, Error> { + let ttl = ttl.or(self.default_in_memory_ttl()); + let encoded = self.redis.get().map(|_| encode(value.clone())); + self.memory.set(key, value, ttl)?; + let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else { + return Ok(ready(())); + }; + let encoded = match encoded { + Ok(encoded) => encoded, + Err(error) => { + redis_failure!(&error, "a cache value did not encode, kept in memory only"); + return Ok(ready(())); + } + }; + if let Some(misses) = send.misses(&remote.client) { + misses.forget(key); + } + Ok(pending(send.batch(&remote.client).add(commands::set( + key, + encoded, + ttl.unwrap_or(remote.default_ttl), + )))) + } +} + +impl DualCache { + /// `DualCache.async_delete_cache_pre_call`: memory now, the `DEL` on the batch `send` picks, + /// after which the request reads the key as missing. + pub(crate) fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { + self.memory.delete(key)?; + let Some(client) = self.redis.client() else { + return Ok(ready(())); + }; + if let Some(misses) = send.misses(client) { + misses.note([key.to_owned()]); + } + let deleted = send.batch(client).add(commands::del(&[key])); + Ok(Box::pin( + async move { Ok(deleted.resolve().await.map(drop)?) }, + )) + } + + /// `DualCache.async_delete_cache_keys`: memory, then one multi-key `DEL` per chunk, in order. + pub(crate) async fn delete_keys(&self, keys: &[String]) -> Result<(), Error> { + keys.iter().try_for_each(|key| self.memory.delete(key))?; + let Some(client) = self.redis.client() else { + return Ok(()); + }; + for chunk in keys.chunks(DELETE_CHUNK) { + let chunk: Vec<&str> = chunk.iter().map(String::as_str).collect(); + client.send(commands::del(&chunk)).await?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/redis-store/src/error.rs b/litellm-rust/crates/redis-store/src/error.rs new file mode 100644 index 00000000000..fdc722855bb --- /dev/null +++ b/litellm-rust/crates/redis-store/src/error.rs @@ -0,0 +1,63 @@ +use crate::DecodeError; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Redis(#[from] litellm_redis_client::Error), + #[error(transparent)] + Memory(#[from] litellm_cache::Error), + #[error(transparent)] + Decode(#[from] DecodeError), + #[error("the value is not JSON serializable: {0}")] + Encode(#[from] litellm_python_compat::Error), +} + +impl Error { + pub fn kind(&self) -> ErrorKind { + match self { + Self::Redis(error) => error.into(), + Self::Memory(_) => ErrorKind::Memory, + Self::Decode(_) => ErrorKind::Decode, + Self::Encode(_) => ErrorKind::Encode, + } + } +} + +/// The failure class a host reports with a store error, `timeout` and `breaker_open` being the ones a +/// caller may drop. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr)] +pub enum ErrorKind { + #[strum(serialize = "breaker_open")] + BreakerOpen, + #[strum(serialize = "timeout")] + Timeout, + #[strum(serialize = "transport")] + Transport, + #[strum(serialize = "rejected")] + Rejected, + #[strum(serialize = "decode")] + Decode, + #[strum(serialize = "encode")] + Encode, + #[strum(serialize = "memory")] + Memory, +} + +impl From<&Error> for ErrorKind { + fn from(error: &Error) -> Self { + error.kind() + } +} + +impl From<&litellm_redis_client::Error> for ErrorKind { + fn from(error: &litellm_redis_client::Error) -> Self { + use litellm_redis_client::Error; + match error { + Error::BreakerOpen => Self::BreakerOpen, + Error::Transport(error) if error.is_timeout() => Self::Timeout, + Error::Transport(_) => Self::Transport, + Error::Server(_) => Self::Rejected, + Error::Decode(_) => Self::Decode, + } + } +} diff --git a/litellm-rust/crates/redis-store/src/failure_log.rs b/litellm-rust/crates/redis-store/src/failure_log.rs new file mode 100644 index 00000000000..e6bf5187cfd --- /dev/null +++ b/litellm-rust/crates/redis-store/src/failure_log.rs @@ -0,0 +1,18 @@ +/// `log_redis_failure`: an open breaker is DEBUG, anything else WARNING. The event's `redis.failure` +/// field names the class, so the host can rate-limit timeout lines with the throttle redis-py shares. +/// A macro so the event carries its caller's target, file and line, as `stacklevel=2` does in Python. +macro_rules! redis_failure { + ($error:expr, $message:expr $(,)?) => {{ + let error = $error; + let kind = $crate::ErrorKind::from(error); + let failure: &'static str = kind.into(); + match kind { + $crate::ErrorKind::BreakerOpen => { + tracing::debug!(redis.failure = failure, "{}: {}", $message, error) + } + _ => tracing::warn!(redis.failure = failure, "{}: {}", $message, error), + } + }}; +} + +pub(crate) use redis_failure; diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs new file mode 100644 index 00000000000..69600aac2a6 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -0,0 +1,24 @@ +pub mod auth_cache; +mod codec; +mod dual_cache; +mod error; +mod failure_log; +mod memory; +mod number; +pub mod rate_limits; +mod remote; +mod request; +pub mod router_cooldowns; +pub mod router_usage; +pub mod sensitive_route_pins; +pub mod spend_counters; +mod value; + +pub use codec::DecodeError; +pub use dual_cache::TierPolicy; +pub use error::{Error, ErrorKind}; +pub use memory::{Clock, system_clock}; +pub use number::Number; +pub use remote::Remote; +pub use request::{Pending, RequestRedis, Send}; +pub use value::{FromValue, IntoValue}; diff --git a/litellm-rust/crates/redis-store/src/memory.rs b/litellm-rust/crates/redis-store/src/memory.rs new file mode 100644 index 00000000000..f89d5926760 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/memory.rs @@ -0,0 +1,122 @@ +use std::{ + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::Error; +use litellm_cache_memory::InMemoryCache; + +use crate::Number; + +pub type Clock = Arc Duration + Send + Sync>; + +pub fn system_clock() -> Clock { + Arc::new(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + }) +} + +/// A store's memory tier: Python's `InMemoryCache`, with its 200 entries and 600 s TTL unless the +/// store names others, plus one lock for every read-modify-write on it. +pub struct Memory { + cache: InMemoryCache, + updates: Mutex<()>, +} + +impl Memory { + pub fn new(entries: Option, now: &Clock) -> Self { + let now = now.clone(); + Self { + cache: InMemoryCache::with_clock(entries, None, move || now()), + updates: Mutex::new(()), + } + } + + pub fn get(&self, key: &str) -> Result, Error> { + self.cache.get_cache(key) + } + + pub fn set(&self, key: &str, value: V, ttl: Option) -> Result<(), Error> { + self.cache.set_cache(key, value, ttl).map(drop) + } + + pub fn delete(&self, key: &str) -> Result<(), Error> { + self.cache.delete_cache(key) + } + + pub fn flush(&self) -> Result<(), Error> { + self.cache.flush_cache() + } + + pub fn len(&self) -> Result { + self.cache.len() + } + + pub fn expires_at(&self, key: &str) -> Result, Error> { + self.cache.expires_at(key) + } + + pub fn entries(&self) -> Result, Error> { + self.cache.entries() + } + + /// `DualCache.update_in_memory_max_size`: `None` restores the 200-entry default. + pub fn update_max_size(&self, entries: Option) { + self.cache.update_max_size_in_memory(entries); + } + + /// Read-modify-write of one entry, atomic against every other update of this tier. + pub fn update( + &self, + key: &str, + ttl: Option, + update: impl FnOnce(Option) -> Result, + ) -> Result { + let _guard = self.guard(); + let value = update(self.get(key)?)?; + self.set(key, value.clone(), ttl)?; + Ok(value) + } + + /// Held across a read and a conditional write that `update` cannot express. + pub fn guard(&self) -> MutexGuard<'_, ()> { + lock(&self.updates) + } +} + +/// A value `InMemoryCache.increment_cache` can add to. +pub(crate) trait Counter: Clone + From { + fn as_number(&self) -> Option; +} + +impl Counter for Number { + fn as_number(&self) -> Option { + Some(*self) + } +} + +impl Memory { + /// `InMemoryCache.increment_cache`: `(current or 0) + amount`, refusing an entry that is not a + /// number and leaving it in place. + pub(crate) fn increment( + &self, + key: &str, + amount: Number, + ttl: Option, + ) -> Result { + let _guard = self.guard(); + let current = match self.get(key)? { + None => Number::Int(0), + Some(value) => value.as_number().ok_or(Error::InvalidEntry(None))?, + }; + let total = current + amount; + self.set(key, V::from(total), ttl)?; + Ok(total) + } +} + +pub(crate) fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} diff --git a/litellm-rust/crates/redis-store/src/number.rs b/litellm-rust/crates/redis-store/src/number.rs new file mode 100644 index 00000000000..7f3e9f4d10b --- /dev/null +++ b/litellm-rust/crates/redis-store/src/number.rs @@ -0,0 +1,83 @@ +use std::{cmp::Ordering, ops::Add}; + +use litellm_python_compat::{BigInt, Value, json, repr}; + +/// A Python `int` or `float`, kept apart because `repr` and `json.dumps` write them differently. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Number { + Int(i64), + Float(f64), +} + +impl From for litellm_redis_client::Amount { + fn from(number: Number) -> Self { + match number { + Number::Int(value) => Self::Int(value), + Number::Float(value) => Self::Float(value), + } + } +} + +impl Number { + pub fn as_f64(self) -> f64 { + match self { + Self::Int(value) => value as f64, + Self::Float(value) => value, + } + } + + /// Python's `int(value)`, which truncates toward zero. + pub fn truncate(self) -> i64 { + match self { + Self::Int(value) => value, + Self::Float(value) => value.trunc() as i64, + } + } + + pub fn to_value(self) -> Value { + match self { + Self::Int(value) => Value::Int(BigInt::from(value)), + Self::Float(value) => Value::Float(value), + } + } + + /// `json.dumps`, which `RedisCache.async_set_cache` writes. + pub fn to_json(self) -> String { + let value = self.to_value(); + json::dumps(&value).unwrap_or_else(|_| repr::to_str(&value)) + } + + pub fn from_value(value: &Value) -> Option { + match value { + Value::Int(value) => i64::try_from(value).ok().map(Self::Int), + Value::Float(value) => Some(Self::Float(*value)), + _ => None, + } + } + + /// Python's `min` over numbers: the first of the smallest, keeping its type. + pub fn min(values: impl IntoIterator) -> Option { + values + .into_iter() + .fold(None, |smallest, value| match smallest { + Some(smallest) + if value.as_f64().partial_cmp(&smallest.as_f64()) != Some(Ordering::Less) => + { + Some(smallest) + } + _ => Some(value), + }) + } +} + +/// Python's `+`: two ints stay an int, anything with a float is a float. +impl Add for Number { + type Output = Self; + + fn add(self, other: Self) -> Self { + match (self, other) { + (Self::Int(left), Self::Int(right)) => Self::Int(left + right), + (left, right) => Self::Float(left.as_f64() + right.as_f64()), + } + } +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/keys.rs b/litellm-rust/crates/redis-store/src/rate_limits/keys.rs new file mode 100644 index 00000000000..7f5714a19ad --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/keys.rs @@ -0,0 +1,74 @@ +/// A rate-limit descriptor, `{descriptor:value}`, the hash tag every key of its window shares. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Scope { + pub descriptor: String, + pub value: String, +} + +impl Scope { + pub(super) fn key(&self, suffix: &str) -> String { + format!("{{{}:{}}}:{suffix}", self.descriptor, self.value) + } + + pub(super) fn window_key(&self) -> String { + self.key(WINDOW) + } + + pub(super) fn gauge_key(&self) -> String { + self.key(GAUGE) + } +} + +const WINDOW: &str = "window"; +const GAUGE: &str = "max_parallel_requests"; + +/// A counter's key suffix, which is also how the Python limiter names it. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)] +pub enum Kind { + #[strum(serialize = "requests")] + Requests, + #[strum(serialize = "tokens")] + Tokens, +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Counter { + pub scope: Scope, + pub kind: Kind, +} + +impl Counter { + pub(super) fn key(&self) -> String { + self.scope.key(self.kind.into()) + } +} + +/// One value of the limiter's keyspace, as both tiers address it. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum Entry { + Window(Scope), + Counter(Counter), + Gauge(Scope), +} + +impl Entry { + /// The entry the Python limiter names by its key suffix: `window`, `max_parallel_requests`, or a + /// counter's `Kind`. + pub fn new(suffix: &str, scope: Scope) -> Result { + match suffix { + WINDOW => Ok(Self::Window(scope)), + GAUGE => Ok(Self::Gauge(scope)), + kind => kind + .parse() + .map(|kind| Self::Counter(Counter { scope, kind })), + } + } + + pub(super) fn key(&self) -> String { + match self { + Self::Window(scope) => scope.window_key(), + Self::Counter(counter) => counter.key(), + Self::Gauge(scope) => scope.gauge_key(), + } + } +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/lua.rs b/litellm-rust/crates/redis-store/src/rate_limits/lua.rs new file mode 100644 index 00000000000..96451212d35 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/lua.rs @@ -0,0 +1,260 @@ +//! The limiter's Lua, byte for byte as `main`'s Python registers it, so both share one script SHA. + +pub(super) const BATCH_RATE_LIMITER: &str = r#" +local results = {} +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) + +-- Process each window/counter pair +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local increment_value = 1 + + -- Check if window exists and is valid + local window_start = redis.call('GET', window_key) + if not window_start or (now - tonumber(window_start)) >= window_size then + -- Reset window and counter + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment_value) + redis.call('EXPIRE', window_key, window_size) + redis.call('EXPIRE', counter_key, window_size) + table.insert(results, tostring(now)) -- window_start + table.insert(results, increment_value) -- counter + else + local counter = redis.call('INCR', counter_key) + -- This happens when window_key exists but counter_key doesn't (e.g., tokens key + -- created after requests key when both share the same window_key) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 then + redis.call('EXPIRE', counter_key, window_size) + end + table.insert(results, window_start) -- window_start + table.insert(results, counter) -- counter + end +end + +return results +"#; + +pub(super) const BATCH_COUNTER_READ: &str = r#"return redis.call('MGET', unpack(KEYS))"#; + +pub(super) const CHECK_AND_INCREMENT_BY_N: &str = r#" +-- Atomic check-and-increment-by-N across one or more descriptors. +-- All-or-nothing: if any descriptor would exceed its limit, no counter is +-- modified. +-- +-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied +-- timestamp so that window resets are deterministic across replicas with +-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the +-- TOCTOU window across multi-replica deployments. +-- +-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. +-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: +-- ARGV[(i-1)*4 + 1] = limit +-- ARGV[(i-1)*4 + 2] = increment +-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) +-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) +-- +-- Return on success: +-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... } +-- Return on over-limit: { 1, descriptor_index, current_counter, limit } +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local descriptor_count = #KEYS / 2 +local reset_windows = {} + +-- Pass 1: read state, validate. Abort without writing if any over limit. +local descriptor_state = {} +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[arg_base]) + local increment = tonumber(ARGV[arg_base + 1]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_start = redis.call('GET', window_key) + local window_expired = (not window_start) or + ((now - tonumber(window_start)) >= window_size) + + local current_counter + if window_expired then + current_counter = 0 + else + current_counter = tonumber(redis.call('GET', counter_key) or 0) + end + + local blocked + if increment > 0 then + blocked = current_counter + increment > limit + else + blocked = current_counter >= limit + end + if blocked then + return { 1, i, current_counter, limit } + end + + descriptor_state[i] = { window_expired, current_counter, window_start } +end + +-- Pass 2: all checks passed. Apply increments. +local results = { 0 } +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_expired = descriptor_state[i][1] + local active_window_start + + if window_expired then + active_window_start = now + if not reset_windows[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset_windows[window_key] = true + end + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + else + active_window_start = tonumber(descriptor_state[i][3]) + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, new_counter) + end + table.insert(results, active_window_start) +end + +return results +"#; + +pub(super) const WINDOW_GUARDED_TOKEN_INCREMENT: &str = r#" +local results = {} +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local arg_base = ((i - 1) / 2) * 3 + 1 + local expected_window_start = ARGV[arg_base] + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local active_window_start = redis.call('GET', window_key) + + if active_window_start and active_window_start == expected_window_start then + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, 1) + table.insert(results, new_counter) + else + table.insert(results, 0) + table.insert(results, tonumber(redis.call('GET', counter_key) or 0)) + end +end +return results +"#; + +pub(super) const PARALLEL_ACQUIRE: &str = r#" +-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge. +-- Each gauge key is a sorted set of per-request slot ids scored by acquire +-- time (Redis server clock). In-flight requests are counted by ZCARD after +-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM +-- counters the gauge is never reset while requests are in flight, a +-- rejected request never occupies a slot, and a slot leaked by a crashed +-- worker self-heals after the slot TTL even under continuous traffic. +-- +-- KEYS: one gauge zset key per descriptor. +-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id). +-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +for i = 1, #KEYS do + local limit = tonumber(ARGV[(i - 1) * 3 + 1]) + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl) + local in_flight = redis.call('ZCARD', KEYS[i]) + if in_flight + 1 > limit then + return { 1, i, in_flight, limit } + end +end +local results = { 0 } +for i = 1, #KEYS do + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + local slot_id = ARGV[(i - 1) * 3 + 3] + redis.call('ZADD', KEYS[i], now, slot_id) + redis.call('EXPIRE', KEYS[i], slot_ttl) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const PARALLEL_RELEASE: &str = r#" +-- Release one slot per gauge key by removing this request's slot id. +-- ZREM of an absent member (or key) is a no-op, so a release without a +-- matching acquire (proxy-side rejection, double-fired callback, slot +-- already expired) can never free a slot owned by another request. +-- KEYS: gauge zset keys. ARGV: per-key slot_id. +-- Returns the remaining in-flight count per key. +local results = {} +for i = 1, #KEYS do + redis.call('ZREM', KEYS[i], ARGV[i]) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const PARALLEL_COUNT: &str = r#" +-- Read the current in-flight count per gauge key (prunes expired slots +-- first so leaked slots do not inflate the reading). +-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local results = {} +for i = 1, #KEYS do + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i])) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const TOKEN_INCREMENT: &str = r#" +local results = {} + +-- Process each key/increment_value/ttl triplet +for i = 1, #KEYS do + local key = KEYS[i] + local increment_value = tonumber(ARGV[i * 2 - 1]) + local ttl_seconds = tonumber(ARGV[i * 2]) + + -- Increment the value + local new_value = redis.call('INCRBYFLOAT', key, increment_value) + + -- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL + -- ttl_seconds can be 0 (no TTL) or positive (set TTL) + if ttl_seconds and ttl_seconds > 0 then + local current_ttl = redis.call('TTL', key) + if current_ttl == -1 then + redis.call('EXPIRE', key, ttl_seconds) + end + end + + table.insert(results, new_value) +end + +return results +"#; diff --git a/litellm-rust/crates/redis-store/src/rate_limits/mod.rs b/litellm-rust/crates/redis-store/src/rate_limits/mod.rs new file mode 100644 index 00000000000..71b67d2ecc0 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/mod.rs @@ -0,0 +1,304 @@ +mod keys; +mod lua; +mod scripts; +mod types; + +use std::{sync::Arc, time::Duration}; + +use litellm_redis_client::{Declared, RedisBatch, RedisClient, Slot}; + +pub use self::keys::{Counter, Entry, Kind, Scope}; +pub use self::types::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, + Increment, Local, WindowCount, +}; +use crate::{ + Error, Number, Pending, RequestRedis, Send, + dual_cache::{DualCache, MemorySettings, PipelineIncrement}, + memory::{Clock, system_clock}, + remote::{Remote, RemoteSlot}, + request::pending, +}; + +/// The v3 rate limiter's keyspace: its memory tier, its Lua scripts and the plain counter writes +/// `DualCache` made for it on `main`. The admission logic stays with the caller. A script needs +/// Redis, so without one it declares nothing and answers `None`. +pub struct RateLimits { + cache: DualCache, +} + +impl RateLimits { + /// `backfill_ttl` is the `default_in_memory_ttl` of the `DualCache` this tier replaces. An + /// increment without a TTL takes the remote's default. + pub fn new(remote: Option, backfill_ttl: Option) -> Self { + Self::with_clock(remote, backfill_ttl, system_clock()) + } + + pub fn with_clock(remote: Option, backfill_ttl: Option, now: Clock) -> Self { + Self { + cache: DualCache::new( + Arc::new(RemoteSlot::new(remote)), + MemorySettings { + entries: None, + ttl: backfill_ttl, + }, + now, + ), + } + } + + pub fn has_redis(&self) -> bool { + self.cache.redis().get().is_some() + } + + /// `_group_keys_by_hash_tag`: indices of `scopes` sharing a cluster slot, in first-seen order, + /// or one group when the client is not a cluster and grouping is not forced. + pub fn slot_groups(&self, scopes: &[Scope], force: bool) -> Vec> { + let by_slot = force + || self + .cache + .redis() + .client() + .is_some_and(RedisClient::is_cluster); + if !by_slot { + return vec![(0..scopes.len()).collect()]; + } + scopes + .iter() + .enumerate() + .fold( + Vec::<(Slot, Vec)>::new(), + |mut groups, (index, scope)| { + let slot = Slot::for_key(scope.window_key()); + match groups.iter_mut().find(|(seen, _)| *seen == slot) { + Some((_, members)) => members.push(index), + None => groups.push((slot, vec![index])), + } + groups + }, + ) + .into_iter() + .map(|(_, members)| members) + .collect() + } + + pub fn get_local(&self, entry: &Entry) -> Result, Error> { + Ok(self.cache.memory().get(&entry.key())?) + } + + pub fn set_local( + &self, + entry: &Entry, + value: Local, + ttl: Option, + ) -> Result<(), Error> { + Ok(self.cache.memory().set(&entry.key(), value, ttl)?) + } + + /// `DualCache.async_get_cache`, backfilling memory for `backfill_ttl` or the store's default. + /// A key `request` already read as missing from Redis is not read again. + pub async fn get( + &self, + request: Option<&RequestRedis>, + entry: &Entry, + backfill_ttl: Option, + ) -> Result, Error> { + self.cache.get(request, &entry.key(), backfill_ttl).await + } + + pub fn sliding_window( + &self, + send: Send<'_>, + counters: &[Counter], + now: i64, + window_size: i64, + ) -> Option>> { + Some(pending(scripts::sliding_window( + &self.batch(send)?, + counters, + now, + window_size, + ))) + } + + pub fn read_counters( + &self, + send: Send<'_>, + counters: &[Counter], + ) -> Option>>> { + Some(pending(scripts::read_counters( + &self.batch(send)?, + counters, + ))) + } + + pub fn check_and_increment( + &self, + send: Send<'_>, + counters: &[AtomicCounter], + ) -> Option> { + Some(pending(scripts::check_and_increment( + &self.batch(send)?, + counters, + ))) + } + + pub fn guarded_increment( + &self, + send: Send<'_>, + increment: &GuardedIncrement, + ) -> Option> { + Some(pending(scripts::guarded_increment( + &self.batch(send)?, + increment, + ))) + } + + pub fn parallel_acquire( + &self, + send: Send<'_>, + gauges: &[Gauge], + slot_ttl: Duration, + slot_id: &str, + ) -> Option> { + Some(pending(scripts::parallel_acquire( + &self.batch(send)?, + gauges, + slot_ttl, + slot_id, + ))) + } + + pub fn parallel_release( + &self, + send: Send<'_>, + gauges: &[Scope], + slot_id: &str, + ) -> Option>> { + Some(pending(self.declare_release(send, gauges, slot_id)?)) + } + + pub fn parallel_count( + &self, + send: Send<'_>, + gauges: &[Scope], + slot_ttl: Duration, + ) -> Option>> { + Some(pending(scripts::parallel_count( + &self.batch(send)?, + gauges, + slot_ttl, + ))) + } + + pub fn token_increment( + &self, + send: Send<'_>, + increments: &[Increment], + ) -> Option>> { + Some(pending(self.declare_token_increment(send, increments)?)) + } + + /// `TOKEN_INCREMENT_SCRIPT` on the request's post-call batch. A failed script falls back to + /// `increment_many` for these increments once the batch has gone out. `false` without Redis. + pub fn token_increment_post_call( + self: &Arc, + request: &RequestRedis, + increments: Vec, + ) -> bool { + let Some(declared) = self.declare_token_increment(Send::PostCall(request), &increments) + else { + return false; + }; + let store = Arc::clone(self); + request.spawn(async move { + let Err(error) = declared.settled().await else { + return; + }; + crate::failure_log::redis_failure!( + &error, + "TTL preservation failed, falling back to regular pipeline" + ); + store.increment_many(&increments).await; + }); + true + } + + /// `PARALLEL_RELEASE_SCRIPT` on the request's post-call batch; the caller has already freed + /// the slot in memory. `false` without Redis. + pub fn parallel_release_post_call( + &self, + request: &RequestRedis, + gauges: &[Scope], + slot_id: &str, + ) -> bool { + let Some(declared) = self.declare_release(Send::PostCall(request), gauges, slot_id) else { + return false; + }; + request.spawn(async move { + if let Err(error) = declared.settled().await { + crate::failure_log::redis_failure!( + &error, + "parallel_release_script failed, the slot stays released in memory only" + ); + } + }); + true + } + + /// `DualCache.async_increment_cache`. A Redis failure answers with the memory count, and an + /// entry memory cannot add to with `None`. + pub async fn increment(&self, increment: &Increment) -> Option { + self.cache + .increment(&increment.counter.key(), increment.amount, increment.ttl) + .await + } + + /// `DualCache.async_increment_cache_pipeline`, whose failures only log. + pub async fn increment_many(&self, increments: &[Increment]) { + self.cache.increment_pipeline(&pipeline(increments)).await; + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, the Redis increments on + /// the request's post-call batch. + pub fn increment_many_post_call(&self, request: &RequestRedis, increments: &[Increment]) { + self.cache + .increment_pipeline_post_call(request, &pipeline(increments)); + } + + fn declare_release( + &self, + send: Send<'_>, + gauges: &[Scope], + slot_id: &str, + ) -> Option>> { + Some(scripts::parallel_release( + &self.batch(send)?, + gauges, + slot_id, + )) + } + + fn declare_token_increment( + &self, + send: Send<'_>, + increments: &[Increment], + ) -> Option>> { + Some(scripts::token_increment(&self.batch(send)?, increments)) + } + + fn batch(&self, send: Send<'_>) -> Option { + self.cache.redis().batch(send) + } +} + +fn pipeline(increments: &[Increment]) -> Vec { + increments + .iter() + .map(|increment| PipelineIncrement { + key: increment.counter.key(), + amount: increment.amount, + ttl: increment.ttl, + }) + .collect() +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs new file mode 100644 index 00000000000..04bdad38268 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs @@ -0,0 +1,303 @@ +use std::{sync::LazyLock, time::Duration}; + +use litellm_python_compat::repr::float_repr; +use litellm_redis_client::{ + Command, Declared, Error, RedisBatch, Script, Value, + commands::{self, unexpected_reply}, +}; + +use super::{ + keys::{Counter, Scope}, + lua, + types::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, + Increment, WindowCount, + }, +}; +use crate::Number; + +static BATCH_RATE_LIMITER: LazyLock