diff --git a/.circleci/config.yml b/.circleci/config.yml index 9b976462b13..8c75bdc5f33 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3051,10 +3051,19 @@ jobs: - ui/litellm-dashboard/node_modules - run: name: Build UI from source + # Prior version used `cp -r out/ ../../litellm/proxy/_experimental/out/`. + # GNU cp (used on CircleCI's Ubuntu image) interprets that as "copy the + # source directory as a child of the destination" when the destination + # already exists — silently creating `_experimental/out/out/` instead of + # replacing the served bundle. The proxy continued serving whatever was + # checked into `_experimental/out/*`, so this job was effectively testing + # the pre-build bundle on every run. Replace-and-move guarantees the + # freshly built bundle is what the proxy actually serves. command: | cd ui/litellm-dashboard npm run build - cp -r out/ ../../litellm/proxy/_experimental/out/ + rm -rf ../../litellm/proxy/_experimental/out + mv out ../../litellm/proxy/_experimental/out # Restructure HTML so extensionless routes work (login.html -> login/index.html) find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html" diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 87e7e17feb7..a631a7c3005 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -31,8 +31,15 @@ jobs: test-path: "tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py" workers: 8 timeout: 20 + # test_proxy_utils.py is large (168+ parametrized tests) — run it on its + # own matrix so --dist=loadscope doesn't pin all of it to a single xdist + # worker and push the "remaining" group past the job timeout. + - test-group: proxy-utils + test-path: "tests/proxy_unit_tests/test_proxy_utils.py" + workers: 8 + timeout: 20 - test-group: remaining - test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py" + test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py --ignore=tests/proxy_unit_tests/test_proxy_utils.py" workers: 8 timeout: 30 uses: ./.github/workflows/_test-unit-services-base.yml diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py index 29101bf9505..feec4046ee1 100644 --- a/ci_cd/run_migration.py +++ b/ci_cd/run_migration.py @@ -1,22 +1,231 @@ +import argparse import os -import subprocess -from pathlib import Path -from datetime import datetime -import testing.postgresql +import re import shutil +import subprocess +import sys +from datetime import datetime +from pathlib import Path + +import testing.postgresql -def create_migration(migration_name: str = None): +DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE) +DEFAULT_BASE_BRANCH = "litellm_internal_staging" + + +def _find_destructive_statements(sql: str) -> list: + """Return SQL lines containing DROP COLUMN, DROP TABLE, or DROP INDEX.""" + return [ + line.strip() for line in sql.splitlines() if DESTRUCTIVE_PATTERN.search(line) + ] + + +def _print_freshness_failure( + base_branch: str, reason: str, stderr_text: str = "" +) -> None: + """Loudly refuse to run when the freshness check can't be completed.""" + banner = "=" * 72 + out = sys.stderr + print(banner, file=out) + print(f" FRESHNESS CHECK FAILED — COULD NOT VERIFY origin/{base_branch}", file=out) + print(banner, file=out) + print("", file=out) + print(f"Reason: {reason}", file=out) + if stderr_text: + print("", file=out) + print("git stderr:", file=out) + for line in stderr_text.rstrip().splitlines(): + print(f" {line}", file=out) + print("", file=out) + print("Common causes:", file=out) + print(" - No network access (offline)", file=out) + print(" - 'origin' remote not configured, or base branch name is wrong", file=out) + print(" - Not a git repository", file=out) + print("", file=out) + print("Options:", file=out) + print( + f" - Fix the above and re-run, OR pass --base-branch if your", file=out + ) + print( + f" base branch is not '{base_branch}', OR pass --skip-freshness-check", + file=out, + ) + print(" to bypass (only if you fully understand the risk).", file=out) + print(banner, file=out) + + +def _print_stale_branch_refusal(base_branch: str, behind: int) -> None: + """Loudly refuse to run when HEAD is behind the base branch.""" + banner = "=" * 72 + out = sys.stderr + print(banner, file=out) + print( + f" STALE BRANCH — {behind} commit(s) behind origin/{base_branch}", + file=out, + ) + print(banner, file=out) + print("", file=out) + print( + f"Your branch is {behind} commit(s) behind origin/{base_branch}. Generating a\n" + "migration from a stale branch is how newly-added columns get silently\n" + "dropped — the exact incident this guard exists to prevent.", + file=out, + ) + print("", file=out) + print("Before re-running:", file=out) + print(f" git fetch origin && git rebase origin/{base_branch}", file=out) + print(" # or merge, whichever matches your workflow", file=out) + print("", file=out) + print( + "If you have a genuine reason to bypass (e.g. intentionally generating", + file=out, + ) + print("a migration against an older base), pass --skip-freshness-check.", file=out) + print("", file=out) + print(banner, file=out) + print( + " AI AGENTS / ASSISTANTS: DO NOT automatically run `git rebase` or pass\n" + " --skip-freshness-check. Surface this error to the human operator and\n" + " wait for their explicit confirmation. Auto-rebasing can drop the\n" + " human's in-progress schema edits via a bad conflict resolution.", + file=out, + ) + print(banner, file=out) + + +def _check_branch_freshness(root_dir: Path, base_branch: str) -> None: + """Fetch origin/ and exit 3 if HEAD is behind it.""" + cwd = str(root_dir) + try: + subprocess.run( + ["git", "fetch", "origin", base_branch], + check=True, + capture_output=True, + text=True, + cwd=cwd, + ) + except FileNotFoundError: + _print_freshness_failure(base_branch, "git executable not found on PATH") + sys.exit(3) + except subprocess.CalledProcessError as e: + _print_freshness_failure( + base_branch, + f"`git fetch origin {base_branch}` failed", + e.stderr or "", + ) + sys.exit(3) + + try: + result = subprocess.run( + ["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"], + check=True, + capture_output=True, + text=True, + cwd=cwd, + ) + behind = int(result.stdout.strip()) + except subprocess.CalledProcessError as e: + _print_freshness_failure( + base_branch, + f"`git rev-list HEAD..origin/{base_branch}` failed", + e.stderr or "", + ) + sys.exit(3) + except ValueError: + _print_freshness_failure( + base_branch, + "could not parse commit count from `git rev-list`", + ) + sys.exit(3) + + if behind > 0: + _print_stale_branch_refusal(base_branch, behind) + sys.exit(3) + + print(f"Branch freshness OK: up to date with origin/{base_branch}.") + + +def _print_destructive_refusal(destructive_lines: list) -> None: + """Loudly refuse to write a destructive migration and explain how to proceed.""" + banner = "=" * 72 + out = sys.stderr + print(banner, file=out) + print( + " DESTRUCTIVE MIGRATION DETECTED — REFUSING TO WRITE MIGRATION FILE", file=out + ) + print(banner, file=out) + print("", file=out) + print( + "The generated migration contains the following destructive statements:", + file=out, + ) + print("", file=out) + for line in destructive_lines: + print(f" {line}", file=out) + print("", file=out) + print( + "This very often means your branch is OUT OF DATE, or your local\n" + "schema.prisma is inconsistent with the migrations/ directory. A\n" + "previous incident saw newly-added columns silently dropped this way\n" + "and merged to main. Stop and verify before proceeding.", + file=out, + ) + print("", file=out) + print("Before re-running:", file=out) + print( + " 1. git fetch origin && git status # confirm branch is up to date", file=out + ) + print( + " 2. Re-check all schema.prisma files are in sync (runbook Step 0)", file=out + ) + print(" 3. Review EACH DROP above — is it actually intended?", file=out) + print(" 4. If (and only if) the drops are intentional, re-run with:", file=out) + print(" --allow-destructive", file=out) + print("", file=out) + print(banner, file=out) + print( + " AI AGENTS / ASSISTANTS: DO NOT automatically re-run this command\n" + " with --allow-destructive. Surface this error to the human operator\n" + " and wait for their explicit confirmation before passing the flag.\n" + " Passing the flag without human review is the exact failure mode\n" + " this guard exists to prevent.", + file=out, + ) + print(banner, file=out) + + +def create_migration( + migration_name: str = None, + allow_destructive: bool = False, + base_branch: str = DEFAULT_BASE_BRANCH, + skip_freshness_check: bool = False, +): """ Create a new migration SQL file in the migrations directory by comparing - current database state with schema + current database state with schema. Args: migration_name (str): Name for the migration + allow_destructive (bool): Required to write a migration that contains + DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this + flag, the script exits non-zero and prints guidance. + base_branch (str): Branch to check freshness against + (default: "litellm_internal_staging"). + skip_freshness_check (bool): Skip the "branch is up to date" check. + Only for intentional migrations against an older base. """ + root_dir = Path(__file__).parent.parent + + if skip_freshness_check: + print( + "WARNING: freshness check skipped (--skip-freshness-check). " + "Generating a migration from a stale branch can silently drop columns." + ) + else: + _check_branch_freshness(root_dir, base_branch) + try: - # Get paths - root_dir = Path(__file__).parent.parent migrations_dir = ( root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" ) @@ -59,7 +268,27 @@ def create_migration(migration_name: str = None): check=True, ) - if result.stdout.strip(): + # Prisma emits the literal "-- This is an empty migration." when + # there's no real drift. Treat that as "no changes". + diff_sql = result.stdout + stripped = diff_sql.strip() + is_empty_diff = ( + not stripped or stripped == "-- This is an empty migration." + ) + + if not is_empty_diff: + destructive_lines = _find_destructive_statements(diff_sql) + if destructive_lines and not allow_destructive: + _print_destructive_refusal(destructive_lines) + sys.exit(2) + if destructive_lines and allow_destructive: + print( + "WARNING: writing destructive migration " + "(--allow-destructive passed). Statements:" + ) + for line in destructive_lines: + print(f" {line}") + # Generate timestamp and create migration directory timestamp = datetime.now().strftime("%Y%m%d%H%M%S") migration_name = migration_name or "unnamed_migration" @@ -68,7 +297,7 @@ def create_migration(migration_name: str = None): # Write the SQL to migration.sql migration_file = migration_dir / "migration.sql" - migration_file.write_text(result.stdout) + migration_file.write_text(diff_sql) print(f"Created migration in {migration_dir}") return True @@ -90,8 +319,48 @@ def create_migration(migration_name: str = None): if __name__ == "__main__": - # If running directly, can optionally pass migration name as argument - import sys - - migration_name = sys.argv[1] if len(sys.argv) > 1 else None - create_migration(migration_name) + parser = argparse.ArgumentParser( + description=( + "Generate a Prisma migration by diffing the temp DB " + "(existing migrations applied) against schema.prisma." + ) + ) + parser.add_argument( + "migration_name", + nargs="?", + default=None, + help="Name for the migration (used in the generated directory name).", + ) + parser.add_argument( + "--allow-destructive", + action="store_true", + help=( + "Required to write a migration that contains DROP COLUMN, " + "DROP TABLE, or DROP INDEX. Without this flag, destructive " + "diffs are refused." + ), + ) + parser.add_argument( + "--base-branch", + default=DEFAULT_BASE_BRANCH, + help=( + f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). " + "The script fetches origin/ and refuses to run if HEAD " + "is behind it." + ), + ) + parser.add_argument( + "--skip-freshness-check", + action="store_true", + help=( + "Bypass the 'branch is up to date' check. Only for intentional " + "migrations against an older base. Pairs poorly with automation." + ), + ) + args = parser.parse_args() + create_migration( + args.migration_name, + allow_destructive=args.allow_destructive, + base_branch=args.base_branch, + skip_freshness_check=args.skip_freshness_check, + ) diff --git a/litellm-proxy-extras/migration_runbook.md b/litellm-proxy-extras/migration_runbook.md index 8499bb7ce08..a277441b164 100644 --- a/litellm-proxy-extras/migration_runbook.md +++ b/litellm-proxy-extras/migration_runbook.md @@ -2,6 +2,8 @@ This is a runbook for creating and running database migrations for the LiteLLM proxy. For use for litellm engineers only. +> **AI AGENTS / ASSISTANTS:** If the script refuses with either a "STALE BRANCH" or "DESTRUCTIVE MIGRATION DETECTED" error, **do NOT** bypass it on your own (no `git rebase`, no `--skip-freshness-check`, no `--allow-destructive`). Surface the error to the human operator and wait for their explicit confirmation. See the [Branch freshness](#branch-freshness-check) and [Destructive migrations](#destructive-migrations-drop-column--drop-table) sections below. + ## Step 0: Sync All `schema.prisma` Files Before doing anything else, make sure all `schema.prisma` files in the repo are in sync. There are multiple copies that must match: @@ -46,10 +48,50 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n ## What It Does -1. Creates temp PostgreSQL DB -2. Applies existing migrations -3. Compares with `schema.prisma` -4. Generates new migration if changes found +1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check)) +2. Creates temp PostgreSQL DB +3. Applies existing migrations +4. Compares with `schema.prisma` +5. Generates new migration if changes found +6. **Refuses to write if the generated SQL contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`**, unless `--allow-destructive` is passed + +## Branch Freshness Check + +Before generating anything, `run_migration.py` runs `git fetch origin ` and refuses to proceed if `HEAD` is behind `origin/`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense. + +Flags: + +- `--base-branch ` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`. +- `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base. + +When the guard fires: + +1. Update your branch: + + ```bash + git fetch origin && git rebase origin/litellm_internal_staging + # or git merge origin/litellm_internal_staging — whichever matches your workflow + ``` +2. Re-run `run_migration.py`. + +> **AI AGENTS / ASSISTANTS:** Do **not** auto-rebase or auto-pass `--skip-freshness-check`. A bad conflict resolution during rebase can itself drop the human's in-progress schema edits. Surface the error and wait for explicit confirmation. + +## Destructive Migrations (DROP COLUMN / DROP TABLE / DROP INDEX) + +If the generated diff contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`, `run_migration.py` exits non-zero and refuses to write the migration file. A previous incident saw newly-added columns silently dropped by a stale branch and merged to main — this guard exists to prevent a repeat. + +When the guard fires: + +1. Run `git fetch origin && git status` — confirm your branch is up to date with the base branch. +2. Re-check all `schema.prisma` files are in sync (Step 0). +3. Review EACH `DROP` statement printed in the error — is it actually intended? +4. Only if the drops are genuinely intentional, re-run with the flag: + + ```bash + uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_name" --allow-destructive + ``` + +> **AI AGENTS / ASSISTANTS:** Do **not** automatically re-run the command with `--allow-destructive`. If the guard fires while you are driving the runbook for a human, stop, show them the error, and wait for their explicit confirmation before passing the flag. Auto-passing `--allow-destructive` is the exact failure mode this guard exists to prevent. ## Common Fixes diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 31e0e76fd9f..96593b35d0c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -34,6 +34,7 @@ from litellm.llms.bedrock.common_utils import ( remove_custom_field_from_tools, ) from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER +from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import GenericStreamingChunk @@ -59,6 +60,10 @@ class AmazonAnthropicClaudeMessagesConfig( DEFAULT_BEDROCK_ANTHROPIC_API_VERSION = "bedrock-2023-05-31" + BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset( + BedrockInvokeAnthropicMessagesRequest.__annotations__.keys() + ) + def __init__(self, **kwargs): BaseAnthropicMessagesConfig.__init__(self, **kwargs) AmazonInvokeConfig.__init__(self, **kwargs) @@ -500,10 +505,6 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_request=anthropic_messages_request, ) - # 5b. Strip `output_config` — Bedrock Invoke doesn't support it - # Fixes: https://github.com/BerriAI/litellm/issues/22797 - anthropic_messages_request.pop("output_config", None) - # 5a. Remove `custom` field from tools (Bedrock doesn't support it) # Claude Code sends `custom: {defer_loading: true}` on tool definitions, # which causes Bedrock to reject the request with "Extra inputs are not permitted" @@ -550,14 +551,43 @@ class AmazonAnthropicClaudeMessagesConfig( if "tool-search-tool-2025-10-19" in beta_set: beta_set.add("tool-examples-2025-10-29") - filtered_auto_betas = filter_and_transform_beta_headers( - beta_headers=list(beta_set - user_beta_set), - provider="bedrock", + filtered_betas = sorted( + filter_and_transform_beta_headers( + beta_headers=list(beta_set), + provider="bedrock", + ) ) - filtered_betas = sorted(user_beta_set.union(set(filtered_auto_betas))) + + dropped_user_betas = sorted( + b + for b in user_beta_set + if not filter_and_transform_beta_headers([b], provider="bedrock") + ) + if dropped_user_betas: + verbose_logger.warning( + "Bedrock Invoke: dropping unsupported anthropic-beta values " + "from client headers: %s. Bedrock has no mapping entry for " + "these; forwarding them would cause a 400.", + dropped_user_betas, + ) + if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas + # 7. Final safety net: filter top-level fields to the Bedrock Invoke allowlist. + # Catches Anthropic-only extensions (context_management, output_config, speed, + # mcp_servers, ...) and any future additions Claude Code may start sending. + allowed = self.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS + stripped = sorted(k for k in anthropic_messages_request if k not in allowed) + if stripped: + verbose_logger.debug( + "Bedrock Invoke: stripping unsupported top-level request fields: %s", + stripped, + ) + anthropic_messages_request = { + k: v for k, v in anthropic_messages_request.items() if k in allowed + } + return anthropic_messages_request def get_async_streaming_response_iterator( diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 6830d95d36f..9ffb52ef88d 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -997,3 +997,47 @@ class BedrockToolBlock(TypedDict, total=False): toolSpec: Optional[ToolSpecBlock] systemTool: Optional[SystemToolBlock] # For Nova grounding cachePoint: Optional[CachePointBlock] + + +class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): + """ + Top-level request body accepted by AWS Bedrock `InvokeModel` / + `InvokeModelWithResponseStream` when calling an Anthropic Claude model with + the Messages API format. The LiteLLM /v1/messages → Bedrock Invoke + transformation filters outgoing requests to the keys of this TypedDict; any + other field (Anthropic-only extension, internal metadata, future addition) + is dropped before signing so Bedrock doesn't 400 with + "Extra inputs are not permitted". + + Reference: + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + + Editing this type is the single source of truth — the runtime allowlist in + `AmazonAnthropicClaudeMessagesConfig.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS` + is derived from `__annotations__`, and a test asserts the resolved set + exactly, so any edit forces a conscious review. + + Value types are intentionally loose (`list`, `dict`) — this type exists to + pin the allowed field names, not to validate nested structure. + """ + + # Required by Bedrock + anthropic_version: str + max_tokens: int + messages: list + + # Documented optional fields + anthropic_beta: List[str] + system: object # str or list[TextBlock] + stop_sequences: List[str] + temperature: float + top_p: float + top_k: int + tools: list + tool_choice: dict + + # `thinking` is required for Opus 4.5 / Sonnet 4 extended thinking, + # `metadata` is part of the common Anthropic Messages API shape. + thinking: dict + metadata: dict diff --git a/pyproject.toml b/pyproject.toml index d5d238473b1..ef06a628fb6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -208,7 +208,7 @@ build-backend = "uv_build" [tool.uv] default-groups = ["dev"] -required-version = "==0.10.9" +required-version = ">=0.10.9" exclude-newer = "3 days" [tool.uv.sources] diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py index 1421700c9a8..a0326f64ed7 100644 --- a/tests/proxy_unit_tests/conftest.py +++ b/tests/proxy_unit_tests/conftest.py @@ -1,48 +1,143 @@ # conftest.py -import importlib +import asyncio +import copy +import inspect import os import sys +import warnings import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path + import litellm +import litellm.proxy.proxy_server + + +# Top-level assignments of these types are the ones importlib.reload(litellm) +# would have effectively reset. We snapshot them at conftest import time and +# deep-copy the snapshot back before every test. +_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes) + + +def _snapshot_mutable_state(module): + """Capture a per-module snapshot of primitive and collection attributes.""" + snapshot = {} + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception as exc: + warnings.warn( + f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}", + stacklevel=2, + ) + continue + if value is None or isinstance(value, _SNAPSHOT_TYPES): + try: + snapshot[attr] = copy.deepcopy(value) + except Exception as exc: + warnings.warn( + f"conftest: could not snapshot {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + return snapshot + + +def _restore_mutable_state(module, snapshot): + for attr, default in snapshot.items(): + try: + setattr(module, attr, copy.deepcopy(default)) + except Exception as exc: + warnings.warn( + f"conftest: could not restore {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +def _collect_flushable_caches(): + """Return (module, attr) pairs whose values expose flush_cache().""" + targets = [] + for module in (litellm, litellm.proxy.proxy_server): + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception: + continue + # Only instances — a class reference has an unbound flush_cache + # that can't be called without a self argument. + if inspect.isclass(value) or inspect.ismodule(value): + continue + if callable(getattr(value, "flush_cache", None)): + targets.append((module, attr)) + return targets + + +def _flush_caches(targets): + for module, attr in targets: + try: + value = getattr(module, attr) + except Exception: + continue + flush = getattr(value, "flush_cache", None) + if callable(flush): + try: + flush() + except Exception as exc: + warnings.warn( + f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +# Snapshot once at conftest import — these are the "clean" module states. +_LITELLM_STATE = _snapshot_mutable_state(litellm) +_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server) +_FLUSHABLE_CACHES = _collect_flushable_caches() @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): + """Reset mutable module state on litellm and proxy_server before each test. + + Replaces a previous importlib.reload(litellm) approach that cost ~17s + per test (re-executing the full litellm __init__ import chain). + + What IS reset: + - Top-level module attributes of type list / dict / set / tuple + / str / int / float / bool / bytes, and None-valued attributes. + These cover callback lists, general_settings, master_key, + premium_user, prisma_client, etc. — anything the old reload() reset + by re-executing the module body. + - Any module-level object instance that exposes flush_cache() (the + DualCache and LLMClientCache family), which handles cache state + that can't round-trip through deepcopy because of internal locks. + + What is NOT reset: + - Class instances without flush_cache() (e.g. ProxyLogging, + JWTHandler, FastAPI routers, loggers). If a test mutates such an + instance in-place (setattr on the instance, appending to one of + its internal lists, etc.), the mutation will leak into later tests. + Use pytest's monkeypatch.setattr() or a local fixture for those + cases — don't rely on this autouse fixture to undo them. """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. - """ - curr_dir = os.getcwd() # Get the current working directory - sys.path.insert( - 0, os.path.abspath("../..") - ) # Adds the project directory to the system path - - import litellm - from litellm import Router - - importlib.reload(litellm) - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - importlib.reload(litellm.proxy.proxy_server) - except Exception as e: - print(f"Error reloading litellm.proxy.proxy_server: {e}") - - import asyncio + _restore_mutable_state(litellm, _LITELLM_STATE) + _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) + _flush_caches(_FLUSHABLE_CACHES) loop = asyncio.get_event_loop_policy().new_event_loop() asyncio.set_event_loop(loop) - print(litellm) - # from litellm import Router, completion, aembedding, acompletion, embedding - yield - - # Teardown code (executes after the yield point) - loop.close() # Close the loop created earlier - asyncio.set_event_loop(None) # Remove the reference to the loop + try: + yield + finally: + loop.close() + asyncio.set_event_loop(None) def pytest_collection_modifyitems(config, items): diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index d3a9c94ea55..7a2a6f56d6f 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -579,6 +579,155 @@ def test_bedrock_messages_strips_output_config_with_output_format(): assert "output_format" not in result +def test_bedrock_messages_strips_context_management(): + """ + Ensure context_management is stripped from the request before sending to + Bedrock Invoke, which doesn't support this Anthropic-specific parameter. + + Claude Code sends context_management on every request; leaving it in the body + causes a 400 "context_management: Extra inputs are not permitted" from Bedrock. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, + } + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + "context_management" not in result + ), "context_management should be stripped — Bedrock Invoke rejects it" + assert result.get("max_tokens") == 4096 + + +def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): + """ + Bedrock Invoke rejects any top-level body field it doesn't recognize with + "Extra inputs are not permitted". Defend against that by filtering the + outgoing body to a Bedrock-supported allowlist — catches Anthropic-only + extensions (speed, mcp_servers, container, ...) and any future additions + Claude Code starts sending before we learn about them. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "temperature": 0.5, + "speed": "fast", + "mcp_servers": [{"type": "url", "url": "https://example.com"}], + "container": {"skills": []}, + "inference_geo": "us", + "output_config": {"effort": "low"}, + "context_management": {"edits": []}, + } + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + for bad in ( + "speed", + "mcp_servers", + "container", + "inference_geo", + "output_config", + "context_management", + "model", + "stream", + ): + assert bad not in result, f"{bad} should be stripped by the allowlist" + + # Supported fields pass through. + assert result["max_tokens"] == 4096 + assert result["temperature"] == 0.5 + assert result["anthropic_version"] == cfg.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION + # Every surviving key is in the allowlist. + assert set(result).issubset(cfg.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS) + + +def test_bedrock_messages_filters_user_provided_unsupported_beta_header(): + """ + In proxy deployments the client (e.g. Claude Code) doesn't know the backend + is Bedrock and may send Anthropic-direct beta headers Bedrock can't handle. + All betas must go through the provider mapping, not just auto-injected ones + — otherwise Bedrock 400s on the unsupported value. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = {"max_tokens": 128} + # `advisor-tool-2026-03-01` has no bedrock mapping entry → must be dropped. + # `context-1m-2025-08-07` does → must pass through. + headers = { + "anthropic-beta": "advisor-tool-2026-03-01,context-1m-2025-08-07", + } + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers=headers, + ) + + betas = result.get("anthropic_beta") or [] + assert ( + "advisor-tool-2026-03-01" not in betas + ), "user-provided beta not in the Bedrock mapping must be dropped" + assert ( + "context-1m-2025-08-07" in betas + ), "user-provided beta that IS in the Bedrock mapping should survive" + + +def test_bedrock_messages_renames_user_provided_aliased_beta_header(): + """ + Bedrock's config maps `advanced-tool-use-2025-11-20` to + `tool-search-tool-2025-10-19`. User-provided betas must go through the + rename too, not be forwarded under their Anthropic-direct spelling. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = {"max_tokens": 128} + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers=headers, + ) + + betas = result.get("anthropic_beta") or [] + assert ( + "advanced-tool-use-2025-11-20" not in betas + ), "Anthropic-direct spelling should be rewritten, not forwarded verbatim" + assert ( + "tool-search-tool-2025-10-19" in betas + ), "user-provided beta should be renamed to the Bedrock-side spelling" + + @pytest.mark.asyncio async def test_promote_message_stop_usage_preserves_message_delta_output_tokens(): """ diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py index 46fbd67902e..a20ec94a99d 100644 --- a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py +++ b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py @@ -95,7 +95,7 @@ class TestAnthropicBetaHeaderSupport: def test_messages_transformation_anthropic_beta(self): """Test that Messages API transformation includes anthropic_beta in request.""" config = AmazonAnthropicClaudeMessagesConfig() - headers = {"anthropic-beta": "output-128k-2025-02-19"} + headers = {"anthropic-beta": "context-1m-2025-08-07"} result = config.transform_anthropic_messages_request( model="anthropic.claude-haiku-4-5-20251001-v1:0", @@ -107,7 +107,7 @@ class TestAnthropicBetaHeaderSupport: assert "anthropic_beta" in result # Sort both arrays before comparing to avoid flakiness from ordering differences - assert sorted(result["anthropic_beta"]) == sorted(["output-128k-2025-02-19"]) + assert sorted(result["anthropic_beta"]) == sorted(["context-1m-2025-08-07"]) def test_converse_computer_use_compatibility(self): """Test that user anthropic_beta headers work with computer use tools."""