From 7b43f5981fa60a0eae3218ffe9b2ce907d64900e Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 20 Apr 2026 16:36:05 -0700
Subject: [PATCH 1/9] [Fix] CI: split test_proxy_utils.py into its own proxy-db
 matrix entry

The "remaining" proxy-db job was consistently timing out at ~98% because
--dist=loadscope pins every test in test_proxy_utils.py (168+ parametrized
tests) to a single xdist worker. 7 workers finished their files in ~15
minutes, then one worker ran alone for another 8+ minutes and hit the
30-minute job cap.

Give test_proxy_utils.py its own matrix entry so its tests spread across
all 8 workers, and add it to the "remaining" ignore list.
---
 .github/workflows/test-unit-proxy-db.yml | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

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

From ccf928361be6c36d8f5e6ac197775067cff442e4 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 20 Apr 2026 17:59:05 -0700
Subject: [PATCH 2/9] [Infra] Speed up proxy unit tests by replacing litellm
 reload with state snapshot
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

tests/proxy_unit_tests/conftest.py was calling importlib.reload(litellm) in an
autouse function-scoped fixture, which cost ~17s per test because it re-ran
the full litellm __init__ import chain. With 400+ proxy unit tests, this was
the single biggest driver of CI wall time — 18 of the top 20 slowest durations
in a typical run were just the 17s fixture setup.

Replace the reload with a snapshot-and-restore approach: snapshot the mutable
lists/dicts/sets on litellm and litellm.proxy.proxy_server once at conftest
import, then deep-copy that snapshot back before each test. Callback lists,
caches, router state, etc. still get reset between tests, but the expensive
import chain only runs once per worker.

Local measurement on test_proxy_utils.py: 188 tests in 3.50s (previously took
~15 minutes of CI wall time on a single worker).
---
 tests/proxy_unit_tests/conftest.py | 81 +++++++++++++++++++++---------
 1 file changed, 56 insertions(+), 25 deletions(-)

diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py
index 1421700c9a8..0cde5bdf28b 100644
--- a/tests/proxy_unit_tests/conftest.py
+++ b/tests/proxy_unit_tests/conftest.py
@@ -1,6 +1,7 @@
 # conftest.py
 
-import importlib
+import asyncio
+import copy
 import os
 import sys
 
@@ -9,40 +10,70 @@ import pytest
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
+
 import litellm
+import litellm.proxy.proxy_server
+
+
+def _snapshot_mutable_state(module):
+    """Deep-copy every list/dict/set module attribute for later restore.
+
+    Classes, functions, submodules and primitives are skipped — only the
+    collections that tests mutate (callbacks, caches, routers, etc.) need
+    per-test isolation.
+    """
+    snapshot = {}
+    for attr in list(vars(module)):
+        if attr.startswith("_"):
+            continue
+        try:
+            value = getattr(module, attr)
+        except Exception:
+            continue
+        if isinstance(value, (list, dict, set)):
+            try:
+                snapshot[attr] = copy.deepcopy(value)
+            except Exception:
+                # Unpickleable collections (e.g. holding open clients) can't
+                # round-trip through deepcopy; skip them rather than crash.
+                pass
+    return snapshot
+
+
+def _restore_mutable_state(module, snapshot):
+    for attr, default in snapshot.items():
+        try:
+            setattr(module, attr, copy.deepcopy(default))
+        except Exception:
+            pass
+
+
+# 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)
 
 
 @pytest.fixture(scope="function", autouse=True)
 def setup_and_teardown():
     """
-    This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained.
+    Reset mutable module state on litellm and proxy_server before every test.
+
+    Replaces a previous importlib.reload(litellm) approach that cost ~17s
+    per test (re-executing the full litellm __init__ import chain). The
+    snapshot-and-restore below only touches collections that actually leak
+    across tests — callbacks, caches, router, etc. — and is effectively
+    instantaneous.
     """
-    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)
 
     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):

From 5411ebedae0f77ed0832289ae877a75a1cca836f Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 20 Apr 2026 21:03:07 -0700
Subject: [PATCH 3/9] [Fix] conftest snapshot: also reset scalar module
 attributes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The previous snapshot only tracked list/dict/set values. Tests mutate
scalar module attrs too — master_key, premium_user, prisma_client — and
importlib.reload used to reset those implicitly. Under the snapshot
approach they were leaking between tests, so test_active_callbacks
failed in CI with "No api key passed in." once an earlier test left
master_key set to sk-1234.

Expand the snapshot to cover primitives (str/int/float/bool/bytes/tuple)
and None-valued attributes. Complex object instances are still skipped
to avoid deepcopy issues.
---
 tests/proxy_unit_tests/conftest.py | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py
index 0cde5bdf28b..544b6a0b421 100644
--- a/tests/proxy_unit_tests/conftest.py
+++ b/tests/proxy_unit_tests/conftest.py
@@ -15,12 +15,17 @@ import litellm
 import litellm.proxy.proxy_server
 
 
-def _snapshot_mutable_state(module):
-    """Deep-copy every list/dict/set module attribute for later restore.
+_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes)
 
-    Classes, functions, submodules and primitives are skipped — only the
-    collections that tests mutate (callbacks, caches, routers, etc.) need
-    per-test isolation.
+
+def _snapshot_mutable_state(module):
+    """Snapshot every module attribute that importlib.reload would have reset.
+
+    Covers the top-level assignments that tests mutate — collections
+    (callbacks, caches, general_settings) plus scalar flags (master_key,
+    premium_user, etc.) that gate auth and feature behavior. Classes,
+    functions, submodules and complex object instances are skipped: those
+    either aren't meant to be reset or can't round-trip through deepcopy.
     """
     snapshot = {}
     for attr in list(vars(module)):
@@ -30,12 +35,12 @@ def _snapshot_mutable_state(module):
             value = getattr(module, attr)
         except Exception:
             continue
-        if isinstance(value, (list, dict, set)):
+        if value is None or isinstance(value, _SNAPSHOT_TYPES):
             try:
                 snapshot[attr] = copy.deepcopy(value)
             except Exception:
-                # Unpickleable collections (e.g. holding open clients) can't
-                # round-trip through deepcopy; skip them rather than crash.
+                # Skip anything that can't round-trip through deepcopy
+                # rather than crash collection.
                 pass
     return snapshot
 

From 0f5d503169aaad1f8bb3b221978c8030449b669b Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 20 Apr 2026 22:09:54 -0700
Subject: [PATCH 4/9] fix(ci): make e2e_ui_testing actually test the freshly
 built UI bundle
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Build UI from source step used:

    cp -r out/ ../../litellm/proxy/_experimental/out/

GNU cp (CircleCI's Ubuntu image, coreutils 8.32) interprets this as
copy the source directory as a CHILD of the destination when the
destination already exists — so the command silently created
litellm/proxy/_experimental/out/out/ instead of replacing the served
bundle at litellm/proxy/_experimental/out/*.

The proxy continued serving whatever bundle was checked in, so every
e2e_ui_testing run between this job's introduction (d09d98a70a,
2026-04-08) and the bundle-rebuild commit (de790fd273, 2026-04-18) was
effectively testing a STALE bundle — not the fresh build. That is why
the double-prefix regression (NEXT_PUBLIC_BASE_URL="ui/" combined with
networking.tsx reading the env var) was never caught in CI even though
the source contained the trigger the whole time: the bundle the proxy
served never picked up the source change.

Replace cp -r with rm + mv so the destination is cleanly swapped.

Verified end-to-end on an Ubuntu 22.04 / GNU coreutils 8.32 container:
- Before fix: fresh build has 9 "ui/" literals in chunks; after cp,
  _experimental/out/*  still has 0 (stale); _experimental/out/out/ is a
  nested dir the proxy does not serve.
- After fix: _experimental/out/*  has 9 "ui/" literals — the proxy now
  serves the freshly built (broken, in this repro) bundle, so
  globalSetup fails at login and every spec is blocked. Removing the
  bug from .env.production and rebuilding brings the count back to 0
  and the suite passes.

No spec changes, no fixtures, no new infrastructure. The existing
Playwright suite already catches this class of regression via the
login flow in globalSetup; it just needs the CI to actually hand it
the freshly built bundle.
---
 .circleci/config.yml | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

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"

From 4b3f5d7f81d38e2019882a63bbe91411b2e31065 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 20 Apr 2026 22:19:36 -0700
Subject: [PATCH 5/9] [Fix] conftest: flush cache instances and warn on silent
 skips
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Addresses review feedback on the snapshot approach:

1. Class-instance mutable state
   The snapshot only covers primitives + collections + None. Class
   instances (DualCache, LLMClientCache) weren't reset between tests,
   so in-place cache mutations could leak. Can't deepcopy these — they
   hold thread locks — but they expose flush_cache(). Collect every
   module attribute whose value implements flush_cache() at conftest
   import, and invoke it per-test alongside the snapshot restore.

2. Silent skips are now warnings
   _snapshot_mutable_state and _restore_mutable_state previously
   swallowed exceptions, so if a future attr gained a property without
   a setter (or other non-round-trippable state), an isolation gap
   would have no signal. Emit warnings.warn on each failure path.

3. Docstring
   Explicitly documents what IS and IS NOT reset, and tells authors to
   use monkeypatch.setattr() for in-place mutations of instances
   without flush_cache() (ProxyLogging, JWTHandler, etc.).
---
 tests/proxy_unit_tests/conftest.py | 99 ++++++++++++++++++++++++------
 1 file changed, 79 insertions(+), 20 deletions(-)

diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py
index 544b6a0b421..a0326f64ed7 100644
--- a/tests/proxy_unit_tests/conftest.py
+++ b/tests/proxy_unit_tests/conftest.py
@@ -2,8 +2,10 @@
 
 import asyncio
 import copy
+import inspect
 import os
 import sys
+import warnings
 
 import pytest
 
@@ -15,33 +17,34 @@ 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):
-    """Snapshot every module attribute that importlib.reload would have reset.
-
-    Covers the top-level assignments that tests mutate — collections
-    (callbacks, caches, general_settings) plus scalar flags (master_key,
-    premium_user, etc.) that gate auth and feature behavior. Classes,
-    functions, submodules and complex object instances are skipped: those
-    either aren't meant to be reset or can't round-trip through deepcopy.
-    """
+    """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:
+        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:
-                # Skip anything that can't round-trip through deepcopy
-                # rather than crash collection.
-                pass
+            except Exception as exc:
+                warnings.warn(
+                    f"conftest: could not snapshot {module.__name__}.{attr}: {exc}",
+                    stacklevel=2,
+                )
     return snapshot
 
 
@@ -49,28 +52,84 @@ 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:
-            pass
+            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 every test.
+    """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). The
-    snapshot-and-restore below only touches collections that actually leak
-    across tests — callbacks, caches, router, etc. — and is effectively
-    instantaneous.
+    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.
     """
     _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)

From 583bdd34a237c4995d1cce94eb63d7cd5a3d1a52 Mon Sep 17 00:00:00 2001
From: SwiftWinds <12981958+SwiftWinds@users.noreply.github.com>
Date: Mon, 20 Apr 2026 20:20:47 -0700
Subject: [PATCH 6/9] fix(bedrock): allowlist Bedrock Invoke body fields and
 filter all anthropic-beta values
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two fail-safes for the /v1/messages → Bedrock Invoke pass-through so new
Anthropic-only extensions Claude Code starts sending can't reach Bedrock
and trigger a 400 "Extra inputs are not permitted":

1. Top-level body fields are filtered to a typed allowlist. New
   `BedrockInvokeAnthropicMessagesRequest` TypedDict (in
   `litellm/types/llms/bedrock.py`) captures the Bedrock Invoke Anthropic
   Messages body schema; the runtime allowlist is derived from its
   `__annotations__` so the type and the filter can't drift. Anchored to
   the AWS reference page in docstrings + transform comment. An
   exact-set test pins the resolved allowlist so any future edit forces
   conscious review.

   Drops context_management, output_config, speed, mcp_servers,
   container, inference_geo, internal litellm_metadata, and any future
   Anthropic addition. output_format stays as an active inline-schema
   conversion (not just a strip).

2. The anthropic-beta header list is filtered + transformed against the
   bedrock mapping for ALL betas, not just auto-injected ones. The
   previous code union'd user-provided betas back in unfiltered, so a
   client on a new Anthropic-direct beta (e.g. advisor-tool-…,
   context-management-…) could still pin the request to fail. In a proxy
   context the client can't know the backend is Bedrock; the provider
   mapping is authoritative. User-provided drops are logged at WARNING
   so intentional overrides leave a breadcrumb.

Updates one existing test that happened to assert on the old buggy
pass-through (it used output-128k-2025-02-19, which is null in the
bedrock mapping and would 400 at runtime); rewrote it against a
bedrock-supported beta.

Scope: messages/invoke only. The same user-beta bypass exists in
chat/invoke but that's a different code path with different
user-expectation trade-offs — follow-up.
---
 .../anthropic_claude3_transformation.py       |  46 +++++-
 litellm/types/llms/bedrock.py                 |  44 ++++++
 .../test_anthropic_claude3_transformation.py  | 149 ++++++++++++++++++
 .../bedrock/test_anthropic_beta_support.py    |   4 +-
 4 files changed, 233 insertions(+), 10 deletions(-)

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/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."""

From 11b776935d4878b513772ee5b0cde1105c943d0d Mon Sep 17 00:00:00 2001
From: SwiftWinds <12981958+SwiftWinds@users.noreply.github.com>
Date: Tue, 21 Apr 2026 11:39:11 -0700
Subject: [PATCH 7/9] chore: make `uv` newer than 0.10 allowable

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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]

From b39f210a6cf5ea0aaa4de86d835b21d12d93590f Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Tue, 21 Apr 2026 12:00:23 -0700
Subject: [PATCH 8/9] [Infra] Add freshness and destructive guards to migration
 workflow

Generating a migration from a stale branch could silently emit DROP
COLUMN for columns the stale branch did not know about, and the
script would write that SQL to a new migration file with no warning.

Adds two guards to ci_cd/run_migration.py:

- Branch freshness check: fetches origin/<base-branch> and exits 3 if
  HEAD is behind. Default base is litellm_internal_staging. New
  flags: --base-branch, --skip-freshness-check.
- Destructive guard: refuses (exit 2) if the generated diff contains
  DROP COLUMN / DROP TABLE / DROP INDEX, unless --allow-destructive
  is passed.

Refusal banners include guidance and an explicit callout instructing
AI agents not to auto-bypass the flags. Also treats Prisma's
"-- This is an empty migration." output as a no-op rather than
writing an empty file.

Updates litellm-proxy-extras/migration_runbook.md with the new
workflow, flag documentation, and agent warnings.
---
 ci_cd/run_migration.py                    | 298 ++++++++++++++++++++--
 litellm-proxy-extras/migration_runbook.md |  50 +++-
 2 files changed, 329 insertions(+), 19 deletions(-)

diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py
index 29101bf9505..1cbe9fb59d5 100644
--- a/ci_cd/run_migration.py
+++ b/ci_cd/run_migration.py
@@ -1,22 +1,230 @@
+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 or DROP TABLE."""
+    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 <name> 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/<base_branch> 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: "main").
+        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 +267,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 +296,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 +318,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/<base-branch> 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 <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. 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 <name>` — 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
 

From 5b007add62e2d2bd31eaad6dbf4e988c7e0246ea Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Tue, 21 Apr 2026 12:07:19 -0700
Subject: [PATCH 9/9] [Docs] Fix docstring inaccuracies in run_migration.py

- _find_destructive_statements: add DROP INDEX to the docstring (the
  regex already detects it; only the docstring lagged).
- create_migration: correct the base_branch default documented in the
  docstring from "main" to "litellm_internal_staging".
---
 ci_cd/run_migration.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py
index 1cbe9fb59d5..feec4046ee1 100644
--- a/ci_cd/run_migration.py
+++ b/ci_cd/run_migration.py
@@ -15,7 +15,7 @@ DEFAULT_BASE_BRANCH = "litellm_internal_staging"
 
 
 def _find_destructive_statements(sql: str) -> list:
-    """Return SQL lines containing DROP COLUMN or DROP TABLE."""
+    """Return SQL lines containing DROP COLUMN, DROP TABLE, or DROP INDEX."""
     return [
         line.strip() for line in sql.splitlines() if DESTRUCTIVE_PATTERN.search(line)
     ]
@@ -210,7 +210,8 @@ def create_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: "main").
+        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.
     """